sana_video model/pipeline review

`sana_video` model/pipeline review

Commit tested: 0f1abc4ae8b0eb2a3b40e82a310507281144c423

Review performed against the repository review rules. Note: .ai/review-rules.md references AGENTS.md, but no AGENTS.md exists in this checkout; the other referenced rule files were applied.

Duplicate-search status: searched GitHub Issues/PRs for sana_video, SanaVideoPipeline, SanaImageToVideoPipeline, SanaVideoTransformer3DModel, num_videos_per_prompt, conditioning_mask, image_latents, attention_mask, cross_attention_dim, and the docs URL typo. I found related merged PRs #12584, #12634, #13229, #12675, and issue #12760, but no exact duplicate for the findings below.

Test coverage status: fast model/pipeline tests exist, but slow Sana Video tests are skipped TODOs; see Issue 5. Direct repro snippets were run with .venv. Pytest collection is currently blocked in this .venv by the installed torch build missing torch._C._distributed_c10d.

Issue 1: I2V fails for `num_videos_per_prompt > 1`

Affected code:

diffusers/src/diffusers/pipelines/sana_video/pipeline_sana_video_i2v.py

Lines 953 to 975 in 0f1abc4

    
           latents = self.prepare_latents( 
        
               image, 
        
               batch_size * num_videos_per_prompt, 
        
               latent_channels, 
        
               height, 
        
               width, 
        
               frames, 
        
               torch.float32, 
        
               device, 
        
               generator, 
        
               latents, 
        
           ) 
        
           conditioning_mask = latents.new_zeros( 
        
               batch_size, 
        
               1, 
        
               latents.shape[2] // self.transformer_temporal_patch_size, 
        
               latents.shape[3] // self.transformer_spatial_patch_size, 
        
               latents.shape[4] // self.transformer_spatial_patch_size, 
        
           ) 
        
           conditioning_mask[:, :, 0] = 1.0 
        
           if self.do_classifier_free_guidance: 
        
               conditioning_mask = torch.cat([conditioning_mask, conditioning_mask])

diffusers/src/diffusers/pipelines/sana_video/pipeline_sana_video_i2v.py

Lines 990 to 1004 in 0f1abc4

    
           latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents 
        
           # broadcast to batch dimension in a way that's compatible with ONNX/Core ML 
        
           timestep = t.expand(conditioning_mask.shape) 
        
           timestep = timestep * (1 - conditioning_mask) 
        
           # predict noise model_output 
        
           noise_pred = self.transformer( 
        
               latent_model_input.to(dtype=transformer_dtype), 
        
               encoder_hidden_states=prompt_embeds.to(dtype=transformer_dtype), 
        
               encoder_attention_mask=prompt_attention_mask, 
        
               timestep=timestep, 
        
               return_dict=False, 
        
               attention_kwargs=self.attention_kwargs, 
        
           )[0]

Problem:
latents are prepared for batch_size * num_videos_per_prompt, but conditioning_mask is created with only batch_size. The timestep tensor is therefore too small for the transformer batch and reshapes to the wrong token length.

Impact:
SanaImageToVideoPipeline cannot generate multiple videos per prompt.

Reproduction:

import torch
from PIL import Image
from diffusers import AutoencoderKLWan, FlowMatchEulerDiscreteScheduler, SanaImageToVideoPipeline, SanaVideoTransformer3DModel

vae = AutoencoderKLWan(base_dim=3, z_dim=16, dim_mult=[1, 1, 1, 1], num_res_blocks=1, temperal_downsample=[False, True, True])
transformer = SanaVideoTransformer3DModel(in_channels=16, out_channels=16, num_attention_heads=2, attention_head_dim=12, num_layers=1, num_cross_attention_heads=2, cross_attention_head_dim=12, cross_attention_dim=24, caption_channels=8, sample_size=8, patch_size=(1, 2, 2), rope_max_seq_len=32)
pipe = SanaImageToVideoPipeline(None, None, vae, transformer, FlowMatchEulerDiscreteScheduler())
pipe.set_progress_bar_config(disable=True)
pipe(
    image=Image.new("RGB", (32, 32)),
    prompt_embeds=torch.randn(1, 16, 8),
    prompt_attention_mask=torch.ones(1, 16, dtype=torch.long),
    height=32, width=32, frames=9, num_inference_steps=1,
    guidance_scale=1.0, num_videos_per_prompt=2,
    output_type="latent", use_resolution_binning=False,
)

