update

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

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

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
 

fastvideo/pipelines/stages/ltx2_denoising.py

@@ -216,13 +216,23 @@ 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
-
-        do_cfg = batch.do_classifier_free_guidance
+        # 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
 
         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,
@@ -231,6 +241,8 @@ class LTX2DenoisingStage(PipelineStage):
             modality_scale_video,
             modality_scale_audio,
             rescale_scale,
+            stg_scale_video,
+            stg_scale_audio,
             tuple(sigmas.shape),
             tuple(latents.shape),
         )
@@ -267,51 +279,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 and neg_prompt_embeds is not None:
+                        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: