huggingface / huggingface/diffusers

sana model/pipeline review

Open
#13,614 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
34.5k
Forks
7.3k
Avg merge
3d 3h
Merged PRs (30d)
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/`.

Contributor guide

Open the contributing guide

Research direction

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.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
backend, machine-learning, testing-qa
Issue type
Bug
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.