`model_infrastructure` model/pipeline review

# `model_infrastructure` model/pipeline review

Commit tested: `0f1abc4ae8b0eb2a3b40e82a310507281144c423`

Review performed against the repository review rules.

All requested `model_infrastructure` files were reviewed:

```text
src/diffusers/models/__init__.py
src/diffusers/models/_modeling_parallel.py
src/diffusers/models/activations.py
src/diffusers/models/attention.py
src/diffusers/models/attention_dispatch.py
src/diffusers/models/attention_flax.py
src/diffusers/models/attention_processor.py
src/diffusers/models/auto_model.py
src/diffusers/models/cache_utils.py
src/diffusers/models/downsampling.py
src/diffusers/models/embeddings.py
src/diffusers/models/embeddings_flax.py
src/diffusers/models/lora.py
src/diffusers/models/model_loading_utils.py
src/diffusers/models/modeling_flax_pytorch_utils.py
src/diffusers/models/modeling_flax_utils.py
src/diffusers/models/modeling_outputs.py
src/diffusers/models/modeling_pytorch_flax_utils.py
src/diffusers/models/modeling_utils.py
src/diffusers/models/normalization.py
src/diffusers/models/resnet.py
src/diffusers/models/resnet_flax.py
src/diffusers/models/upsampling.py
src/diffusers/models/vae_flax.py
src/diffusers/models/vq_model.py
```

Duplicate-search status: searched existing GitHub Issues and PRs for `model_infrastructure`, affected file/function names, and the specific failures below. No direct duplicates were found. Related but not duplicate: `huggingface/diffusers#12409` and `#12533` around distributed/context-parallel availability.

Test coverage status: fast/unit coverage exists for parts of AutoModel, layer helpers, cache utilities, attention backends, and parallelism helpers, but the slow/integration coverage gaps listed in Issue 9 remain.

## Issue 1: `AutoModel.from_pretrained` rejects `PathLike` pipeline roots

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/auto_model.py#L343-L345

Problem:
`from_pretrained()` accepts `Union[str, os.PathLike]`, but the model-index path builds `_diffusers_load_id` with `"|".join(parts)` while `parts` may contain a `Path` object. Passing a `Path` root with `subfolder` fails before loading.

Impact:
Documented/local loading behavior breaks for users who pass `pathlib.Path`, and the failure is a low-level `TypeError` rather than a loading error.

Reproduction:
```python
import json
import tempfile
from pathlib import Path
from diffusers import AutoModel, UNet2DModel

with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as d:
root = Path(d)
unet_dir = root / "unet"
UNet2DModel(
sample_size=4, in_channels=1, out_channels=1, layers_per_block=1,
block_out_channels=(4,), down_block_types=("DownBlock2D",),
up_block_types=("UpBlock2D",), norm_num_groups=1,
).save_pretrained(unet_dir, safe_serialization=False)
(root / "model_index.json").write_text(json.dumps({"_class_name": "DummyPipeline", "unet": ["diffusers", "UNet2DModel"]}))

try:
AutoModel.from_pretrained(root, subfolder="unet", use_safetensors=False)
except Exception as e:
print(type(e).__name__)
print(str(e).splitlines()[0])
```

Relevant precedent:
Other loading paths normalize path-like inputs before string operations.

Suggested fix:
```python
load_id = "|".join("null" if p is None else str(p) for p in parts)
```

## Issue 2: `AutoModel.from_pretrained` leaks `config_name` after model-index loading

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/auto_model.py#L278-L293

Problem:
`from_pretrained()` mutates the class attribute `cls.config_name` to `"model_index.json"` and does not restore it after a successful model-index load. A later `AutoModel.from_config(model_dir)` then looks for `model_index.json` inside a plain model directory instead of `config.json`.

Impact:
One successful `AutoModel.from_pretrained()` call can change later behavior process-wide.

Reproduction:
```python
import json
import tempfile
from pathlib import Path
from diffusers import AutoModel, UNet2DModel

with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as d:
root = Path(d)
unet_dir = root / "unet"
UNet2DModel(
sample_size=4, in_channels=1, out_channels=1, layers_per_block=1,
block_out_channels=(4,), down_block_types=("DownBlock2D",),
up_block_types=("UpBlock2D",), norm_num_groups=1,
).save_pretrained(unet_dir, safe_serialization=False)
(root / "model_index.json").write_text(json.dumps({"_class_name": "DummyPipeline", "unet": ["diffusers", "UNet2DModel"]}))

AutoModel.from_pretrained(str(root), subfolder="unet", use_safetensors=False)
print("leaked config_name:", AutoModel.config_name)
try:
AutoModel.from_config(str(unet_dir))
except Exception as e:
print(type(e).__name__, str(e).splitlines()[0])
```

Relevant precedent:
Config name overrides should be local to the load attempt, not stored on the public class.

