nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[bugfix] Added ltx2 guidance missing modulation term (#1100)

Browse Source 
Co-authored-by: SolitaryThinker <wlsaidhi@gmail.com>
This commit is contained in:
Shao Duan
2026-02-23 14:22:31 -08:00
committed by GitHub
co-authored by SolitaryThinker
parent 128d0b7fc5
commit 37cab18780
4 changed files with 145 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/configs/sample/ltx2.py
+6 -1
View File
@@ -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
+43 -13
View File
@@ -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
fastvideo/pipelines/pipeline_batch_info.py
+9 -2
View File
@@ -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 = []
fastvideo/pipelines/stages/ltx2_denoising.py
+87 -37
View File
@@ -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: