Skip to content

Fix SanaTransformer2DModel crash when guidance_embeds=True with non-guidance pipelines#13517

Closed
Ricardo-M-L wants to merge 1 commit intohuggingface:mainfrom
Ricardo-M-L:fix-sana-guidance-embeds
Closed

Fix SanaTransformer2DModel crash when guidance_embeds=True with non-guidance pipelines#13517
Ricardo-M-L wants to merge 1 commit intohuggingface:mainfrom
Ricardo-M-L:fix-sana-guidance-embeds

Conversation

@Ricardo-M-L
Copy link
Copy Markdown
Contributor

What does this PR do?

Fixes #12540.

When `SanaTransformer2DModel` is built with `guidance_embeds=True` and then driven by a pipeline that does not thread `guidance` through (`SanaPipeline`, `SanaPAGPipeline`), the transformer forward crashes with:

```
TypeError: SanaCombinedTimestepGuidanceEmbeddings.forward() got an unexpected keyword argument 'batch_size'
```

Root cause

`SanaTransformer2DModel.forward` routed the time-embedding call on whether `guidance` was passed:

```python
if guidance is not None:
timestep, embedded_timestep = self.time_embed(timestep, guidance=guidance, hidden_dtype=...)
else:
timestep, embedded_timestep = self.time_embed(timestep, batch_size=batch_size, hidden_dtype=...)
```

But the two embedding classes have incompatible signatures:

  • `AdaLayerNormSingle.forward` takes `batch_size` and does not accept `guidance`
  • `SanaCombinedTimestepGuidanceEmbeddings.forward` takes `guidance` and does not accept `batch_size`

So with `guidance_embeds=True` and no guidance kwarg, the `batch_size` branch is hit against the wrong embedder and blows up.

Fix

Dispatch on the embedding type instead of on the input. If the model was configured with guidance embeddings but the caller omitted `guidance`, raise a clear `ValueError` naming the misconfiguration rather than surfacing a cryptic kwargs error from deeper in the embedder.

Reproducer (from #12540)

```python
from diffusers import SanaTransformer2DModel, SanaPAGPipeline
import torch

config = SanaTransformer2DModel.load_config('Efficient-Large-Model/Sana_600M_512px_diffusers', subfolder='transformer')
config['guidance_embeds'] = True
new_transformer = SanaTransformer2DModel.from_config(config)

pipe = SanaPAGPipeline.from_pretrained(
'Efficient-Large-Model/Sana_1600M_512px_MultiLing_diffusers',
pag_applied_layers=['transformer_blocks.8'],
transformer=new_transformer,
torch_dtype=torch.bfloat16,
vae=None,
).to('cuda')
pipe(prompt='cat', output_type='latent', height=768, width=512, pag_scale=1.0, guidance_scale=4.0, num_inference_steps=20)
```

Before: `TypeError: SanaCombinedTimestepGuidanceEmbeddings.forward() got an unexpected keyword argument 'batch_size'`
After: `ValueError: SanaTransformer2DModel was configured with guidance_embeds=True, but guidance was not provided. Either pass guidance to the transformer or rebuild the model with guidance_embeds=False.`

Testing

Manually built tiny SanaTransformer2DModel instances and confirmed:

  1. `guidance_embeds=True` without `guidance` → clear ValueError
  2. `guidance_embeds=True` with `guidance` → works
  3. `guidance_embeds=False` (default, AdaLayerNormSingle) → unchanged, works

Before submitting

Who can review?

@yiyixuxu @DN6 @dg845

…nce presence

When SanaTransformer2DModel was built with guidance_embeds=True, the
forward pass routed on `guidance is not None`. Pipelines that do not
thread guidance through (SanaPipeline, SanaPAGPipeline) therefore fell
into the `batch_size` branch and crashed with
`TypeError: SanaCombinedTimestepGuidanceEmbeddings.forward() got an
unexpected keyword argument 'batch_size'` (see huggingface#12540).

Dispatch on `isinstance(self.time_embed, SanaCombinedTimestepGuidanceEmbeddings)`
instead: the embedder type is what dictates the call signature. If the
model was configured with guidance embeddings but the pipeline omits
guidance, raise a clear ValueError naming the misconfiguration instead
of surfacing a confusing kwargs error from the embedder.

Fixes huggingface#12540
@github-actions github-actions Bot added models size/S PR with diff < 50 LOC labels Apr 21, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

models size/S PR with diff < 50 LOC

Projects

None yet

Development

Successfully merging this pull request may close these issues.

SanaCombinedTimestepGuidanceEmbeddings do not work with SanaPipeline

1 participant