Suggested fix:
```python
old_config_name = cls.config_name
try:
cls.config_name = "model_index.json"
model_index, kwargs = cls.load_config(...)

# existing model-index handling
finally:
cls.config_name = old_config_name
```

## Issue 3: `CacheMixin.cache_context` leaves stale hook state after exceptions

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/cache_utils.py#L154-L164

Problem:
`cache_context()` clears the hook context only after the `yield`. If the wrapped forward pass raises, `_current_context` remains set on stateful hooks.

Impact:
A failed denoising step can leak cached/stateful hook state into later calls, and `get_state()` can incorrectly succeed outside a context.

Reproduction:
```python
import torch
from diffusers.hooks.hooks import BaseState, HookRegistry, ModelHook, StateManager
from diffusers.models.cache_utils import CacheMixin

class State(BaseState):
def reset(self):
pass

class StatefulHook(ModelHook):
_is_stateful = True
def __init__(self):
super().__init__()
self.state_manager = StateManager(State)
def reset_state(self, module):
self.state_manager.reset()

class Model(torch.nn.Module, CacheMixin):
def forward(self, x):
return x

model = Model()
hook = StatefulHook()
HookRegistry.check_if_exists_or_initialize(model).register_hook(hook, "stateful")

try:
with model.cache_context("failed-call"):
raise RuntimeError("simulate an interrupted denoise step")
except RuntimeError:
pass

print(hook.state_manager._current_context)
print(type(hook.state_manager.get_state()).__name__)
```

Relevant precedent:
Context managers that mutate global or hook state should restore that state in `finally`.

Suggested fix:
```python
registry._set_context(name)
try:
yield
finally:
registry._set_context(None)
```

## Issue 4: FIR up/downsampling fails on `bfloat16` and `float16` inputs

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/downsampling.py#L217-L239
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/downsampling.py#L386-L397
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/upsampling.py#L261-L312
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/upsampling.py#L502-L513

Problem:
The FIR helpers create kernels as default `float32` tensors and pass them to convolution against low-precision hidden states.

Impact:
Low-precision model execution can fail in FIR downsample/upsample paths with dtype mismatch.

Reproduction:
```python
import torch
from diffusers.models.downsampling import downsample_2d
from diffusers.models.upsampling import upsample_2d

x = torch.randn(1, 1, 4, 4, dtype=torch.bfloat16)
for fn in (downsample_2d, upsample_2d):
try:
print(fn.__name__, fn(x).dtype)
except Exception as e:
print(fn.__name__, type(e).__name__, str(e).splitlines()[0])
```

Relevant precedent:
`Upsample2D` already has low-precision test coverage; FIR helper/module paths need equivalent dtype handling.

Suggested fix:
```python
kernel = torch.as_tensor(kernel, device=hidden_states.device, dtype=torch.float32)
if kernel.ndim == 1:
kernel = torch.outer(kernel, kernel)
kernel = kernel / kernel.sum() * gain
kernel = kernel.to(dtype=hidden_states.dtype)
```

## Issue 5: `enable_parallelism` uses the wrong distributed guard

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/modeling_utils.py#L1595-L1598

Problem:
The guard uses `and` instead of `or`: `not is_available() and not is_initialized()`. When distributed is unavailable, this may call `is_initialized()` anyway; when distributed is available but uninitialized, the guard can fail to raise.

Impact:
Users get backend-specific errors instead of the intended clear `RuntimeError`, or parallelism proceeds before `torch.distributed` is initialized.

Reproduction:
```python
from diffusers import ContextParallelConfig, UNet2DModel

model = UNet2DModel(
sample_size=4, in_channels=1, out_channels=1, layers_per_block=1,
block_out_channels=(4,), down_block_types=("DownBlock2D",),
up_block_types=("UpBlock2D",), norm_num_groups=1,
)
try:
model.enable_parallelism(config=ContextParallelConfig(ring_degree=2))
except Exception as e:
print(type(e).__name__)
print(str(e).splitlines()[0])
```

Relevant precedent:
Distributed APIs normally require both availability and initialization before model wrapping.

Suggested fix:
```python
if not torch.distributed.is_available() or not torch.distributed.is_initialized():
raise RuntimeError("torch.distributed must be available and initialized before calling `enable_parallelism`.")
```

## Issue 6: Legacy `Attention.set_use_xla_flash_attention` checks the function object instead of calling it

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/attention_processor.py#L311-L342

Problem:
The method checks `if is_torch_xla_available:` instead of `if is_torch_xla_available():`, so environments without `torch_xla` enter the XLA version checks and can raise `InvalidVersion`. The method also raises strings in two branches.

Impact:
Users enabling XLA flash attention get confusing exceptions instead of a clear dependency/version error.

Reproduction:
```python
from diffusers.models.attention_processor import Attention

attn = Attention(query_dim=4, heads=1, dim_head=4)
try:
attn.set_use_xla_flash_attention(True)
except Exception as e:
print(type(e).__name__)
print(str(e).splitlines()[0])
```

