huggingface / huggingface/diffusers

sana model/pipeline review

Aperta
#13,614 0 commenti 0 reazioni 0 assegnatari Vedi su GitHub
Lingua principale
Python
Stelle
34.5k
Fork
7.3k
Merge medio
3g 3h
PR unite (30g)
91

Descrizione

# `sana` model/pipeline review

Commit tested: `0f1abc4ae8b0eb2a3b40e82a310507281144c423`

Review performed against the repository review rules.

Duplicate search: searched GitHub Issues and PRs for `sana`, affected class/file names, and failure terms. No likely duplicates found except Issue 2, which is already tracked.

## Issue 1: Sana Sprint rejects documented 1/3/4-step inference by default

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/sana/pipeline_sana_sprint.py#L436-L437
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/sana/pipeline_sana_sprint.py#L620-L623
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/sana/pipeline_sana_sprint_img2img.py#L463-L464
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/sana/pipeline_sana_sprint_img2img.py#L691-L694

Problem:
`intermediate_timesteps` defaults to `1.3`, but `check_inputs()` rejects any non-`None` value unless `num_inference_steps == 2`. As a result, `num_inference_steps=1`, `3`, or `4` fails unless users know to pass `intermediate_timesteps=None`.

Impact:
SANA-Sprint is documented as a 1-4 step model, but the pipeline blocks the one-step path by default.

Reproduction:
```python
from diffusers import SanaSprintImg2ImgPipeline, SanaSprintPipeline

common = dict(
prompt="cat",
height=1024,
width=1024,
num_inference_steps=1,
timesteps=None,
max_timesteps=1.5708,
intermediate_timesteps=1.3,
callback_on_step_end_tensor_inputs=None,
prompt_embeds=None,
prompt_attention_mask=None,
)

for cls, extra in [(SanaSprintPipeline, {}), (SanaSprintImg2ImgPipeline, {"strength": 0.5})]:
try:
cls.check_inputs(None, **common, **extra)
except Exception as e:
print(cls.__name__, type(e).__name__, e)
```

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/docs/source/en/api/pipelines/sana_sprint.md#L25

Suggested fix:
```python
# In both Sprint pipeline __call__ signatures:
intermediate_timesteps: float | None = None,

# Before retrieve_timesteps:
if num_inference_steps == 2 and intermediate_timesteps is None:
intermediate_timesteps = 1.3
```

## Issue 2: Known duplicate: `guidance_embeds=True` crashes without `guidance`

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/sana_transformer.py#L460-L466

Problem:
`SanaTransformer2DModel.forward()` dispatches the time embedding call based on whether `guidance` was passed, not on which time embedding module was configured. With `guidance_embeds=True` and no `guidance`, it calls `SanaCombinedTimestepGuidanceEmbeddings.forward(..., batch_size=...)`, which is not accepted.

Impact:
A model configured with guidance embeddings cannot be used by non-guidance Sana pipelines; users get a low-level `TypeError`.

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

model = SanaTransformer2DModel(
in_channels=4, out_channels=4, num_attention_heads=2, attention_head_dim=4,
num_layers=1, num_cross_attention_heads=2, cross_attention_head_dim=4,
cross_attention_dim=8, caption_channels=8, sample_size=4, patch_size=1,
guidance_embeds=True,
)

model(
hidden_states=torch.randn(1, 4, 4, 4),
encoder_hidden_states=torch.randn(1, 3, 8),
timestep=torch.tensor([1.0]),
)
```

Relevant precedent:
Duplicate: https://github.com/huggingface/diffusers/issues/12540
Related PRs: https://github.com/huggingface/diffusers/pull/13109 and closed-unmerged https://github.com/huggingface/diffusers/pull/13517

Suggested fix:
Route by embedding type/configuration instead of `guidance is not None`, and raise a clear `ValueError` when `guidance_embeds=True` but `guidance` is absent.

## Issue 3: `cross_attention_dim=None` constructs a broken block

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/sana_transformer.py#L226-L281
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/controlnets/controlnet_sana.py#L56-L99

Problem:
`cross_attention_dim` is annotated as optional, but `SanaTransformerBlock` only defines `self.attn2` and `self.norm2` inside `if cross_attention_dim is not None`. `forward()` always reads them.

Impact:
Valid-looking configs fail at runtime, including `SanaControlNetModel`, which reuses the same block.

Reproduction:
```python
import torch
from diffusers import SanaControlNetModel, SanaTransformer2DModel

for cls in (SanaTransformer2DModel, SanaControlNetModel):
model = cls(
in_channels=4, out_channels=4, num_attention_heads=2, attention_head_dim=4,
num_layers=1, num_cross_attention_heads=2, cross_attention_head_dim=4,
cross_attention_dim=None, caption_channels=8, sample_size=4, patch_size=1,
)
kwargs = dict(
hidden_states=torch.randn(1, 4, 4, 4),
encoder_hidden_states=torch.randn(1, 3, 8),
timestep=torch.tensor([1.0]),
)
if cls is SanaControlNetModel:
kwargs["controlnet_cond"] = torch.randn(1, 4, 4, 4)

