model_infrastructure model/pipeline review

Commit tested: 0f1abc4ae8b0eb2a3b40e82a310507281144c423

Review performed against the repository review rules.

All requested model_infrastructure files were reviewed:

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:

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:

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:

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:

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:

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:

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:

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:

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:

registry._set_context(name)
try:
    yield
finally:
    registry._set_context(None)

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

Affected code:

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:

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:

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:

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:

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:

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:

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:

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:

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:

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:

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:

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:

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:

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:

if dtype is None:
    dtype = torch.float32

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

Affected code:

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:

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.