Relevant precedent:
`AttentionModuleMixin.set_use_xla_flash_attention()` in `src/diffusers/models/attention.py` calls `is_torch_xla_available()` and raises `ImportError`.

Suggested fix:
```python
if use_xla_flash_attention:
if not is_torch_xla_available():
raise ImportError("torch_xla is not available")
if is_torch_xla_version("<", "2.3"):
raise ImportError("flash attention pallas kernel is supported from torch_xla version 2.3")
```

## Issue 7: `FlaxModelMixin.from_pretrained(config=...)` can reference `unused_kwargs` before assignment

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/modeling_flax_utils.py#L307-L322

Problem:
`unused_kwargs` is assigned only in the `config is None` branch, but later used unconditionally. Passing a preloaded config can therefore raise `UnboundLocalError`. This was statically verified; the local `.venv` does not include Flax.

Impact:
Callers that pass an already loaded Flax config cannot reliably use `from_pretrained(config=...)`.

Reproduction:
```python
from diffusers import FlaxAutoencoderKL

config = {
"in_channels": 3,
"out_channels": 3,
"down_block_types": ("DownEncoderBlock2D",),
"up_block_types": ("UpDecoderBlock2D",),
"block_out_channels": (32,),
"layers_per_block": 1,
"act_fn": "silu",
"latent_channels": 4,
"norm_num_groups": 32,
"sample_size": 32,
"scaling_factor": 0.18215,
}

try:
FlaxAutoencoderKL.from_pretrained("does-not-matter", config=config, local_files_only=True)
except Exception as e:
print(type(e).__name__, str(e).splitlines()[0])
```

Relevant precedent:
The PyTorch loading path preserves extra kwargs regardless of whether config is loaded internally or supplied by the caller.

Suggested fix:
```python
if config is None:
config, unused_kwargs = cls.load_config(..., **kwargs)
else:
unused_kwargs = kwargs
```

## Issue 8: Sinusoidal embedding helper defaults to `float64` on non-MPS devices

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/embeddings.py#L321-L356

Problem:
For `output_type="pt"`, the helper defaults `dtype` to `torch.float64` unless the device is MPS. The review rules call out NPU float64 limitations, but NPU takes the same non-MPS branch.

Impact:
Embedding creation can fail on NPU, and CPU/GPU callers get an unexpectedly high-precision tensor unless they override dtype.

Reproduction:
```python
import torch
from diffusers.models.embeddings import get_1d_sincos_pos_embed_from_grid

pos = torch.arange(4, dtype=torch.float32)
emb = get_1d_sincos_pos_embed_from_grid(8, pos, output_type="pt")
print(emb.dtype)

# On NPU this same non-MPS branch attempts float64 tensor creation:
# pos = pos.to("npu")
# get_1d_sincos_pos_embed_from_grid(8, pos, output_type="pt")
```

Relevant precedent:
NPU-safe code paths in the repository avoid implicit float64 tensors.

Suggested fix:
```python
if dtype is None:
dtype = torch.float32
```

## Issue 9: Slow/integration coverage is missing for shared model infrastructure regressions

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/models/test_models_auto.py#L15-L82
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/models/test_layers_utils.py#L112-L327
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/models/testing_utils/cache.py#L194-L221

Problem:
Fast coverage exists, but it does not cover the integration-style failures found above: AutoModel `PathLike`/state leakage across calls, cache cleanup after exceptions, FIR low-precision paths, and unavailable/uninitialized distributed backends. No dedicated slow test covers these shared infrastructure behaviors through a tiny pipeline/model load.

Impact:
Regressions in shared model infrastructure can affect many model and pipeline families without being caught by family-specific fast tests.

Reproduction:
```python
from pathlib import Path

interesting = [
Path("tests/models/test_models_auto.py"),
Path("tests/models/test_layers_utils.py"),
Path("tests/models/testing_utils/cache.py"),
Path("tests/others/test_attention_backends.py"),
]
for path in interesting:
text = path.read_text()
print(path, "@slow" in text or "slow(" in text)
```

Relevant precedent:
Other model/pipeline families combine focused fast tests with slow tests that exercise actual loading/runtime behavior.

Suggested fix:
Add fast regression tests for Issues 1-8. Add at least one slow/integration test using a tiny Hub fixture or saved tiny local pipeline to exercise `AutoModel.from_pretrained`, cache/offload/attention setup, and shared dtype/device behavior through public APIs.

Start with the affected files and line references under each numbered issue, especially auto_model.py, cache_utils.py, downsampling.py, upsampling.py, modeling_utils.py, attention_processor.py, and modeling_flax_utils.py. Run the supplied reproductions and inspect the existing fast/unit coverage for AutoModel, cache utilities, attention backends, and parallelism helpers. Done means the reported failures are fixed and corresponding regression coverage is added.

Vom Indexierungsmodell aus dem Issue-Text verfasst.