try:
model(**kwargs)
except Exception as e:
print(cls.__name__, type(e).__name__, e)
```

Relevant precedent:
Standard transformer blocks either define the optional attention attributes unconditionally or reject unsupported configs early.

Suggested fix:
```python
self.norm2 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps)
self.attn2 = None
if cross_attention_dim is not None:
self.attn2 = Attention(...)
```

## Issue 4: Sana attention bypasses backend dispatch and silently ignores self-attention masks

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/sana_transformer.py#L122-L172
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/sana_transformer.py#L246-L268
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/sana_transformer.py#L424-L451

Problem:
`SanaAttnProcessor2_0` calls `F.scaled_dot_product_attention` directly and has no `_attention_backend` / `_parallel_config`, so `set_attention_backend()` cannot configure it. Separately, the public `attention_mask` is normalized and passed into blocks, but self-attention calls `self.attn1(norm_hidden_states)` without the mask.

Impact:
Backend selection and context-parallel attention support do not behave like newer transformer families. Passing `attention_mask` gives users a false signal because it is ignored.

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

model = SanaTransformer2DModel(
in_channels=4, out_channels=4, num_attention_heads=2, attention_head_dim=4,
num_layers=1, num_cross_attention_heads=2, cross_attention_head_dim=4,
cross_attention_dim=8, caption_channels=8, sample_size=4, patch_size=1,
).eval()

print({name: hasattr(proc, "_attention_backend") for name, proc in model.attn_processors.items()})
model.set_attention_backend("_native_math")
print({name: getattr(proc, "_attention_backend", None) for name, proc in model.attn_processors.items()})

inputs = dict(
hidden_states=torch.randn(1, 4, 4, 4),
encoder_hidden_states=torch.randn(1, 3, 8),
timestep=torch.tensor([1.0]),
)

with torch.no_grad():
a = model(**inputs).sample
b = model(**inputs, attention_mask=torch.zeros(1, 16)).sample

print((a - b).abs().max().item()) # 0.0: mask had no effect
```

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_sana_video.py#L277-L335

Suggested fix:
Port the Sana Video processor pattern: add `_attention_backend` / `_parallel_config` and call `dispatch_attention_fn()` for cross-attention. For self-attention, either implement mask handling in `SanaLinearAttnProcessor2_0` or remove/reject the unsupported public `attention_mask`.

## Issue 5: `SanaPipelineOutput` is not exported from the Sana subpackage

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/sana/__init__.py#L14-L28
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/sana/pipeline_output.py#L10

Problem:
`pipeline_output.py` defines `SanaPipelineOutput`, and docs autodoc it, but `src/diffusers/pipelines/sana/__init__.py` never adds `pipeline_output` to `_import_structure`.

Impact:
The expected subpackage import fails while similar pipeline families expose their output classes through lazy imports.

Reproduction:
```python
from diffusers.pipelines.sana import SanaPipelineOutput
```

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/flux/__init__.py#L13-L15

Suggested fix:
```python
_import_structure = {"pipeline_output": ["SanaPipelineOutput"]}

# in TYPE_CHECKING / slow import branch
from .pipeline_output import SanaPipelineOutput
```

## Issue 6: Test coverage gaps for Sana ControlNet and Sprint variants

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/sana/test_sana.py#L313-L315
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/sana/test_sana_controlnet.py#L39-L40
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/sana/test_sana_sprint.py#L32-L33
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/sana/test_sana_sprint_img2img.py#L37-L38

Problem:
Fast pipeline tests exist for all four pipelines, and slow tests exist for base `SanaPipeline` only. There are no slow tests for `SanaControlNetPipeline`, `SanaSprintPipeline`, or `SanaSprintImg2ImgPipeline`. There is also no `tests/models/controlnets/test_models_controlnet_sana.py`, so `SanaControlNetModel` lacks direct model-mixin coverage.

Impact:
Real checkpoint loading, expected-output slices, ControlNet serialization/model behavior, and Sprint 1/3/4-step behavior are not covered.

Reproduction:
```python
from pathlib import Path

for path in sorted(Path("tests/pipelines/sana").glob("test_*.py")):
text = path.read_text()
print(path, "@slow" in text)

print("Sana ControlNet model tests:", list(Path("tests/models/controlnets").glob("*sana*.py")))
```

Relevant precedent:
Base Sana has slow integration tests:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/sana/test_sana.py#L313-L378

Suggested fix:
Add slow tests with small output slices for the public ControlNet and Sprint checkpoints, add `num_inference_steps` coverage for Sprint `1`, `2`, `3`, and `4`, and add a `ModelTesterMixin`-style `SanaControlNetModel` test file under `tests/models/controlnets/`.

Guida per i contributori

Apri la guida per i contributori

Direzione di ricerca

Start with the affected Sana pipeline and transformer files, then run the reproductions and inspect the existing tests under tests/pipelines/sana. Compare the attention implementation with transformer_sana_video.py and the export pattern in pipelines/flux/__init__.py. Done means the reported Sprint, transformer, attention, import, and coverage gaps are addressed with focused regression tests.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Valutazione

Stack tecnologico
python, pytorch
Ambito
backend, 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

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.