`model_unets_shared` model/pipeline review

# `model_unets_shared` model/pipeline review

Commit tested: `0f1abc4ae8b0eb2a3b40e82a310507281144c423`

Review performed against the repository review rules.

Files reviewed: `unet_1d.py`, `unet_1d_blocks.py`, `unet_2d.py`, `unet_2d_blocks.py`, `unet_2d_blocks_flax.py`, `unet_3d_blocks.py`, `uvit_2d.py`, `unet_2d_condition.py`, `unet_2d_condition_flax.py`, `unet_3d_condition.py`.

Duplicate search: searched GitHub Issues/PRs for `model_unets_shared`, affected class/file names, and the specific failure modes. Exact duplicate found only for the UViT checkpointing failure: https://github.com/huggingface/diffusers/issues/11214. No exact duplicate found for the scalar timestep truncation, `UNet3DConditionModel` mask drop, or missing coverage; old issue #1890 is related to general attention masking but not this 3D UNet path.

## Issue 1: `UNet3DConditionModel` accepts `attention_mask` but drops it

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/unets/unet_3d_condition.py#L483-L547
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/unets/unet_3d_blocks.py#L501-L532
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/unets/unet_3d_blocks.py#L728-L778

Problem:
`UNet3DConditionModel.forward()` documents and prepares `attention_mask`, but `CrossAttnDownBlock3D` and `CrossAttnUpBlock3D` explicitly do not use it. The mid block also declares `attention_mask` and does not pass it to the spatial transformer. Padding masks therefore silently have no effect.

Impact:
Batched prompts with padding can attend to discarded text tokens in video UNets. The API suggests masking is honored, so users get silently incorrect conditioning.

Reproduction:
```python
import torch
from diffusers import UNet3DConditionModel

torch.manual_seed(0)
model = UNet3DConditionModel(
block_out_channels=(8, 16),
norm_num_groups=4,
down_block_types=("CrossAttnDownBlock3D", "DownBlock3D"),
up_block_types=("UpBlock3D", "CrossAttnUpBlock3D"),
cross_attention_dim=8,
attention_head_dim=2,
out_channels=4,
in_channels=4,
layers_per_block=1,
sample_size=16,
).eval()

sample = torch.randn(1, 4, 2, 16, 16)
encoder_hidden_states = torch.randn(1, 4, 8)

with torch.no_grad():
keep = model(sample, torch.tensor([10]), encoder_hidden_states, attention_mask=torch.ones(1, 4)).sample
drop = model(sample, torch.tensor([10]), encoder_hidden_states, attention_mask=torch.zeros(1, 4)).sample

print((keep - drop).abs().max().item()) # 0.0: mask is ignored
```

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/unets/unet_2d_condition.py#L1074-L1076
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/unets/unet_2d_blocks.py#L1264-L1275

Suggested fix:
```python
# In UNet3DConditionModel.forward(), after num_frames is known:
if attention_mask is not None:
attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0
attention_mask = attention_mask.unsqueeze(1)
attention_mask = attention_mask.repeat_interleave(
num_frames, dim=0, output_size=attention_mask.shape[0] * num_frames
)

# In 3D cross-attn blocks, pass it to the spatial Transformer2DModel cross-attn mask:
hidden_states = attn(
hidden_states,
encoder_hidden_states=encoder_hidden_states,
encoder_attention_mask=attention_mask,
cross_attention_kwargs=cross_attention_kwargs,
return_dict=False,
)[0]
```

## Issue 2: scalar float timesteps are truncated in unconditional UNets

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/unets/unet_1d.py#L228-L234
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/unets/unet_2d.py#L278-L286
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/unets/unet_2d_condition_flax.py#L371-L376

Problem:
`UNet1DModel` and `UNet2DModel` type their public `timestep` as `torch.Tensor | float | int`, but scalar Python floats are converted with integer dtype. `UNet2DModel` then divides Fourier outputs by the integer timestep, so `1e-4` becomes `0` and produces non-finite output. The Flax model has the same scalar-float-to-`int32` pattern.

Impact:
Calling the public API with a scalar float timestep changes semantics versus passing `torch.tensor([float])`; VE/NCSN-style small sigma timesteps can produce NaNs.

Reproduction:
```python
import torch
from diffusers import UNet1DModel, UNet2DModel

unet2d = UNet2DModel(
sample_size=8,
in_channels=3,
out_channels=3,
block_out_channels=(8,),
layers_per_block=1,
down_block_types=("DownBlock2D",),
up_block_types=("UpBlock2D",),
norm_num_groups=4,
time_embedding_type="fourier",
).eval()

sample2d = torch.randn(1, 3, 8, 8)
with torch.no_grad():
print(torch.isfinite(unet2d(sample2d, 1e-4).sample).all().item()) # False
print(torch.isfinite(unet2d(sample2d, torch.tensor([1e-4])).sample).all().item()) # True

unet1d = UNet1DModel(
sample_size=16,
in_channels=4,
out_channels=4,
block_out_channels=(8, 8),
down_block_types=("DownResnetBlock1D", "DownResnetBlock1D"),
up_block_types=("UpResnetBlock1D",),
mid_block_type="MidResTemporalBlock1D",
out_block_type="OutConv1DBlock",
time_embedding_type="positional",
use_timestep_embedding=True,
norm_num_groups=4,
act_fn="swish",
).eval()

sample1d = torch.randn(2, 4, 16)
with torch.no_grad():
a = unet1d(sample1d, 0.5).sample
b = unet1d(sample1d, torch.tensor([0.5])).sample
print((a - b).abs().max().item()) # non-zero: scalar path was truncated
```

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/unets/unet_2d_condition.py#L851-L873
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/unets/unet_3d_condition.py#L548-L569