Relevant precedent:
SanaVideoPipeline expands timestep from latent_model_input.shape[0]:

diffusers/src/diffusers/pipelines/sana_video/pipeline_sana_video.py

Lines 944 to 947 in 0f1abc4

    
           latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents 
        
           # broadcast to batch dimension in a way that's compatible with ONNX/Core ML 
        
           timestep = t.expand(latent_model_input.shape[0])

Suggested fix:

effective_batch_size = batch_size * num_videos_per_prompt
conditioning_mask = latents.new_zeros(
    effective_batch_size,
    1,
    latents.shape[2] // self.transformer_temporal_patch_size,
    latents.shape[3] // self.transformer_spatial_patch_size,
    latents.shape[4] // self.transformer_spatial_patch_size,
)
conditioning_mask[:, :, 0] = 1.0
if self.do_classifier_free_guidance:
    conditioning_mask = torch.cat([conditioning_mask, conditioning_mask])

Issue 2: I2V batched images are broken

Affected code:

diffusers/src/diffusers/pipelines/sana_video/pipeline_sana_video_i2v.py

Lines 465 to 466 in 0f1abc4

    
           if image is not None and not isinstance(image, torch.Tensor) and not isinstance(image, PIL.Image.Image): 
        
               raise ValueError(f"`image` has to be of type `torch.Tensor` or `PIL.Image.Image` but is {type(image)}")

diffusers/src/diffusers/pipelines/sana_video/pipeline_sana_video_i2v.py

Lines 694 to 715 in 0f1abc4

    
           if isinstance(generator, list): 
        
               image_latents = [retrieve_latents(self.vae.encode(image), sample_mode="argmax") for _ in generator] 
        
               image_latents = torch.cat(image_latents) 
        
           else: 
        
               image_latents = retrieve_latents(self.vae.encode(image), sample_mode="argmax") 
        
               image_latents = image_latents.repeat(batch_size, 1, 1, 1, 1) 
        
           if isinstance(self.vae, AutoencoderKLLTX2Video): 
        
               _latents_mean = self.vae.latents_mean 
        
               _latents_std = self.vae.latents_std 
        
           elif isinstance(self.vae, AutoencoderKLWan): 
        
               _latents_mean = torch.tensor(self.vae.config.latents_mean) 
        
               _latents_std = torch.tensor(self.vae.config.latents_std) 
        
           else: 
        
               _latents_mean = torch.zeros(image_latents.shape[1], device=image_latents.device, dtype=image_latents.dtype) 
        
               _latents_std = torch.ones(image_latents.shape[1], device=image_latents.device, dtype=image_latents.dtype) 
        
           latents_mean = _latents_mean.view(1, -1, 1, 1, 1).to(image_latents.device, image_latents.dtype) 
        
           latents_std = 1.0 / _latents_std.view(1, -1, 1, 1, 1).to(image_latents.device, image_latents.dtype) 
        
           image_latents = (image_latents - latents_mean) * latents_std 
        
           latents[:, :, 0:1] = image_latents.to(dtype)

Problem:
The signature accepts PipelineImageInput, but validation rejects image lists. Batched tensor images pass validation, then prepare_latents repeats encoded image latents by the effective batch size, producing too many image latents.

Impact:
Batched I2V inference fails for normal batched tensor inputs and cannot accept list-style image batches.

Reproduction:

import torch
from diffusers import AutoencoderKLWan, SanaImageToVideoPipeline

pipe = SanaImageToVideoPipeline.__new__(SanaImageToVideoPipeline)
pipe.vae = AutoencoderKLWan(base_dim=3, z_dim=16, dim_mult=[1, 1, 1, 1], num_res_blocks=1, temperal_downsample=[False, True, True])
pipe.vae_scale_factor_temporal = pipe.vae.config.scale_factor_temporal
pipe.vae_scale_factor_spatial = pipe.vae.config.scale_factor_spatial

pipe.prepare_latents(
    image=torch.zeros(2, 3, 32, 32),
    batch_size=2,
    num_channels_latents=16,
    height=32,
    width=32,
    num_frames=9,
    dtype=torch.float32,
    device=torch.device("cpu"),
)

Relevant precedent:
CogVideoX I2V accepts list inputs in validation:

diffusers/src/diffusers/pipelines/cogvideo/pipeline_cogvideox_image2video.py

