From d09c44bfa547a94835fd7de458c462a2203bb5ef Mon Sep 17 00:00:00 2001 From: Shao Duan Date: Fri, 13 Feb 2026 22:26:14 +0000 Subject: [PATCH] aligned guidance modulation term --- fastvideo/configs/sample/ltx2.py | 7 ++- fastvideo/models/dits/ltx2.py | 54 +++++++++++++++----- fastvideo/pipelines/pipeline_batch_info.py | 5 ++ fastvideo/pipelines/stages/ltx2_denoising.py | 36 +++++++++++++ 4 files changed, 88 insertions(+), 14 deletions(-) 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 db9b9a21c..ce7285a99 100644 --- a/fastvideo/models/dits/ltx2.py +++ b/fastvideo/models/dits/ltx2.py @@ -1535,6 +1535,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. @@ -1545,6 +1547,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 @@ -1558,12 +1564,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), @@ -1575,12 +1582,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), @@ -2010,14 +2018,22 @@ 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: + for idx, block in enumerate(self.transformer_blocks): + skip_v_sa = (skip_video_self_attn_blocks is not None + and idx in skip_video_self_attn_blocks) + skip_a_sa = (skip_audio_self_attn_blocks is not None + and idx in skip_audio_self_attn_blocks) 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 @@ -2045,6 +2061,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. @@ -2056,6 +2074,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( @@ -2080,6 +2102,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 = ( @@ -2249,6 +2273,8 @@ class LTX2Transformer3DModel(CachableDiT): 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): @@ -2406,6 +2432,8 @@ class LTX2Transformer3DModel(CachableDiT): 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 3ec7340ed..bc51a3802 100644 --- a/fastvideo/pipelines/pipeline_batch_info.py +++ b/fastvideo/pipelines/pipeline_batch_info.py @@ -181,6 +181,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 diff --git a/fastvideo/pipelines/stages/ltx2_denoising.py b/fastvideo/pipelines/stages/ltx2_denoising.py index 707d8ae82..ff3baf812 100644 --- a/fastvideo/pipelines/stages/ltx2_denoising.py +++ b/fastvideo/pipelines/stages/ltx2_denoising.py @@ -216,6 +216,12 @@ 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 = stg_scale_video != 0.0 or stg_scale_audio != 0.0 do_cfg = batch.do_classifier_free_guidance @@ -223,6 +229,7 @@ class LTX2DenoisingStage(PipelineStage): "[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, @@ -231,6 +238,8 @@ class LTX2DenoisingStage(PipelineStage): modality_scale_video, modality_scale_audio, rescale_scale, + stg_scale_video, + stg_scale_audio, tuple(sigmas.shape), tuple(latents.shape), ) @@ -301,17 +310,44 @@ class LTX2DenoisingStage(PipelineStage): mod_denoised = mod_outputs mod_audio = None + # Pass 4: STG perturbed (skip self-attn in + # specified blocks) + ptb_denoised = None + ptb_audio = None + 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), + skip_audio_self_attn_blocks=(stg_blocks_audio), + ) + 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)) + if ptb_denoised is not None: + vid = vid + 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)) + if ptb_audio is not None: + aud = aud + stg_scale_audio * (pos_audio - + ptb_audio) # Guidance rescaling (prevents saturation). if rescale_scale > 0: