huggingface / huggingface/diffusers
`model_transformers_shared` model/pipeline review
- Lingua principale
- Python
- Stelle
- 34.5k
- Fork
- 7.3k
- Merge medio
- 3g 3h
- PR unite (30g)
- 91
Descrizione
# `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.
Guida per i contributori
Apri la guida per i contributori
Direzione di ricerca
Start with the four transformer files listed in the review and reproduce the reported failures, then inspect their existing tests in tests/models/test_layers_utils.py and tests/models/transformers/test_models_transformer_temporal.py. Compare import wiring in src/diffusers/__init__.py and src/diffusers/models/__init__.py with the existing shared classes. Done means the five reported areas have agreed behavior and regression coverage, including dual-transformer calls, exports, temporal channels, non-fp64 execution, and missing model or pipeline coverage.
Scritto dal modello di indicizzazione a partire dal testo della issue.
Valutazione
- Stack tecnologico
- python, pytorch
- Ambito
- machine-learning, testing-qa
- Tipo di issue
- Bug
- Difficoltà
- 5/5
- Tempo stimato
- Più di una settimana
- Stato di attività
- Tranquilla
- Chiarezza
- Abbastanza chiara
- Idoneità per principianti
- 35/100