huggingface / huggingface/diffusers

sana model/pipeline review

Ouverte
#13,614 0 commentaires 0 réactions 0 personnes assignées Voir sur GitHub
Langage dominant
Python
Étoiles
34.5k
Forks
7.3k
Merge moyen
3 j 3 h
PR mergées (30 j)
91

Description

# `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/`.

Guide de contribution

Ouvrir le guide de contribution

Piste de recherche

Commencez par les fichiers affectés de Sana pipeline et transformer, puis exécutez les reproductions et examinez les tests existants dans tests/pipelines/sana. Comparez l’implémentation de l’attention avec transformer_sana_video.py et le modèle d’exportation dans pipelines/flux/__init__.py. Le travail est terminé lorsque les lacunes signalées concernant Sprint, transformer, attention, import et coverage sont corrigées par des tests de régression ciblés.

Rédigé par le modèle d'indexation à partir du texte de l'issue.

Évaluation

Stack technique
python, pytorch
Domaine
backend, machine-learning, testing-qa
Type d'issue
Bug
Difficulté
5/5
Temps estimé
Plus d'une semaine
Activité
Calme
Clarté
Plutôt claire
Accessibilité débutants
35/100

Recevez les nouvelles issues par e-mail

Un résumé court des issues GitHub adaptées aux débutants.