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:

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:

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:

Suggested fix:

# 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:

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:

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: #12540
Related PRs: #13109 and closed-unmerged #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:

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:

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:

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:

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:

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:

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:

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:

from diffusers.pipelines.sana import SanaPipelineOutput

Relevant precedent:

Suggested fix:

_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:

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:

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:

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