Lines 467 to 474 in 0f1abc4

    
           if ( 
        
               not isinstance(image, torch.Tensor) 
        
               and not isinstance(image, PIL.Image.Image) 
        
               and not isinstance(image, list) 
        
           ): 
        
               raise ValueError( 
        
                   "`image` has to be of type `torch.Tensor` or `PIL.Image.Image` or `list[PIL.Image.Image]` but is" 
        
                   f" {type(image)}"

Suggested fix:

image_latents = retrieve_latents(self.vae.encode(image), sample_mode="argmax")
if image_latents.shape[0] != batch_size:
    if batch_size % image_latents.shape[0] != 0:
        raise ValueError("Image batch size must divide the effective prompt batch size.")
    image_latents = image_latents.repeat_interleave(batch_size // image_latents.shape[0], dim=0)

Also update validation to accept valid PipelineImageInput lists/arrays.

Issue 3: `attention_mask` is accepted but ignored by self-attention

Affected code:

diffusers/src/diffusers/models/transformers/transformer_sana_video.py

Lines 596 to 607 in 0f1abc4

    
           if attention_mask is not None and attention_mask.ndim == 2: 
        
               # assume that mask is expressed as: 
        
               #   (1 = keep,      0 = discard) 
        
               # convert mask into a bias that can be added to attention scores: 
        
               #       (keep = +0,     discard = -10000.0) 
        
               attention_mask = (1 - attention_mask.to(hidden_states.dtype)) * -10000.0 
        
               attention_mask = attention_mask.unsqueeze(1) 
        
           # convert encoder_attention_mask to a bias the same way we do for attention_mask 
        
           if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2: 
        
               encoder_attention_mask = (1 - encoder_attention_mask.to(hidden_states.dtype)) * -10000.0 
        
               encoder_attention_mask = encoder_attention_mask.unsqueeze(1)

diffusers/src/diffusers/models/transformers/transformer_sana_video.py

Lines 638 to 668 in 0f1abc4

    
           # 2. Transformer blocks 
        
           if torch.is_grad_enabled() and self.gradient_checkpointing: 
        
               for index_block, block in enumerate(self.transformer_blocks): 
        
                   hidden_states = self._gradient_checkpointing_func( 
        
                       block, 
        
                       hidden_states, 
        
                       attention_mask, 
        
                       encoder_hidden_states, 
        
                       encoder_attention_mask, 
        
                       timestep, 
        
                       post_patch_num_frames, 
        
                       post_patch_height, 
        
                       post_patch_width, 
        
                       rotary_emb, 
        
                   ) 
        
                   if controlnet_block_samples is not None and 0 < index_block <= len(controlnet_block_samples): 
        
                       hidden_states = hidden_states + controlnet_block_samples[index_block - 1] 
        
           else: 
        
               for index_block, block in enumerate(self.transformer_blocks): 
        
                   hidden_states = block( 
        
                       hidden_states, 
        
                       attention_mask, 
        
                       encoder_hidden_states, 
        
                       encoder_attention_mask, 
        
                       timestep, 
        
                       post_patch_num_frames, 
        
                       post_patch_height, 
        
                       post_patch_width, 
        
                       rotary_emb, 
        
                   )

diffusers/src/diffusers/models/transformers/transformer_sana_video.py

Lines 433 to 441 in 0f1abc4

    
           attn_output = self.attn1(norm_hidden_states, rotary_emb=rotary_emb) 
        
           hidden_states = hidden_states + gate_msa * attn_output 
        
           # 3. Cross Attention 
        
           if self.attn2 is not None: 
        
               attn_output = self.attn2( 
        
                   hidden_states, 
        
                   encoder_hidden_states=encoder_hidden_states, 
        
                   attention_mask=encoder_attention_mask,

Problem:
The model converts attention_mask and passes it into each block, but SanaVideoTransformerBlock never forwards it to attn1; only encoder_attention_mask is used for cross-attention.

Impact:
Callers can pass latent-token masks and get silently unmasked output.

Reproduction:

import torch
from diffusers import SanaVideoTransformer3DModel

torch.manual_seed(0)
model = SanaVideoTransformer3DModel(in_channels=16, out_channels=16, num_attention_heads=2, attention_head_dim=12, num_layers=1, num_cross_attention_heads=2, cross_attention_head_dim=12, cross_attention_dim=24, caption_channels=8, sample_size=8, patch_size=(1, 2, 2), rope_max_seq_len=32).eval()
hidden_states = torch.randn(1, 16, 2, 8, 8)
encoder_hidden_states = torch.randn(1, 12, 8)
timestep = torch.tensor([1])
mask = torch.zeros(1, 32, dtype=torch.long)

with torch.no_grad():
    a = model(hidden_states, encoder_hidden_states, timestep, return_dict=False)[0]
    b = model(hidden_states, encoder_hidden_states, timestep, attention_mask=mask, return_dict=False)[0]
print((a - b).abs().max().item())  # 0.0

Relevant precedent:
The model review rules say declared masks must be honored or omitted.

Suggested fix:
Route attention_mask into attn1 and implement padding-mask support in SanaLinearAttnProcessor3_0, preferably keeping boolean masks until the processor. If self-attention masks are not supported, remove attention_mask from the public forward signatures.

Issue 4: `cross_attention_dim=None` is accepted but crashes

Affected code:

diffusers/src/diffusers/models/transformers/transformer_sana_video.py

Lines 510 to 512 in 0f1abc4

    
           num_cross_attention_heads: int | None = 20, 
        
           cross_attention_head_dim: int | None = 112, 
        
           cross_attention_dim: int | None = 2240,

diffusers/src/diffusers/models/transformers/transformer_sana_video.py

Lines 389 to 446 in 0f1abc4

    
               if cross_attention_dim is not None: 
        
                   self.norm2 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) 
        
                   self.attn2 = Attention( 
        
                       query_dim=dim, 
        
                       qk_norm=qk_norm, 
        
                       kv_heads=num_cross_attention_heads if qk_norm is not None else None, 
        
                       cross_attention_dim=cross_attention_dim, 
        
                       heads=num_cross_attention_heads, 
        
                       dim_head=cross_attention_head_dim, 
        
                       dropout=dropout, 
        
                       bias=True, 
        
                       out_bias=attention_out_bias, 
        
                       processor=SanaAttnProcessor2_0(), 
        
                   ) 
        
               # 3. Feed-forward 
        
               self.ff = GLUMBTempConv(dim, dim, mlp_ratio, norm_type=None, residual_connection=False) 
        
               self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5) 
        
           def forward( 
        
               self, 
        
               hidden_states: torch.Tensor, 
        
               attention_mask: torch.Tensor | None = None, 
        
               encoder_hidden_states: torch.Tensor | None = None, 
        
               encoder_attention_mask: torch.Tensor | None = None, 
        
               timestep: torch.LongTensor | None = None, 
        
               frames: int = None, 
        
               height: int = None, 
        
               width: int = None, 
        
               rotary_emb: torch.Tensor | None = None, 
        
           ) -> torch.Tensor: 
        
               batch_size = hidden_states.shape[0] 
        
               # 1. Modulation 
        
               shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( 
        
                   self.scale_shift_table[None, None] + timestep.reshape(batch_size, timestep.shape[1], 6, -1) 
        
               ).unbind(dim=2) 
        
               # 2. Self Attention 
        
               norm_hidden_states = self.norm1(hidden_states) 
        
               norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa 
        
               norm_hidden_states = norm_hidden_states.to(hidden_states.dtype) 
        
               attn_output = self.attn1(norm_hidden_states, rotary_emb=rotary_emb) 
        
               hidden_states = hidden_states + gate_msa * attn_output 
        
               # 3. Cross Attention 
        
               if self.attn2 is not None: 
        
                   attn_output = self.attn2( 
        
                       hidden_states, 
        
                       encoder_hidden_states=encoder_hidden_states, 
        
                       attention_mask=encoder_attention_mask, 
        
                   ) 
        
                   hidden_states = attn_output + hidden_states 
        
               # 4. Feed-forward 
        
               norm_hidden_states = self.norm2(hidden_states)

Problem:
cross_attention_dim is typed as optional, but when it is None, the block never defines attn2 or norm2. Forward then accesses both unconditionally.

Impact:
A serialized config with cross_attention_dim: null loads but cannot run.

Reproduction:

import torch
from diffusers import SanaVideoTransformer3DModel

model = SanaVideoTransformer3DModel(in_channels=16, out_channels=16, num_attention_heads=2, attention_head_dim=12, num_layers=1, cross_attention_dim=None, caption_channels=8, sample_size=8, patch_size=(1, 2, 2), rope_max_seq_len=32)
model(
    hidden_states=torch.randn(1, 16, 2, 8, 8),
    encoder_hidden_states=torch.randn(1, 12, 8),
    timestep=torch.tensor([1]),
    return_dict=False,
)

Relevant precedent:
Other transformer blocks either make cross-attention required or initialize optional attention attributes to None.

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 5: Slow tests are skipped TODOs

Affected code:

diffusers/tests/pipelines/sana_video/test_sana_video.py

Lines 208 to 225 in 0f1abc4

    
           @slow 
        
           @require_torch_accelerator 
        
           class SanaVideoPipelineIntegrationTests(unittest.TestCase): 
        
               prompt = "Evening, backlight, side lighting, soft light, high contrast, mid-shot, centered composition, clean solo shot, warm color. A young Caucasian man stands in a forest." 
        
               def setUp(self): 
        
                   super().setUp() 
        
                   gc.collect() 
        
                   backend_empty_cache(torch_device) 
        
               def tearDown(self): 
        
                   super().tearDown() 
        
                   gc.collect() 
        
                   backend_empty_cache(torch_device) 
        
               @unittest.skip("TODO: test needs to be implemented") 
        
               def test_sana_video_480p(self): 
        
                   pass

diffusers/tests/pipelines/sana_video/test_sana_video_i2v.py

Lines 221 to 238 in 0f1abc4

    
           @slow 
        
           @require_torch_accelerator 
        
           class SanaVideoPipelineIntegrationTests(unittest.TestCase): 
        
               prompt = "Evening, backlight, side lighting, soft light, high contrast, mid-shot, centered composition, clean solo shot, warm color. A young Caucasian man stands in a forest." 
        
               def setUp(self): 
        
                   super().setUp() 
        
                   gc.collect() 
        
                   backend_empty_cache(torch_device) 
        
               def tearDown(self): 
        
                   super().tearDown() 
        
                   gc.collect() 
        
                   backend_empty_cache(torch_device) 
        
               @unittest.skip("TODO: test needs to be implemented") 
        
               def test_sana_video_480p(self): 
        
                   pass

Problem:
Both slow integration tests are present only as @unittest.skip("TODO: test needs to be implemented").

Impact:
There is no real slow coverage for published Sana Video checkpoints, scheduler/config loading, or end-to-end output stability.

Reproduction:

from pathlib import Path

for path in [
    "tests/pipelines/sana_video/test_sana_video.py",
    "tests/pipelines/sana_video/test_sana_video_i2v.py",
]:
    text = Path(path).read_text()
    print(path, '@unittest.skip("TODO: test needs to be implemented")' in text)

Relevant precedent:
Wan video pipelines include slow tests with expected output slices, e.g. tests/pipelines/wan/test_wan.py.

Suggested fix:
Implement T2V and I2V slow tests against the published Efficient-Large-Model/SANA-Video_2B_480p_diffusers checkpoint with fixed seeds, low step counts, and expected tensor/video slices.

Issue 6: Docs model link points to a non-existent repo

Affected code:

diffusers/docs/source/en/api/pipelines/sana_video.md

Line 34 in 0f1abc4

    
           | [`Efficient-Large-Model/SANA-Video_2B_480p_diffusers`](https://huggingface.co/Efficient-Large-Model/ANA-Video_2B_480p_diffusers) | `torch.bfloat16` |

Problem:
The link target is Efficient-Large-Model/ANA-Video_2B_480p_diffusers, missing the leading S.

Impact:
Users following the docs model table hit the wrong Hugging Face URL.

Reproduction:

from pathlib import Path

text = Path("docs/source/en/api/pipelines/sana_video.md").read_text()
assert "https://huggingface.co/Efficient-Large-Model/ANA-Video_2B_480p_diffusers" in text

Relevant precedent:
The examples in the pipeline code use Efficient-Large-Model/SANA-Video_2B_480p_diffusers.

Suggested fix:

| [`Efficient-Large-Model/SANA-Video_2B_480p_diffusers`](https://huggingface.co/Efficient-Large-Model/SANA-Video_2B_480p_diffusers) | `torch.bfloat16` |

Navigation Menu

Repository navigation

sana_video model/pipeline review #13606

`sana_video` model/pipeline review

Issue 1: I2V fails for `num_videos_per_prompt > 1`

Issue 2: I2V batched images are broken

Issue 3: `attention_mask` is accepted but ignored by self-attention

Issue 4: `cross_attention_dim=None` is accepted but crashes

Issue 5: Slow tests are skipped TODOs

Issue 6: Docs model link points to a non-existent repo

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

	latents = self.prepare_latents(
	image,
	batch_size * num_videos_per_prompt,
	latent_channels,
	height,
	width,
	frames,
	torch.float32,
	device,
	generator,
	latents,
	)

	conditioning_mask = latents.new_zeros(
	batch_size,
	1,
	latents.shape[2] // self.transformer_temporal_patch_size,
	latents.shape[3] // self.transformer_spatial_patch_size,
	latents.shape[4] // self.transformer_spatial_patch_size,
	)
	conditioning_mask[:, :, 0] = 1.0
	if self.do_classifier_free_guidance:
	conditioning_mask = torch.cat([conditioning_mask, conditioning_mask])

	latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents

	# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
	timestep = t.expand(conditioning_mask.shape)
	timestep = timestep * (1 - conditioning_mask)

	# predict noise model_output
	noise_pred = self.transformer(
	latent_model_input.to(dtype=transformer_dtype),
	encoder_hidden_states=prompt_embeds.to(dtype=transformer_dtype),
	encoder_attention_mask=prompt_attention_mask,
	timestep=timestep,
	return_dict=False,
	attention_kwargs=self.attention_kwargs,
	)[0]

	latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents

	# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
	timestep = t.expand(latent_model_input.shape[0])

	if image is not None and not isinstance(image, torch.Tensor) and not isinstance(image, PIL.Image.Image):
	raise ValueError(f"`image` has to be of type `torch.Tensor` or `PIL.Image.Image` but is {type(image)}")

	if isinstance(generator, list):
	image_latents = [retrieve_latents(self.vae.encode(image), sample_mode="argmax") for _ in generator]
	image_latents = torch.cat(image_latents)
	else:
	image_latents = retrieve_latents(self.vae.encode(image), sample_mode="argmax")
	image_latents = image_latents.repeat(batch_size, 1, 1, 1, 1)

	if isinstance(self.vae, AutoencoderKLLTX2Video):
	_latents_mean = self.vae.latents_mean
	_latents_std = self.vae.latents_std
	elif isinstance(self.vae, AutoencoderKLWan):
	_latents_mean = torch.tensor(self.vae.config.latents_mean)
	_latents_std = torch.tensor(self.vae.config.latents_std)
	else:
	_latents_mean = torch.zeros(image_latents.shape[1], device=image_latents.device, dtype=image_latents.dtype)
	_latents_std = torch.ones(image_latents.shape[1], device=image_latents.device, dtype=image_latents.dtype)

	latents_mean = _latents_mean.view(1, -1, 1, 1, 1).to(image_latents.device, image_latents.dtype)
	latents_std = 1.0 / _latents_std.view(1, -1, 1, 1, 1).to(image_latents.device, image_latents.dtype)
	image_latents = (image_latents - latents_mean) * latents_std

	latents[:, :, 0:1] = image_latents.to(dtype)

	if (
	not isinstance(image, torch.Tensor)
	and not isinstance(image, PIL.Image.Image)
	and not isinstance(image, list)
	):
	raise ValueError(
	"`image` has to be of type `torch.Tensor` or `PIL.Image.Image` or `list[PIL.Image.Image]` but is"
	f" {type(image)}"

	if attention_mask is not None and attention_mask.ndim == 2:
	# assume that mask is expressed as:
	# (1 = keep, 0 = discard)
	# convert mask into a bias that can be added to attention scores:
	# (keep = +0, discard = -10000.0)
	attention_mask = (1 - attention_mask.to(hidden_states.dtype)) * -10000.0
	attention_mask = attention_mask.unsqueeze(1)

	# convert encoder_attention_mask to a bias the same way we do for attention_mask
	if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:
	encoder_attention_mask = (1 - encoder_attention_mask.to(hidden_states.dtype)) * -10000.0
	encoder_attention_mask = encoder_attention_mask.unsqueeze(1)

	# 2. Transformer blocks
	if torch.is_grad_enabled() and self.gradient_checkpointing:
	for index_block, block in enumerate(self.transformer_blocks):
	hidden_states = self._gradient_checkpointing_func(
	block,
	hidden_states,
	attention_mask,
	encoder_hidden_states,
	encoder_attention_mask,
	timestep,
	post_patch_num_frames,
	post_patch_height,
	post_patch_width,
	rotary_emb,
	)
	if controlnet_block_samples is not None and 0 < index_block <= len(controlnet_block_samples):
	hidden_states = hidden_states + controlnet_block_samples[index_block - 1]

	else:
	for index_block, block in enumerate(self.transformer_blocks):
	hidden_states = block(
	hidden_states,
	attention_mask,
	encoder_hidden_states,
	encoder_attention_mask,
	timestep,
	post_patch_num_frames,
	post_patch_height,
	post_patch_width,
	rotary_emb,
	)

	attn_output = self.attn1(norm_hidden_states, rotary_emb=rotary_emb)
	hidden_states = hidden_states + gate_msa * attn_output

	# 3. Cross Attention
	if self.attn2 is not None:
	attn_output = self.attn2(
	hidden_states,
	encoder_hidden_states=encoder_hidden_states,
	attention_mask=encoder_attention_mask,

	num_cross_attention_heads: int \| None = 20,
	cross_attention_head_dim: int \| None = 112,
	cross_attention_dim: int \| None = 2240,

	if cross_attention_dim is not None:
	self.norm2 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps)
	self.attn2 = Attention(
	query_dim=dim,
	qk_norm=qk_norm,
	kv_heads=num_cross_attention_heads if qk_norm is not None else None,
	cross_attention_dim=cross_attention_dim,
	heads=num_cross_attention_heads,
	dim_head=cross_attention_head_dim,
	dropout=dropout,
	bias=True,
	out_bias=attention_out_bias,
	processor=SanaAttnProcessor2_0(),
	)

	# 3. Feed-forward
	self.ff = GLUMBTempConv(dim, dim, mlp_ratio, norm_type=None, residual_connection=False)

	self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5)

	def forward(
	self,
	hidden_states: torch.Tensor,
	attention_mask: torch.Tensor \| None = None,
	encoder_hidden_states: torch.Tensor \| None = None,
	encoder_attention_mask: torch.Tensor \| None = None,
	timestep: torch.LongTensor \| None = None,
	frames: int = None,
	height: int = None,
	width: int = None,
	rotary_emb: torch.Tensor \| None = None,
	) -> torch.Tensor:
	batch_size = hidden_states.shape[0]

	# 1. Modulation
	shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
	self.scale_shift_table[None, None] + timestep.reshape(batch_size, timestep.shape[1], 6, -1)
	).unbind(dim=2)

	# 2. Self Attention
	norm_hidden_states = self.norm1(hidden_states)
	norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa
	norm_hidden_states = norm_hidden_states.to(hidden_states.dtype)

	attn_output = self.attn1(norm_hidden_states, rotary_emb=rotary_emb)
	hidden_states = hidden_states + gate_msa * attn_output

	# 3. Cross Attention
	if self.attn2 is not None:
	attn_output = self.attn2(
	hidden_states,
	encoder_hidden_states=encoder_hidden_states,
	attention_mask=encoder_attention_mask,
	)
	hidden_states = attn_output + hidden_states

	# 4. Feed-forward
	norm_hidden_states = self.norm2(hidden_states)

	@slow
	@require_torch_accelerator
	class SanaVideoPipelineIntegrationTests(unittest.TestCase):
	prompt = "Evening, backlight, side lighting, soft light, high contrast, mid-shot, centered composition, clean solo shot, warm color. A young Caucasian man stands in a forest."

	def setUp(self):
	super().setUp()
	gc.collect()
	backend_empty_cache(torch_device)

	def tearDown(self):
	super().tearDown()
	gc.collect()
	backend_empty_cache(torch_device)

	@unittest.skip("TODO: test needs to be implemented")
	def test_sana_video_480p(self):
	pass

sana_video model/pipeline review #13606

Description

sana_video model/pipeline review

Issue 1: I2V fails for num_videos_per_prompt > 1

Issue 2: I2V batched images are broken

Issue 3: attention_mask is accepted but ignored by self-attention

Issue 4: cross_attention_dim=None is accepted but crashes

Issue 5: Slow tests are skipped TODOs

Issue 6: Docs model link points to a non-existent repo

Activity

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

`sana_video` model/pipeline review

Issue 1: I2V fails for `num_videos_per_prompt > 1`

Issue 3: `attention_mask` is accepted but ignored by self-attention

Issue 4: `cross_attention_dim=None` is accepted but crashes