Suggested fix:
```python
if not torch.is_tensor(timesteps):
is_mps = sample.device.type == "mps"
is_npu = sample.device.type == "npu"
if isinstance(timesteps, float):
dtype = torch.float32 if (is_mps or is_npu) else torch.float64
else:
dtype = torch.int32 if (is_mps or is_npu) else torch.int64
timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
```

## Issue 3: Existing duplicate: `UVit2DModel` gradient checkpointing crashes

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/unets/uvit_2d.py#L180-L195

Problem:
Duplicate of https://github.com/huggingface/diffusers/issues/11214. `UVit2DModel` sets `_supports_gradient_checkpointing = True`, but its checkpoint wrapper only accepts `*args` and is then called with keyword arguments.

Impact:
Training or fine-tuning with `enable_gradient_checkpointing()` fails immediately.

Reproduction:
```python
import torch
from diffusers import UVit2DModel

model = UVit2DModel(
hidden_size=8,
cond_embed_dim=4,
micro_cond_encode_dim=2,
micro_cond_embed_dim=4,
encoder_hidden_size=6,
vocab_size=16,
codebook_size=15,
in_channels=4,
block_out_channels=4,
num_res_blocks=1,
block_num_heads=1,
num_hidden_layers=1,
num_attention_heads=1,
intermediate_size=16,
sample_size=2,
)
model.enable_gradient_checkpointing()
model(torch.randint(0, 16, (1, 2, 2)), torch.randn(1, 3, 6), torch.randn(1, 4), torch.tensor([[1.0, 2.0]]))
```

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_flux.py#L715-L725

Suggested fix:
```python
if torch.is_grad_enabled() and self.gradient_checkpointing:
hidden_states = self._gradient_checkpointing_func(
layer,
hidden_states,
None,
encoder_hidden_states,
None,
None,
cross_attention_kwargs,
None,
{"pooled_text_emb": pooled_text_emb},
)
else:
hidden_states = layer(
hidden_states,
encoder_hidden_states=encoder_hidden_states,
cross_attention_kwargs=cross_attention_kwargs,
added_cond_kwargs={"pooled_text_emb": pooled_text_emb},
)
```

## Issue 4: Missing fast/slow coverage for parts of the family

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/unets/uvit_2d.py#L38-L39
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/unets/unet_2d_condition_flax.py#L51-L124
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/models/unets/test_models_unet_3d_condition.py#L35-L35

Problem:
`UVit2DModel` has no model or pipeline tests under `tests/`. `FlaxUNet2DConditionModel` is only covered by dummy-object checks, not fast or slow behavior tests. `UNet3DConditionModel` has fast model tests, but no `@slow` integration coverage.

Impact:
The UViT checkpointing regression above is currently unguarded, Flax UNet serialization/runtime can regress unnoticed while still publicly exported, and 3D UNet published-checkpoint behavior is not covered by slow tests.

Reproduction:
```python
from pathlib import Path

for needle in ["UVit2DModel", "FlaxUNet2DConditionModel", "UNet3DConditionModel"]:
hits = [str(p) for p in Path("tests").rglob("*.py") if needle in p.read_text(encoding="utf-8")]
print(needle, hits)

text = Path("tests/models/unets/test_models_unet_3d_condition.py").read_text(encoding="utf-8")
print("UNet3D slow marker?", "@slow" in text or "pytest.mark.slow" in text)
```

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/models/unets/test_models_unet_1d.py#L141-L141
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/models/unets/test_models_unet_2d.py#L330-L344
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/models/unets/test_models_unet_2d_condition.py#L1161-L1162

Suggested fix:
```python
# Add focused fast tests:
# - tests/models/unets/test_models_uvit_2d.py with a tiny UVit2DModel forward,
# save/load, and gradient-checkpointing regression.
# - tests/models/unets/test_models_unet_2d_condition_flax.py guarded by require_flax,
# or explicitly remove/limit public Flax coverage if deprecated support is no longer maintained.
# - Add at least one @slow UNet3DConditionModel checkpoint or deprecated text-to-video pipeline regression.
```

Verification: ran the import checks and all reproduction snippets above with `.venv`. JAX/Flax is not installed in this `.venv`, so Flax runtime behavior was code-reviewed but not executed.

Start by running the provided reproduction snippets, then inspect the affected files under src/diffusers/models/unets/ and the existing tests/models/unets/ coverage named in the report. Separate the duplicate UVit2DModel issue from the remaining findings; done means the 3D attention-mask and scalar-timestep regressions are covered by passing tests and the identified coverage gaps are addressed.

Written by the indexing model from the issue text.