diff --git a/fastvideo/configs/sample/ltx2.py b/fastvideo/configs/sample/ltx2.py index 4fae60709..30d932721 100644 --- a/fastvideo/configs/sample/ltx2.py +++ b/fastvideo/configs/sample/ltx2.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -from dataclasses import dataclass +from dataclasses import dataclass, field from fastvideo.configs.sample.base import SamplingParam @@ -44,6 +44,11 @@ class LTX2BaseSamplingParam(SamplingParam): ltx2_modality_scale_video: float = 3.0 ltx2_modality_scale_audio: float = 3.0 ltx2_rescale_scale: float = 0.7 + # STG (Spatio-Temporal Guidance) defaults from official LTX-2. + ltx2_stg_scale_video: float = 1.0 + ltx2_stg_scale_audio: float = 1.0 + ltx2_stg_blocks_video: list[int] = field(default_factory=lambda: [29]) + ltx2_stg_blocks_audio: list[int] = field(default_factory=lambda: [29]) @dataclass diff --git a/fastvideo/models/dits/ltx2.py b/fastvideo/models/dits/ltx2.py index bb582f3d3..55abc4fd5 100644 --- a/fastvideo/models/dits/ltx2.py +++ b/fastvideo/models/dits/ltx2.py @@ -1636,6 +1636,8 @@ class BasicAVTransformerBlock(torch.nn.Module): video_attention_mask: torch.Tensor | None = None, audio_attention_mask: torch.Tensor | None = None, skip_cross_modal_attn: bool = False, + skip_video_self_attn: bool = False, + skip_audio_self_attn: bool = False, ) -> tuple[TransformerArgs | None, TransformerArgs | None]: """Forward pass for transformer block. @@ -1646,6 +1648,10 @@ class BasicAVTransformerBlock(torch.nn.Module): audio_attention_mask: SP padding attention mask for audio [B, padded_seq_len] skip_cross_modal_attn: If True, skip A2V and V2A cross-modal attention (used for the modality-isolated CFG pass). + skip_video_self_attn: If True, skip video self-attention + in this block (used for STG perturbed pass). + skip_audio_self_attn: If True, skip audio self-attention + in this block (used for STG perturbed pass). """ vx = video.x if video is not None else None ax = audio.x if audio is not None else None @@ -1659,12 +1665,13 @@ class BasicAVTransformerBlock(torch.nn.Module): vshift_msa, vscale_msa, vgate_msa = self.get_ada_values( self.scale_shift_table, vx.shape[0], video.timesteps, slice(0, 3) ) - norm_vx = torch.nn.functional.rms_norm(vx, (vx.shape[-1],), eps=self.norm_eps) * (1 + vscale_msa) + vshift_msa - # Self-attention: pass SP attention mask for distributed attention - if self.use_distributed_attention: - vx = vx + self.attn1(norm_vx, pe=video.positional_embeddings, attention_mask=video_attention_mask) * vgate_msa - else: - vx = vx + self.attn1(norm_vx, pe=video.positional_embeddings) * vgate_msa + if not skip_video_self_attn: + norm_vx = torch.nn.functional.rms_norm(vx, (vx.shape[-1],), eps=self.norm_eps) * (1 + vscale_msa) + vshift_msa + # Self-attention: pass SP attention mask for distributed attention + if self.use_distributed_attention: + vx = vx + self.attn1(norm_vx, pe=video.positional_embeddings, attention_mask=video_attention_mask) * vgate_msa + else: + vx = vx + self.attn1(norm_vx, pe=video.positional_embeddings) * vgate_msa # Text cross-attention: no SP mask needed (text is replicated, uses local attention) vx = vx + self.attn2( torch.nn.functional.rms_norm(vx, (vx.shape[-1],), eps=self.norm_eps), @@ -1676,12 +1683,13 @@ class BasicAVTransformerBlock(torch.nn.Module): ashift_msa, ascale_msa, agate_msa = self.get_ada_values( self.audio_scale_shift_table, ax.shape[0], audio.timesteps, slice(0, 3) ) - norm_ax = torch.nn.functional.rms_norm(ax, (ax.shape[-1],), eps=self.norm_eps) * (1 + ascale_msa) + ashift_msa - # Self-attention: pass SP attention mask for distributed attention - if self.use_distributed_attention: - ax = ax + self.audio_attn1(norm_ax, pe=audio.positional_embeddings, attention_mask=audio_attention_mask) * agate_msa - else: - ax = ax + self.audio_attn1(norm_ax, pe=audio.positional_embeddings) * agate_msa + if not skip_audio_self_attn: + norm_ax = torch.nn.functional.rms_norm(ax, (ax.shape[-1],), eps=self.norm_eps) * (1 + ascale_msa) + ashift_msa + # Self-attention: pass SP attention mask for distributed attention + if self.use_distributed_attention: + ax = ax + self.audio_attn1(norm_ax, pe=audio.positional_embeddings, attention_mask=audio_attention_mask) * agate_msa + else: + ax = ax + self.audio_attn1(norm_ax, pe=audio.positional_embeddings) * agate_msa # Text cross-attention: no SP mask needed (text is replicated, uses local attention) ax = ax + self.audio_attn2( torch.nn.functional.rms_norm(ax, (ax.shape[-1],), eps=self.norm_eps), @@ -2111,14 +2119,24 @@ class LTXModel(torch.nn.Module): video_attention_mask: torch.Tensor | None = None, audio_attention_mask: torch.Tensor | None = None, skip_cross_modal_attn: bool = False, + skip_video_self_attn_blocks: list[int] | None = None, + skip_audio_self_attn_blocks: list[int] | None = None, ) -> tuple[TransformerArgs | None, TransformerArgs | None]: - for block in self.transformer_blocks: + # Convert once so per-block membership checks stay O(1). + skip_video_self_attn_block_set = set(skip_video_self_attn_blocks or []) + skip_audio_self_attn_block_set = set(skip_audio_self_attn_blocks or []) + + for idx, block in enumerate(self.transformer_blocks): + skip_v_sa = idx in skip_video_self_attn_block_set + skip_a_sa = idx in skip_audio_self_attn_block_set video, audio = block( video=video, audio=audio, video_attention_mask=video_attention_mask, audio_attention_mask=audio_attention_mask, skip_cross_modal_attn=skip_cross_modal_attn, + skip_video_self_attn=skip_v_sa, + skip_audio_self_attn=skip_a_sa, ) return video, audio @@ -2146,6 +2164,8 @@ class LTXModel(torch.nn.Module): video_attention_mask: torch.Tensor | None = None, audio_attention_mask: torch.Tensor | None = None, skip_cross_modal_attn: bool = False, + skip_video_self_attn_blocks: list[int] | None = None, + skip_audio_self_attn_blocks: list[int] | None = None, ) -> tuple[torch.Tensor | None, torch.Tensor | None]: """Forward pass through the LTX model. @@ -2157,6 +2177,10 @@ class LTXModel(torch.nn.Module): skip_cross_modal_attn: If True, skip A2V and V2A cross-modal attention in all transformer blocks (modality-isolated CFG pass). + skip_video_self_attn_blocks: Block indices where video + self-attention is skipped (STG perturbed pass). + skip_audio_self_attn_blocks: Block indices where audio + self-attention is skipped (STG perturbed pass). """ if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") == "1": _debug_block_log_line( @@ -2181,6 +2205,8 @@ class LTXModel(torch.nn.Module): video_attention_mask=video_attention_mask, audio_attention_mask=audio_attention_mask, skip_cross_modal_attn=skip_cross_modal_attn, + skip_video_self_attn_blocks=skip_video_self_attn_blocks, + skip_audio_self_attn_blocks=skip_audio_self_attn_blocks, ) vx = ( @@ -2356,6 +2382,8 @@ class LTX2Transformer3DModel(BaseDiT): audio_timestep: torch.Tensor | None = None, audio_encoder_attention_mask: torch.Tensor | None = None, skip_cross_modal_attn: bool = False, + skip_video_self_attn_blocks: list[int] | None = None, + skip_audio_self_attn_blocks: list[int] | None = None, **kwargs, ) -> torch.Tensor: if isinstance(encoder_hidden_states, list): @@ -2513,6 +2541,8 @@ class LTX2Transformer3DModel(BaseDiT): video_attention_mask=video_attention_mask, audio_attention_mask=audio_attention_mask, skip_cross_modal_attn=skip_cross_modal_attn, + skip_video_self_attn_blocks=skip_video_self_attn_blocks, + skip_audio_self_attn_blocks=skip_audio_self_attn_blocks, ) # Denoised prediction diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/pipelines/pipeline_batch_info.py index ea97f2e8b..33bc8fc1d 100644 --- a/fastvideo/pipelines/pipeline_batch_info.py +++ b/fastvideo/pipelines/pipeline_batch_info.py @@ -180,6 +180,11 @@ class ForwardBatch: ltx2_modality_scale_video: float = 1.0 ltx2_modality_scale_audio: float = 1.0 ltx2_rescale_scale: float = 0.0 + # STG (Spatio-Temporal Guidance) parameters + ltx2_stg_scale_video: float = 0.0 + ltx2_stg_scale_audio: float = 0.0 + ltx2_stg_blocks_video: list[int] = field(default_factory=list) + ltx2_stg_blocks_audio: list[int] = field(default_factory=list) n_tokens: int | None = None @@ -220,8 +225,10 @@ class ForwardBatch: def __post_init__(self): """Initialize dependent fields after dataclass initialization.""" - # Set do_classifier_free_guidance based on guidance scale and negative prompt - if self.guidance_scale > 1.0: + # Enable CFG for standard guidance_scale and LTX-2 text CFG scales. + ltx2_text_cfg_enabled = (self.ltx2_cfg_scale_video != 1.0 + or self.ltx2_cfg_scale_audio != 1.0) + if self.guidance_scale > 1.0 or ltx2_text_cfg_enabled: self.do_classifier_free_guidance = True if self.negative_prompt_embeds is None: self.negative_prompt_embeds = [] diff --git a/fastvideo/pipelines/stages/ltx2_denoising.py b/fastvideo/pipelines/stages/ltx2_denoising.py index 6ac573f49..fab71d5e2 100644 --- a/fastvideo/pipelines/stages/ltx2_denoising.py +++ b/fastvideo/pipelines/stages/ltx2_denoising.py @@ -109,8 +109,9 @@ class LTX2DenoisingStage(PipelineStage): neg_prompt_mask = None # Only load negative prompts if CFG is actually enabled if batch.do_classifier_free_guidance: - assert batch.negative_prompt_embeds is not None, ( - "CFG is enabled but negative_prompt_embeds is None") + if not batch.negative_prompt_embeds: + raise ValueError( + "CFG is enabled but negative_prompt_embeds is empty") neg_prompt_embeds = batch.negative_prompt_embeds[0] # Ensure text conditioning is on the same device as latents. @@ -225,13 +226,28 @@ class LTX2DenoisingStage(PipelineStage): modality_scale_video = batch.ltx2_modality_scale_video modality_scale_audio = batch.ltx2_modality_scale_audio rescale_scale = batch.ltx2_rescale_scale + # STG (Spatio-Temporal Guidance) parameters. + stg_scale_video = batch.ltx2_stg_scale_video + stg_scale_audio = batch.ltx2_stg_scale_audio + stg_blocks_video = batch.ltx2_stg_blocks_video + stg_blocks_audio = batch.ltx2_stg_blocks_audio + do_stg_video = stg_scale_video != 0.0 + do_stg_audio = stg_scale_audio != 0.0 + do_stg = do_stg_video or do_stg_audio + do_cfg_text = (cfg_scale_video != 1.0 or cfg_scale_audio != 1.0) + do_mod = (modality_scale_video != 1.0 or modality_scale_audio != 1.0) + do_guidance = do_cfg_text or do_mod or do_stg - do_cfg = batch.do_classifier_free_guidance + if do_cfg_text and neg_prompt_embeds is None: + raise ValueError("LTX-2 text CFG is enabled " + "(ltx2_cfg_scale_video/audio != 1.0), " + "but negative prompt embeddings are missing") logger.info( "[LTX2] Denoising start: steps=%d dtype=%s " "cfg_video=%.1f cfg_audio=%.1f mod_video=%.1f " "mod_audio=%.1f rescale=%.2f " + "stg_video=%.1f stg_audio=%.1f " "sigmas_shape=%s latents_shape=%s", batch.num_inference_steps, target_dtype, @@ -240,6 +256,8 @@ class LTX2DenoisingStage(PipelineStage): modality_scale_video, modality_scale_audio, rescale_scale, + stg_scale_video, + stg_scale_audio, tuple(sigmas.shape), tuple(latents.shape), ) @@ -290,51 +308,83 @@ class LTX2DenoisingStage(PipelineStage): pos_denoised = pos_outputs pos_audio = None - if do_cfg: - # Pass 2: Unconditioned text (negative prompt) - neg_outputs = self.transformer( - hidden_states=latents.to(target_dtype), - encoder_hidden_states=neg_prompt_embeds, - encoder_attention_mask=neg_prompt_mask, - timestep=timestep, - audio_hidden_states=audio_latents, - audio_encoder_hidden_states=audio_context_n, - audio_timestep=audio_timestep, - ) - if isinstance(neg_outputs, tuple): - neg_denoised, neg_audio = neg_outputs - else: - neg_denoised = neg_outputs - neg_audio = None + if do_guidance: + # Defaults: (pos - pos) = 0 under each scale. + neg_denoised = pos_denoised + neg_audio = pos_audio + mod_denoised = pos_denoised + mod_audio = pos_audio + ptb_denoised = pos_denoised + ptb_audio = pos_audio - # Pass 3: Modality-isolated (skip cross-modal attn) - mod_outputs = self.transformer( - hidden_states=latents.to(target_dtype), - encoder_hidden_states=prompt_embeds, - encoder_attention_mask=prompt_mask, - timestep=timestep, - audio_hidden_states=audio_latents, - audio_encoder_hidden_states=audio_context_p, - audio_timestep=audio_timestep, - skip_cross_modal_attn=True, - ) - if isinstance(mod_outputs, tuple): - mod_denoised, mod_audio = mod_outputs - else: - mod_denoised = mod_outputs - mod_audio = None + # Pass 2: text CFG (negative prompt) + if do_cfg_text: + neg_outputs = self.transformer( + hidden_states=latents.to(target_dtype), + encoder_hidden_states=neg_prompt_embeds, + encoder_attention_mask=neg_prompt_mask, + timestep=timestep, + audio_hidden_states=audio_latents, + audio_encoder_hidden_states=audio_context_n, + audio_timestep=audio_timestep, + ) + if isinstance(neg_outputs, tuple): + neg_denoised, neg_audio = neg_outputs + else: + neg_denoised = neg_outputs + + # Pass 3: Modality-isolated (skip cross-modal + # attn) + if do_mod: + mod_outputs = self.transformer( + hidden_states=latents.to(target_dtype), + encoder_hidden_states=prompt_embeds, + encoder_attention_mask=prompt_mask, + timestep=timestep, + audio_hidden_states=audio_latents, + audio_encoder_hidden_states=audio_context_p, + audio_timestep=audio_timestep, + skip_cross_modal_attn=True, + ) + if isinstance(mod_outputs, tuple): + mod_denoised, mod_audio = mod_outputs + else: + mod_denoised = mod_outputs + + # Pass 4: STG perturbed (skip self-attn in + # specified blocks) + if do_stg: + ptb_outputs = self.transformer( + hidden_states=latents.to(target_dtype), + encoder_hidden_states=prompt_embeds, + encoder_attention_mask=prompt_mask, + timestep=timestep, + audio_hidden_states=audio_latents, + audio_encoder_hidden_states=(audio_context_p), + audio_timestep=audio_timestep, + skip_video_self_attn_blocks=( + stg_blocks_video if do_stg_video else None), + skip_audio_self_attn_blocks=( + stg_blocks_audio if do_stg_audio else None), + ) + if isinstance(ptb_outputs, tuple): + ptb_denoised, ptb_audio = ptb_outputs + else: + ptb_denoised = ptb_outputs # Multi-modal guidance formula per stream. vid = (pos_denoised + (cfg_scale_video - 1) * (pos_denoised - neg_denoised) + (modality_scale_video - 1) * - (pos_denoised - mod_denoised)) + (pos_denoised - mod_denoised) + stg_scale_video * + (pos_denoised - ptb_denoised)) aud = None if pos_audio is not None: aud = (pos_audio + (cfg_scale_audio - 1) * (pos_audio - neg_audio) + (modality_scale_audio - 1) * - (pos_audio - mod_audio)) + (pos_audio - mod_audio) + stg_scale_audio * + (pos_audio - ptb_audio)) # Guidance rescaling (prevents saturation). if rescale_scale > 0: