huggingface / huggingface/diffusers

`model_transformers_shared` model/pipeline review

Offen
#13,651 0 Kommentare 0 Reaktionen 0 zugewiesene Personen Auf GitHub ansehen
Vorherrschende Sprache
Python
Sterne
34.5k
Forks
7.3k
Ø Merge
3 T. 3 Std.
Gemergte PRs (30 T.)
91

Beschreibung

# `model_transformers_shared` model/pipeline review

Commit tested: `0f1abc4ae8b0eb2a3b40e82a310507281144c423`

Review performed against the repository review rules. `AGENTS.md` is referenced by `.ai/review-rules.md` but is not present in this checkout; all available referenced rule files were read.

Files reviewed:
- `src/diffusers/models/transformers/dual_transformer_2d.py`
- `src/diffusers/models/transformers/prior_transformer.py`
- `src/diffusers/models/transformers/transformer_2d.py`
- `src/diffusers/models/transformers/transformer_temporal.py`

Duplicate-search status: checked GitHub Issues/PRs for `model_transformers_shared`, all target class/file names, and the specific failures below. No likely duplicate found for Issues 1-4. Shap-E slow-test coverage is already reported in https://github.com/huggingface/diffusers/issues/13593 and is called out in Issue 5.

## Issue 1: `DualTransformer2DModel` rejects UNet encoder-mask calls and drops masks

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/dual_transformer_2d.py#L96-L145

Problem:
`DualTransformer2DModel.forward()` does not accept `encoder_attention_mask`, but `CrossAttnDownBlock2D` and related UNet blocks call transformer modules with that keyword. The same method also declares `attention_mask` but never forwards it to either child transformer.

Impact:
Any UNet block configured with `dual_cross_attention=True` fails before inference/training. Masked text/image conditioning is also silently ignored if callers invoke the wrapper directly.

Reproduction:
```python
import torch
from diffusers.models.unets.unet_2d_blocks import CrossAttnDownBlock2D

block = CrossAttnDownBlock2D(
in_channels=4, out_channels=4, temb_channels=8, num_layers=1,
transformer_layers_per_block=1, num_attention_heads=1,
cross_attention_dim=8, dual_cross_attention=True, resnet_groups=1,
)

block(
torch.randn(1, 4, 4, 4),
temb=torch.randn(1, 8),
encoder_hidden_states=torch.randn(1, 77 + 257, 8),
)
```

Relevant precedent:
`Transformer2DModel.forward()` accepts and routes `encoder_attention_mask`:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_2d.py#L333
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_2d.py#L431

Suggested fix:
```python
def forward(..., attention_mask=None, encoder_attention_mask=None, ...):
...
condition_mask = None
if encoder_attention_mask is not None:
condition_mask = encoder_attention_mask[..., tokens_start : tokens_start + self.condition_lengths[i]]

encoded_state = self.transformers[transformer_index](
input_states,
encoder_hidden_states=condition_state,
timestep=timestep,
attention_mask=attention_mask,
encoder_attention_mask=condition_mask,
cross_attention_kwargs=cross_attention_kwargs,
return_dict=False,
)[0]
```

## Issue 2: `DualTransformer2DModel` is not available from the top-level lazy import

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/__init__.py#L86
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/__init__.py#L194-L306

Problem:
`DualTransformer2DModel` is exported from `diffusers.models` and `diffusers.models.transformers`, but not from top-level `diffusers`. The local review rules require model classes to be wired through both subpackage and top-level lazy imports.

Impact:
Users can import related shared transformer classes from `diffusers`, but this one raises `ImportError`.

Reproduction:
```python
from diffusers.models import DualTransformer2DModel
print(DualTransformer2DModel.__name__)

from diffusers import DualTransformer2DModel
```

Relevant precedent:
The same top-level list already exports `PriorTransformer`, `Transformer2DModel`, and `TransformerTemporalModel`:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/__init__.py#L273-L292

Suggested fix:
```python
# src/diffusers/__init__.py
_import_structure["models"].extend([
...
"DualTransformer2DModel",
...
])

if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
from .models import (
...
DualTransformer2DModel,
...
)
```

Also regenerate/add the matching `dummy_pt_objects.py` entry for no-torch imports.

## Issue 3: Temporal transformer `out_channels` is serialized but ignored

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_temporal.py#L78-L121
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_temporal.py#L225-L276

Problem:
Both `TransformerTemporalModel` and `TransformerSpatioTemporalModel` accept `out_channels`, and the spatio-temporal class stores it, but `proj_out` is hardwired to `in_channels`. The model config can say `out_channels=8` while the output still has 4 channels.

Impact:
Saved configs misrepresent the architecture. Users cannot rely on the public constructor/config contract, and future checkpoint conversion can silently produce the wrong projection shape.

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

model = TransformerTemporalModel(
num_attention_heads=1, attention_head_dim=4,
in_channels=4, out_channels=8, num_layers=1, norm_num_groups=1,
)
print(model.config.out_channels) # 8
print(model.proj_out.out_features) # 4

