`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.

Written by the indexing model from the issue text.