x = torch.randn(2, 4, 4, 4)
print(model(x, num_frames=2).sample.shape) # torch.Size([2, 4, 4, 4])
```

Relevant precedent:
SD3 wires `self.out_channels` into its output projection:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_sd3.py#L139-L170

Suggested fix:
If different output channels are unsupported, reject them explicitly:
```python
if out_channels is not None and out_channels != in_channels:
raise ValueError("`out_channels` must be None or equal to `in_channels` for this temporal transformer.")
```

If support is intended, wire `proj_out` to `self.out_channels` and handle the residual path when channel counts differ.

## Issue 4: `Transformer2DModel` discrete output hardcasts through float64

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_2d.py#L514-L520
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/models/test_layers_utils.py#L435-L436

Problem:
The vectorized/discrete path computes `F.log_softmax(logits.double(), dim=1).float()`. The local model rules explicitly prohibit unconditional float64 in model forwards because MPS and several NPU backends do not support it. The only discrete-path test is gated behind `require_torch_accelerator_with_fp64`, so unsupported devices are skipped instead of protected.

Impact:
Discrete `Transformer2DModel` inference can fail on devices without float64 support, and the test suite encodes that limitation rather than catching it.

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

device = torch.device("mps") # or another backend without float64 forward support
model = Transformer2DModel(
num_attention_heads=1,
attention_head_dim=32,
num_vector_embeds=8,
sample_size=2,
).to(device)

sample = torch.randint(0, 8, (1, 4), device=device)
model(sample)
```

Relevant precedent:
The review rules require avoiding unconditional float64 in model code.

Suggested fix:
```python
output = F.log_softmax(logits.float(), dim=1)
```

This preserves the current float32 output contract without routing through float64.

## Issue 5: Slow/direct coverage is incomplete for the shared transformer family

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/models/test_layers_utils.py#L324-L436
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/models/transformers/test_models_transformer_temporal.py#L32-L33
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/kandinsky/test_kandinsky_prior.py#L171
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/stable_unclip/test_stable_unclip.py#L202-L204

Problem:
`DualTransformer2DModel` has no direct test references. `Transformer2DModel` tests live in `test_layers_utils.py` with no slow coverage. `TransformerTemporalModel` has a fast model test but no slow coverage. Prior model slow coverage exists, but the prior-family pipelines are uneven: Kandinsky prior tests are fast-only, Stable UnCLIP is nightly-only, and Shap-E slow coverage is already tracked separately in https://github.com/huggingface/diffusers/issues/13593.

Impact:
The exact regressions above are not covered: dual UNet call compatibility, top-level import parity, temporal `out_channels`, and non-fp64 discrete transformer execution.

Reproduction:
```python
from pathlib import Path

for path in [
Path("tests/models/test_layers_utils.py"),
Path("tests/models/transformers/test_models_prior.py"),
Path("tests/models/transformers/test_models_transformer_temporal.py"),
]:
text = path.read_text()
print(path, "@slow" in text, "@nightly" in text)

print(
"DualTransformer2DModel test refs:",
sum("DualTransformer2DModel" in p.read_text(errors="ignore") for p in Path("tests").rglob("test_*.py")),
)
```

Relevant precedent:
PriorTransformer already has direct slow model coverage:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/models/transformers/test_models_prior.py#L142-L174

Suggested fix:
Add a dedicated dual-transformer fast test that runs through a `dual_cross_attention=True` UNet block with masks, add import/export assertions for all public shared classes, add temporal `out_channels` validation coverage, and add slow or integration coverage for `TransformerTemporalModel` and the prior-family pipelines that are currently fast-only or nightly-only.

Beitragsleitfaden

Beitragsleitfaden öffnen

Rechercherichtung

Beginne mit den vier in der Review aufgeführten Transformer-Dateien und reproduziere die gemeldeten Fehler. Untersuche anschließend deren bestehende Tests in tests/models/test_layers_utils.py und tests/models/transformers/test_models_transformer_temporal.py. Vergleiche die Import-Verdrahtung in src/diffusers/__init__.py und src/diffusers/models/__init__.py mit den vorhandenen gemeinsamen Klassen. Als abgeschlossen gilt, dass die fünf gemeldeten Bereiche einheitliches Verhalten und Regressionstestabdeckung aufweisen, einschließlich Aufrufen dualer Transformer, Exports, temporaler Kanäle, Ausführung ohne fp64 und fehlender Abdeckung für Modelle oder Pipelines.

Vom Indexierungsmodell aus dem Issue-Text verfasst.

Bewertung

Tech-Stack
python, pytorch
Bereich
machine-learning, testing-qa
Issue-Typ
Bug
Schwierigkeit
5/5
Geschätzter Aufwand
Über eine Woche
Aktivitätsstatus
Ruhig
Klarheit
Größtenteils klar
Anfängerfreundlichkeit
35/100

Neue Issues direkt in Ihr Postfach

Eine kurze Übersicht über anfängerfreundliche GitHub-Issues.