From 6c7e720dd8e4821461fdd4cb1d5ea40282f1535b Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Fri, 6 Mar 2026 05:02:02 +0100
Subject: [PATCH 01/41] Initial implementation of perturbed attn processor for
 LTX 2.3

---
 .../models/transformers/transformer_ltx2.py   | 108 +++++++++++++++++-
 1 file changed, 107 insertions(+), 1 deletion(-)

diff --git a/src/diffusers/models/transformers/transformer_ltx2.py b/src/diffusers/models/transformers/transformer_ltx2.py
index db949ca34a1f..d29c390f9d93 100644
--- a/src/diffusers/models/transformers/transformer_ltx2.py
+++ b/src/diffusers/models/transformers/transformer_ltx2.py
@@ -217,6 +217,105 @@ def __call__(
         return hidden_states
 
 
+class LTX2PerturbedAttnProcessor:
+    r"""
+    Processor which implements attention with perturbation masking and per-head gating for LTX-2.X models.
+    """
+
+    _attention_backend = None
+    _parallel_config = None
+
+    def __init__(self):
+        if is_torch_version("<", "2.0"):
+            raise ValueError(
+                "LTX attention processors require a minimum PyTorch version of 2.0. Please upgrade your PyTorch installation."
+            )
+
+    def __call__(
+        self,
+        attn: "LTX2Attention",
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        query_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+        key_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+        perturbation_mask: torch.Tensor | None = None,
+        all_perturbed: bool | None = None,
+    ) -> torch.Tensor:
+        batch_size, sequence_length, _ = (
+            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+        )
+
+        if attention_mask is not None:
+            attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+            attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+
+        if attn.to_gate_logits is not None:
+            # Calculate gate logits on original hidden_states
+            gate_logits = attn.to_gate_logits(hidden_states)
+
+        value = attn.to_v(encoder_hidden_states)
+        if all_perturbed is None:
+            all_perturbed = torch.all(perturbation_mask == 0) if perturbation_mask is not None else False
+
+        if all_perturbed:
+            # Skip attention, use the value projection value
+            hidden_states = value
+        else:
+            query = attn.to_q(hidden_states)
+            key = attn.to_k(encoder_hidden_states)
+
+            query = attn.norm_q(query)
+            key = attn.norm_k(key)
+
+            if query_rotary_emb is not None:
+                if attn.rope_type == "interleaved":
+                    query = apply_interleaved_rotary_emb(query, query_rotary_emb)
+                    key = apply_interleaved_rotary_emb(
+                        key, key_rotary_emb if key_rotary_emb is not None else query_rotary_emb
+                    )
+                elif attn.rope_type == "split":
+                    query = apply_split_rotary_emb(query, query_rotary_emb)
+                    key = apply_split_rotary_emb(
+                        key, key_rotary_emb if key_rotary_emb is not None else query_rotary_emb
+                    )
+
+            query = query.unflatten(2, (attn.heads, -1))
+            key = key.unflatten(2, (attn.heads, -1))
+            value = value.unflatten(2, (attn.heads, -1))
+
+            hidden_states = dispatch_attention_fn(
+                query,
+                key,
+                value,
+                attn_mask=attention_mask,
+                dropout_p=0.0,
+                is_causal=False,
+                backend=self._attention_backend,
+                parallel_config=self._parallel_config,
+            )
+            hidden_states = hidden_states.flatten(2, 3)
+            hidden_states = hidden_states.to(query.dtype)
+
+            if perturbation_mask is not None:
+                value = value.flatten(2, 3)
+                hidden_states = torch.lerp(value, hidden_states, perturbation_mask)
+
+        if attn.to_gate_logits is not None:
+            hidden_states = hidden_states.unflatten(2, (attn.heads, -1))  # [B, T, H, D]
+            # The factor of 2.0 is so that if the gates logits are zero-initialized the initial gates are all 1
+            gates = 2.0 * torch.sigmoid(gate_logits)  # [B, T, H]
+            hidden_states = hidden_states * gates.unsqueeze(-1)
+            hidden_states = hidden_states.flatten(2, 3)
+
+        hidden_states = attn.to_out[0](hidden_states)
+        hidden_states = attn.to_out[1](hidden_states)
+        return hidden_states
+
+
 class LTX2Attention(torch.nn.Module, AttentionModuleMixin):
     r"""
     Attention class for all LTX-2.0 attention layers. Compared to LTX-1.0, this supports specifying the query and key
@@ -224,7 +323,7 @@ class LTX2Attention(torch.nn.Module, AttentionModuleMixin):
     """
 
     _default_processor_cls = LTX2AudioVideoAttnProcessor
-    _available_processors = [LTX2AudioVideoAttnProcessor]
+    _available_processors = [LTX2AudioVideoAttnProcessor, LTX2PerturbedAttnProcessor]
 
     def __init__(
         self,
@@ -240,6 +339,7 @@ def __init__(
         norm_eps: float = 1e-6,
         norm_elementwise_affine: bool = True,
         rope_type: str = "interleaved",
+        apply_gated_attention: bool = False,
         processor=None,
     ):
         super().__init__()
@@ -266,6 +366,12 @@ def __init__(
         self.to_out.append(torch.nn.Linear(self.inner_dim, self.out_dim, bias=out_bias))
         self.to_out.append(torch.nn.Dropout(dropout))
 
+        if apply_gated_attention:
+            # Per head gate values
+            self.to_gate_logits = torch.nn.Linear(query_dim, heads, bias=True)
+        else:
+            self.to_gate_logits = None
+
         if processor is None:
             processor = self._default_processor_cls()
         self.set_processor(processor)

From e90b90a3cc64f87c086dc949f88f684efe157b40 Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Sat, 7 Mar 2026 03:32:03 +0100
Subject: [PATCH 02/41] Update DiT block for LTX 2.3 + add self_attention_mask

---
 .../models/transformers/transformer_ltx2.py   | 158 +++++++++++++-----
 1 file changed, 118 insertions(+), 40 deletions(-)

diff --git a/src/diffusers/models/transformers/transformer_ltx2.py b/src/diffusers/models/transformers/transformer_ltx2.py
index d29c390f9d93..625b26bac6f4 100644
--- a/src/diffusers/models/transformers/transformer_ltx2.py
+++ b/src/diffusers/models/transformers/transformer_ltx2.py
@@ -178,6 +178,10 @@ def __call__(
         if encoder_hidden_states is None:
             encoder_hidden_states = hidden_states
 
+        if attn.to_gate_logits is not None:
+            # Calculate gate logits on original hidden_states
+            gate_logits = attn.to_gate_logits(hidden_states)
+
         query = attn.to_q(hidden_states)
         key = attn.to_k(encoder_hidden_states)
         value = attn.to_v(encoder_hidden_states)
@@ -212,6 +216,13 @@ def __call__(
         hidden_states = hidden_states.flatten(2, 3)
         hidden_states = hidden_states.to(query.dtype)
 
+        if attn.to_gate_logits is not None:
+            hidden_states = hidden_states.unflatten(2, (attn.heads, -1))  # [B, T, H, D]
+            # The factor of 2.0 is so that if the gates logits are zero-initialized the initial gates are all 1
+            gates = 2.0 * torch.sigmoid(gate_logits)  # [B, T, H]
+            hidden_states = hidden_states * gates.unsqueeze(-1)
+            hidden_states = hidden_states.flatten(2, 3)
+
         hidden_states = attn.to_out[0](hidden_states)
         hidden_states = attn.to_out[1](hidden_states)
         return hidden_states
@@ -427,6 +438,10 @@ def __init__(
         audio_num_attention_heads: int,
         audio_attention_head_dim,
         audio_cross_attention_dim: int,
+        video_gated_attn: bool = False,
+        video_cross_attn_adaln: bool = False,
+        audio_gated_attn: bool = False,
+        audio_cross_attn_adaln: bool = False,
         qk_norm: str = "rms_norm_across_heads",
         activation_fn: str = "gelu-approximate",
         attention_bias: bool = True,
@@ -449,6 +464,8 @@ def __init__(
             out_bias=attention_out_bias,
             qk_norm=qk_norm,
             rope_type=rope_type,
+            apply_gated_attention=video_gated_attn,
+            processor=LTX2AudioVideoAttnProcessor(),
         )
 
         self.audio_norm1 = RMSNorm(audio_dim, eps=eps, elementwise_affine=elementwise_affine)
@@ -462,6 +479,8 @@ def __init__(
             out_bias=attention_out_bias,
             qk_norm=qk_norm,
             rope_type=rope_type,
+            apply_gated_attention=audio_gated_attn,
+            processor=LTX2AudioVideoAttnProcessor(),
         )
 
         # 2. Prompt Cross-Attention
@@ -476,6 +495,8 @@ def __init__(
             out_bias=attention_out_bias,
             qk_norm=qk_norm,
             rope_type=rope_type,
+            apply_gated_attention=video_gated_attn,
+            processor=LTX2AudioVideoAttnProcessor(),
         )
 
         self.audio_norm2 = RMSNorm(audio_dim, eps=eps, elementwise_affine=elementwise_affine)
@@ -489,6 +510,8 @@ def __init__(
             out_bias=attention_out_bias,
             qk_norm=qk_norm,
             rope_type=rope_type,
+            apply_gated_attention=audio_gated_attn,
+            processor=LTX2AudioVideoAttnProcessor(),
         )
 
         # 3. Audio-to-Video (a2v) and Video-to-Audio (v2a) Cross-Attention
@@ -504,6 +527,8 @@ def __init__(
             out_bias=attention_out_bias,
             qk_norm=qk_norm,
             rope_type=rope_type,
+            apply_gated_attention=video_gated_attn,
+            processor=LTX2AudioVideoAttnProcessor(),
         )
 
         # Video-to-Audio (v2a) Attention --> Q: Audio; K,V: Video
@@ -518,6 +543,8 @@ def __init__(
             out_bias=attention_out_bias,
             qk_norm=qk_norm,
             rope_type=rope_type,
+            apply_gated_attention=audio_gated_attn,
+            processor=LTX2AudioVideoAttnProcessor(),
         )
 
         # 4. Feedforward layers
@@ -528,14 +555,37 @@ def __init__(
         self.audio_ff = FeedForward(audio_dim, activation_fn=activation_fn)
 
         # 5. Per-Layer Modulation Parameters
-        # Self-Attention / Feedforward AdaLayerNorm-Zero mod params
-        self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5)
-        self.audio_scale_shift_table = nn.Parameter(torch.randn(6, audio_dim) / audio_dim**0.5)
+        # Self-Attention (attn1) / Feedforward AdaLayerNorm-Zero mod params
+        # 6 base mod params for text cross-attn K,V; if cross_attn_adaln, also has mod params for Q
+        self.video_cross_attn_adaln = video_cross_attn_adaln
+        self.audio_cross_attn_adaln = audio_cross_attn_adaln
+        video_mod_param_num = 9 if self.video_cross_attn_adaln else 6
+        audio_mod_param_num = 9 if self.audio_cross_attn_adaln else 6
+        self.scale_shift_table = nn.Parameter(torch.randn(video_mod_param_num, dim) / dim**0.5)
+        self.audio_scale_shift_table = nn.Parameter(torch.randn(audio_mod_param_num, audio_dim) / audio_dim**0.5)
+
+        # Prompt cross-attn (attn2) additional modulation params
+        self.cross_attn_adaln = video_cross_attn_adaln or audio_cross_attn_adaln
+        if self.cross_attn_adaln:
+            self.prompt_scale_shift_table = nn.Parameter(torch.randn(2, dim))
+            self.audio_prompt_scale_shift_table = nn.Parameter(torch.randn(2, dim))
 
         # Per-layer a2v, v2a Cross-Attention mod params
         self.video_a2v_cross_attn_scale_shift_table = nn.Parameter(torch.randn(5, dim))
         self.audio_a2v_cross_attn_scale_shift_table = nn.Parameter(torch.randn(5, audio_dim))
 
+    @staticmethod
+    def get_mod_params(
+        scale_shift_table: torch.Tensor, temb: torch.Tensor, batch_size: int
+    ) -> tuple[torch.Tensor, ...]:
+        num_ada_params = scale_shift_table.shape[0]
+        ada_values = (
+            scale_shift_table[None, None].to(temb.device)
+            + temb.reshape(batch_size, temb.shape[1], num_ada_params, -1)
+        )
+        ada_params = ada_values.unbind(dim=2)
+        return ada_params
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -548,6 +598,8 @@ def forward(
         temb_ca_audio_scale_shift: torch.Tensor,
         temb_ca_gate: torch.Tensor,
         temb_ca_audio_gate: torch.Tensor,
+        temb_prompt: torch.Tensor | None = None,
+        temb_prompt_audio: torch.Tensor | None = None,
         video_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
         audio_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
         ca_video_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
@@ -560,13 +612,13 @@ def forward(
         batch_size = hidden_states.size(0)
 
         # 1. Video and Audio Self-Attention
-        norm_hidden_states = self.norm1(hidden_states)
+        # 1.1. Video Self-Attention
+        video_ada_params = self.get_mod_params(self.scale_shift_table, temb, batch_size)
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = video_ada_params[:6]
+        if self.video_cross_attn_adaln:
+            shift_text_q, scale_text_q, gate_text_q = video_ada_params[6:9]
 
-        num_ada_params = self.scale_shift_table.shape[0]
-        ada_values = self.scale_shift_table[None, None].to(temb.device) + temb.reshape(
-            batch_size, temb.size(1), num_ada_params, -1
-        )
-        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ada_values.unbind(dim=2)
+        norm_hidden_states = self.norm1(hidden_states)
         norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa
 
         attn_hidden_states = self.attn1(
@@ -576,15 +628,15 @@ def forward(
         )
         hidden_states = hidden_states + attn_hidden_states * gate_msa
 
-        norm_audio_hidden_states = self.audio_norm1(audio_hidden_states)
-
-        num_audio_ada_params = self.audio_scale_shift_table.shape[0]
-        audio_ada_values = self.audio_scale_shift_table[None, None].to(temb_audio.device) + temb_audio.reshape(
-            batch_size, temb_audio.size(1), num_audio_ada_params, -1
-        )
+        # 1.2. Audio Self-Attention
+        audio_ada_params = self.get_mod_params(self.audio_scale_shift_table, temb_audio, batch_size)
         audio_shift_msa, audio_scale_msa, audio_gate_msa, audio_shift_mlp, audio_scale_mlp, audio_gate_mlp = (
-            audio_ada_values.unbind(dim=2)
+            audio_ada_params[:6]
         )
+        if self.audio_cross_attn_adaln:
+            audio_shift_text_q, audio_scale_text_q, audio_gate_text_q = audio_ada_params[6:9]
+
+        norm_audio_hidden_states = self.audio_norm1(audio_hidden_states)
         norm_audio_hidden_states = norm_audio_hidden_states * (1 + audio_scale_msa) + audio_shift_msa
 
         attn_audio_hidden_states = self.audio_attn1(
@@ -594,63 +646,74 @@ def forward(
         )
         audio_hidden_states = audio_hidden_states + attn_audio_hidden_states * audio_gate_msa
 
-        # 2. Video and Audio Cross-Attention with the text embeddings
+        # 2. Video and Audio Cross-Attention with the text embeddings (Q: Video or Audio; K,V: Text)
+        if self.cross_attn_adaln:
+            video_prompt_ada_params = self.get_mod_params(self.prompt_scale_shift_table, temb_prompt, batch_size)
+            shift_text_kv, scale_text_kv = video_prompt_ada_params
+
+            audio_prompt_ada_params = self.get_mod_params(self.audio_prompt_scale_shift_table, temb_prompt_audio, batch_size)
+            audio_shift_text_kv, audio_scale_text_kv = audio_prompt_ada_params
+
+        # 2.1. Video-Text Cross-Attention (Q: Video; K,V: Test)
         norm_hidden_states = self.norm2(hidden_states)
+        if self.video_cross_attn_adaln:
+            norm_hidden_states = norm_hidden_states * (1 + scale_text_q) + shift_text_q
+        if self.cross_attn_adaln:
+            encoder_hidden_states = encoder_hidden_states * (1 + scale_text_kv) + shift_text_kv
+
         attn_hidden_states = self.attn2(
             norm_hidden_states,
             encoder_hidden_states=encoder_hidden_states,
             query_rotary_emb=None,
             attention_mask=encoder_attention_mask,
         )
+        if self.video_cross_attn_adaln:
+            attn_hidden_states = attn_hidden_states * gate_text_q
         hidden_states = hidden_states + attn_hidden_states
 
+        # 2.2. Audio-Text Cross-Attention
         norm_audio_hidden_states = self.audio_norm2(audio_hidden_states)
+        if self.audio_cross_attn_adaln:
+            norm_audio_hidden_states = norm_audio_hidden_states * (1 + audio_scale_text_q) + audio_shift_text_q
+        if self.cross_attn_adaln:
+            audio_encoder_hidden_states = audio_encoder_hidden_states * (1 + audio_scale_text_kv) + audio_shift_text_kv
+
         attn_audio_hidden_states = self.audio_attn2(
             norm_audio_hidden_states,
             encoder_hidden_states=audio_encoder_hidden_states,
             query_rotary_emb=None,
             attention_mask=audio_encoder_attention_mask,
         )
+        if self.audio_cross_attn_adaln:
+            attn_audio_hidden_states = attn_audio_hidden_states * audio_gate_text_q
         audio_hidden_states = audio_hidden_states + attn_audio_hidden_states
 
         # 3. Audio-to-Video (a2v) and Video-to-Audio (v2a) Cross-Attention
         norm_hidden_states = self.audio_to_video_norm(hidden_states)
         norm_audio_hidden_states = self.video_to_audio_norm(audio_hidden_states)
 
-        # Combine global and per-layer cross attention modulation parameters
+        # 3.1. Combine global and per-layer cross attention modulation parameters
         # Video
         video_per_layer_ca_scale_shift = self.video_a2v_cross_attn_scale_shift_table[:4, :]
         video_per_layer_ca_gate = self.video_a2v_cross_attn_scale_shift_table[4:, :]
 
-        video_ca_scale_shift_table = (
-            video_per_layer_ca_scale_shift[:, :, ...].to(temb_ca_scale_shift.dtype)
-            + temb_ca_scale_shift.reshape(batch_size, temb_ca_scale_shift.shape[1], 4, -1)
-        ).unbind(dim=2)
-        video_ca_gate = (
-            video_per_layer_ca_gate[:, :, ...].to(temb_ca_gate.dtype)
-            + temb_ca_gate.reshape(batch_size, temb_ca_gate.shape[1], 1, -1)
-        ).unbind(dim=2)
+        video_ca_ada_params = self.get_mod_params(video_per_layer_ca_scale_shift, temb_ca_scale_shift, batch_size)
+        video_ca_gate_param = self.get_mod_params(video_per_layer_ca_gate, temb_ca_gate, batch_size)
 
-        video_a2v_ca_scale, video_a2v_ca_shift, video_v2a_ca_scale, video_v2a_ca_shift = video_ca_scale_shift_table
-        a2v_gate = video_ca_gate[0].squeeze(2)
+        video_a2v_ca_scale, video_a2v_ca_shift, video_v2a_ca_scale, video_v2a_ca_shift = video_ca_ada_params
+        a2v_gate = video_ca_gate_param[0].squeeze(2)
 
         # Audio
         audio_per_layer_ca_scale_shift = self.audio_a2v_cross_attn_scale_shift_table[:4, :]
         audio_per_layer_ca_gate = self.audio_a2v_cross_attn_scale_shift_table[4:, :]
 
-        audio_ca_scale_shift_table = (
-            audio_per_layer_ca_scale_shift[:, :, ...].to(temb_ca_audio_scale_shift.dtype)
-            + temb_ca_audio_scale_shift.reshape(batch_size, temb_ca_audio_scale_shift.shape[1], 4, -1)
-        ).unbind(dim=2)
-        audio_ca_gate = (
-            audio_per_layer_ca_gate[:, :, ...].to(temb_ca_audio_gate.dtype)
-            + temb_ca_audio_gate.reshape(batch_size, temb_ca_audio_gate.shape[1], 1, -1)
-        ).unbind(dim=2)
+        audio_ca_ada_params = self.get_mod_params(audio_per_layer_ca_scale_shift, temb_ca_audio_scale_shift, batch_size)
+        audio_ca_gate_param = self.get_mod_params(audio_per_layer_ca_gate, temb_ca_audio_gate, batch_size)
 
-        audio_a2v_ca_scale, audio_a2v_ca_shift, audio_v2a_ca_scale, audio_v2a_ca_shift = audio_ca_scale_shift_table
-        v2a_gate = audio_ca_gate[0].squeeze(2)
+        audio_a2v_ca_scale, audio_a2v_ca_shift, audio_v2a_ca_scale, audio_v2a_ca_shift = audio_ca_ada_params
+        v2a_gate = audio_ca_gate_param[0].squeeze(2)
 
-        # Audio-to-Video Cross Attention: Q: Video; K,V: Audio
+        # 3.2. Audio-to-Video Cross Attention: Q: Video; K,V: Audio
         mod_norm_hidden_states = norm_hidden_states * (1 + video_a2v_ca_scale.squeeze(2)) + video_a2v_ca_shift.squeeze(
             2
         )
@@ -668,7 +731,7 @@ def forward(
 
         hidden_states = hidden_states + a2v_gate * a2v_attn_hidden_states
 
-        # Video-to-Audio Cross Attention: Q: Audio; K,V: Video
+        # 3.3. Video-to-Audio Cross Attention: Q: Audio; K,V: Video
         mod_norm_hidden_states = norm_hidden_states * (1 + video_v2a_ca_scale.squeeze(2)) + video_v2a_ca_shift.squeeze(
             2
         )
@@ -1209,6 +1272,7 @@ def forward(
         audio_timestep: torch.LongTensor | None = None,
         encoder_attention_mask: torch.Tensor | None = None,
         audio_encoder_attention_mask: torch.Tensor | None = None,
+        self_attention_mask: torch.Tensor | None = None,
         num_frames: int | None = None,
         height: int | None = None,
         width: int | None = None,
@@ -1241,6 +1305,8 @@ def forward(
                 Optional multiplicative text attention mask of shape `(batch_size, text_seq_len)`.
             audio_encoder_attention_mask (`torch.Tensor`, *optional*):
                 Optional multiplicative text attention mask of shape `(batch_size, text_seq_len)` for audio modeling.
+            self_attention_mask (`torch.Tensor`, *optional*):
+                Optional multiplicative self-attention mask of shape `(batch_size, seq_len, seq_len)`.
             num_frames (`int`, *optional*):
                 The number of latent video frames. Used if calculating the video coordinates for RoPE.
             height (`int`, *optional*):
@@ -1281,6 +1347,18 @@ def forward(
             audio_encoder_attention_mask = (1 - audio_encoder_attention_mask.to(audio_hidden_states.dtype)) * -10000.0
             audio_encoder_attention_mask = audio_encoder_attention_mask.unsqueeze(1)
 
+        if self_attention_mask is not None and self_attention_mask.ndim == 3:
+            # Convert to additive attention mask in log-space where 0 (masked) values get mapped to a large negative
+            # number and positive values are mapped to their logarithm.
+            dtype_finfo = torch.finfo(hidden_states.dtype)
+            additive_self_attn_mask = torch.full_like(self_attention_mask, dtype_finfo.min, dtype=hidden_states.dtype)
+            unmasked_entries = self_attention_mask > 0
+            if torch.any(unmasked_entries):
+                additive_self_attn_mask[unmasked_entries] = torch.log(
+                    self_attention_mask[unmasked_entries].clamp(min=dtype_finfo.tiny)
+                ).to(hidden_states.dtype)
+            self_attention_mask = additive_self_attn_mask.unsqueeze(1)  # [batch_size, 1, seq_len, seq_len]
+
         batch_size = hidden_states.size(0)
 
         # 1. Prepare RoPE positional embeddings

From f768f8dae8d5a425e8e7f0963a0da020a57cbeab Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Sun, 8 Mar 2026 03:09:38 +0100
Subject: [PATCH 03/41] Add flag to control using perturbed attn processor for
 now

---
 .../models/transformers/transformer_ltx2.py    | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/src/diffusers/models/transformers/transformer_ltx2.py b/src/diffusers/models/transformers/transformer_ltx2.py
index 625b26bac6f4..8c230c507670 100644
--- a/src/diffusers/models/transformers/transformer_ltx2.py
+++ b/src/diffusers/models/transformers/transformer_ltx2.py
@@ -449,9 +449,15 @@ def __init__(
         eps: float = 1e-6,
         elementwise_affine: bool = False,
         rope_type: str = "interleaved",
+        perturbed_attn: bool = False,
     ):
         super().__init__()
 
+        if perturbed_attn:
+            attn_processor_cls = LTX2PerturbedAttnProcessor
+        else:
+            attn_processor_cls = LTX2AudioVideoAttnProcessor
+
         # 1. Self-Attention (video and audio)
         self.norm1 = RMSNorm(dim, eps=eps, elementwise_affine=elementwise_affine)
         self.attn1 = LTX2Attention(
@@ -465,7 +471,7 @@ def __init__(
             qk_norm=qk_norm,
             rope_type=rope_type,
             apply_gated_attention=video_gated_attn,
-            processor=LTX2AudioVideoAttnProcessor(),
+            processor=attn_processor_cls(),
         )
 
         self.audio_norm1 = RMSNorm(audio_dim, eps=eps, elementwise_affine=elementwise_affine)
@@ -480,7 +486,7 @@ def __init__(
             qk_norm=qk_norm,
             rope_type=rope_type,
             apply_gated_attention=audio_gated_attn,
-            processor=LTX2AudioVideoAttnProcessor(),
+            processor=attn_processor_cls(),
         )
 
         # 2. Prompt Cross-Attention
@@ -496,7 +502,7 @@ def __init__(
             qk_norm=qk_norm,
             rope_type=rope_type,
             apply_gated_attention=video_gated_attn,
-            processor=LTX2AudioVideoAttnProcessor(),
+            processor=attn_processor_cls(),
         )
 
         self.audio_norm2 = RMSNorm(audio_dim, eps=eps, elementwise_affine=elementwise_affine)
@@ -511,7 +517,7 @@ def __init__(
             qk_norm=qk_norm,
             rope_type=rope_type,
             apply_gated_attention=audio_gated_attn,
-            processor=LTX2AudioVideoAttnProcessor(),
+            processor=attn_processor_cls(),
         )
 
         # 3. Audio-to-Video (a2v) and Video-to-Audio (v2a) Cross-Attention
@@ -528,7 +534,7 @@ def __init__(
             qk_norm=qk_norm,
             rope_type=rope_type,
             apply_gated_attention=video_gated_attn,
-            processor=LTX2AudioVideoAttnProcessor(),
+            processor=attn_processor_cls(),
         )
 
         # Video-to-Audio (v2a) Attention --> Q: Audio; K,V: Video
@@ -544,7 +550,7 @@ def __init__(
             qk_norm=qk_norm,
             rope_type=rope_type,
             apply_gated_attention=audio_gated_attn,
-            processor=LTX2AudioVideoAttnProcessor(),
+            processor=attn_processor_cls(),
         )
 
         # 4. Feedforward layers

From cde67486cfab167e9c97afc021ae31e816dce1a9 Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Sun, 8 Mar 2026 04:32:58 +0100
Subject: [PATCH 04/41] Add support for new video upsampling blocks used by
 LTX-2.3

---
 .../autoencoders/autoencoder_kl_ltx2.py       | 52 ++++++++++++++-----
 1 file changed, 40 insertions(+), 12 deletions(-)

diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_ltx2.py b/src/diffusers/models/autoencoders/autoencoder_kl_ltx2.py
index 7c04bd715c25..775a83651c98 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_ltx2.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_ltx2.py
@@ -237,7 +237,7 @@ def forward(
 
 
 # Like LTX 1.0 LTXVideoDownsampler3d, but uses new causal Conv3d
-class LTXVideoDownsampler3d(nn.Module):
+class LTX2VideoDownsampler3d(nn.Module):
     def __init__(
         self,
         in_channels: int,
@@ -285,10 +285,11 @@ def forward(self, hidden_states: torch.Tensor, causal: bool = True) -> torch.Ten
 
 
 # Like LTX 1.0 LTXVideoUpsampler3d, but uses new causal Conv3d
-class LTXVideoUpsampler3d(nn.Module):
+class LTX2VideoUpsampler3d(nn.Module):
     def __init__(
         self,
         in_channels: int,
+        out_channels: int | None = None,
         stride: int | tuple[int, int, int] = 1,
         residual: bool = False,
         upscale_factor: int = 1,
@@ -300,7 +301,8 @@ def __init__(
         self.residual = residual
         self.upscale_factor = upscale_factor
 
-        out_channels = (in_channels * stride[0] * stride[1] * stride[2]) // upscale_factor
+        out_channels = out_channels or in_channels
+        out_channels = (out_channels * stride[0] * stride[1] * stride[2]) // upscale_factor
 
         self.conv = LTX2VideoCausalConv3d(
             in_channels=in_channels,
@@ -408,7 +410,7 @@ def __init__(
                 )
             elif downsample_type == "spatial":
                 self.downsamplers.append(
-                    LTXVideoDownsampler3d(
+                    LTX2VideoDownsampler3d(
                         in_channels=in_channels,
                         out_channels=out_channels,
                         stride=(1, 2, 2),
@@ -417,7 +419,7 @@ def __init__(
                 )
             elif downsample_type == "temporal":
                 self.downsamplers.append(
-                    LTXVideoDownsampler3d(
+                    LTX2VideoDownsampler3d(
                         in_channels=in_channels,
                         out_channels=out_channels,
                         stride=(2, 1, 1),
@@ -426,7 +428,7 @@ def __init__(
                 )
             elif downsample_type == "spatiotemporal":
                 self.downsamplers.append(
-                    LTXVideoDownsampler3d(
+                    LTX2VideoDownsampler3d(
                         in_channels=in_channels,
                         out_channels=out_channels,
                         stride=(2, 2, 2),
@@ -580,6 +582,7 @@ def __init__(
         resnet_eps: float = 1e-6,
         resnet_act_fn: str = "swish",
         spatio_temporal_scale: bool = True,
+        upsample_type: str = "spatiotemporal",
         inject_noise: bool = False,
         timestep_conditioning: bool = False,
         upsample_residual: bool = False,
@@ -609,17 +612,38 @@ def __init__(
 
         self.upsamplers = None
         if spatio_temporal_scale:
-            self.upsamplers = nn.ModuleList(
-                [
-                    LTXVideoUpsampler3d(
-                        out_channels * upscale_factor,
+            self.upsamplers = nn.ModuleList()
+
+            if upsample_type == "spatial":
+                self.upsamplers.append(
+                    LTX2VideoUpsampler3d(
+                        in_channels=out_channels * upscale_factor,
+                        stride=(1, 2, 2),
+                        residual=upsample_residual,
+                        upscale_factor=upscale_factor,
+                        spatial_padding_mode=spatial_padding_mode,
+                    )
+                )
+            elif upsample_type == "temporal":
+                self.upsamplers.append(
+                    LTX2VideoUpsampler3d(
+                        in_channels=out_channels * upscale_factor,
+                        stride=(2, 1, 1),
+                        residual=upsample_residual,
+                        upscale_factor=upscale_factor,
+                        spatial_padding_mode=spatial_padding_mode,
+                    )
+                )
+            elif upsample_type == "spatiotemporal":
+                self.upsamplers.append(
+                    LTX2VideoUpsampler3d(
+                        in_channels=out_channels * upscale_factor,
                         stride=(2, 2, 2),
                         residual=upsample_residual,
                         upscale_factor=upscale_factor,
                         spatial_padding_mode=spatial_padding_mode,
                     )
-                ]
-            )
+                )
 
         resnets = []
         for _ in range(num_layers):
@@ -862,6 +886,7 @@ def __init__(
         block_out_channels: tuple[int, ...] = (256, 512, 1024),
         spatio_temporal_scaling: tuple[bool, ...] = (True, True, True),
         layers_per_block: tuple[int, ...] = (5, 5, 5, 5),
+        upsample_type: tuple[str, ...] = ("spatiotemporal", "spatiotemporal", "spatiotemporal"),
         patch_size: int = 4,
         patch_size_t: int = 1,
         resnet_norm_eps: float = 1e-6,
@@ -917,6 +942,7 @@ def __init__(
                 num_layers=layers_per_block[i + 1],
                 resnet_eps=resnet_norm_eps,
                 spatio_temporal_scale=spatio_temporal_scaling[i],
+                upsample_type=upsample_type[i],
                 inject_noise=inject_noise[i + 1],
                 timestep_conditioning=timestep_conditioning,
                 upsample_residual=upsample_residual[i],
@@ -1062,6 +1088,7 @@ def __init__(
         decoder_spatio_temporal_scaling: tuple[bool, ...] = (True, True, True),
         decoder_inject_noise: tuple[bool, ...] = (False, False, False, False),
         downsample_type: tuple[str, ...] = ("spatial", "temporal", "spatiotemporal", "spatiotemporal"),
+        upsample_type: tuple[str, ...] = ("spatiotemporal", "spatiotemporal", "spatiotemporal"),
         upsample_residual: tuple[bool, ...] = (True, True, True),
         upsample_factor: tuple[int, ...] = (2, 2, 2),
         timestep_conditioning: bool = False,
@@ -1098,6 +1125,7 @@ def __init__(
             block_out_channels=decoder_block_out_channels,
             spatio_temporal_scaling=decoder_spatio_temporal_scaling,
             layers_per_block=decoder_layers_per_block,
+            upsample_type=upsample_type,
             patch_size=patch_size,
             patch_size_t=patch_size_t,
             resnet_norm_eps=resnet_norm_eps,

From 236eb8db64b342748494ec302dc389b0a54ec4b7 Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Sun, 8 Mar 2026 06:47:56 +0100
Subject: [PATCH 05/41] Support LTX-2.3 Big-VGAN V2-style vocoder

---
 src/diffusers/pipelines/ltx2/vocoder.py | 259 ++++++++++++++++++++++--
 1 file changed, 247 insertions(+), 12 deletions(-)

diff --git a/src/diffusers/pipelines/ltx2/vocoder.py b/src/diffusers/pipelines/ltx2/vocoder.py
index 551c3ac5980f..03b583511e74 100644
--- a/src/diffusers/pipelines/ltx2/vocoder.py
+++ b/src/diffusers/pipelines/ltx2/vocoder.py
@@ -8,6 +8,173 @@
 from ...models.modeling_utils import ModelMixin
 
 
+def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) -> torch.Tensor:
+    """
+    Creates a Kaiser sinc kernel for low-pass filtering.
+
+    Args:
+        cutoff (`float`):
+            Normalized frequency cutoff (relative to the sampling rate). Must be between 0 and 0.5 (the Nyquist
+            frequency).
+        half_width (`float`):
+            Used to determine the Kaiser window's beta parameter.
+        kernel_size:
+            Size of the Kaiser window (and ultimately the Kaiser sinc kernel).
+
+    Returns:
+        `torch.Tensor` of shape `(kernel_size,)`:
+            The Kaiser sinc kernel.
+    """
+    delta_f = 4 * half_width
+    amplitude = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95
+    if amplitude > 50.0:
+        beta = 0.1102 * (amplitude - 8.7)
+    elif amplitude >= 21.0:
+        beta = 0.5842 * (amplitude - 21) ** 0.4 + 0.07886 * (amplitude - 21.0)
+    else:
+        beta = 0.0
+
+    window = torch.kaiser_window(kernel_size, beta=beta, periodic=False)
+
+    even = kernel_size % 2 == 0
+    half_size = kernel_size // 2
+    time = torch.arange(-half_size, half_size) + 0.5 if even else torch.arange(kernel_size) - half_size
+
+    if cutoff == 0.0:
+        filter = torch.zeros_like(time)
+    else:
+        time = 2 * cutoff * time
+        sinc = torch.where(
+            time == 0,
+            torch.ones_like(time),
+            torch.sin(math.pi * time) / math.pi / time,
+        )
+        filter = 2 * cutoff * window * sinc
+        filter = filter / filter.sum()
+    return filter
+
+
+class DownSample1d(nn.Module):
+    """1D low-pass filter for antialias downsampling."""
+
+    def __init__(
+        self,
+        ratio: int = 2,
+        kernel_size: int | None = None,
+        use_padding: bool = True,
+        padding_mode: str = "replicate",
+    ):
+        super().__init__()
+        self.ratio = ratio
+        self.kernel_size = kernel_size or int(6 * ratio // 2) * 2
+        self.pad_left = self.kernel_size // 2 + (self.kernel_size % 2) - 1
+        self.pad_right = self.kernel_size // 2
+        self.use_padding = use_padding
+        self.padding_mode = padding_mode
+
+        cutoff = 0.5 / ratio
+        half_width = 0.6 / ratio
+        low_pass_filter = kaiser_sinc_filter1d(cutoff, half_width, self.kernel_size)
+        self.register_buffer("filter", low_pass_filter.view(1, 1, self,kernel_size))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x expected shape: [batch_size, num_channels, hidden_dim]
+        num_channels = x.shape[1]
+        if self.use_padding:
+            x = F.pad(x, (self.pad_left, self.pad_right), mode=self.padding_mode)
+        x_filtered = F.conv1d(x, self.filter.expand(num_channels, -1, -1), stride=self.ratio, groups=num_channels)
+        return x_filtered
+
+
+class UpSample1d(nn.Module):
+    def __init__(
+        self,
+        ratio: int = 2,
+        kernel_size: int | None  = None,
+        window_type: str = "kaiser",
+        padding_mode: str = "replicate",
+    ):
+        super().__init__()
+        self.ratio = ratio
+        self.padding_mode = padding_mode
+
+        if window_type == "hann":
+            rolloff = 0.99
+            lowpass_filter_width = 6
+            width = math.ceil(lowpass_filter_width / rolloff)
+            self.kernel_size = 2 * width * ratio + 1
+            self.pad = width
+            self.pad_left = 2 * width * ratio
+            self.pad_right = self.kernel_size - ratio
+
+            time_axis = (torch.arange(self.kernel_size) / ratio - width) * rolloff
+            time_clamped = time_axis.clamp(-lowpass_filter_width, lowpass_filter_width)
+            window = torch.cos(time_clamped * math.pi / lowpass_filter_width / 2) ** 2
+            sinc_filter = (torch.sinc(time_axis) * window * rolloff / ratio).view(1, 1, -1)
+        else:
+            # Kaiser sinc filter is BigVGAN default
+            self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+            self.pad = self.kernel_size // ratio - 1
+            self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2
+            self.pad_right = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2
+
+            sinc_filter = kaiser_sinc_filter1d(
+                cutoff=0.5 / ratio,
+                half_width=0.6 / ratio,
+                kernel_size=self.kernel_size,
+            )
+
+        self.register_buffer("filter", sinc_filter, persistent=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x expected shape: [batch_size, num_channels, hidden_dim]
+        num_channels = x.shape[1]
+        x = F.pad(x, (self.pad, self.pad), mode=self.padding_mode)
+        low_pass_filter = self.filter.to(dtype=x.dtype, device=x.device).expand(num_channels, -1, -1)
+        x = self.ratio * F.conv_transpose1d(x, low_pass_filter, stride=self.ratio, groups=num_channels)
+        return x[..., self.pad_left:-self.pad_right]
+
+
+class SnakeBeta(nn.Module):
+    """
+    Implements the Snake and SnakeBeta activations, which help with learning periodic patterns.
+    """
+    def __init__(
+        self,
+        channels: int,
+        alpha: float = 1.0,
+        eps: float = 1e-9,
+        trainable_params: bool = True,
+        logscale: bool = True,
+        use_beta: bool = True,
+    ):
+        self.eps = eps
+        self.logscale = logscale
+        self.use_beta = use_beta
+
+        self.alpha = nn.Parameter(torch.zeros(channels) if self.logscale else torch.ones(channels) * alpha)
+        self.alpha.requires_grad = trainable_params
+        if use_beta:
+            self.beta = nn.Parameter(torch.zeros(channels) if self.logscale else torch.ones(channels) * alpha)
+            self.beta.requires_grad = trainable_params
+
+    def forward(self, hidden_states: torch.Tensor, channel_dim: int = 1) -> torch.Tensor:
+        broadcast_shape = [1] * hidden_states.ndim
+        broadcast_shape[channel_dim] = -1
+        alpha = self.alpha.view(broadcast_shape)
+        if self.use_beta:
+            beta = self.beta.view(broadcast_shape)
+
+        if self.logscale:
+            alpha = torch.exp(alpha)
+            if self.use_beta:
+                beta = torch.exp(beta)
+
+        amplitude = beta if self.use_beta else alpha
+        hidden_states = hidden_states + (1.0 / (amplitude + self.eps)) * torch.sin(hidden_states * alpha).pow(2)
+        return hidden_states
+
+
 class ResBlock(nn.Module):
     def __init__(
         self,
@@ -15,12 +182,15 @@ def __init__(
         kernel_size: int = 3,
         stride: int = 1,
         dilations: tuple[int, ...] = (1, 3, 5),
+        act_fn: str = "leaky_relu",
         leaky_relu_negative_slope: float = 0.1,
+        antialias: bool = False,
+        antialias_ratio: int = 2,
+        antialias_kernel_size: int = 12,
         padding_mode: str = "same",
     ):
         super().__init__()
         self.dilations = dilations
-        self.negative_slope = leaky_relu_negative_slope
 
         self.convs1 = nn.ModuleList(
             [
@@ -28,6 +198,22 @@ def __init__(
                 for dilation in dilations
             ]
         )
+        self.acts1 = nn.ModuleList()
+        for _ in range(len(self.convs1)):
+            if act_fn == "snakebeta":
+                act = SnakeBeta(channels, use_beta=True)
+            elif act_fn == "snake":
+                act = SnakeBeta(channels, use_beta=False)
+            else:
+                act_fn = nn.LeakyReLU(negative_slope=leaky_relu_negative_slope)
+
+            if antialias:
+                act = nn.Sequential(
+                    UpSample1d(ratio=antialias_ratio, kernel_size=antialias_kernel_size),
+                    act,
+                    DownSample1d(ratio=antialias_ratio, kernel_size=antialias_kernel_size),
+                )
+            self.acts1.append(act)
 
         self.convs2 = nn.ModuleList(
             [
@@ -35,12 +221,28 @@ def __init__(
                 for _ in range(len(dilations))
             ]
         )
+        self.acts2 = nn.ModuleList()
+        for _ in range(len(self.convs2)):
+            if act_fn == "snakebeta":
+                act = SnakeBeta(channels, use_beta=True)
+            elif act_fn == "snake":
+                act = SnakeBeta(channels, use_beta=False)
+            else:
+                act_fn = nn.LeakyReLU(negative_slope=leaky_relu_negative_slope)
+
+            if antialias:
+                act = nn.Sequential(
+                    UpSample1d(ratio=antialias_ratio, kernel_size=antialias_kernel_size),
+                    act,
+                    DownSample1d(ratio=antialias_ratio, kernel_size=antialias_kernel_size),
+                )
+            self.acts2.append(act)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        for conv1, conv2 in zip(self.convs1, self.convs2):
-            xt = F.leaky_relu(x, negative_slope=self.negative_slope)
+        for act1, conv1, act2, conv2 in zip(self.acts1, self.convs1, self.acts2, self.convs2):
+            xt = act1(xt)
             xt = conv1(xt)
-            xt = F.leaky_relu(xt, negative_slope=self.negative_slope)
+            xt = act2(xt)
             xt = conv2(xt)
             x = x + xt
         return x
@@ -61,7 +263,13 @@ def __init__(
         upsample_factors: list[int] = [6, 5, 2, 2, 2],
         resnet_kernel_sizes: list[int] = [3, 7, 11],
         resnet_dilations: list[list[int]] = [[1, 3, 5], [1, 3, 5], [1, 3, 5]],
+        act_fn: str = "leaky_relu",
         leaky_relu_negative_slope: float = 0.1,
+        antialias: bool = False,
+        antialias_ratio: int = 2,
+        antialias_kernel_size: int = 12,
+        final_act_fn: str | None = "tanh",
+        final_bias: bool = True,
         output_sampling_rate: int = 24000,
     ):
         super().__init__()
@@ -69,7 +277,9 @@ def __init__(
         self.resnets_per_upsample = len(resnet_kernel_sizes)
         self.out_channels = out_channels
         self.total_upsample_factor = math.prod(upsample_factors)
+        self.act_fn = act_fn
         self.negative_slope = leaky_relu_negative_slope
+        self.final_act_fn = final_act_fn
 
         if self.num_upsample_layers != len(upsample_factors):
             raise ValueError(
@@ -83,6 +293,13 @@ def __init__(
                 f" {len(self.resnets_per_upsample)} and {len(resnet_dilations)}, respectively."
             )
 
+        supported_act_fns = ["snakebeta", "snake", "leaky_relu"]
+        if self.act_fn not in supported_act_fns:
+            raise ValueError(
+                f"Unsupported activation function: {self.act_fn}. Currently supported values of `act_fn` are "
+                f"{supported_act_fns}."
+            )
+
         self.conv_in = nn.Conv1d(in_channels, hidden_channels, kernel_size=7, stride=1, padding=3)
 
         self.upsamplers = nn.ModuleList()
@@ -103,15 +320,30 @@ def __init__(
             for kernel_size, dilations in zip(resnet_kernel_sizes, resnet_dilations):
                 self.resnets.append(
                     ResBlock(
-                        output_channels,
-                        kernel_size,
+                        channels=output_channels,
+                        kernel_size=kernel_size,
                         dilations=dilations,
+                        act_fn=act_fn,
                         leaky_relu_negative_slope=leaky_relu_negative_slope,
+                        antialias=antialias,
+                        antialias_ratio=antialias_ratio,
+                        antialias_kernel_size=antialias_kernel_size,
                     )
                 )
             input_channels = output_channels
 
-        self.conv_out = nn.Conv1d(output_channels, out_channels, 7, stride=1, padding=3)
+        if act_fn == "snakebeta" or act_fn == "snake":
+            # Always use antialiasing
+            self.act_out = nn.Sequential(
+                UpSample1d(ratio=antialias_ratio, kernel_size=antialias_kernel_size),
+                SnakeBeta(channels=out_channels, use_beta=True),
+                DownSample1d(ratio=antialias_ratio, kernel_size=antialias_kernel_size),
+            )
+        elif act_fn == "leaky_relu":
+            # NOTE: does NOT use self.negative_slope, following the original code
+            self.act_out = nn.LeakyReLU()
+
+        self.conv_out = nn.Conv1d(output_channels, out_channels, 7, stride=1, padding=3, bias=final_bias)
 
     def forward(self, hidden_states: torch.Tensor, time_last: bool = False) -> torch.Tensor:
         r"""
@@ -139,7 +371,9 @@ def forward(self, hidden_states: torch.Tensor, time_last: bool = False) -> torch
         hidden_states = self.conv_in(hidden_states)
 
         for i in range(self.num_upsample_layers):
-            hidden_states = F.leaky_relu(hidden_states, negative_slope=self.negative_slope)
+            if self.act_fn == "leaky_relu":
+                # Other activations are inside each upsampling block
+                hidden_states = F.leaky_relu(hidden_states, negative_slope=self.negative_slope)
             hidden_states = self.upsamplers[i](hidden_states)
 
             # Run all resnets in parallel on hidden_states
@@ -149,10 +383,11 @@ def forward(self, hidden_states: torch.Tensor, time_last: bool = False) -> torch
 
             hidden_states = torch.mean(resnet_outputs, dim=0)
 
-        # NOTE: unlike the first leaky ReLU, this leaky ReLU is set to use the default F.leaky_relu negative slope of
-        # 0.01 (whereas the others usually use a slope of 0.1). Not sure if this is intended
-        hidden_states = F.leaky_relu(hidden_states, negative_slope=0.01)
+        hidden_states = self.act_out(hidden_states)
         hidden_states = self.conv_out(hidden_states)
-        hidden_states = torch.tanh(hidden_states)
+        if self.final_act_fn == "tanh":
+            hidden_states = torch.tanh(hidden_states)
+        else:
+            hidden_states = torch.clamp(hidden_states, -1, 1)
 
         return hidden_states

From 1e89cb3652942f2cf1dc1799f8ebca4e97301924 Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Sun, 8 Mar 2026 08:04:42 +0100
Subject: [PATCH 06/41] Initial implementation of LTX-2.3 vocoder with
 bandwidth extender

---
 src/diffusers/pipelines/ltx2/__init__.py |   2 +-
 src/diffusers/pipelines/ltx2/vocoder.py  | 186 ++++++++++++++++++++++-
 2 files changed, 185 insertions(+), 3 deletions(-)

diff --git a/src/diffusers/pipelines/ltx2/__init__.py b/src/diffusers/pipelines/ltx2/__init__.py
index d6a408d5c546..0c95a17f3e31 100644
--- a/src/diffusers/pipelines/ltx2/__init__.py
+++ b/src/diffusers/pipelines/ltx2/__init__.py
@@ -44,7 +44,7 @@
         from .pipeline_ltx2_condition import LTX2ConditionPipeline
         from .pipeline_ltx2_image2video import LTX2ImageToVideoPipeline
         from .pipeline_ltx2_latent_upsample import LTX2LatentUpsamplePipeline
-        from .vocoder import LTX2Vocoder
+        from .vocoder import LTX2Vocoder, LTX2VocoderWithBWE
 
 else:
     import sys
diff --git a/src/diffusers/pipelines/ltx2/vocoder.py b/src/diffusers/pipelines/ltx2/vocoder.py
index 03b583511e74..0248b8bc6cd5 100644
--- a/src/diffusers/pipelines/ltx2/vocoder.py
+++ b/src/diffusers/pipelines/ltx2/vocoder.py
@@ -63,6 +63,7 @@ def __init__(
         kernel_size: int | None = None,
         use_padding: bool = True,
         padding_mode: str = "replicate",
+        persistent: bool = True,
     ):
         super().__init__()
         self.ratio = ratio
@@ -75,7 +76,7 @@ def __init__(
         cutoff = 0.5 / ratio
         half_width = 0.6 / ratio
         low_pass_filter = kaiser_sinc_filter1d(cutoff, half_width, self.kernel_size)
-        self.register_buffer("filter", low_pass_filter.view(1, 1, self,kernel_size))
+        self.register_buffer("filter", low_pass_filter.view(1, 1, self,kernel_size), persistent=persistent)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         # x expected shape: [batch_size, num_channels, hidden_dim]
@@ -93,6 +94,7 @@ def __init__(
         kernel_size: int | None  = None,
         window_type: str = "kaiser",
         padding_mode: str = "replicate",
+        persistent: bool = True,
     ):
         super().__init__()
         self.ratio = ratio
@@ -124,7 +126,7 @@ def __init__(
                 kernel_size=self.kernel_size,
             )
 
-        self.register_buffer("filter", sinc_filter, persistent=True)
+        self.register_buffer("filter", sinc_filter, persistent=persistent)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         # x expected shape: [batch_size, num_channels, hidden_dim]
@@ -391,3 +393,183 @@ def forward(self, hidden_states: torch.Tensor, time_last: bool = False) -> torch
             hidden_states = torch.clamp(hidden_states, -1, 1)
 
         return hidden_states
+
+
+class CausalSTFT(nn.Module):
+    """
+    Performs a causal short-time Fourier transform (STFT) using causal Hann windows on a waveform. The DFT bases
+    multiplied by the Hann windows are pre-calculated and stored as buffers. For exact parity with training, the
+    exact buffers should be loaded from the checkpoint in bfloat16.
+    """
+
+    def __init__(self, filter_length: int = 512, hop_length: int = 80, window_length: int = 512):
+        super().__init__()
+        self.hop_length = hop_length
+        self.window_length = window_length
+        n_freqs = filter_length // 2 + 1
+
+        self.register_buffer("forward_basis", torch.zeros(n_freqs * 2, 1, filter_length), persistent=True)
+        self.register_buffer("inverse_basis", torch.zeros(n_freqs * 2, 1, filter_length), persistent=True)
+
+    def forward(self, waveform: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        if waveform.ndim == 2:
+            waveform = waveform.unsqueeze(1)  # [B, num_channels, num_samples]
+
+        left_pad = max(0, self.window_length - self.hop_length)  # causal: left-only
+        waveform = F.pad(waveform, (left_pad, 0))
+
+        spec = F.conv1d(waveform, self.forward_basis, stride=self.hop_length, padding=0)
+        n_freqs = spec.shape[1] // 2
+        real, imag = spec[:, :n_freqs], spec[:, n_freqs:]
+        magnitude = torch.sqrt(real**2 + imag**2)
+        phase = torch.atan2(imag.float(), real.float()).to(dtype=real.dtype)
+        return magnitude, phase
+
+
+class MelSTFT(nn.Module):
+    """
+    Calculates a causal log-mel spectrogram from a waveform. Uses a pre-calculated mel filterbank, which should be
+    loaded from the checkpoint in bfloat16.
+    """
+
+    def __init__(
+        self,
+        filter_length: int = 512,
+        hop_length: int = 80,
+        window_length: int = 512,
+        num_mel_channels: int = 64,
+    ):
+        super().__init__()
+        self.stft_fn = CausalSTFT(filter_length, hop_length, window_length)
+
+        num_freqs = filter_length // 2 + 1
+        self.register_buffer("mel_basis", torch.zeros(num_mel_channels, num_freqs), persistent=True)
+
+    def forward(self, waveform: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        magnitude, phase = self.stft_fn(waveform)
+        energy = torch.norm(magnitude, dim=1)
+        mel = torch.matmul(self.mel_basis.to(magnitude.dtype), magnitude)
+        log_mel = torch.log(torch.clamp(mel, min=1e-5))
+        return log_mel, magnitude, phase, energy
+
+
+class LTX2VocoderWithBWE(ModelMixin, ConfigMixin):
+    """
+    LTX-2.X vocoder with bandwidth extension (BWE) upsampling. The vocoder and the BWE module run in sequence, with the
+    BWE module upsampling the vocoder output waveform to a higher sampling rate. The BWE module itself has the same
+    architecture as the original vocoder.
+    """
+
+    @register_to_config
+    def __init__(
+        self,
+        in_channels: int = 128,
+        hidden_channels: int = 1536,
+        out_channels: int = 2,
+        upsample_kernel_sizes: list[int] = [11, 4, 4, 4, 4, 4],
+        upsample_factors: list[int] = [5, 2, 2, 2, 2, 2],
+        resnet_kernel_sizes: list[int] = [3, 7, 11],
+        resnet_dilations: list[list[int]] = [[1, 3, 5], [1, 3, 5], [1, 3, 5]],
+        act_fn: str = "snakebeta",
+        leaky_relu_negative_slope: float = 0.1,
+        antialias: bool = True,
+        antialias_ratio: int = 2,
+        antialias_kernel_size: int = 12,
+        final_act_fn: str | None = None,
+        final_bias: bool = False,
+        bwe_in_channels: int = 128,
+        bwe_hidden_channels: int = 512,
+        bwe_out_channels: int = 2,
+        bwe_upsample_kernel_sizes: list[int] = [12, 11, 8, 4, 4],
+        bwe_upsample_factors: list[int] = [6, 5, 2, 2, 2],
+        bwe_resnet_kernel_sizes: list[int] = [3, 7, 11],
+        bwe_resnet_dilations: list[list[int]] = [[1, 3, 5], [1, 3, 5], [1, 3, 5]],
+        bwe_act_fn: str = "snakebeta",
+        bwe_leaky_relu_negative_slope: float = 0.1,
+        bwe_antialias: bool = True,
+        bwe_antialias_ratio: int = 2,
+        bwe_antialias_kernel_size: int = 12,
+        bwe_final_act_fn: str | None = None,
+        bwe_final_bias: bool = False,
+        filter_length: int = 512,
+        hop_length: int = 80,
+        window_length: int = 512,
+        num_mel_channels: int = 64,
+        input_sampling_rate: int = 16000,
+        output_sampling_rate: int = 48000,
+    ):
+        super().__init__()
+
+        self.vocoder = LTX2Vocoder(
+            in_channels=in_channels,
+            hidden_channels=hidden_channels,
+            out_channels=out_channels,
+            upsample_kernel_sizes=upsample_kernel_sizes,
+            upsample_factors=upsample_factors,
+            resnet_kernel_sizes=resnet_kernel_sizes,
+            resnet_dilations=resnet_dilations,
+            act_fn=act_fn,
+            leaky_relu_negative_slope=leaky_relu_negative_slope,
+            antialias=antialias,
+            antialias_ratio=antialias_ratio,
+            antialias_kernel_size=antialias_kernel_size,
+            final_act_fn=final_act_fn,
+            final_bias=final_bias,
+            output_sampling_rate=input_sampling_rate,
+        )
+        self.bwe_generator = LTX2Vocoder(
+            in_channels=bwe_in_channels,
+            hidden_channels=bwe_hidden_channels,
+            out_channels=bwe_out_channels,
+            upsample_kernel_sizes=bwe_upsample_kernel_sizes,
+            upsample_factors=bwe_upsample_factors,
+            resnet_kernel_sizes=bwe_resnet_kernel_sizes,
+            resnet_dilations=bwe_resnet_dilations,
+            act_fn=bwe_act_fn,
+            leaky_relu_negative_slope=bwe_leaky_relu_negative_slope,
+            antialias=bwe_antialias,
+            antialias_ratio=bwe_antialias_ratio,
+            antialias_kernel_size=bwe_antialias_kernel_size,
+            final_act_fn=bwe_final_act_fn,
+            final_bias=bwe_final_bias,
+            output_sampling_rate=output_sampling_rate,
+        )
+
+        self.mel_stft = MelSTFT(
+            filter_length=filter_length,
+            hop_length=hop_length,
+            window_length=window_length,
+            num_mel_channels=num_mel_channels,
+        )
+
+        self.resampler = UpSample1d(
+            ratio=output_sampling_rate // input_sampling_rate,
+            window_type="hann",
+            persistent=False,
+        )
+
+    def forward(self, mel_spec: torch.Tensor) -> torch.Tensor:
+        # 1. Run stage 1 vocoder to get low sampling rate waveform
+        x = self.vocoder(mel_spec)
+        batch_size, num_channels, num_samples = x.shape
+
+        # Pad to exact multiple of hop_length for exact mel frame count
+        remainder = num_samples % self.config.hop_length
+        if remainder != 0:
+            x = F.pad(x, (0, self.hop_length - remainder))
+
+        # 2. Compute mel spectrogram on vocoder output
+        x = x.flatten(0, 1)
+        mel, _, _, _ = self.mel_stft(x)
+        mel = mel.unflatten(0, (-1, num_channels))
+
+        # 3. Run bandwidth extender (BWE) on new mel spectrogram
+        mel_for_bwe = mel.transpose(2, 3)  # [B, C, num_mel_bins, num_frames] --> [B, C, num_frames, num_mel_bins]
+        residual = self.bwe_generator(mel_for_bwe)
+
+        # 4. Residual connection with resampler
+        skip = self.resampler(x)
+        waveform = torch.clamp(residual + skip, -1, 1)
+        output_samples = num_samples * self.config.output_sampling_rate // self.config.input_sampling_rate
+        waveform = waveform[..., :output_samples]
+        return waveform

From 5a44adb0b068e13675f95ff2b3feed6da2ae895a Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Mon, 9 Mar 2026 02:17:23 +0100
Subject: [PATCH 07/41] Initial support for LTX-2.3 per-modality feature
 extractor

---
 src/diffusers/pipelines/ltx2/connectors.py | 93 +++++++++++++++++-----
 1 file changed, 73 insertions(+), 20 deletions(-)

diff --git a/src/diffusers/pipelines/ltx2/connectors.py b/src/diffusers/pipelines/ltx2/connectors.py
index 4b2a81a9dc2c..e0826053cd52 100644
--- a/src/diffusers/pipelines/ltx2/connectors.py
+++ b/src/diffusers/pipelines/ltx2/connectors.py
@@ -1,3 +1,5 @@
+import math
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -9,6 +11,12 @@
 from ...models.transformers.transformer_ltx2 import LTX2Attention, LTX2AudioVideoAttnProcessor
 
 
+def per_token_rms_norm(text_encoder_hidden_states: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
+    variance = torch.mean(text_encoder_hidden_states ** 2, dim=2, keepdim=2)
+    norm_text_encoder_hidden_states = text_encoder_hidden_states + torch.rsqrt(variance + eps)
+    return norm_text_encoder_hidden_states
+
+
 class LTX2RotaryPosEmbed1d(nn.Module):
     """
     1D rotary positional embeddings (RoPE) for the LTX 2.0 text encoder connectors.
@@ -260,24 +268,34 @@ class LTX2TextConnectors(ModelMixin, PeftAdapterMixin, ConfigMixin):
     @register_to_config
     def __init__(
         self,
-        caption_channels: int,
-        text_proj_in_factor: int,
-        video_connector_num_attention_heads: int,
-        video_connector_attention_head_dim: int,
-        video_connector_num_layers: int,
-        video_connector_num_learnable_registers: int | None,
-        audio_connector_num_attention_heads: int,
-        audio_connector_attention_head_dim: int,
-        audio_connector_num_layers: int,
-        audio_connector_num_learnable_registers: int | None,
-        connector_rope_base_seq_len: int,
-        rope_theta: float,
-        rope_double_precision: bool,
-        causal_temporal_positioning: bool,
+        caption_channels: int = 3840,  # default Gemma-3-12B text encoder hidden_size
+        text_proj_in_factor: int = 49,  # num_layers + 1 for embedding layer = 48 + 1 for Gemma-3-12B
+        video_connector_num_attention_heads: int = 30,
+        video_connector_attention_head_dim: int = 128,
+        video_connector_num_layers: int = 2,
+        video_connector_num_learnable_registers: int | None = 128,
+        audio_connector_num_attention_heads: int = 30,
+        audio_connector_attention_head_dim: int = 128,
+        audio_connector_num_layers: int = 2,
+        audio_connector_num_learnable_registers: int | None = 128,
+        connector_rope_base_seq_len: int = 4096,
+        rope_theta: float = 10000.0,
+        rope_double_precision: bool = True,
+        causal_temporal_positioning: bool = False,
         rope_type: str = "interleaved",
+        per_modality_projections: bool = True,
+        video_hidden_dim: int = 4096,
+        audio_hidden_dim: int = 2048,
+        proj_bias: bool = False,
     ):
         super().__init__()
-        self.text_proj_in = nn.Linear(caption_channels * text_proj_in_factor, caption_channels, bias=False)
+        text_encoder_dim = caption_channels * text_proj_in_factor
+        if per_modality_projections:
+            self.video_text_proj_in = nn.Linear(text_encoder_dim, video_hidden_dim, bias=proj_bias)
+            self.audio_text_proj_in = nn.Linear(text_encoder_dim, audio_hidden_dim, bias=proj_bias)
+        else:
+            self.text_proj_in = nn.Linear(text_encoder_dim, caption_channels, bias=proj_bias)
+
         self.video_connector = LTX2ConnectorTransformer1d(
             num_attention_heads=video_connector_num_attention_heads,
             attention_head_dim=video_connector_attention_head_dim,
@@ -303,22 +321,57 @@ def __init__(
 
     def forward(
         self, text_encoder_hidden_states: torch.Tensor, attention_mask: torch.Tensor, additive_mask: bool = False
-    ):
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Given per-layer text encoder hidden_states, extracts features and runs per-modality connectors to get text
+        embeddings for the LTX-2.X DiT models.
+
+        Args:
+            text_encoder_hidden_states (`torch.Tensor` of shape `(batch_size, seq_len, caption_channels, text_proj_in_factor)`):
+                Per-layer text encoder hidden_states.
+            attention_mask (`torch.Tensor` of shape `(batch_size, seq_len)`):
+                Multiplicative binary attention mask where 1s indicate unmasked positions and 0s indicate masked
+                positions.
+        """
         # Convert to additive attention mask, if necessary
         if not additive_mask:
             text_dtype = text_encoder_hidden_states.dtype
             attention_mask = (attention_mask - 1).reshape(attention_mask.shape[0], 1, -1, attention_mask.shape[-1])
             attention_mask = attention_mask.to(text_dtype) * torch.finfo(text_dtype).max
 
-        text_encoder_hidden_states = self.text_proj_in(text_encoder_hidden_states)
-
-        video_text_embedding, new_attn_mask = self.video_connector(text_encoder_hidden_states, attention_mask)
+        if self.config.per_modality_projections:
+            norm_text_encoder_hidden_states = per_token_rms_norm(text_encoder_hidden_states)
+
+            norm_text_encoder_hidden_states = norm_text_encoder_hidden_states.flatten(2, 3)
+            bool_mask = attention_mask.bool().unsqueeze(-1)
+            norm_text_encoder_hidden_states = torch.where(
+                bool_mask, norm_text_encoder_hidden_states, torch.zeros_like(norm_text_encoder_hidden_states)
+            )
+
+            # Rescale norms for video and audio dims
+            video_scale_factor = math.sqrt(self.config.video_hidden_dim / self.config.caption_channels)
+            video_norm_text_emb = norm_text_encoder_hidden_states * video_scale_factor
+            audio_scale_factor = math.sqrt(self.config.audio_hidden_dim / self.config.caption_channels)
+            audio_norm_text_emb = norm_text_encoder_hidden_states * audio_scale_factor
+
+            # Per-Modality Feature extractors
+            video_text_emb_proj = self.video_text_proj_in(video_norm_text_emb)
+            audio_text_emb_proj = self.audio_text_proj_in(audio_norm_text_emb)
+        else:
+            # Convert to additive attention mask, if necessary
+            if not additive_mask:
+                text_dtype = text_encoder_hidden_states.dtype
+                attention_mask = (attention_mask - 1).reshape(attention_mask.shape[0], 1, -1, attention_mask.shape[-1])
+                attention_mask = attention_mask.to(text_dtype) * torch.finfo(text_dtype).max
+            video_text_emb_proj, audio_text_emb_proj = self.text_proj_in(text_encoder_hidden_states)
+
+        video_text_embedding, new_attn_mask = self.video_connector(video_text_emb_proj, attention_mask)
 
         attn_mask = (new_attn_mask < 1e-6).to(torch.int64)
         attn_mask = attn_mask.reshape(video_text_embedding.shape[0], video_text_embedding.shape[1], 1)
         video_text_embedding = video_text_embedding * attn_mask
         new_attn_mask = attn_mask.squeeze(-1)
 
-        audio_text_embedding, _ = self.audio_connector(text_encoder_hidden_states, attention_mask)
+        audio_text_embedding, _ = self.audio_connector(audio_text_emb_proj, attention_mask)
 
         return video_text_embedding, audio_text_embedding, new_attn_mask

From 4ff31688c779a382cb007421a2e2f4ddaea9eaa3 Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Mon, 9 Mar 2026 03:08:24 +0100
Subject: [PATCH 08/41] Refactor so that text connectors own all text encoder
 hidden_states normalization logic

---
 src/diffusers/pipelines/ltx2/connectors.py    | 133 +++++++++++++++---
 src/diffusers/pipelines/ltx2/pipeline_ltx2.py |  84 +----------
 .../pipelines/ltx2/pipeline_ltx2_condition.py |  85 +----------
 .../ltx2/pipeline_ltx2_image2video.py         |  85 +----------
 4 files changed, 126 insertions(+), 261 deletions(-)

diff --git a/src/diffusers/pipelines/ltx2/connectors.py b/src/diffusers/pipelines/ltx2/connectors.py
index e0826053cd52..7c6b35a1ce2c 100644
--- a/src/diffusers/pipelines/ltx2/connectors.py
+++ b/src/diffusers/pipelines/ltx2/connectors.py
@@ -11,6 +11,73 @@
 from ...models.transformers.transformer_ltx2 import LTX2Attention, LTX2AudioVideoAttnProcessor
 
 
+def per_batch_per_layer_mean_norm(
+    text_hidden_states: torch.Tensor,
+    sequence_lengths: torch.Tensor,
+    device: str | torch.device,
+    padding_side: str = "left",
+    scale_factor: int = 8,
+    eps: float = 1e-6,
+):
+    """
+    Performs per-batch per-layer normalization using a masked mean and range on per-layer text encoder hidden_states.
+    Respects the padding of the hidden states.
+
+    Args:
+        text_hidden_states (`torch.Tensor` of shape `(batch_size, seq_len, hidden_dim, num_layers)`):
+            Per-layer hidden_states from a text encoder (e.g. `Gemma3ForConditionalGeneration`).
+        sequence_lengths (`torch.Tensor of shape `(batch_size,)`):
+            The number of valid (non-padded) tokens for each batch instance.
+        device: (`str` or `torch.device`, *optional*):
+            torch device to place the resulting embeddings on
+        padding_side: (`str`, *optional*, defaults to `"left"`):
+            Whether the text tokenizer performs padding on the `"left"` or `"right"`.
+        scale_factor (`int`, *optional*, defaults to `8`):
+            Scaling factor to multiply the normalized hidden states by.
+        eps (`float`, *optional*, defaults to `1e-6`):
+            A small positive value for numerical stability when performing normalization.
+
+    Returns:
+        `torch.Tensor` of shape `(batch_size, seq_len, hidden_dim * num_layers)`:
+            Normed and flattened text encoder hidden states.
+    """
+    batch_size, seq_len, hidden_dim, num_layers = text_hidden_states.shape
+    original_dtype = text_hidden_states.dtype
+
+    # Create padding mask
+    token_indices = torch.arange(seq_len, device=device).unsqueeze(0)
+    if padding_side == "right":
+        # For right padding, valid tokens are from 0 to sequence_length-1
+        mask = token_indices < sequence_lengths[:, None]  # [batch_size, seq_len]
+    elif padding_side == "left":
+        # For left padding, valid tokens are from (T - sequence_length) to T-1
+        start_indices = seq_len - sequence_lengths[:, None]  # [batch_size, 1]
+        mask = token_indices >= start_indices  # [B, T]
+    else:
+        raise ValueError(f"padding_side must be 'left' or 'right', got {padding_side}")
+    mask = mask[:, :, None, None]  # [batch_size, seq_len] --> [batch_size, seq_len, 1, 1]
+
+    # Compute masked mean over non-padding positions of shape (batch_size, 1, 1, seq_len)
+    masked_text_hidden_states = text_hidden_states.masked_fill(~mask, 0.0)
+    num_valid_positions = (sequence_lengths * hidden_dim).view(batch_size, 1, 1, 1)
+    masked_mean = masked_text_hidden_states.sum(dim=(1, 2), keepdim=True) / (num_valid_positions + eps)
+
+    # Compute min/max over non-padding positions of shape (batch_size, 1, 1 seq_len)
+    x_min = text_hidden_states.masked_fill(~mask, float("inf")).amin(dim=(1, 2), keepdim=True)
+    x_max = text_hidden_states.masked_fill(~mask, float("-inf")).amax(dim=(1, 2), keepdim=True)
+
+    # Normalization
+    normalized_hidden_states = (text_hidden_states - masked_mean) / (x_max - x_min + eps)
+    normalized_hidden_states = normalized_hidden_states * scale_factor
+
+    # Pack the hidden states to a 3D tensor (batch_size, seq_len, hidden_dim * num_layers)
+    normalized_hidden_states = normalized_hidden_states.flatten(2)
+    mask_flat = mask.squeeze(-1).expand(-1, -1, hidden_dim * num_layers)
+    normalized_hidden_states = normalized_hidden_states.masked_fill(~mask_flat, 0.0)
+    normalized_hidden_states = normalized_hidden_states.to(dtype=original_dtype)
+    return normalized_hidden_states
+
+
 def per_token_rms_norm(text_encoder_hidden_states: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
     variance = torch.mean(text_encoder_hidden_states ** 2, dim=2, keepdim=2)
     norm_text_encoder_hidden_states = text_encoder_hidden_states + torch.rsqrt(variance + eps)
@@ -320,26 +387,37 @@ def __init__(
         )
 
     def forward(
-        self, text_encoder_hidden_states: torch.Tensor, attention_mask: torch.Tensor, additive_mask: bool = False
+        self,
+        text_encoder_hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        padding_side: str = "left",
+        scale_factor: int = 8,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """
         Given per-layer text encoder hidden_states, extracts features and runs per-modality connectors to get text
         embeddings for the LTX-2.X DiT models.
 
         Args:
-            text_encoder_hidden_states (`torch.Tensor` of shape `(batch_size, seq_len, caption_channels, text_proj_in_factor)`):
-                Per-layer text encoder hidden_states.
+            text_encoder_hidden_states (`torch.Tensor`)):
+                Per-layer text encoder hidden_states. Can either be 4D with shape `(batch_size, seq_len,
+                caption_channels, text_proj_in_factor) or 3D with the last two dimensions flattened.
             attention_mask (`torch.Tensor` of shape `(batch_size, seq_len)`):
                 Multiplicative binary attention mask where 1s indicate unmasked positions and 0s indicate masked
                 positions.
+            padding_side (`str`, *optional*, defaults to `"left"`):
+                The padding side used by the text encoder's text encoder (either `"left"` or `"right"`). Defaults to
+                `"left"` as this is what the default Gemma3-12B text encoder uses. Only used if
+                `per_modality_projections` is `False` (LTX-2.0 models).
+            scale_factor (`int`, *optional*, defaults to `8`):
+                Scale factor for masked mean/range normalization. Only used if `per_modality_projections` is `False`
+                (LTX-2.0 models).
         """
-        # Convert to additive attention mask, if necessary
-        if not additive_mask:
-            text_dtype = text_encoder_hidden_states.dtype
-            attention_mask = (attention_mask - 1).reshape(attention_mask.shape[0], 1, -1, attention_mask.shape[-1])
-            attention_mask = attention_mask.to(text_dtype) * torch.finfo(text_dtype).max
+        if text_encoder_hidden_states.ndim == 3:
+            # Ensure shape is [batch_size, seq_len, caption_channels, text_proj_in_factor]
+            text_encoder_hidden_states = text_encoder_hidden_states.unflatten(2, (self.config.caption_channels, -1))
 
         if self.config.per_modality_projections:
+            # LTX-2.3
             norm_text_encoder_hidden_states = per_token_rms_norm(text_encoder_hidden_states)
 
             norm_text_encoder_hidden_states = norm_text_encoder_hidden_states.flatten(2, 3)
@@ -348,7 +426,7 @@ def forward(
                 bool_mask, norm_text_encoder_hidden_states, torch.zeros_like(norm_text_encoder_hidden_states)
             )
 
-            # Rescale norms for video and audio dims
+            # Rescale norms with respect to video and audio dims for feature extractors
             video_scale_factor = math.sqrt(self.config.video_hidden_dim / self.config.caption_channels)
             video_norm_text_emb = norm_text_encoder_hidden_states * video_scale_factor
             audio_scale_factor = math.sqrt(self.config.audio_hidden_dim / self.config.caption_channels)
@@ -358,20 +436,31 @@ def forward(
             video_text_emb_proj = self.video_text_proj_in(video_norm_text_emb)
             audio_text_emb_proj = self.audio_text_proj_in(audio_norm_text_emb)
         else:
-            # Convert to additive attention mask, if necessary
-            if not additive_mask:
-                text_dtype = text_encoder_hidden_states.dtype
-                attention_mask = (attention_mask - 1).reshape(attention_mask.shape[0], 1, -1, attention_mask.shape[-1])
-                attention_mask = attention_mask.to(text_dtype) * torch.finfo(text_dtype).max
-            video_text_emb_proj, audio_text_emb_proj = self.text_proj_in(text_encoder_hidden_states)
+            # LTX-2.0
+            sequence_lengths = attention_mask.sum(dim=-1)
+            norm_text_encoder_hidden_states = per_batch_per_layer_mean_norm(
+                text_hidden_states=text_encoder_hidden_states,
+                sequence_lengths=sequence_lengths,
+                device=text_encoder_hidden_states.device,
+                padding_side=padding_side,
+                scale_factor=scale_factor,
+            )
+
+            video_text_emb_proj, audio_text_emb_proj = self.text_proj_in(norm_text_encoder_hidden_states)
+
+        # Convert to additive attention mask for connectors
+        text_dtype = video_text_emb_proj.dtype
+        attention_mask = (attention_mask.to(torch.int64) - 1).to(text_dtype)
+        attention_mask = attention_mask.reshape(attention_mask.shape[0], 1, -1, attention_mask.shape[-1])
+        add_attn_mask = attention_mask * torch.finfo(text_dtype).max
 
-        video_text_embedding, new_attn_mask = self.video_connector(video_text_emb_proj, attention_mask)
+        video_text_embedding, video_attn_mask = self.video_connector(video_text_emb_proj, add_attn_mask)
 
-        attn_mask = (new_attn_mask < 1e-6).to(torch.int64)
-        attn_mask = attn_mask.reshape(video_text_embedding.shape[0], video_text_embedding.shape[1], 1)
-        video_text_embedding = video_text_embedding * attn_mask
-        new_attn_mask = attn_mask.squeeze(-1)
+        # Convert video attn mask to binary (multiplicative) mask and mask video text embedding
+        binary_attn_mask = (video_attn_mask < 1e-6).to(torch.int64)
+        binary_attn_mask = binary_attn_mask.reshape(video_text_embedding.shape[0], video_text_embedding.shape[1], 1)
+        video_text_embedding = video_text_embedding * binary_attn_mask
 
-        audio_text_embedding, _ = self.audio_connector(audio_text_emb_proj, attention_mask)
+        audio_text_embedding, _ = self.audio_connector(audio_text_emb_proj, add_attn_mask)
 
-        return video_text_embedding, audio_text_embedding, new_attn_mask
+        return video_text_embedding, audio_text_embedding, binary_attn_mask.squeeze(-1)
diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
index 037840360137..4b350d3a4883 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
@@ -268,73 +268,6 @@ def __init__(
             self.tokenizer.model_max_length if getattr(self, "tokenizer", None) is not None else 1024
         )
 
-    @staticmethod
-    def _pack_text_embeds(
-        text_hidden_states: torch.Tensor,
-        sequence_lengths: torch.Tensor,
-        device: str | torch.device,
-        padding_side: str = "left",
-        scale_factor: int = 8,
-        eps: float = 1e-6,
-    ) -> torch.Tensor:
-        """
-        Packs and normalizes text encoder hidden states, respecting padding. Normalization is performed per-batch and
-        per-layer in a masked fashion (only over non-padded positions).
-
-        Args:
-            text_hidden_states (`torch.Tensor` of shape `(batch_size, seq_len, hidden_dim, num_layers)`):
-                Per-layer hidden_states from a text encoder (e.g. `Gemma3ForConditionalGeneration`).
-            sequence_lengths (`torch.Tensor of shape `(batch_size,)`):
-                The number of valid (non-padded) tokens for each batch instance.
-            device: (`str` or `torch.device`, *optional*):
-                torch device to place the resulting embeddings on
-            padding_side: (`str`, *optional*, defaults to `"left"`):
-                Whether the text tokenizer performs padding on the `"left"` or `"right"`.
-            scale_factor (`int`, *optional*, defaults to `8`):
-                Scaling factor to multiply the normalized hidden states by.
-            eps (`float`, *optional*, defaults to `1e-6`):
-                A small positive value for numerical stability when performing normalization.
-
-        Returns:
-            `torch.Tensor` of shape `(batch_size, seq_len, hidden_dim * num_layers)`:
-                Normed and flattened text encoder hidden states.
-        """
-        batch_size, seq_len, hidden_dim, num_layers = text_hidden_states.shape
-        original_dtype = text_hidden_states.dtype
-
-        # Create padding mask
-        token_indices = torch.arange(seq_len, device=device).unsqueeze(0)
-        if padding_side == "right":
-            # For right padding, valid tokens are from 0 to sequence_length-1
-            mask = token_indices < sequence_lengths[:, None]  # [batch_size, seq_len]
-        elif padding_side == "left":
-            # For left padding, valid tokens are from (T - sequence_length) to T-1
-            start_indices = seq_len - sequence_lengths[:, None]  # [batch_size, 1]
-            mask = token_indices >= start_indices  # [B, T]
-        else:
-            raise ValueError(f"padding_side must be 'left' or 'right', got {padding_side}")
-        mask = mask[:, :, None, None]  # [batch_size, seq_len] --> [batch_size, seq_len, 1, 1]
-
-        # Compute masked mean over non-padding positions of shape (batch_size, 1, 1, seq_len)
-        masked_text_hidden_states = text_hidden_states.masked_fill(~mask, 0.0)
-        num_valid_positions = (sequence_lengths * hidden_dim).view(batch_size, 1, 1, 1)
-        masked_mean = masked_text_hidden_states.sum(dim=(1, 2), keepdim=True) / (num_valid_positions + eps)
-
-        # Compute min/max over non-padding positions of shape (batch_size, 1, 1 seq_len)
-        x_min = text_hidden_states.masked_fill(~mask, float("inf")).amin(dim=(1, 2), keepdim=True)
-        x_max = text_hidden_states.masked_fill(~mask, float("-inf")).amax(dim=(1, 2), keepdim=True)
-
-        # Normalization
-        normalized_hidden_states = (text_hidden_states - masked_mean) / (x_max - x_min + eps)
-        normalized_hidden_states = normalized_hidden_states * scale_factor
-
-        # Pack the hidden states to a 3D tensor (batch_size, seq_len, hidden_dim * num_layers)
-        normalized_hidden_states = normalized_hidden_states.flatten(2)
-        mask_flat = mask.squeeze(-1).expand(-1, -1, hidden_dim * num_layers)
-        normalized_hidden_states = normalized_hidden_states.masked_fill(~mask_flat, 0.0)
-        normalized_hidden_states = normalized_hidden_states.to(dtype=original_dtype)
-        return normalized_hidden_states
-
     def _get_gemma_prompt_embeds(
         self,
         prompt: str | list[str],
@@ -387,16 +320,7 @@ def _get_gemma_prompt_embeds(
         )
         text_encoder_hidden_states = text_encoder_outputs.hidden_states
         text_encoder_hidden_states = torch.stack(text_encoder_hidden_states, dim=-1)
-        sequence_lengths = prompt_attention_mask.sum(dim=-1)
-
-        prompt_embeds = self._pack_text_embeds(
-            text_encoder_hidden_states,
-            sequence_lengths,
-            device=device,
-            padding_side=self.tokenizer.padding_side,
-            scale_factor=scale_factor,
-        )
-        prompt_embeds = prompt_embeds.to(dtype=dtype)
+        prompt_embeds = text_encoder_hidden_states.flatten(2, 3).to(dtype=dtype)  # Pack to 3D
 
         # duplicate text embeddings for each generation per prompt, using mps friendly method
         _, seq_len, _ = prompt_embeds.shape
@@ -960,9 +884,11 @@ def __call__(
             prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
             prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0)
 
-        additive_attention_mask = (1 - prompt_attention_mask.to(prompt_embeds.dtype)) * -1000000.0
+        tokenizer_padding_side = "left"  # Padding side for default Gemma3-12B text encoder
+        if getattr(self, "tokenizer", None) is not None:
+            tokenizer_padding_side = getattr(self.tokenizer, "padding_side", "left")
         connector_prompt_embeds, connector_audio_prompt_embeds, connector_attention_mask = self.connectors(
-            prompt_embeds, additive_attention_mask, additive_mask=True
+            prompt_embeds, prompt_attention_mask, padding_side=tokenizer_padding_side
         )
 
         # 4. Prepare latent variables
diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py
index 4c451330f439..d66cd272c2c3 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py
@@ -300,74 +300,6 @@ def __init__(
             self.tokenizer.model_max_length if getattr(self, "tokenizer", None) is not None else 1024
         )
 
-    @staticmethod
-    # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._pack_text_embeds
-    def _pack_text_embeds(
-        text_hidden_states: torch.Tensor,
-        sequence_lengths: torch.Tensor,
-        device: str | torch.device,
-        padding_side: str = "left",
-        scale_factor: int = 8,
-        eps: float = 1e-6,
-    ) -> torch.Tensor:
-        """
-        Packs and normalizes text encoder hidden states, respecting padding. Normalization is performed per-batch and
-        per-layer in a masked fashion (only over non-padded positions).
-
-        Args:
-            text_hidden_states (`torch.Tensor` of shape `(batch_size, seq_len, hidden_dim, num_layers)`):
-                Per-layer hidden_states from a text encoder (e.g. `Gemma3ForConditionalGeneration`).
-            sequence_lengths (`torch.Tensor of shape `(batch_size,)`):
-                The number of valid (non-padded) tokens for each batch instance.
-            device: (`str` or `torch.device`, *optional*):
-                torch device to place the resulting embeddings on
-            padding_side: (`str`, *optional*, defaults to `"left"`):
-                Whether the text tokenizer performs padding on the `"left"` or `"right"`.
-            scale_factor (`int`, *optional*, defaults to `8`):
-                Scaling factor to multiply the normalized hidden states by.
-            eps (`float`, *optional*, defaults to `1e-6`):
-                A small positive value for numerical stability when performing normalization.
-
-        Returns:
-            `torch.Tensor` of shape `(batch_size, seq_len, hidden_dim * num_layers)`:
-                Normed and flattened text encoder hidden states.
-        """
-        batch_size, seq_len, hidden_dim, num_layers = text_hidden_states.shape
-        original_dtype = text_hidden_states.dtype
-
-        # Create padding mask
-        token_indices = torch.arange(seq_len, device=device).unsqueeze(0)
-        if padding_side == "right":
-            # For right padding, valid tokens are from 0 to sequence_length-1
-            mask = token_indices < sequence_lengths[:, None]  # [batch_size, seq_len]
-        elif padding_side == "left":
-            # For left padding, valid tokens are from (T - sequence_length) to T-1
-            start_indices = seq_len - sequence_lengths[:, None]  # [batch_size, 1]
-            mask = token_indices >= start_indices  # [B, T]
-        else:
-            raise ValueError(f"padding_side must be 'left' or 'right', got {padding_side}")
-        mask = mask[:, :, None, None]  # [batch_size, seq_len] --> [batch_size, seq_len, 1, 1]
-
-        # Compute masked mean over non-padding positions of shape (batch_size, 1, 1, seq_len)
-        masked_text_hidden_states = text_hidden_states.masked_fill(~mask, 0.0)
-        num_valid_positions = (sequence_lengths * hidden_dim).view(batch_size, 1, 1, 1)
-        masked_mean = masked_text_hidden_states.sum(dim=(1, 2), keepdim=True) / (num_valid_positions + eps)
-
-        # Compute min/max over non-padding positions of shape (batch_size, 1, 1 seq_len)
-        x_min = text_hidden_states.masked_fill(~mask, float("inf")).amin(dim=(1, 2), keepdim=True)
-        x_max = text_hidden_states.masked_fill(~mask, float("-inf")).amax(dim=(1, 2), keepdim=True)
-
-        # Normalization
-        normalized_hidden_states = (text_hidden_states - masked_mean) / (x_max - x_min + eps)
-        normalized_hidden_states = normalized_hidden_states * scale_factor
-
-        # Pack the hidden states to a 3D tensor (batch_size, seq_len, hidden_dim * num_layers)
-        normalized_hidden_states = normalized_hidden_states.flatten(2)
-        mask_flat = mask.squeeze(-1).expand(-1, -1, hidden_dim * num_layers)
-        normalized_hidden_states = normalized_hidden_states.masked_fill(~mask_flat, 0.0)
-        normalized_hidden_states = normalized_hidden_states.to(dtype=original_dtype)
-        return normalized_hidden_states
-
     # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._get_gemma_prompt_embeds
     def _get_gemma_prompt_embeds(
         self,
@@ -421,16 +353,7 @@ def _get_gemma_prompt_embeds(
         )
         text_encoder_hidden_states = text_encoder_outputs.hidden_states
         text_encoder_hidden_states = torch.stack(text_encoder_hidden_states, dim=-1)
-        sequence_lengths = prompt_attention_mask.sum(dim=-1)
-
-        prompt_embeds = self._pack_text_embeds(
-            text_encoder_hidden_states,
-            sequence_lengths,
-            device=device,
-            padding_side=self.tokenizer.padding_side,
-            scale_factor=scale_factor,
-        )
-        prompt_embeds = prompt_embeds.to(dtype=dtype)
+        prompt_embeds = text_encoder_hidden_states.flatten(2, 3).to(dtype=dtype)  # Pack to 3D
 
         # duplicate text embeddings for each generation per prompt, using mps friendly method
         _, seq_len, _ = prompt_embeds.shape
@@ -1208,9 +1131,11 @@ def __call__(
             prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
             prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0)
 
-        additive_attention_mask = (1 - prompt_attention_mask.to(prompt_embeds.dtype)) * -1000000.0
+        tokenizer_padding_side = "left"  # Padding side for default Gemma3-12B text encoder
+        if getattr(self, "tokenizer", None) is not None:
+            tokenizer_padding_side = getattr(self.tokenizer, "padding_side", "left")
         connector_prompt_embeds, connector_audio_prompt_embeds, connector_attention_mask = self.connectors(
-            prompt_embeds, additive_attention_mask, additive_mask=True
+            prompt_embeds, prompt_attention_mask, padding_side=tokenizer_padding_side
         )
 
         # 4. Prepare latent variables
diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
index 83ba2cd7c685..26329685f2b3 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
@@ -271,74 +271,6 @@ def __init__(
             self.tokenizer.model_max_length if getattr(self, "tokenizer", None) is not None else 1024
         )
 
-    @staticmethod
-    # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._pack_text_embeds
-    def _pack_text_embeds(
-        text_hidden_states: torch.Tensor,
-        sequence_lengths: torch.Tensor,
-        device: str | torch.device,
-        padding_side: str = "left",
-        scale_factor: int = 8,
-        eps: float = 1e-6,
-    ) -> torch.Tensor:
-        """
-        Packs and normalizes text encoder hidden states, respecting padding. Normalization is performed per-batch and
-        per-layer in a masked fashion (only over non-padded positions).
-
-        Args:
-            text_hidden_states (`torch.Tensor` of shape `(batch_size, seq_len, hidden_dim, num_layers)`):
-                Per-layer hidden_states from a text encoder (e.g. `Gemma3ForConditionalGeneration`).
-            sequence_lengths (`torch.Tensor of shape `(batch_size,)`):
-                The number of valid (non-padded) tokens for each batch instance.
-            device: (`str` or `torch.device`, *optional*):
-                torch device to place the resulting embeddings on
-            padding_side: (`str`, *optional*, defaults to `"left"`):
-                Whether the text tokenizer performs padding on the `"left"` or `"right"`.
-            scale_factor (`int`, *optional*, defaults to `8`):
-                Scaling factor to multiply the normalized hidden states by.
-            eps (`float`, *optional*, defaults to `1e-6`):
-                A small positive value for numerical stability when performing normalization.
-
-        Returns:
-            `torch.Tensor` of shape `(batch_size, seq_len, hidden_dim * num_layers)`:
-                Normed and flattened text encoder hidden states.
-        """
-        batch_size, seq_len, hidden_dim, num_layers = text_hidden_states.shape
-        original_dtype = text_hidden_states.dtype
-
-        # Create padding mask
-        token_indices = torch.arange(seq_len, device=device).unsqueeze(0)
-        if padding_side == "right":
-            # For right padding, valid tokens are from 0 to sequence_length-1
-            mask = token_indices < sequence_lengths[:, None]  # [batch_size, seq_len]
-        elif padding_side == "left":
-            # For left padding, valid tokens are from (T - sequence_length) to T-1
-            start_indices = seq_len - sequence_lengths[:, None]  # [batch_size, 1]
-            mask = token_indices >= start_indices  # [B, T]
-        else:
-            raise ValueError(f"padding_side must be 'left' or 'right', got {padding_side}")
-        mask = mask[:, :, None, None]  # [batch_size, seq_len] --> [batch_size, seq_len, 1, 1]
-
-        # Compute masked mean over non-padding positions of shape (batch_size, 1, 1, seq_len)
-        masked_text_hidden_states = text_hidden_states.masked_fill(~mask, 0.0)
-        num_valid_positions = (sequence_lengths * hidden_dim).view(batch_size, 1, 1, 1)
-        masked_mean = masked_text_hidden_states.sum(dim=(1, 2), keepdim=True) / (num_valid_positions + eps)
-
-        # Compute min/max over non-padding positions of shape (batch_size, 1, 1 seq_len)
-        x_min = text_hidden_states.masked_fill(~mask, float("inf")).amin(dim=(1, 2), keepdim=True)
-        x_max = text_hidden_states.masked_fill(~mask, float("-inf")).amax(dim=(1, 2), keepdim=True)
-
-        # Normalization
-        normalized_hidden_states = (text_hidden_states - masked_mean) / (x_max - x_min + eps)
-        normalized_hidden_states = normalized_hidden_states * scale_factor
-
-        # Pack the hidden states to a 3D tensor (batch_size, seq_len, hidden_dim * num_layers)
-        normalized_hidden_states = normalized_hidden_states.flatten(2)
-        mask_flat = mask.squeeze(-1).expand(-1, -1, hidden_dim * num_layers)
-        normalized_hidden_states = normalized_hidden_states.masked_fill(~mask_flat, 0.0)
-        normalized_hidden_states = normalized_hidden_states.to(dtype=original_dtype)
-        return normalized_hidden_states
-
     # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._get_gemma_prompt_embeds
     def _get_gemma_prompt_embeds(
         self,
@@ -392,16 +324,7 @@ def _get_gemma_prompt_embeds(
         )
         text_encoder_hidden_states = text_encoder_outputs.hidden_states
         text_encoder_hidden_states = torch.stack(text_encoder_hidden_states, dim=-1)
-        sequence_lengths = prompt_attention_mask.sum(dim=-1)
-
-        prompt_embeds = self._pack_text_embeds(
-            text_encoder_hidden_states,
-            sequence_lengths,
-            device=device,
-            padding_side=self.tokenizer.padding_side,
-            scale_factor=scale_factor,
-        )
-        prompt_embeds = prompt_embeds.to(dtype=dtype)
+        prompt_embeds = text_encoder_hidden_states.flatten(2, 3).to(dtype=dtype)  # Pack to 3D
 
         # duplicate text embeddings for each generation per prompt, using mps friendly method
         _, seq_len, _ = prompt_embeds.shape
@@ -1017,9 +940,11 @@ def __call__(
             prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
             prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0)
 
-        additive_attention_mask = (1 - prompt_attention_mask.to(prompt_embeds.dtype)) * -1000000.0
+        tokenizer_padding_side = "left"  # Padding side for default Gemma3-12B text encoder
+        if getattr(self, "tokenizer", None) is not None:
+            tokenizer_padding_side = getattr(self.tokenizer, "padding_side", "left")
         connector_prompt_embeds, connector_audio_prompt_embeds, connector_attention_mask = self.connectors(
-            prompt_embeds, additive_attention_mask, additive_mask=True
+            prompt_embeds, prompt_attention_mask, padding_side=tokenizer_padding_side
         )
 
         # 4. Prepare latent variables

From 835bed615d88b37c8f30eba1106c3025e33d256b Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Mon, 9 Mar 2026 04:00:42 +0100
Subject: [PATCH 09/41] Fix some bugs for inference

---
 src/diffusers/pipelines/ltx2/connectors.py | 6 ++++--
 src/diffusers/pipelines/ltx2/vocoder.py    | 4 ++--
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/src/diffusers/pipelines/ltx2/connectors.py b/src/diffusers/pipelines/ltx2/connectors.py
index 7c6b35a1ce2c..e054decafff0 100644
--- a/src/diffusers/pipelines/ltx2/connectors.py
+++ b/src/diffusers/pipelines/ltx2/connectors.py
@@ -350,7 +350,7 @@ def __init__(
         rope_double_precision: bool = True,
         causal_temporal_positioning: bool = False,
         rope_type: str = "interleaved",
-        per_modality_projections: bool = True,
+        per_modality_projections: bool = False,
         video_hidden_dim: int = 4096,
         audio_hidden_dim: int = 2048,
         proj_bias: bool = False,
@@ -446,7 +446,9 @@ def forward(
                 scale_factor=scale_factor,
             )
 
-            video_text_emb_proj, audio_text_emb_proj = self.text_proj_in(norm_text_encoder_hidden_states)
+            text_emb_proj = self.text_proj_in(norm_text_encoder_hidden_states)
+            video_text_emb_proj = text_emb_proj
+            audio_text_emb_proj = text_emb_proj
 
         # Convert to additive attention mask for connectors
         text_dtype = video_text_emb_proj.dtype
diff --git a/src/diffusers/pipelines/ltx2/vocoder.py b/src/diffusers/pipelines/ltx2/vocoder.py
index 0248b8bc6cd5..2abbf9ac0182 100644
--- a/src/diffusers/pipelines/ltx2/vocoder.py
+++ b/src/diffusers/pipelines/ltx2/vocoder.py
@@ -207,7 +207,7 @@ def __init__(
             elif act_fn == "snake":
                 act = SnakeBeta(channels, use_beta=False)
             else:
-                act_fn = nn.LeakyReLU(negative_slope=leaky_relu_negative_slope)
+                act = nn.LeakyReLU(negative_slope=leaky_relu_negative_slope)
 
             if antialias:
                 act = nn.Sequential(
@@ -242,7 +242,7 @@ def __init__(
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         for act1, conv1, act2, conv2 in zip(self.acts1, self.convs1, self.acts2, self.convs2):
-            xt = act1(xt)
+            xt = act1(x)
             xt = conv1(xt)
             xt = act2(xt)
             xt = conv2(xt)

From 19004efc6c62f8533d5aa52b47d01c491de83244 Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Mon, 9 Mar 2026 07:29:08 +0100
Subject: [PATCH 10/41] Fix LTX-2.X DiT block forward pass

---
 src/diffusers/models/transformers/transformer_ltx2.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/diffusers/models/transformers/transformer_ltx2.py b/src/diffusers/models/transformers/transformer_ltx2.py
index 8c230c507670..84db3f6eac28 100644
--- a/src/diffusers/models/transformers/transformer_ltx2.py
+++ b/src/diffusers/models/transformers/transformer_ltx2.py
@@ -1466,6 +1466,8 @@ def forward(
                     audio_cross_attn_scale_shift,
                     video_cross_attn_a2v_gate,
                     audio_cross_attn_v2a_gate,
+                    None,  # temb_prompt
+                    None,  # temb_prompt_audio
                     video_rotary_emb,
                     audio_rotary_emb,
                     video_cross_attn_rotary_emb,
@@ -1485,6 +1487,8 @@ def forward(
                     temb_ca_audio_scale_shift=audio_cross_attn_scale_shift,
                     temb_ca_gate=video_cross_attn_a2v_gate,
                     temb_ca_audio_gate=audio_cross_attn_v2a_gate,
+                    temb_prompt=None,
+                    temb_prompt_audio=None,
                     video_rotary_emb=video_rotary_emb,
                     audio_rotary_emb=audio_rotary_emb,
                     ca_video_rotary_emb=video_cross_attn_rotary_emb,

From 4dfcfeb3aaebcffe23e9f33ffa51602eb16e28c1 Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Mon, 9 Mar 2026 08:28:24 +0100
Subject: [PATCH 11/41] Support prompt timestep embeds and prompt cross attn
 modulation

---
 .../models/transformers/transformer_ltx2.py   | 47 +++++++++++++++++--
 1 file changed, 43 insertions(+), 4 deletions(-)

diff --git a/src/diffusers/models/transformers/transformer_ltx2.py b/src/diffusers/models/transformers/transformer_ltx2.py
index 84db3f6eac28..c00bcf2489ad 100644
--- a/src/diffusers/models/transformers/transformer_ltx2.py
+++ b/src/diffusers/models/transformers/transformer_ltx2.py
@@ -1093,6 +1093,8 @@ def __init__(
         pos_embed_max_pos: int = 20,
         base_height: int = 2048,
         base_width: int = 2048,
+        gated_attn: bool = False,
+        cross_attn_mod: bool = False,
         audio_in_channels: int = 128,  # Audio Arguments
         audio_out_channels: int | None = 128,
         audio_patch_size: int = 1,
@@ -1104,6 +1106,8 @@ def __init__(
         audio_pos_embed_max_pos: int = 20,
         audio_sampling_rate: int = 16000,
         audio_hop_length: int = 160,
+        audio_gated_attn: bool = False,
+        audio_cross_attn_mod: bool = False,
         num_layers: int = 48,  # Shared arguments
         activation_fn: str = "gelu-approximate",
         qk_norm: str = "rms_norm_across_heads",
@@ -1118,6 +1122,7 @@ def __init__(
         timestep_scale_multiplier: int = 1000,
         cross_attn_timestep_scale_multiplier: int = 1000,
         rope_type: str = "interleaved",
+        prompt_modulation: bool = False,
     ) -> None:
         super().__init__()
 
@@ -1170,6 +1175,13 @@ def __init__(
         self.scale_shift_table = nn.Parameter(torch.randn(2, inner_dim) / inner_dim**0.5)
         self.audio_scale_shift_table = nn.Parameter(torch.randn(2, audio_inner_dim) / audio_inner_dim**0.5)
 
+        # 3.4. Prompt Scale/Shift Modulation parameters (LTX-2.3)
+        if prompt_modulation:
+            self.prompt_adaln = LTX2AdaLayerNormSingle(inner_dim, num_mod_params=2, use_additional_conditions=False)
+            self.audio_prompt_adaln = LTX2AdaLayerNormSingle(
+                inner_dim, num_mod_params=2, use_additional_conditions=False
+            )
+
         # 4. Rotary Positional Embeddings (RoPE)
         # Self-Attention
         self.rope = LTX2AudioVideoRotaryPosEmbed(
@@ -1246,6 +1258,10 @@ def __init__(
                     audio_num_attention_heads=audio_num_attention_heads,
                     audio_attention_head_dim=audio_attention_head_dim,
                     audio_cross_attention_dim=audio_cross_attention_dim,
+                    video_gated_attn=gated_attn,
+                    video_cross_attn_adaln=cross_attn_mod,
+                    audio_gated_attn=audio_gated_attn,
+                    audio_cross_attn_adaln=audio_cross_attn_mod,
                     qk_norm=qk_norm,
                     activation_fn=activation_fn,
                     attention_bias=attention_bias,
@@ -1276,6 +1292,8 @@ def forward(
         audio_encoder_hidden_states: torch.Tensor,
         timestep: torch.LongTensor,
         audio_timestep: torch.LongTensor | None = None,
+        sigma: torch.Tensor | None = None,
+        audio_sigma: torch.Tensor | None = None,
         encoder_attention_mask: torch.Tensor | None = None,
         audio_encoder_attention_mask: torch.Tensor | None = None,
         self_attention_mask: torch.Tensor | None = None,
@@ -1307,6 +1325,13 @@ def forward(
             audio_timestep (`torch.Tensor`, *optional*):
                 Input timestep of shape `(batch_size,)` or `(batch_size, num_audio_tokens)` for audio modulation
                 params. This is only used by certain pipelines such as the I2V pipeline.
+            sigma (`torch.Tensor`, *optional*):
+                Input scaled timestep of shape (batch_size,). Used for video prompt cross attention modulation in
+                models such as LTX-2.3.
+            audio_sigma (`torch.Tensor`, *optional*):
+                Input scaled timestep of shape (batch_size,). Used for audio prompt cross attention modulation in
+                models such as LTX-2.3. If `sigma` is supplied but `audio_sigma` is not, `audio_sigma` will be set to
+                the provided `sigma` value.
             encoder_attention_mask (`torch.Tensor`, *optional*):
                 Optional multiplicative text attention mask of shape `(batch_size, text_seq_len)`.
             audio_encoder_attention_mask (`torch.Tensor`, *optional*):
@@ -1343,6 +1368,7 @@ def forward(
         """
         # Determine timestep for audio.
         audio_timestep = audio_timestep if audio_timestep is not None else timestep
+        audio_sigma = audio_sigma if audio_sigma is not None else sigma
 
         # 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:
@@ -1413,6 +1439,19 @@ def forward(
         temb_audio = temb_audio.view(batch_size, -1, temb_audio.size(-1))
         audio_embedded_timestep = audio_embedded_timestep.view(batch_size, -1, audio_embedded_timestep.size(-1))
 
+        if self.config.prompt_modulation:
+            # LTX-2.3
+            temb_prompt, _ = self.prompt_adaln(
+                sigma.flatten(), batch_size=batch_size, hidden_dtype=hidden_states.dtype
+            )
+            temb_prompt_audio, _ = self.audio_prompt_adaln(
+                audio_sigma.flatten(), batch_size=batch_size, hidden_dtype=audio_hidden_states.dtype
+            )
+            temb_prompt = temb_prompt.view(batch_size, -1, temb_prompt.size(-1))
+            temb_prompt_audio = temb_prompt_audio.view(batch_size, -1, temb_prompt_audio.size(-1))
+        else:
+            temb_prompt = temb_prompt_audio = None
+
         # 3.2. Prepare global modality cross attention modulation parameters
         video_cross_attn_scale_shift, _ = self.av_cross_attn_video_scale_shift(
             timestep.flatten(),
@@ -1466,8 +1505,8 @@ def forward(
                     audio_cross_attn_scale_shift,
                     video_cross_attn_a2v_gate,
                     audio_cross_attn_v2a_gate,
-                    None,  # temb_prompt
-                    None,  # temb_prompt_audio
+                    temb_prompt,
+                    temb_prompt_audio,
                     video_rotary_emb,
                     audio_rotary_emb,
                     video_cross_attn_rotary_emb,
@@ -1487,8 +1526,8 @@ def forward(
                     temb_ca_audio_scale_shift=audio_cross_attn_scale_shift,
                     temb_ca_gate=video_cross_attn_a2v_gate,
                     temb_ca_audio_gate=audio_cross_attn_v2a_gate,
-                    temb_prompt=None,
-                    temb_prompt_audio=None,
+                    temb_prompt=temb_prompt,
+                    temb_prompt_audio=temb_prompt_audio,
                     video_rotary_emb=video_rotary_emb,
                     audio_rotary_emb=audio_rotary_emb,
                     ca_video_rotary_emb=video_cross_attn_rotary_emb,

From 13292dde4d1b250c3d2dfd7b68e55177a6cb3edf Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Tue, 10 Mar 2026 02:44:36 +0100
Subject: [PATCH 12/41] Add LTX-2.3 configs to conversion script

---
 scripts/convert_ltx2_to_diffusers.py          | 205 +++++++++++++++++-
 .../autoencoders/autoencoder_kl_ltx2.py       |  36 ++-
 .../models/transformers/transformer_ltx2.py   |   8 +-
 src/diffusers/pipelines/ltx2/connectors.py    |  10 +-
 src/diffusers/pipelines/ltx2/vocoder.py       |   4 +-
 5 files changed, 244 insertions(+), 19 deletions(-)

diff --git a/scripts/convert_ltx2_to_diffusers.py b/scripts/convert_ltx2_to_diffusers.py
index 9dee954af6d0..9232497f75d1 100644
--- a/scripts/convert_ltx2_to_diffusers.py
+++ b/scripts/convert_ltx2_to_diffusers.py
@@ -225,7 +225,7 @@ def get_ltx2_transformer_config(version: str) -> tuple[dict[str, Any], dict[str,
         special_keys_remap = LTX_2_0_TRANSFORMER_SPECIAL_KEYS_REMAP
     elif version == "2.0":
         config = {
-            "model_id": "diffusers-internal-dev/new-ltx-model",
+            "model_id": "Lightricks/LTX-2",
             "diffusers_config": {
                 "in_channels": 128,
                 "out_channels": 128,
@@ -238,6 +238,8 @@ def get_ltx2_transformer_config(version: str) -> tuple[dict[str, Any], dict[str,
                 "pos_embed_max_pos": 20,
                 "base_height": 2048,
                 "base_width": 2048,
+                "gated_attn": False,
+                "cross_attn_mod": False,
                 "audio_in_channels": 128,
                 "audio_out_channels": 128,
                 "audio_patch_size": 1,
@@ -249,6 +251,8 @@ def get_ltx2_transformer_config(version: str) -> tuple[dict[str, Any], dict[str,
                 "audio_pos_embed_max_pos": 20,
                 "audio_sampling_rate": 16000,
                 "audio_hop_length": 160,
+                "audio_gated_attn": False,
+                "audio_cross_attn_mod": False,
                 "num_layers": 48,
                 "activation_fn": "gelu-approximate",
                 "qk_norm": "rms_norm_across_heads",
@@ -263,6 +267,56 @@ def get_ltx2_transformer_config(version: str) -> tuple[dict[str, Any], dict[str,
                 "timestep_scale_multiplier": 1000,
                 "cross_attn_timestep_scale_multiplier": 1000,
                 "rope_type": "split",
+                "perturbed_attn": False,
+            },
+        }
+        rename_dict = LTX_2_0_TRANSFORMER_KEYS_RENAME_DICT
+        special_keys_remap = LTX_2_0_TRANSFORMER_SPECIAL_KEYS_REMAP
+    elif version == "2.3":
+        config = {
+            "model_id": "Lightricks/LTX-2.3",
+            "diffusers_config": {
+                "in_channels": 128,
+                "out_channels": 128,
+                "patch_size": 1,
+                "patch_size_t": 1,
+                "num_attention_heads": 32,
+                "attention_head_dim": 128,
+                "cross_attention_dim": 4096,
+                "vae_scale_factors": (8, 32, 32),
+                "pos_embed_max_pos": 20,
+                "base_height": 2048,
+                "base_width": 2048,
+                "gated_attn": True,
+                "cross_attn_mod": True,
+                "audio_in_channels": 128,
+                "audio_out_channels": 128,
+                "audio_patch_size": 1,
+                "audio_patch_size_t": 1,
+                "audio_num_attention_heads": 32,
+                "audio_attention_head_dim": 64,
+                "audio_cross_attention_dim": 2048,
+                "audio_scale_factor": 4,
+                "audio_pos_embed_max_pos": 20,
+                "audio_sampling_rate": 16000,
+                "audio_hop_length": 160,
+                "audio_gated_attn": False,
+                "audio_cross_attn_mod": False,
+                "num_layers": 48,
+                "activation_fn": "gelu-approximate",
+                "qk_norm": "rms_norm_across_heads",
+                "norm_elementwise_affine": False,
+                "norm_eps": 1e-6,
+                "caption_channels": 3840,
+                "attention_bias": True,
+                "attention_out_bias": True,
+                "rope_theta": 10000.0,
+                "rope_double_precision": True,
+                "causal_offset": 1,
+                "timestep_scale_multiplier": 1000,
+                "cross_attn_timestep_scale_multiplier": 1000,
+                "rope_type": "split",
+                "perturbed_attn": True,
             },
         }
         rename_dict = LTX_2_0_TRANSFORMER_KEYS_RENAME_DICT
@@ -293,7 +347,7 @@ def get_ltx2_connectors_config(version: str) -> tuple[dict[str, Any], dict[str,
         }
     elif version == "2.0":
         config = {
-            "model_id": "diffusers-internal-dev/new-ltx-model",
+            "model_id": "Lightricks/LTX-2",
             "diffusers_config": {
                 "caption_channels": 3840,
                 "text_proj_in_factor": 49,
@@ -301,15 +355,46 @@ def get_ltx2_connectors_config(version: str) -> tuple[dict[str, Any], dict[str,
                 "video_connector_attention_head_dim": 128,
                 "video_connector_num_layers": 2,
                 "video_connector_num_learnable_registers": 128,
+                "video_gated_attn": False,
                 "audio_connector_num_attention_heads": 30,
                 "audio_connector_attention_head_dim": 128,
                 "audio_connector_num_layers": 2,
                 "audio_connector_num_learnable_registers": 128,
+                "audio_gated_attn": False,
                 "connector_rope_base_seq_len": 4096,
                 "rope_theta": 10000.0,
                 "rope_double_precision": True,
                 "causal_temporal_positioning": False,
                 "rope_type": "split",
+                "per_modality_projections": False,
+                "proj_bias": False,
+            },
+        }
+    elif version == "2.3":
+        config = {
+            "model_id": "Lightricks/LTX-2.3",
+            "diffusers_config": {
+                "caption_channels": 3840,
+                "text_proj_in_factor": 49,
+                "video_connector_num_attention_heads": 32,
+                "video_connector_attention_head_dim": 128,
+                "video_connector_num_layers": 8,
+                "video_connector_num_learnable_registers": 128,
+                "video_gated_attn": True,
+                "audio_connector_num_attention_heads": 32,
+                "audio_connector_attention_head_dim": 64,
+                "audio_connector_num_layers": 8,
+                "audio_connector_num_learnable_registers": 128,
+                "audio_gated_attn": True,
+                "connector_rope_base_seq_len": 4096,
+                "rope_theta": 10000.0,
+                "rope_double_precision": True,
+                "causal_temporal_positioning": False,
+                "rope_type": "split",
+                "per_modality_projections": True,
+                "video_hidden_dim": 4096,
+                "audio_hidden_dim": 2048,
+                "proj_bias": True,
             },
         }
 
@@ -416,7 +501,7 @@ def get_ltx2_video_vae_config(
         special_keys_remap = LTX_2_0_VAE_SPECIAL_KEYS_REMAP
     elif version == "2.0":
         config = {
-            "model_id": "diffusers-internal-dev/dummy-ltx2",
+            "model_id": "Lightricks/LTX-2",
             "diffusers_config": {
                 "in_channels": 3,
                 "out_channels": 3,
@@ -435,6 +520,7 @@ def get_ltx2_video_vae_config(
                 "decoder_spatio_temporal_scaling": (True, True, True),
                 "decoder_inject_noise": (False, False, False, False),
                 "downsample_type": ("spatial", "temporal", "spatiotemporal", "spatiotemporal"),
+                "upsample_type": ("spatiotemporal", "spatiotemporal", "spatiotemporal"),
                 "upsample_residual": (True, True, True),
                 "upsample_factor": (2, 2, 2),
                 "timestep_conditioning": timestep_conditioning,
@@ -451,6 +537,44 @@ def get_ltx2_video_vae_config(
         }
         rename_dict = LTX_2_0_VIDEO_VAE_RENAME_DICT
         special_keys_remap = LTX_2_0_VAE_SPECIAL_KEYS_REMAP
+    elif version == "2.3":
+        config = {
+            "model_id": "Lightricks/LTX-2.3",
+            "diffusers_config": {
+                "in_channels": 3,
+                "out_channels": 3,
+                "latent_channels": 128,
+                "block_out_channels": (256, 512, 1024, 1024),
+                "down_block_types": (
+                    "LTX2VideoDownBlock3D",
+                    "LTX2VideoDownBlock3D",
+                    "LTX2VideoDownBlock3D",
+                    "LTX2VideoDownBlock3D",
+                ),
+                "decoder_block_out_channels": (256, 512, 512, 1024),
+                "layers_per_block": (4, 6, 4, 2, 2),
+                "decoder_layers_per_block": (4, 6, 4, 2, 2),
+                "spatio_temporal_scaling": (True, True, True, True),
+                "decoder_spatio_temporal_scaling": (True, True, True, True),
+                "decoder_inject_noise": (False, False, False, False, False),
+                "downsample_type": ("spatial", "temporal", "spatiotemporal", "spatiotemporal"),
+                "upsample_type": ("spatial", "temporal", "spatiotemporal", "spatiotemporal"),
+                "upsample_residual": (True, True, True, True),
+                "upsample_factor": (2, 2, 1, 2),
+                "timestep_conditioning": timestep_conditioning,
+                "patch_size": 4,
+                "patch_size_t": 1,
+                "resnet_norm_eps": 1e-6,
+                "encoder_causal": True,
+                "decoder_causal": False,
+                "encoder_spatial_padding_mode": "zeros",
+                "decoder_spatial_padding_mode": "zeros",
+                "spatial_compression_ratio": 32,
+                "temporal_compression_ratio": 8,
+            },
+        }
+        rename_dict = LTX_2_0_VIDEO_VAE_RENAME_DICT
+        special_keys_remap = LTX_2_0_VAE_SPECIAL_KEYS_REMAP
     return config, rename_dict, special_keys_remap
 
 
@@ -485,7 +609,7 @@ def convert_ltx2_video_vae(
 def get_ltx2_audio_vae_config(version: str) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]:
     if version == "2.0":
         config = {
-            "model_id": "diffusers-internal-dev/new-ltx-model",
+            "model_id": "Lightricks/LTX-2",
             "diffusers_config": {
                 "base_channels": 128,
                 "output_channels": 2,
@@ -508,6 +632,31 @@ def get_ltx2_audio_vae_config(version: str) -> tuple[dict[str, Any], dict[str, A
         }
         rename_dict = LTX_2_0_AUDIO_VAE_RENAME_DICT
         special_keys_remap = LTX_2_0_AUDIO_VAE_SPECIAL_KEYS_REMAP
+    elif version == "2.3":
+        config = {
+            "model_id": "Lightricks/LTX-2.3",
+            "diffusers_config": {
+                "base_channels": 128,
+                "output_channels": 2,
+                "ch_mult": (1, 2, 4),
+                "num_res_blocks": 2,
+                "attn_resolutions": None,
+                "in_channels": 2,
+                "resolution": 256,
+                "latent_channels": 8,
+                "norm_type": "pixel",
+                "causality_axis": "height",
+                "dropout": 0.0,
+                "mid_block_add_attention": False,
+                "sample_rate": 16000,
+                "mel_hop_length": 160,
+                "is_causal": True,
+                "mel_bins": 64,
+                "double_z": True,
+            },  # Same config as LTX-2.0
+        }
+        rename_dict = LTX_2_0_AUDIO_VAE_RENAME_DICT
+        special_keys_remap = LTX_2_0_AUDIO_VAE_SPECIAL_KEYS_REMAP
     return config, rename_dict, special_keys_remap
 
 
@@ -540,7 +689,7 @@ def convert_ltx2_audio_vae(original_state_dict: dict[str, Any], version: str) ->
 def get_ltx2_vocoder_config(version: str) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]:
     if version == "2.0":
         config = {
-            "model_id": "diffusers-internal-dev/new-ltx-model",
+            "model_id": "Lightricks/LTX-2",
             "diffusers_config": {
                 "in_channels": 128,
                 "hidden_channels": 1024,
@@ -549,12 +698,58 @@ def get_ltx2_vocoder_config(version: str) -> tuple[dict[str, Any], dict[str, Any
                 "upsample_factors": [6, 5, 2, 2, 2],
                 "resnet_kernel_sizes": [3, 7, 11],
                 "resnet_dilations": [[1, 3, 5], [1, 3, 5], [1, 3, 5]],
+                "act_fn": "leaky_relu",
                 "leaky_relu_negative_slope": 0.1,
+                "antialias": False,
+                "final_act_fn": "tanh",
+                "final_bias": True, 
                 "output_sampling_rate": 24000,
             },
         }
         rename_dict = LTX_2_0_VOCODER_RENAME_DICT
         special_keys_remap = LTX_2_0_VOCODER_SPECIAL_KEYS_REMAP
+    elif version == "2.3":
+        config = {
+            "model_id": "Lightricks/LTX-2.3",
+            "diffusers_config": {
+                "in_channels": 128,
+                "hidden_channels": 1024,
+                "out_channels": 2,
+                "upsample_kernel_sizes": [11, 4, 4, 4, 4, 4],
+                "upsample_factors": [5, 2, 2, 2, 2, 2],
+                "resnet_kernel_sizes": [3, 7, 11],
+                "resnet_dilations": [[1, 3, 5], [1, 3, 5], [1, 3, 5]],
+                "act_fn": "snakebeta",
+                "leaky_relu_negative_slope": 0.1,
+                "antialias": True,
+                "antialias_ratio": 2,
+                "antialias_kernel_size": 12,
+                "final_act_fn": None,
+                "final_bias": False,
+                "bwe_in_channels": 128,
+                "bwe_hidden_channels": 512,
+                "bwe_out_channels": 2,
+                "bwe_upsample_kernel_sizes": [12, 11, 8, 4, 4],
+                "bwe_upsample_factors": [6, 5, 2, 2, 2],
+                "bwe_resnet_kernel_sizes": [3, 7, 11],
+                "bwe_resnet_dilations": [[1, 3, 5], [1, 3, 5], [1, 3, 5]],
+                "bwe_act_fn": "snakebeta",
+                "bwe_leaky_relu_negative_slope": 0.1,
+                "bwe_antialias": True,
+                "bwe_antialias_ratio": 2,
+                "bwe_antialias_kernel_size": 12,
+                "bwe_final_act_fn": None,
+                "bwe_final_bias": False,
+                "filter_length": 512,
+                "hop_length": 80,
+                "window_length": 512,
+                "num_mel_channels": 64,
+                "input_sampling_rate": 16000,
+                "output_sampling_rate": 48000,
+            },
+        }
+        rename_dict = LTX_2_0_VOCODER_RENAME_DICT
+        special_keys_remap = LTX_2_0_VOCODER_SPECIAL_KEYS_REMAP
     return config, rename_dict, special_keys_remap
 
 
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_ltx2.py b/src/diffusers/models/autoencoders/autoencoder_kl_ltx2.py
index 775a83651c98..36cf73e4bb2a 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_ltx2.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_ltx2.py
@@ -740,7 +740,7 @@ def __init__(
             "LTX2VideoDownBlock3D",
             "LTX2VideoDownBlock3D",
         ),
-        spatio_temporal_scaling: tuple[bool, ...] = (True, True, True, True),
+        spatio_temporal_scaling: bool | tuple[bool, ...] = (True, True, True, True),
         layers_per_block: tuple[int, ...] = (4, 6, 6, 2, 2),
         downsample_type: tuple[str, ...] = ("spatial", "temporal", "spatiotemporal", "spatiotemporal"),
         patch_size: int = 4,
@@ -750,6 +750,9 @@ def __init__(
         spatial_padding_mode: str = "zeros",
     ):
         super().__init__()
+        num_encoder_blocks = len(layers_per_block)
+        if isinstance(spatio_temporal_scaling, bool):
+            spatio_temporal_scaling = (spatio_temporal_scaling,) * (num_encoder_blocks - 1)
 
         self.patch_size = patch_size
         self.patch_size_t = patch_size_t
@@ -884,20 +887,27 @@ def __init__(
         in_channels: int = 128,
         out_channels: int = 3,
         block_out_channels: tuple[int, ...] = (256, 512, 1024),
-        spatio_temporal_scaling: tuple[bool, ...] = (True, True, True),
+        spatio_temporal_scaling: bool | tuple[bool, ...] = (True, True, True),
         layers_per_block: tuple[int, ...] = (5, 5, 5, 5),
         upsample_type: tuple[str, ...] = ("spatiotemporal", "spatiotemporal", "spatiotemporal"),
         patch_size: int = 4,
         patch_size_t: int = 1,
         resnet_norm_eps: float = 1e-6,
         is_causal: bool = False,
-        inject_noise: tuple[bool, ...] = (False, False, False),
+        inject_noise: bool | tuple[bool, ...] = (False, False, False),
         timestep_conditioning: bool = False,
-        upsample_residual: tuple[bool, ...] = (True, True, True),
+        upsample_residual: bool | tuple[bool, ...] = (True, True, True),
         upsample_factor: tuple[bool, ...] = (2, 2, 2),
         spatial_padding_mode: str = "reflect",
     ) -> None:
         super().__init__()
+        num_decoder_blocks = len(layers_per_block)
+        if isinstance(spatio_temporal_scaling, bool):
+            spatio_temporal_scaling = (spatio_temporal_scaling,) * (num_decoder_blocks - 1)
+        if isinstance(inject_noise, bool):
+            inject_noise = (inject_noise,) * num_decoder_blocks
+        if isinstance(upsample_residual, bool):
+            upsample_residual = (upsample_residual,) * (num_decoder_blocks - 1)
 
         self.patch_size = patch_size
         self.patch_size_t = patch_size_t
@@ -1084,12 +1094,12 @@ def __init__(
         decoder_block_out_channels: tuple[int, ...] = (256, 512, 1024),
         layers_per_block: tuple[int, ...] = (4, 6, 6, 2, 2),
         decoder_layers_per_block: tuple[int, ...] = (5, 5, 5, 5),
-        spatio_temporal_scaling: tuple[bool, ...] = (True, True, True, True),
-        decoder_spatio_temporal_scaling: tuple[bool, ...] = (True, True, True),
-        decoder_inject_noise: tuple[bool, ...] = (False, False, False, False),
+        spatio_temporal_scaling: bool | tuple[bool, ...] = (True, True, True, True),
+        decoder_spatio_temporal_scaling: bool | tuple[bool, ...] = (True, True, True),
+        decoder_inject_noise: bool | tuple[bool, ...] = (False, False, False, False),
         downsample_type: tuple[str, ...] = ("spatial", "temporal", "spatiotemporal", "spatiotemporal"),
         upsample_type: tuple[str, ...] = ("spatiotemporal", "spatiotemporal", "spatiotemporal"),
-        upsample_residual: tuple[bool, ...] = (True, True, True),
+        upsample_residual: bool | tuple[bool, ...] = (True, True, True),
         upsample_factor: tuple[int, ...] = (2, 2, 2),
         timestep_conditioning: bool = False,
         patch_size: int = 4,
@@ -1104,6 +1114,16 @@ def __init__(
         temporal_compression_ratio: int = None,
     ) -> None:
         super().__init__()
+        num_encoder_blocks = len(layers_per_block)
+        num_decoder_blocks = len(decoder_layers_per_block)
+        if isinstance(spatio_temporal_scaling, bool):
+            spatio_temporal_scaling = (spatio_temporal_scaling,) * (num_encoder_blocks - 1)
+        if isinstance(decoder_spatio_temporal_scaling, bool):
+            decoder_spatio_temporal_scaling = (decoder_spatio_temporal_scaling,) * (num_decoder_blocks - 1)
+        if isinstance(decoder_inject_noise, bool):
+            decoder_inject_noise = (decoder_inject_noise,) * num_decoder_blocks
+        if isinstance(upsample_residual, bool):
+            upsample_residual = (upsample_residual,) * (num_decoder_blocks - 1)
 
         self.encoder = LTX2VideoEncoder3d(
             in_channels=in_channels,
diff --git a/src/diffusers/models/transformers/transformer_ltx2.py b/src/diffusers/models/transformers/transformer_ltx2.py
index c00bcf2489ad..951220cf671e 100644
--- a/src/diffusers/models/transformers/transformer_ltx2.py
+++ b/src/diffusers/models/transformers/transformer_ltx2.py
@@ -1122,7 +1122,7 @@ def __init__(
         timestep_scale_multiplier: int = 1000,
         cross_attn_timestep_scale_multiplier: int = 1000,
         rope_type: str = "interleaved",
-        prompt_modulation: bool = False,
+        perturbed_attn: bool = False,
     ) -> None:
         super().__init__()
 
@@ -1176,7 +1176,8 @@ def __init__(
         self.audio_scale_shift_table = nn.Parameter(torch.randn(2, audio_inner_dim) / audio_inner_dim**0.5)
 
         # 3.4. Prompt Scale/Shift Modulation parameters (LTX-2.3)
-        if prompt_modulation:
+        self.prompt_modulation = cross_attn_mod or audio_cross_attn_mod
+        if self.prompt_modulation:
             self.prompt_adaln = LTX2AdaLayerNormSingle(inner_dim, num_mod_params=2, use_additional_conditions=False)
             self.audio_prompt_adaln = LTX2AdaLayerNormSingle(
                 inner_dim, num_mod_params=2, use_additional_conditions=False
@@ -1269,6 +1270,7 @@ def __init__(
                     eps=norm_eps,
                     elementwise_affine=norm_elementwise_affine,
                     rope_type=rope_type,
+                    perturbed_attn=perturbed_attn,
                 )
                 for _ in range(num_layers)
             ]
@@ -1439,7 +1441,7 @@ def forward(
         temb_audio = temb_audio.view(batch_size, -1, temb_audio.size(-1))
         audio_embedded_timestep = audio_embedded_timestep.view(batch_size, -1, audio_embedded_timestep.size(-1))
 
-        if self.config.prompt_modulation:
+        if self.prompt_modulation:
             # LTX-2.3
             temb_prompt, _ = self.prompt_adaln(
                 sigma.flatten(), batch_size=batch_size, hidden_dtype=hidden_states.dtype
diff --git a/src/diffusers/pipelines/ltx2/connectors.py b/src/diffusers/pipelines/ltx2/connectors.py
index e054decafff0..b75444692b08 100644
--- a/src/diffusers/pipelines/ltx2/connectors.py
+++ b/src/diffusers/pipelines/ltx2/connectors.py
@@ -181,6 +181,7 @@ def __init__(
         activation_fn: str = "gelu-approximate",
         eps: float = 1e-6,
         rope_type: str = "interleaved",
+        apply_gated_attention: bool = False,
     ):
         super().__init__()
 
@@ -190,8 +191,9 @@ def __init__(
             heads=num_attention_heads,
             kv_heads=num_attention_heads,
             dim_head=attention_head_dim,
-            processor=LTX2AudioVideoAttnProcessor(),
             rope_type=rope_type,
+            apply_gated_attention=apply_gated_attention,
+            processor=LTX2AudioVideoAttnProcessor(),
         )
 
         self.norm2 = torch.nn.RMSNorm(dim, eps=eps, elementwise_affine=False)
@@ -235,6 +237,7 @@ def __init__(
         eps: float = 1e-6,
         causal_temporal_positioning: bool = False,
         rope_type: str = "interleaved",
+        gated_attention: bool = False,
     ):
         super().__init__()
         self.num_attention_heads = num_attention_heads
@@ -263,6 +266,7 @@ def __init__(
                     num_attention_heads=num_attention_heads,
                     attention_head_dim=attention_head_dim,
                     rope_type=rope_type,
+                    apply_gated_attention=gated_attention,
                 )
                 for _ in range(num_layers)
             ]
@@ -341,10 +345,12 @@ def __init__(
         video_connector_attention_head_dim: int = 128,
         video_connector_num_layers: int = 2,
         video_connector_num_learnable_registers: int | None = 128,
+        video_gated_attn: bool = False,
         audio_connector_num_attention_heads: int = 30,
         audio_connector_attention_head_dim: int = 128,
         audio_connector_num_layers: int = 2,
         audio_connector_num_learnable_registers: int | None = 128,
+        audio_gated_attn: bool = False,
         connector_rope_base_seq_len: int = 4096,
         rope_theta: float = 10000.0,
         rope_double_precision: bool = True,
@@ -373,6 +379,7 @@ def __init__(
             rope_double_precision=rope_double_precision,
             causal_temporal_positioning=causal_temporal_positioning,
             rope_type=rope_type,
+            gated_attention=video_gated_attn,
         )
         self.audio_connector = LTX2ConnectorTransformer1d(
             num_attention_heads=audio_connector_num_attention_heads,
@@ -384,6 +391,7 @@ def __init__(
             rope_double_precision=rope_double_precision,
             causal_temporal_positioning=causal_temporal_positioning,
             rope_type=rope_type,
+            gated_attention=audio_gated_attn,
         )
 
     def forward(
diff --git a/src/diffusers/pipelines/ltx2/vocoder.py b/src/diffusers/pipelines/ltx2/vocoder.py
index 2abbf9ac0182..22c8f71cf923 100644
--- a/src/diffusers/pipelines/ltx2/vocoder.py
+++ b/src/diffusers/pipelines/ltx2/vocoder.py
@@ -270,7 +270,7 @@ def __init__(
         antialias: bool = False,
         antialias_ratio: int = 2,
         antialias_kernel_size: int = 12,
-        final_act_fn: str | None = "tanh",
+        final_act_fn: str | None = "tanh",  # tanh, clamp, None
         final_bias: bool = True,
         output_sampling_rate: int = 24000,
     ):
@@ -389,7 +389,7 @@ def forward(self, hidden_states: torch.Tensor, time_last: bool = False) -> torch
         hidden_states = self.conv_out(hidden_states)
         if self.final_act_fn == "tanh":
             hidden_states = torch.tanh(hidden_states)
-        else:
+        elif self.final_act_fn == "clamp":
             hidden_states = torch.clamp(hidden_states, -1, 1)
 
         return hidden_states

From 0528fde41dbdd1ab469de2ccce0b79683463e04f Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Tue, 10 Mar 2026 05:50:29 +0100
Subject: [PATCH 13/41] Support converting LTX-2.3 DiT checkpoints

---
 scripts/convert_ltx2_to_diffusers.py          | 16 ++++++--
 .../models/transformers/transformer_ltx2.py   | 37 ++++++++++++-------
 2 files changed, 35 insertions(+), 18 deletions(-)

diff --git a/scripts/convert_ltx2_to_diffusers.py b/scripts/convert_ltx2_to_diffusers.py
index 9232497f75d1..28fdd162b337 100644
--- a/scripts/convert_ltx2_to_diffusers.py
+++ b/scripts/convert_ltx2_to_diffusers.py
@@ -44,6 +44,12 @@
     "k_norm": "norm_k",
 }
 
+LTX_2_3_TRANSFORMER_KEYS_RENAME_DICT = {
+    **LTX_2_0_TRANSFORMER_KEYS_RENAME_DICT,
+    "audio_prompt_adaln_single": "audio_prompt_adaln",
+    "prompt_adaln_single": "prompt_adaln",
+}
+
 LTX_2_0_VIDEO_VAE_RENAME_DICT = {
     # Encoder
     "down_blocks.0": "down_blocks.0",
@@ -267,6 +273,7 @@ def get_ltx2_transformer_config(version: str) -> tuple[dict[str, Any], dict[str,
                 "timestep_scale_multiplier": 1000,
                 "cross_attn_timestep_scale_multiplier": 1000,
                 "rope_type": "split",
+                "use_prompt_embeddings": True,
                 "perturbed_attn": False,
             },
         }
@@ -300,8 +307,8 @@ def get_ltx2_transformer_config(version: str) -> tuple[dict[str, Any], dict[str,
                 "audio_pos_embed_max_pos": 20,
                 "audio_sampling_rate": 16000,
                 "audio_hop_length": 160,
-                "audio_gated_attn": False,
-                "audio_cross_attn_mod": False,
+                "audio_gated_attn": True,
+                "audio_cross_attn_mod": True,
                 "num_layers": 48,
                 "activation_fn": "gelu-approximate",
                 "qk_norm": "rms_norm_across_heads",
@@ -316,10 +323,11 @@ def get_ltx2_transformer_config(version: str) -> tuple[dict[str, Any], dict[str,
                 "timestep_scale_multiplier": 1000,
                 "cross_attn_timestep_scale_multiplier": 1000,
                 "rope_type": "split",
+                "use_prompt_embeddings": False,
                 "perturbed_attn": True,
             },
         }
-        rename_dict = LTX_2_0_TRANSFORMER_KEYS_RENAME_DICT
+        rename_dict = LTX_2_3_TRANSFORMER_KEYS_RENAME_DICT
         special_keys_remap = LTX_2_0_TRANSFORMER_SPECIAL_KEYS_REMAP
     return config, rename_dict, special_keys_remap
 
@@ -881,7 +889,7 @@ def none_or_str(value: str):
         "--version",
         type=str,
         default="2.0",
-        choices=["test", "2.0"],
+        choices=["test", "2.0", "2.3"],
         help="Version of the LTX 2.0 model",
     )
 
diff --git a/src/diffusers/models/transformers/transformer_ltx2.py b/src/diffusers/models/transformers/transformer_ltx2.py
index 951220cf671e..ce4d9f964d29 100644
--- a/src/diffusers/models/transformers/transformer_ltx2.py
+++ b/src/diffusers/models/transformers/transformer_ltx2.py
@@ -574,7 +574,7 @@ def __init__(
         self.cross_attn_adaln = video_cross_attn_adaln or audio_cross_attn_adaln
         if self.cross_attn_adaln:
             self.prompt_scale_shift_table = nn.Parameter(torch.randn(2, dim))
-            self.audio_prompt_scale_shift_table = nn.Parameter(torch.randn(2, dim))
+            self.audio_prompt_scale_shift_table = nn.Parameter(torch.randn(2, audio_dim))
 
         # Per-layer a2v, v2a Cross-Attention mod params
         self.video_a2v_cross_attn_scale_shift_table = nn.Parameter(torch.randn(5, dim))
@@ -1122,6 +1122,7 @@ def __init__(
         timestep_scale_multiplier: int = 1000,
         cross_attn_timestep_scale_multiplier: int = 1000,
         rope_type: str = "interleaved",
+        use_prompt_embeddings=True,
         perturbed_attn: bool = False,
     ) -> None:
         super().__init__()
@@ -1136,17 +1137,25 @@ def __init__(
         self.audio_proj_in = nn.Linear(audio_in_channels, audio_inner_dim)
 
         # 2. Prompt embeddings
-        self.caption_projection = PixArtAlphaTextProjection(in_features=caption_channels, hidden_size=inner_dim)
-        self.audio_caption_projection = PixArtAlphaTextProjection(
-            in_features=caption_channels, hidden_size=audio_inner_dim
-        )
+        if use_prompt_embeddings:
+            # LTX-2.0; LTX-2.3 uses per-modality feature projections in the connector instead
+            self.caption_projection = PixArtAlphaTextProjection(in_features=caption_channels, hidden_size=inner_dim)
+            self.audio_caption_projection = PixArtAlphaTextProjection(
+                in_features=caption_channels, hidden_size=audio_inner_dim
+            )
 
         # 3. Timestep Modulation Params and Embedding
+        self.prompt_modulation = cross_attn_mod or audio_cross_attn_mod  # used by LTX-2.3
+
         # 3.1. Global Timestep Modulation Parameters (except for cross-attention) and timestep + size embedding
         # time_embed and audio_time_embed calculate both the timestep embedding and (global) modulation parameters
-        self.time_embed = LTX2AdaLayerNormSingle(inner_dim, num_mod_params=6, use_additional_conditions=False)
+        video_time_emb_mod_params = 9 if cross_attn_mod else 6
+        audio_time_emb_mod_params = 9 if audio_cross_attn_mod else 6
+        self.time_embed = LTX2AdaLayerNormSingle(
+            inner_dim, num_mod_params=video_time_emb_mod_params, use_additional_conditions=False
+        )
         self.audio_time_embed = LTX2AdaLayerNormSingle(
-            audio_inner_dim, num_mod_params=6, use_additional_conditions=False
+            audio_inner_dim, num_mod_params=audio_time_emb_mod_params, use_additional_conditions=False
         )
 
         # 3.2. Global Cross Attention Modulation Parameters
@@ -1176,11 +1185,10 @@ def __init__(
         self.audio_scale_shift_table = nn.Parameter(torch.randn(2, audio_inner_dim) / audio_inner_dim**0.5)
 
         # 3.4. Prompt Scale/Shift Modulation parameters (LTX-2.3)
-        self.prompt_modulation = cross_attn_mod or audio_cross_attn_mod
         if self.prompt_modulation:
             self.prompt_adaln = LTX2AdaLayerNormSingle(inner_dim, num_mod_params=2, use_additional_conditions=False)
             self.audio_prompt_adaln = LTX2AdaLayerNormSingle(
-                inner_dim, num_mod_params=2, use_additional_conditions=False
+                audio_inner_dim, num_mod_params=2, use_additional_conditions=False
             )
 
         # 4. Rotary Positional Embeddings (RoPE)
@@ -1485,12 +1493,13 @@ def forward(
         )
         audio_cross_attn_v2a_gate = audio_cross_attn_v2a_gate.view(batch_size, -1, audio_cross_attn_v2a_gate.shape[-1])
 
-        # 4. Prepare prompt embeddings
-        encoder_hidden_states = self.caption_projection(encoder_hidden_states)
-        encoder_hidden_states = encoder_hidden_states.view(batch_size, -1, hidden_states.size(-1))
+        # 4. Prepare prompt embeddings (LTX-2.0)
+        if self.config.use_prompt_embeddings:
+            encoder_hidden_states = self.caption_projection(encoder_hidden_states)
+            encoder_hidden_states = encoder_hidden_states.view(batch_size, -1, hidden_states.size(-1))
 
-        audio_encoder_hidden_states = self.audio_caption_projection(audio_encoder_hidden_states)
-        audio_encoder_hidden_states = audio_encoder_hidden_states.view(batch_size, -1, audio_hidden_states.size(-1))
+            audio_encoder_hidden_states = self.audio_caption_projection(audio_encoder_hidden_states)
+            audio_encoder_hidden_states = audio_encoder_hidden_states.view(batch_size, -1, audio_hidden_states.size(-1))
 
         # 5. Run transformer blocks
         for block in self.transformer_blocks:

From c5e1fcc4b7e52a81f8c69aed0dcd5cbbab351cb6 Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Tue, 10 Mar 2026 06:50:18 +0100
Subject: [PATCH 14/41] Support converting LTX-2.3 Video VAE checkpoints

---
 scripts/convert_ltx2_to_diffusers.py | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/scripts/convert_ltx2_to_diffusers.py b/scripts/convert_ltx2_to_diffusers.py
index 28fdd162b337..df2f75f89c32 100644
--- a/scripts/convert_ltx2_to_diffusers.py
+++ b/scripts/convert_ltx2_to_diffusers.py
@@ -78,6 +78,13 @@
     "per_channel_statistics.std-of-means": "latents_std",
 }
 
+LTX_2_3_VIDEO_VAE_RENAME_DICT = {
+    **LTX_2_0_VIDEO_VAE_RENAME_DICT,
+    # Decoder extra blocks
+    "up_blocks.7": "up_blocks.3.upsamplers.0",
+    "up_blocks.8": "up_blocks.3",
+}
+
 LTX_2_0_AUDIO_VAE_RENAME_DICT = {
     "per_channel_statistics.mean-of-means": "latents_mean",
     "per_channel_statistics.std-of-means": "latents_std",
@@ -566,7 +573,7 @@ def get_ltx2_video_vae_config(
                 "decoder_spatio_temporal_scaling": (True, True, True, True),
                 "decoder_inject_noise": (False, False, False, False, False),
                 "downsample_type": ("spatial", "temporal", "spatiotemporal", "spatiotemporal"),
-                "upsample_type": ("spatial", "temporal", "spatiotemporal", "spatiotemporal"),
+                "upsample_type": ("spatiotemporal", "spatiotemporal", "temporal", "spatial"),
                 "upsample_residual": (True, True, True, True),
                 "upsample_factor": (2, 2, 1, 2),
                 "timestep_conditioning": timestep_conditioning,
@@ -581,7 +588,7 @@ def get_ltx2_video_vae_config(
                 "temporal_compression_ratio": 8,
             },
         }
-        rename_dict = LTX_2_0_VIDEO_VAE_RENAME_DICT
+        rename_dict = LTX_2_3_VIDEO_VAE_RENAME_DICT
         special_keys_remap = LTX_2_0_VAE_SPECIAL_KEYS_REMAP
     return config, rename_dict, special_keys_remap
 

From 50da4df0bad649414057be0f0c6298b882d3c165 Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Tue, 10 Mar 2026 08:13:02 +0100
Subject: [PATCH 15/41] Support converting LTX-2.3 Vocoder with bandwidth
 extender

---
 scripts/convert_ltx2_to_diffusers.py     | 43 +++++++++++++---
 src/diffusers/pipelines/ltx2/__init__.py |  2 +-
 src/diffusers/pipelines/ltx2/vocoder.py  | 63 ++++++++++++++++--------
 3 files changed, 79 insertions(+), 29 deletions(-)

diff --git a/scripts/convert_ltx2_to_diffusers.py b/scripts/convert_ltx2_to_diffusers.py
index df2f75f89c32..d8d7a2273dc2 100644
--- a/scripts/convert_ltx2_to_diffusers.py
+++ b/scripts/convert_ltx2_to_diffusers.py
@@ -17,7 +17,7 @@
     LTX2Pipeline,
     LTX2VideoTransformer3DModel,
 )
-from diffusers.pipelines.ltx2 import LTX2LatentUpsamplerModel, LTX2TextConnectors, LTX2Vocoder
+from diffusers.pipelines.ltx2 import LTX2LatentUpsamplerModel, LTX2TextConnectors, LTX2Vocoder, LTX2VocoderWithBWE
 from diffusers.utils.import_utils import is_accelerate_available
 
 
@@ -97,6 +97,15 @@
     "conv_post": "conv_out",
 }
 
+LTX_2_3_VOCODER_RENAME_DICT= {
+    # Handle upsamplers ("ups" --> "upsamplers") due to name clash
+    "resblocks": "resnets",
+    "conv_pre": "conv_in",
+    "conv_post": "conv_out",
+    "act_post": "act_out",
+    "downsample.lowpass": "downsample",
+}
+
 LTX_2_0_TEXT_ENCODER_RENAME_DICT = {
     "video_embeddings_connector": "video_connector",
     "audio_embeddings_connector": "audio_connector",
@@ -142,6 +151,18 @@ def convert_ltx2_audio_vae_per_channel_statistics(key: str, state_dict: dict[str
     return
 
 
+def convert_ltx2_3_vocoder_upsamplers(key: str, state_dict: dict[str, Any]) -> None:
+    # Skip if not a weight, bias
+    if ".weight" not in key and ".bias" not in key:
+        return
+
+    if ".ups." in key:
+        new_key = key.replace(".ups.", ".upsamplers.")
+        param = state_dict.pop(key)
+        state_dict[new_key] = param
+    return
+
+
 LTX_2_0_TRANSFORMER_SPECIAL_KEYS_REMAP = {
     "video_embeddings_connector": remove_keys_inplace,
     "audio_embeddings_connector": remove_keys_inplace,
@@ -168,6 +189,10 @@ def convert_ltx2_audio_vae_per_channel_statistics(key: str, state_dict: dict[str
 
 LTX_2_0_VOCODER_SPECIAL_KEYS_REMAP = {}
 
+LTX_2_3_VOCODER_SPECIAL_KEYS_REMAP = {
+    ".ups.": convert_ltx2_3_vocoder_upsamplers,
+}
+
 
 def split_transformer_and_connector_state_dict(state_dict: dict[str, Any]) -> tuple[dict[str, Any], dict[str, Any]]:
     connector_prefixes = (
@@ -728,7 +753,7 @@ def get_ltx2_vocoder_config(version: str) -> tuple[dict[str, Any], dict[str, Any
             "model_id": "Lightricks/LTX-2.3",
             "diffusers_config": {
                 "in_channels": 128,
-                "hidden_channels": 1024,
+                "hidden_channels": 1536,
                 "out_channels": 2,
                 "upsample_kernel_sizes": [11, 4, 4, 4, 4, 4],
                 "upsample_factors": [5, 2, 2, 2, 2, 2],
@@ -744,7 +769,7 @@ def get_ltx2_vocoder_config(version: str) -> tuple[dict[str, Any], dict[str, Any
                 "bwe_in_channels": 128,
                 "bwe_hidden_channels": 512,
                 "bwe_out_channels": 2,
-                "bwe_upsample_kernel_sizes": [12, 11, 8, 4, 4],
+                "bwe_upsample_kernel_sizes": [12, 11, 4, 4, 4],
                 "bwe_upsample_factors": [6, 5, 2, 2, 2],
                 "bwe_resnet_kernel_sizes": [3, 7, 11],
                 "bwe_resnet_dilations": [[1, 3, 5], [1, 3, 5], [1, 3, 5]],
@@ -763,17 +788,21 @@ def get_ltx2_vocoder_config(version: str) -> tuple[dict[str, Any], dict[str, Any
                 "output_sampling_rate": 48000,
             },
         }
-        rename_dict = LTX_2_0_VOCODER_RENAME_DICT
-        special_keys_remap = LTX_2_0_VOCODER_SPECIAL_KEYS_REMAP
+        rename_dict = LTX_2_3_VOCODER_RENAME_DICT
+        special_keys_remap = LTX_2_3_VOCODER_SPECIAL_KEYS_REMAP
     return config, rename_dict, special_keys_remap
 
 
 def convert_ltx2_vocoder(original_state_dict: dict[str, Any], version: str) -> dict[str, Any]:
     config, rename_dict, special_keys_remap = get_ltx2_vocoder_config(version)
     diffusers_config = config["diffusers_config"]
+    if version == "2.3":
+        vocoder_cls = LTX2VocoderWithBWE
+    else:
+        vocoder_cls = LTX2Vocoder
 
     with init_empty_weights():
-        vocoder = LTX2Vocoder.from_config(diffusers_config)
+        vocoder = vocoder_cls.from_config(diffusers_config)
 
     # Handle official code --> diffusers key remapping via the remap dict
     for key in list(original_state_dict.keys()):
@@ -861,7 +890,7 @@ def get_model_state_dict_from_combined_ckpt(combined_ckpt: dict[str, Any], prefi
     model_state_dict = {}
     for param_name, param in combined_ckpt.items():
         if param_name.startswith(prefix):
-            model_state_dict[param_name.replace(prefix, "")] = param
+            model_state_dict[param_name.removeprefix(prefix)] = param
 
     if prefix == "model.diffusion_model.":
         # Some checkpoints store the text connector projection outside the diffusion model prefix.
diff --git a/src/diffusers/pipelines/ltx2/__init__.py b/src/diffusers/pipelines/ltx2/__init__.py
index 0c95a17f3e31..7177faaf3486 100644
--- a/src/diffusers/pipelines/ltx2/__init__.py
+++ b/src/diffusers/pipelines/ltx2/__init__.py
@@ -28,7 +28,7 @@
     _import_structure["pipeline_ltx2_condition"] = ["LTX2ConditionPipeline"]
     _import_structure["pipeline_ltx2_image2video"] = ["LTX2ImageToVideoPipeline"]
     _import_structure["pipeline_ltx2_latent_upsample"] = ["LTX2LatentUpsamplePipeline"]
-    _import_structure["vocoder"] = ["LTX2Vocoder"]
+    _import_structure["vocoder"] = ["LTX2Vocoder", "LTX2VocoderWithBWE"]
 
 if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
     try:
diff --git a/src/diffusers/pipelines/ltx2/vocoder.py b/src/diffusers/pipelines/ltx2/vocoder.py
index 22c8f71cf923..f46cb249836e 100644
--- a/src/diffusers/pipelines/ltx2/vocoder.py
+++ b/src/diffusers/pipelines/ltx2/vocoder.py
@@ -26,6 +26,7 @@ def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) ->
             The Kaiser sinc kernel.
     """
     delta_f = 4 * half_width
+    half_size = kernel_size // 2
     amplitude = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95
     if amplitude > 50.0:
         beta = 0.1102 * (amplitude - 8.7)
@@ -37,7 +38,6 @@ def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) ->
     window = torch.kaiser_window(kernel_size, beta=beta, periodic=False)
 
     even = kernel_size % 2 == 0
-    half_size = kernel_size // 2
     time = torch.arange(-half_size, half_size) + 0.5 if even else torch.arange(kernel_size) - half_size
 
     if cutoff == 0.0:
@@ -76,7 +76,7 @@ def __init__(
         cutoff = 0.5 / ratio
         half_width = 0.6 / ratio
         low_pass_filter = kaiser_sinc_filter1d(cutoff, half_width, self.kernel_size)
-        self.register_buffer("filter", low_pass_filter.view(1, 1, self,kernel_size), persistent=persistent)
+        self.register_buffer("filter", low_pass_filter.view(1, 1, self.kernel_size), persistent=persistent)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         # x expected shape: [batch_size, num_channels, hidden_dim]
@@ -117,8 +117,8 @@ def __init__(
             # Kaiser sinc filter is BigVGAN default
             self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
             self.pad = self.kernel_size // ratio - 1
-            self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2
-            self.pad_right = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2
+            self.pad_left = self.pad * self.ratio + (self.kernel_size - self.ratio) // 2
+            self.pad_right = self.pad * self.ratio + (self.kernel_size - self.ratio + 1) // 2
 
             sinc_filter = kaiser_sinc_filter1d(
                 cutoff=0.5 / ratio,
@@ -126,7 +126,7 @@ def __init__(
                 kernel_size=self.kernel_size,
             )
 
-        self.register_buffer("filter", sinc_filter, persistent=persistent)
+        self.register_buffer("filter", sinc_filter.view(1, 1, self.kernel_size), persistent=persistent)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         # x expected shape: [batch_size, num_channels, hidden_dim]
@@ -137,6 +137,37 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return x[..., self.pad_left:-self.pad_right]
 
 
+class AntiAliasAct1d(nn.Module):
+    """
+    Antialiasing activation for a 1D signal: upsamples, applies an activation (usually snakebeta), and then downsamples
+    to avoid aliasing.
+    """
+    def __init__(
+        self,
+        act_fn: str | nn.Module,
+        ratio: int = 2,
+        kernel_size: int = 12,
+        **kwargs,
+    ):
+        super().__init__()
+        self.upsample = UpSample1d(ratio=ratio, kernel_size=kernel_size)
+        if isinstance(act_fn, str):
+            if act_fn == "snakebeta":
+                act_fn = SnakeBeta(**kwargs)
+            elif act_fn == "snake":
+                act_fn = SnakeBeta(**kwargs)
+            else:
+                act_fn = nn.LeakyReLU(**kwargs)
+        self.act = act_fn
+        self.downsample = DownSample1d(ratio=ratio, kernel_size=kernel_size)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.upsample(x)
+        x = self.act(x)
+        x = self.downsample(x)
+        return x
+
+
 class SnakeBeta(nn.Module):
     """
     Implements the Snake and SnakeBeta activations, which help with learning periodic patterns.
@@ -150,6 +181,7 @@ def __init__(
         logscale: bool = True,
         use_beta: bool = True,
     ):
+        super().__init__()
         self.eps = eps
         self.logscale = logscale
         self.use_beta = use_beta
@@ -210,11 +242,7 @@ def __init__(
                 act = nn.LeakyReLU(negative_slope=leaky_relu_negative_slope)
 
             if antialias:
-                act = nn.Sequential(
-                    UpSample1d(ratio=antialias_ratio, kernel_size=antialias_kernel_size),
-                    act,
-                    DownSample1d(ratio=antialias_ratio, kernel_size=antialias_kernel_size),
-                )
+                act = AntiAliasAct1d(act, ratio=antialias_ratio, kernel_size=antialias_kernel_size)
             self.acts1.append(act)
 
         self.convs2 = nn.ModuleList(
@@ -233,11 +261,7 @@ def __init__(
                 act_fn = nn.LeakyReLU(negative_slope=leaky_relu_negative_slope)
 
             if antialias:
-                act = nn.Sequential(
-                    UpSample1d(ratio=antialias_ratio, kernel_size=antialias_kernel_size),
-                    act,
-                    DownSample1d(ratio=antialias_ratio, kernel_size=antialias_kernel_size),
-                )
+                act = AntiAliasAct1d(act, ratio=antialias_ratio, kernel_size=antialias_kernel_size)
             self.acts2.append(act)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -336,11 +360,8 @@ def __init__(
 
         if act_fn == "snakebeta" or act_fn == "snake":
             # Always use antialiasing
-            self.act_out = nn.Sequential(
-                UpSample1d(ratio=antialias_ratio, kernel_size=antialias_kernel_size),
-                SnakeBeta(channels=out_channels, use_beta=True),
-                DownSample1d(ratio=antialias_ratio, kernel_size=antialias_kernel_size),
-            )
+            act_out = SnakeBeta(channels=output_channels, use_beta=True)
+            self.act_out = AntiAliasAct1d(act_out, ratio=antialias_ratio, kernel_size=antialias_kernel_size)
         elif act_fn == "leaky_relu":
             # NOTE: does NOT use self.negative_slope, following the original code
             self.act_out = nn.LeakyReLU()
@@ -480,7 +501,7 @@ def __init__(
         bwe_in_channels: int = 128,
         bwe_hidden_channels: int = 512,
         bwe_out_channels: int = 2,
-        bwe_upsample_kernel_sizes: list[int] = [12, 11, 8, 4, 4],
+        bwe_upsample_kernel_sizes: list[int] = [12, 11, 4, 4, 4],
         bwe_upsample_factors: list[int] = [6, 5, 2, 2, 2],
         bwe_resnet_kernel_sizes: list[int] = [3, 7, 11],
         bwe_resnet_dilations: list[list[int]] = [[1, 3, 5], [1, 3, 5], [1, 3, 5]],

From 420628039afaf24c0f5681f969ce5e38d39191e0 Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Tue, 10 Mar 2026 08:39:51 +0100
Subject: [PATCH 16/41] Support converting LTX-2.3 text connectors

---
 scripts/convert_ltx2_to_diffusers.py       | 51 +++++++++++++---------
 src/diffusers/pipelines/ltx2/connectors.py |  4 +-
 2 files changed, 33 insertions(+), 22 deletions(-)

diff --git a/scripts/convert_ltx2_to_diffusers.py b/scripts/convert_ltx2_to_diffusers.py
index d8d7a2273dc2..63fb6486c137 100644
--- a/scripts/convert_ltx2_to_diffusers.py
+++ b/scripts/convert_ltx2_to_diffusers.py
@@ -106,10 +106,25 @@
     "downsample.lowpass": "downsample",
 }
 
-LTX_2_0_TEXT_ENCODER_RENAME_DICT = {
+LTX_2_0_CONNECTORS_KEYS_RENAME_DICT = {
+    "connectors.": "",
     "video_embeddings_connector": "video_connector",
     "audio_embeddings_connector": "audio_connector",
     "transformer_1d_blocks": "transformer_blocks",
+    "text_embedding_projection.aggregate_embed": "text_proj_in",
+    # Attention QK Norms
+    "q_norm": "norm_q",
+    "k_norm": "norm_k",
+}
+
+LTX_2_3_CONNECTORS_KEYS_RENAME_DICT = {
+    "connectors.": "",
+    "video_embeddings_connector": "video_connector",
+    "audio_embeddings_connector": "audio_connector",
+    "transformer_1d_blocks": "transformer_blocks",
+    # LTX-2.3 uses per-modality embedding projections
+    "text_embedding_projection.audio_aggregate_embed": "audio_text_proj_in",
+    "text_embedding_projection.video_aggregate_embed": "video_text_proj_in",
     # Attention QK Norms
     "q_norm": "norm_q",
     "k_norm": "norm_k",
@@ -169,17 +184,6 @@ def convert_ltx2_3_vocoder_upsamplers(key: str, state_dict: dict[str, Any]) -> N
     "adaln_single": convert_ltx2_transformer_adaln_single,
 }
 
-LTX_2_0_CONNECTORS_KEYS_RENAME_DICT = {
-    "connectors.": "",
-    "video_embeddings_connector": "video_connector",
-    "audio_embeddings_connector": "audio_connector",
-    "transformer_1d_blocks": "transformer_blocks",
-    "text_embedding_projection.aggregate_embed": "text_proj_in",
-    # Attention QK Norms
-    "q_norm": "norm_q",
-    "k_norm": "norm_k",
-}
-
 LTX_2_0_VAE_SPECIAL_KEYS_REMAP = {
     "per_channel_statistics.channel": remove_keys_inplace,
     "per_channel_statistics.mean-of-stds": remove_keys_inplace,
@@ -193,13 +197,15 @@ def convert_ltx2_3_vocoder_upsamplers(key: str, state_dict: dict[str, Any]) -> N
     ".ups.": convert_ltx2_3_vocoder_upsamplers,
 }
 
+LTX_2_0_CONNECTORS_SPECIAL_KEYS_REMAP = {}
+
 
 def split_transformer_and_connector_state_dict(state_dict: dict[str, Any]) -> tuple[dict[str, Any], dict[str, Any]]:
     connector_prefixes = (
         "video_embeddings_connector",
         "audio_embeddings_connector",
         "transformer_1d_blocks",
-        "text_embedding_projection.aggregate_embed",
+        "text_embedding_projection",
         "connectors.",
         "video_connector",
         "audio_connector",
@@ -410,6 +416,8 @@ def get_ltx2_connectors_config(version: str) -> tuple[dict[str, Any], dict[str,
                 "proj_bias": False,
             },
         }
+        rename_dict = LTX_2_0_CONNECTORS_KEYS_RENAME_DICT
+        special_keys_remap = LTX_2_0_CONNECTORS_SPECIAL_KEYS_REMAP
     elif version == "2.3":
         config = {
             "model_id": "Lightricks/LTX-2.3",
@@ -437,9 +445,8 @@ def get_ltx2_connectors_config(version: str) -> tuple[dict[str, Any], dict[str,
                 "proj_bias": True,
             },
         }
-
-    rename_dict = LTX_2_0_CONNECTORS_KEYS_RENAME_DICT
-    special_keys_remap = {}
+        rename_dict = LTX_2_3_CONNECTORS_KEYS_RENAME_DICT
+        special_keys_remap = LTX_2_0_CONNECTORS_SPECIAL_KEYS_REMAP
 
     return config, rename_dict, special_keys_remap
 
@@ -894,9 +901,13 @@ def get_model_state_dict_from_combined_ckpt(combined_ckpt: dict[str, Any], prefi
 
     if prefix == "model.diffusion_model.":
         # Some checkpoints store the text connector projection outside the diffusion model prefix.
-        connector_key = "text_embedding_projection.aggregate_embed.weight"
-        if connector_key in combined_ckpt and connector_key not in model_state_dict:
-            model_state_dict[connector_key] = combined_ckpt[connector_key]
+        connector_prefixes = ["text_embedding_projection"]
+        for param_name, param in combined_ckpt.items():
+            for prefix in connector_prefixes:
+                if param_name.startswith(prefix):
+                    # Check to make sure we're not overwriting an existing key
+                    if param_name not in model_state_dict:
+                        model_state_dict[param_name] = combined_ckpt[param_name]
 
     return model_state_dict
 
@@ -1026,7 +1037,7 @@ def main(args):
             args.audio_vae,
             args.dit,
             args.vocoder,
-            args.text_encoder,
+            args.connectors,
             args.full_pipeline,
             args.upsample_pipeline,
         ]
diff --git a/src/diffusers/pipelines/ltx2/connectors.py b/src/diffusers/pipelines/ltx2/connectors.py
index b75444692b08..40cf14769015 100644
--- a/src/diffusers/pipelines/ltx2/connectors.py
+++ b/src/diffusers/pipelines/ltx2/connectors.py
@@ -11,7 +11,7 @@
 from ...models.transformers.transformer_ltx2 import LTX2Attention, LTX2AudioVideoAttnProcessor
 
 
-def per_batch_per_layer_mean_norm(
+def per_layer_masked_mean_norm(
     text_hidden_states: torch.Tensor,
     sequence_lengths: torch.Tensor,
     device: str | torch.device,
@@ -446,7 +446,7 @@ def forward(
         else:
             # LTX-2.0
             sequence_lengths = attention_mask.sum(dim=-1)
-            norm_text_encoder_hidden_states = per_batch_per_layer_mean_norm(
+            norm_text_encoder_hidden_states = per_layer_masked_mean_norm(
                 text_hidden_states=text_encoder_hidden_states,
                 sequence_lengths=sequence_lengths,
                 device=text_encoder_hidden_states.device,

From e719d32c63872aa80ab3b3a96bc561c5f9cd95b8 Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Tue, 10 Mar 2026 09:03:44 +0100
Subject: [PATCH 17/41] Don't convert any upsamplers for now

---
 scripts/convert_ltx2_to_diffusers.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/scripts/convert_ltx2_to_diffusers.py b/scripts/convert_ltx2_to_diffusers.py
index 63fb6486c137..b5fca2d33998 100644
--- a/scripts/convert_ltx2_to_diffusers.py
+++ b/scripts/convert_ltx2_to_diffusers.py
@@ -1102,7 +1102,7 @@ def main(args):
         if not args.full_pipeline:
             tokenizer.save_pretrained(os.path.join(args.output_path, "tokenizer"))
 
-    if args.latent_upsampler or args.full_pipeline or args.upsample_pipeline:
+    if args.latent_upsampler or args.upsample_pipeline:
         original_latent_upsampler_ckpt = load_hub_or_local_checkpoint(
             repo_id=args.original_state_dict_repo_id, filename=args.latent_upsampler_filename
         )

From fbb50d964d6d1dd75ded16641bcf0c6dcc8f14b9 Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Tue, 10 Mar 2026 09:50:08 +0100
Subject: [PATCH 18/41] Support self attention mask for LTX2Pipeline

---
 .../models/transformers/transformer_ltx2.py    | 18 ++++++++++++++++++
 src/diffusers/pipelines/ltx2/pipeline_ltx2.py  |  7 +++++--
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/src/diffusers/models/transformers/transformer_ltx2.py b/src/diffusers/models/transformers/transformer_ltx2.py
index ce4d9f964d29..ebb63f093af9 100644
--- a/src/diffusers/models/transformers/transformer_ltx2.py
+++ b/src/diffusers/models/transformers/transformer_ltx2.py
@@ -612,8 +612,11 @@ def forward(
         ca_audio_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
         encoder_attention_mask: torch.Tensor | None = None,
         audio_encoder_attention_mask: torch.Tensor | None = None,
+        self_attention_mask: torch.Tensor | None = None,
+        audio_self_attention_mask: torch.Tensor | None = None,
         a2v_cross_attention_mask: torch.Tensor | None = None,
         v2a_cross_attention_mask: torch.Tensor | None = None,
+        perturbation_mask: torch.Tensor | None = None,
     ) -> torch.Tensor:
         batch_size = hidden_states.size(0)
 
@@ -631,6 +634,7 @@ def forward(
             hidden_states=norm_hidden_states,
             encoder_hidden_states=None,
             query_rotary_emb=video_rotary_emb,
+            attention_mask=self_attention_mask,
         )
         hidden_states = hidden_states + attn_hidden_states * gate_msa
 
@@ -649,6 +653,7 @@ def forward(
             hidden_states=norm_audio_hidden_states,
             encoder_hidden_states=None,
             query_rotary_emb=audio_rotary_emb,
+            attention_mask=audio_self_attention_mask,
         )
         audio_hidden_states = audio_hidden_states + attn_audio_hidden_states * audio_gate_msa
 
@@ -1307,6 +1312,7 @@ def forward(
         encoder_attention_mask: torch.Tensor | None = None,
         audio_encoder_attention_mask: torch.Tensor | None = None,
         self_attention_mask: torch.Tensor | None = None,
+        audio_self_attention_mask: torch.Tensor | None = None,
         num_frames: int | None = None,
         height: int | None = None,
         width: int | None = None,
@@ -1401,6 +1407,18 @@ def forward(
                 ).to(hidden_states.dtype)
             self_attention_mask = additive_self_attn_mask.unsqueeze(1)  # [batch_size, 1, seq_len, seq_len]
 
+        if audio_self_attention_mask is not None and audio_self_attention_mask.ndim == 3:
+            # Convert to additive attention mask in log-space where 0 (masked) values get mapped to a large negative
+            # number and positive values are mapped to their logarithm.
+            dtype_finfo = torch.finfo(hidden_states.dtype)
+            additive_self_attn_mask = torch.full_like(audio_self_attention_mask, dtype_finfo.min, dtype=hidden_states.dtype)
+            unmasked_entries = audio_self_attention_mask > 0
+            if torch.any(unmasked_entries):
+                additive_self_attn_mask[unmasked_entries] = torch.log(
+                    audio_self_attention_mask[unmasked_entries].clamp(min=dtype_finfo.tiny)
+                ).to(hidden_states.dtype)
+            audio_self_attention_mask = additive_self_attn_mask.unsqueeze(1)  # [batch_size, 1, seq_len, seq_len]
+
         batch_size = hidden_states.size(0)
 
         # 1. Prepare RoPE positional embeddings
diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
index 4b350d3a4883..9791c1e9d80e 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
@@ -31,7 +31,7 @@
 from ..pipeline_utils import DiffusionPipeline
 from .connectors import LTX2TextConnectors
 from .pipeline_output import LTX2PipelineOutput
-from .vocoder import LTX2Vocoder
+from .vocoder import LTX2Vocoder, LTX2VocoderWithBWE
 
 
 if is_torch_xla_available():
@@ -221,7 +221,7 @@ def __init__(
         tokenizer: GemmaTokenizer | GemmaTokenizerFast,
         connectors: LTX2TextConnectors,
         transformer: LTX2VideoTransformer3DModel,
-        vocoder: LTX2Vocoder,
+        vocoder: LTX2Vocoder | LTX2VocoderWithBWE,
     ):
         super().__init__()
 
@@ -1037,8 +1037,11 @@ def __call__(
                         encoder_hidden_states=connector_prompt_embeds,
                         audio_encoder_hidden_states=connector_audio_prompt_embeds,
                         timestep=timestep,
+                        sigma=timestep,  # Used by LTX-2.3
                         encoder_attention_mask=connector_attention_mask,
                         audio_encoder_attention_mask=connector_attention_mask,
+                        self_attention_mask=None,
+                        audio_self_attention_mask=None,
                         num_frames=latent_num_frames,
                         height=latent_height,
                         width=latent_width,

From de3f869b5ca3f313285a1d52411095860a366be0 Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Tue, 10 Mar 2026 10:43:50 +0100
Subject: [PATCH 19/41] Fix some inference bugs

---
 src/diffusers/pipelines/ltx2/connectors.py | 2 +-
 src/diffusers/pipelines/ltx2/vocoder.py    | 3 +--
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/src/diffusers/pipelines/ltx2/connectors.py b/src/diffusers/pipelines/ltx2/connectors.py
index 40cf14769015..f945200962d2 100644
--- a/src/diffusers/pipelines/ltx2/connectors.py
+++ b/src/diffusers/pipelines/ltx2/connectors.py
@@ -79,7 +79,7 @@ def per_layer_masked_mean_norm(
 
 
 def per_token_rms_norm(text_encoder_hidden_states: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
-    variance = torch.mean(text_encoder_hidden_states ** 2, dim=2, keepdim=2)
+    variance = torch.mean(text_encoder_hidden_states ** 2, dim=2, keepdim=True)
     norm_text_encoder_hidden_states = text_encoder_hidden_states + torch.rsqrt(variance + eps)
     return norm_text_encoder_hidden_states
 
diff --git a/src/diffusers/pipelines/ltx2/vocoder.py b/src/diffusers/pipelines/ltx2/vocoder.py
index f46cb249836e..c17bfb786c9d 100644
--- a/src/diffusers/pipelines/ltx2/vocoder.py
+++ b/src/diffusers/pipelines/ltx2/vocoder.py
@@ -580,8 +580,7 @@ def forward(self, mel_spec: torch.Tensor) -> torch.Tensor:
             x = F.pad(x, (0, self.hop_length - remainder))
 
         # 2. Compute mel spectrogram on vocoder output
-        x = x.flatten(0, 1)
-        mel, _, _, _ = self.mel_stft(x)
+        mel, _, _, _ = self.mel_stft(x.flatten(0, 1))
         mel = mel.unflatten(0, (-1, num_channels))
 
         # 3. Run bandwidth extender (BWE) on new mel spectrogram

From 5056aa8203b198cb225e0bef7151471649da1de7 Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Wed, 11 Mar 2026 01:12:04 +0100
Subject: [PATCH 20/41] Support self attn mask and sigmas for LTX-2.3 I2V, Cond
 pipelines

---
 src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py   | 7 +++++--
 src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py | 7 +++++--
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py
index d66cd272c2c3..058ed27146c5 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py
@@ -33,7 +33,7 @@
 from ..pipeline_utils import DiffusionPipeline
 from .connectors import LTX2TextConnectors
 from .pipeline_output import LTX2PipelineOutput
-from .vocoder import LTX2Vocoder
+from .vocoder import LTX2Vocoder, LTX2VocoderWithBWE
 
 
 if is_torch_xla_available():
@@ -254,7 +254,7 @@ def __init__(
         tokenizer: GemmaTokenizer | GemmaTokenizerFast,
         connectors: LTX2TextConnectors,
         transformer: LTX2VideoTransformer3DModel,
-        vocoder: LTX2Vocoder,
+        vocoder: LTX2Vocoder | LTX2VocoderWithBWE,
     ):
         super().__init__()
 
@@ -1269,8 +1269,11 @@ def __call__(
                         audio_encoder_hidden_states=connector_audio_prompt_embeds,
                         timestep=video_timestep,
                         audio_timestep=timestep,
+                        sigma=timestep,  # Used by LTX-2.3
                         encoder_attention_mask=connector_attention_mask,
                         audio_encoder_attention_mask=connector_attention_mask,
+                        self_attention_mask=None,
+                        audio_self_attention_mask=None,
                         num_frames=latent_num_frames,
                         height=latent_height,
                         width=latent_width,
diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
index 26329685f2b3..b5371260d3bb 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
@@ -32,7 +32,7 @@
 from ..pipeline_utils import DiffusionPipeline
 from .connectors import LTX2TextConnectors
 from .pipeline_output import LTX2PipelineOutput
-from .vocoder import LTX2Vocoder
+from .vocoder import LTX2Vocoder, LTX2VocoderWithBWE
 
 
 if is_torch_xla_available():
@@ -224,7 +224,7 @@ def __init__(
         tokenizer: GemmaTokenizer | GemmaTokenizerFast,
         connectors: LTX2TextConnectors,
         transformer: LTX2VideoTransformer3DModel,
-        vocoder: LTX2Vocoder,
+        vocoder: LTX2Vocoder | LTX2VocoderWithBWE,
     ):
         super().__init__()
 
@@ -1102,8 +1102,11 @@ def __call__(
                         audio_encoder_hidden_states=connector_audio_prompt_embeds,
                         timestep=video_timestep,
                         audio_timestep=timestep,
+                        sigma=timestep,  # Used by LTX-2.3
                         encoder_attention_mask=connector_attention_mask,
                         audio_encoder_attention_mask=connector_attention_mask,
+                        self_attention_mask=None,
+                        audio_self_attention_mask=None,
                         num_frames=latent_num_frames,
                         height=latent_height,
                         width=latent_width,

From f875031d2f3e2f1d0b2aeef41ffb25ba3d69db33 Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Wed, 11 Mar 2026 06:41:34 +0100
Subject: [PATCH 21/41] Support STG and modality isolation guidance for LTX-2.3

---
 .../models/transformers/transformer_ltx2.py   | 183 ++++++++-----
 src/diffusers/pipelines/ltx2/pipeline_ltx2.py | 241 +++++++++++++++--
 .../pipelines/ltx2/pipeline_ltx2_condition.py | 245 ++++++++++++++++--
 .../ltx2/pipeline_ltx2_image2video.py         | 245 ++++++++++++++++--
 4 files changed, 777 insertions(+), 137 deletions(-)

diff --git a/src/diffusers/models/transformers/transformer_ltx2.py b/src/diffusers/models/transformers/transformer_ltx2.py
index ebb63f093af9..0aed1124fb4a 100644
--- a/src/diffusers/models/transformers/transformer_ltx2.py
+++ b/src/diffusers/models/transformers/transformer_ltx2.py
@@ -453,6 +453,7 @@ def __init__(
     ):
         super().__init__()
 
+        self.perturbed_attn = perturbed_attn
         if perturbed_attn:
             attn_processor_cls = LTX2PerturbedAttnProcessor
         else:
@@ -616,7 +617,10 @@ def forward(
         audio_self_attention_mask: torch.Tensor | None = None,
         a2v_cross_attention_mask: torch.Tensor | None = None,
         v2a_cross_attention_mask: torch.Tensor | None = None,
+        use_a2v_cross_attention: bool = True,
+        use_v2a_cross_attention: bool = True,
         perturbation_mask: torch.Tensor | None = None,
+        all_perturbed: bool | None = None,
     ) -> torch.Tensor:
         batch_size = hidden_states.size(0)
 
@@ -630,12 +634,17 @@ def forward(
         norm_hidden_states = self.norm1(hidden_states)
         norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa
 
-        attn_hidden_states = self.attn1(
-            hidden_states=norm_hidden_states,
-            encoder_hidden_states=None,
-            query_rotary_emb=video_rotary_emb,
-            attention_mask=self_attention_mask,
-        )
+        video_self_attn_args = {
+            "hidden_states": norm_hidden_states,
+            "encoder_hidden_states": None,
+            "query_rotary_emb": video_rotary_emb,
+            "attention_mask": self_attention_mask,
+        }
+        if self.perturbed_attn:
+            video_self_attn_args["perturbation_mask"] = perturbation_mask
+            video_self_attn_args["all_perturbed"] = all_perturbed
+
+        attn_hidden_states = self.attn1(**video_self_attn_args)
         hidden_states = hidden_states + attn_hidden_states * gate_msa
 
         # 1.2. Audio Self-Attention
@@ -649,12 +658,17 @@ def forward(
         norm_audio_hidden_states = self.audio_norm1(audio_hidden_states)
         norm_audio_hidden_states = norm_audio_hidden_states * (1 + audio_scale_msa) + audio_shift_msa
 
-        attn_audio_hidden_states = self.audio_attn1(
-            hidden_states=norm_audio_hidden_states,
-            encoder_hidden_states=None,
-            query_rotary_emb=audio_rotary_emb,
-            attention_mask=audio_self_attention_mask,
-        )
+        audio_self_attn_args = {
+            "hidden_states": norm_audio_hidden_states,
+            "encoder_hidden_states": None,
+            "query_rotary_emb": audio_rotary_emb,
+            "attention_mask": audio_self_attention_mask,
+        }
+        if self.perturbed_attn:
+            audio_self_attn_args["perturbation_mask"] = perturbation_mask
+            audio_self_attn_args["all_perturbed"] = all_perturbed
+
+        attn_audio_hidden_states = self.audio_attn1(**audio_self_attn_args)
         audio_hidden_states = audio_hidden_states + attn_audio_hidden_states * audio_gate_msa
 
         # 2. Video and Audio Cross-Attention with the text embeddings (Q: Video or Audio; K,V: Text)
@@ -700,65 +714,68 @@ def forward(
         audio_hidden_states = audio_hidden_states + attn_audio_hidden_states
 
         # 3. Audio-to-Video (a2v) and Video-to-Audio (v2a) Cross-Attention
-        norm_hidden_states = self.audio_to_video_norm(hidden_states)
-        norm_audio_hidden_states = self.video_to_audio_norm(audio_hidden_states)
+        if use_a2v_cross_attention or use_v2a_cross_attention:
+            norm_hidden_states = self.audio_to_video_norm(hidden_states)
+            norm_audio_hidden_states = self.video_to_audio_norm(audio_hidden_states)
 
-        # 3.1. Combine global and per-layer cross attention modulation parameters
-        # Video
-        video_per_layer_ca_scale_shift = self.video_a2v_cross_attn_scale_shift_table[:4, :]
-        video_per_layer_ca_gate = self.video_a2v_cross_attn_scale_shift_table[4:, :]
+            # 3.1. Combine global and per-layer cross attention modulation parameters
+            # Video
+            video_per_layer_ca_scale_shift = self.video_a2v_cross_attn_scale_shift_table[:4, :]
+            video_per_layer_ca_gate = self.video_a2v_cross_attn_scale_shift_table[4:, :]
 
-        video_ca_ada_params = self.get_mod_params(video_per_layer_ca_scale_shift, temb_ca_scale_shift, batch_size)
-        video_ca_gate_param = self.get_mod_params(video_per_layer_ca_gate, temb_ca_gate, batch_size)
+            video_ca_ada_params = self.get_mod_params(video_per_layer_ca_scale_shift, temb_ca_scale_shift, batch_size)
+            video_ca_gate_param = self.get_mod_params(video_per_layer_ca_gate, temb_ca_gate, batch_size)
 
-        video_a2v_ca_scale, video_a2v_ca_shift, video_v2a_ca_scale, video_v2a_ca_shift = video_ca_ada_params
-        a2v_gate = video_ca_gate_param[0].squeeze(2)
+            video_a2v_ca_scale, video_a2v_ca_shift, video_v2a_ca_scale, video_v2a_ca_shift = video_ca_ada_params
+            a2v_gate = video_ca_gate_param[0].squeeze(2)
 
-        # Audio
-        audio_per_layer_ca_scale_shift = self.audio_a2v_cross_attn_scale_shift_table[:4, :]
-        audio_per_layer_ca_gate = self.audio_a2v_cross_attn_scale_shift_table[4:, :]
+            # Audio
+            audio_per_layer_ca_scale_shift = self.audio_a2v_cross_attn_scale_shift_table[:4, :]
+            audio_per_layer_ca_gate = self.audio_a2v_cross_attn_scale_shift_table[4:, :]
 
-        audio_ca_ada_params = self.get_mod_params(audio_per_layer_ca_scale_shift, temb_ca_audio_scale_shift, batch_size)
-        audio_ca_gate_param = self.get_mod_params(audio_per_layer_ca_gate, temb_ca_audio_gate, batch_size)
+            audio_ca_ada_params = self.get_mod_params(audio_per_layer_ca_scale_shift, temb_ca_audio_scale_shift, batch_size)
+            audio_ca_gate_param = self.get_mod_params(audio_per_layer_ca_gate, temb_ca_audio_gate, batch_size)
 
-        audio_a2v_ca_scale, audio_a2v_ca_shift, audio_v2a_ca_scale, audio_v2a_ca_shift = audio_ca_ada_params
-        v2a_gate = audio_ca_gate_param[0].squeeze(2)
+            audio_a2v_ca_scale, audio_a2v_ca_shift, audio_v2a_ca_scale, audio_v2a_ca_shift = audio_ca_ada_params
+            v2a_gate = audio_ca_gate_param[0].squeeze(2)
 
-        # 3.2. Audio-to-Video Cross Attention: Q: Video; K,V: Audio
-        mod_norm_hidden_states = norm_hidden_states * (1 + video_a2v_ca_scale.squeeze(2)) + video_a2v_ca_shift.squeeze(
-            2
-        )
-        mod_norm_audio_hidden_states = norm_audio_hidden_states * (
-            1 + audio_a2v_ca_scale.squeeze(2)
-        ) + audio_a2v_ca_shift.squeeze(2)
-
-        a2v_attn_hidden_states = self.audio_to_video_attn(
-            mod_norm_hidden_states,
-            encoder_hidden_states=mod_norm_audio_hidden_states,
-            query_rotary_emb=ca_video_rotary_emb,
-            key_rotary_emb=ca_audio_rotary_emb,
-            attention_mask=a2v_cross_attention_mask,
-        )
+            # 3.2. Audio-to-Video Cross Attention: Q: Video; K,V: Audio
+            if use_a2v_cross_attention:
+                mod_norm_hidden_states = norm_hidden_states * (1 + video_a2v_ca_scale.squeeze(2)) + video_a2v_ca_shift.squeeze(
+                    2
+                )
+                mod_norm_audio_hidden_states = norm_audio_hidden_states * (
+                    1 + audio_a2v_ca_scale.squeeze(2)
+                ) + audio_a2v_ca_shift.squeeze(2)
+
+                a2v_attn_hidden_states = self.audio_to_video_attn(
+                    mod_norm_hidden_states,
+                    encoder_hidden_states=mod_norm_audio_hidden_states,
+                    query_rotary_emb=ca_video_rotary_emb,
+                    key_rotary_emb=ca_audio_rotary_emb,
+                    attention_mask=a2v_cross_attention_mask,
+                )
 
-        hidden_states = hidden_states + a2v_gate * a2v_attn_hidden_states
+                hidden_states = hidden_states + a2v_gate * a2v_attn_hidden_states
 
-        # 3.3. Video-to-Audio Cross Attention: Q: Audio; K,V: Video
-        mod_norm_hidden_states = norm_hidden_states * (1 + video_v2a_ca_scale.squeeze(2)) + video_v2a_ca_shift.squeeze(
-            2
-        )
-        mod_norm_audio_hidden_states = norm_audio_hidden_states * (
-            1 + audio_v2a_ca_scale.squeeze(2)
-        ) + audio_v2a_ca_shift.squeeze(2)
-
-        v2a_attn_hidden_states = self.video_to_audio_attn(
-            mod_norm_audio_hidden_states,
-            encoder_hidden_states=mod_norm_hidden_states,
-            query_rotary_emb=ca_audio_rotary_emb,
-            key_rotary_emb=ca_video_rotary_emb,
-            attention_mask=v2a_cross_attention_mask,
-        )
+            # 3.3. Video-to-Audio Cross Attention: Q: Audio; K,V: Video
+            if use_v2a_cross_attention:
+                mod_norm_hidden_states = norm_hidden_states * (1 + video_v2a_ca_scale.squeeze(2)) + video_v2a_ca_shift.squeeze(
+                    2
+                )
+                mod_norm_audio_hidden_states = norm_audio_hidden_states * (
+                    1 + audio_v2a_ca_scale.squeeze(2)
+                ) + audio_v2a_ca_shift.squeeze(2)
+
+                v2a_attn_hidden_states = self.video_to_audio_attn(
+                    mod_norm_audio_hidden_states,
+                    encoder_hidden_states=mod_norm_hidden_states,
+                    query_rotary_emb=ca_audio_rotary_emb,
+                    key_rotary_emb=ca_video_rotary_emb,
+                    attention_mask=v2a_cross_attention_mask,
+                )
 
-        audio_hidden_states = audio_hidden_states + v2a_gate * v2a_attn_hidden_states
+                audio_hidden_states = audio_hidden_states + v2a_gate * v2a_attn_hidden_states
 
         # 4. Feedforward
         norm_hidden_states = self.norm3(hidden_states) * (1 + scale_mlp) + shift_mlp
@@ -1320,6 +1337,9 @@ def forward(
         audio_num_frames: int | None = None,
         video_coords: torch.Tensor | None = None,
         audio_coords: torch.Tensor | None = None,
+        isolate_modalities: bool = False,
+        spatio_temporal_guidance_blocks: list[int] | None = None,
+        perturbation_mask: torch.Tensor | None = None,
         attention_kwargs: dict[str, Any] | None = None,
         return_dict: bool = True,
     ) -> torch.Tensor:
@@ -1371,6 +1391,17 @@ def forward(
             audio_coords (`torch.Tensor`, *optional*):
                 The audio coordinates to be used when calculating the rotary positional embeddings (RoPE) of shape
                 `(batch_size, 1, num_audio_tokens, 2)`. If not supplied, this will be calculated inside `forward`.
+            isolate_modalities (`bool`, *optional*, defaults to `False`):
+                Whether to isolate each modality by turning off cross-modality (audio-to-video and video-to-audio)
+                cross attention (for all blocks). Use for modality guidance in LTX-2.3.
+            spatio_temporal_guidance_blocks (`list[int]`, *optional*, defaults to `None`):
+                The transformer block indices at which to apply spatio-temporal guidance (STG), which shortcuts the
+                self-attention operations by simply using the values rather than the full scaled dot-product attention
+                (SDPA) operation. If `None` or empty, STG will not be applied to any block.
+            perturbation_mask (`torch.Tensor`, *optional*):
+                Perturbation mask for STG of shape `(batch_size,)` or `(batch_size, 1, 1)`. Should be 0 at batch
+                elements where STG should be applied and 1 elsewhere. If STG is being used but `peturbation_mask` is
+                not supplied, will default to applying STG (perturbing) all batch elements.
             attention_kwargs (`dict[str, Any]`, *optional*):
                 Optional dict of keyword args to be passed to the attention processor.
             return_dict (`bool`, *optional*, defaults to `True`):
@@ -1520,7 +1551,19 @@ def forward(
             audio_encoder_hidden_states = audio_encoder_hidden_states.view(batch_size, -1, audio_hidden_states.size(-1))
 
         # 5. Run transformer blocks
-        for block in self.transformer_blocks:
+        spatio_temporal_guidance_blocks = spatio_temporal_guidance_blocks or []
+        if len(spatio_temporal_guidance_blocks) > 0 and perturbation_mask is None:
+            # If STG is being used and perturbation_mask is not supplied, default to perturbing all batch elements.
+            perturbation_mask = torch.zeros((batch_size,))
+        if perturbation_mask is not None and perturbation_mask.ndim == 1:
+            perturbation_mask = perturbation_mask[:, None, None]  # unsqueeze to 3D to broadcast with hidden_states
+        all_perturbed = torch.all(perturbation_mask == 0) if perturbation_mask is not None else False
+        stg_blocks = set(spatio_temporal_guidance_blocks)
+
+        for block_idx, block in enumerate(self.transformer_blocks):
+            block_perturbation_mask = perturbation_mask if block_idx in stg_blocks else None
+            block_all_perturbed = all_perturbed if block_idx in stg_blocks else False
+
             if torch.is_grad_enabled() and self.gradient_checkpointing:
                 hidden_states, audio_hidden_states = self._gradient_checkpointing_func(
                     block,
@@ -1542,6 +1585,14 @@ def forward(
                     audio_cross_attn_rotary_emb,
                     encoder_attention_mask,
                     audio_encoder_attention_mask,
+                    self_attention_mask,
+                    audio_self_attention_mask,
+                    None,  # a2v_cross_attention_mask
+                    None,  # v2a_cross_attention_mask
+                    not isolate_modalities,  # use_a2v_cross_attention
+                    not isolate_modalities,  # use_v2a_cross_attention
+                    block_perturbation_mask,
+                    block_all_perturbed,
                 )
             else:
                 hidden_states, audio_hidden_states = block(
@@ -1563,6 +1614,14 @@ def forward(
                     ca_audio_rotary_emb=audio_cross_attn_rotary_emb,
                     encoder_attention_mask=encoder_attention_mask,
                     audio_encoder_attention_mask=audio_encoder_attention_mask,
+                    self_attention_mask=self_attention_mask,
+                    audio_self_attention_mask=audio_self_attention_mask,
+                    a2v_cross_attention_mask=None,
+                    v2a_cross_attention_mask=None,
+                    use_a2v_cross_attention=not isolate_modalities,
+                    use_v2a_cross_attention=not isolate_modalities,
+                    perturbation_mask=block_perturbation_mask,
+                    all_perturbed=block_all_perturbed,
                 )
 
         # 6. Output layers (including unpatchification)
diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
index 9791c1e9d80e..5f3a0fd00bb3 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
@@ -155,12 +155,12 @@ def retrieve_timesteps(
     return timesteps, num_inference_steps
 
 
-# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.rescale_noise_cfg
-def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
+def rescale_noise_cfg_delta(noise_cfg, noise_pred_text, guidance_rescale=0.0):
     r"""
     Rescales `noise_cfg` tensor based on `guidance_rescale` to improve image quality and fix overexposure. Based on
     Section 3.4 from [Common Diffusion Noise Schedules and Sample Steps are
-    Flawed](https://huggingface.co/papers/2305.08891).
+    Flawed](https://huggingface.co/papers/2305.08891). Returns a delta value with respect to noise_pred_text, which is
+    useful when there are multiple guidance terms.
 
     Args:
         noise_cfg (`torch.Tensor`):
@@ -179,7 +179,9 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
     noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
     # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images
     noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg
-    return noise_cfg
+    # Return as a delta with respect to noise_pred_text
+    rescale_delta = noise_cfg - noise_pred_text
+    return rescale_delta
 
 
 class LTX2Pipeline(DiffusionPipeline, FromSingleFileMixin, LTX2LoraLoaderMixin):
@@ -428,6 +430,9 @@ def check_inputs(
         negative_prompt_embeds=None,
         prompt_attention_mask=None,
         negative_prompt_attention_mask=None,
+        spatio_temporal_guidance_blocks=None,
+        stg_scale=None,
+        audio_stg_scale=None,
     ):
         if height % 32 != 0 or width % 32 != 0:
             raise ValueError(f"`height` and `width` have to be divisible by 32 but are {height} and {width}.")
@@ -471,6 +476,12 @@ def check_inputs(
                     f" {negative_prompt_attention_mask.shape}."
                 )
 
+        if ((stg_scale > 0.0) or (audio_stg_scale > 0.0)) and not spatio_temporal_guidance_blocks:
+            raise ValueError(
+                "Spatio-Temporal Guidance (STG) is specified but no STG blocks are supplied. Please supply a list of"
+                "block indices at which to apply STG in `spatio_temporal_guidance_blocks`"
+            )
+
     @staticmethod
     def _pack_latents(latents: torch.Tensor, patch_size: int = 1, patch_size_t: int = 1) -> torch.Tensor:
         # Unpacked latents of shape are [B, C, F, H, W] are patched into tokens of shape [B, C, F // p_t, p_t, H // p, p, W // p, p].
@@ -681,9 +692,41 @@ def guidance_scale(self):
     def guidance_rescale(self):
         return self._guidance_rescale
 
+    @property
+    def stg_scale(self):
+        return self._stg_scale
+
+    @property
+    def modality_scale(self):
+        return self._modality_scale
+
+    @property
+    def audio_guidance_scale(self):
+        return self._audio_guidance_scale
+
+    @property
+    def audio_guidance_rescale(self):
+        return self._audio_guidance_rescale
+
+    @property
+    def audio_stg_scale(self):
+        return self._audio_stg_scale
+
+    @property
+    def audio_modality_scale(self):
+        return self._audio_modality_scale
+
     @property
     def do_classifier_free_guidance(self):
-        return self._guidance_scale > 1.0
+        return (self._guidance_scale > 1.0) or (self._audio_guidance_scale > 1.0)
+
+    @property
+    def do_spatio_temporal_guidance(self):
+        return (self._stg_scale > 0.0) or (self._audio_stg_scale > 0.0)
+
+    @property
+    def do_modality_isolation_guidance(self):
+        return (self._modality_scale > 1.0) or (self._audio_modality_scale > 1.0)
 
     @property
     def num_timesteps(self):
@@ -715,7 +758,14 @@ def __call__(
         sigmas: list[float] | None = None,
         timesteps: list[int] = None,
         guidance_scale: float = 4.0,
+        stg_scale: float = 0.0,
+        modality_scale: float = 1.0,
         guidance_rescale: float = 0.0,
+        audio_guidance_scale: float = 4.0,
+        audio_stg_scale: float = 0.0,
+        audio_modality_scale: float = 1.0,
+        audio_guidance_rescale: float = 0.0,
+        spatio_temporal_guidance_blocks: list[int] | None = None,
         noise_scale: float = 0.0,
         num_videos_per_prompt: int = 1,
         generator: torch.Generator | list[torch.Generator] | None = None,
@@ -765,13 +815,44 @@ def __call__(
                 Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2.
                 of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
                 `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
-                the text `prompt`, usually at the expense of lower image quality.
+                the text `prompt`, usually at the expense of lower image quality. Used for the video modality (there is
+                a separate value `audio_guidance_scale` for the audio modality).
+            stg_scale (`float`, *optional*, defaults to `0.0`):
+                Video guidance scale for Spatio-Temporal Guidance (STG), proposed in [Spatiotemporal Skip Guidance for
+                Enhanced Video Diffusion Sampling](https://arxiv.org/abs/2411.18664). STG uses a CFG-like estimate
+                where we move the sample away from a weak sample from a perturbed version of the denoising model.
+                Enabling STG will result in an additional denoising model forward pass; the default value of `0.0`
+                means that STG is disabled.
+            modality_scale (`float`, *optional*, defaults to `1.0`):
+                Video guidance scale for LTX-2.X modality isolation guidance, where we move the sample away from a
+                weaker sample generated by the denoising model withy cross-modality (audio-to-video and video-to-audio)
+                cross attention disabled using a CFG-like estimate. Enabling modality guidance will result in an
+                additional denoising model forward pass; the default value of `1.0` means that modality guidance is
+                disabled.
             guidance_rescale (`float`, *optional*, defaults to 0.0):
                 Guidance rescale factor proposed by [Common Diffusion Noise Schedules and Sample Steps are
                 Flawed](https://huggingface.co/papers/2305.08891) `guidance_scale` is defined as `φ` in equation 16. of
                 [Common Diffusion Noise Schedules and Sample Steps are
                 Flawed](https://huggingface.co/papers/2305.08891). Guidance rescale factor should fix overexposure when
-                using zero terminal SNR.
+                using zero terminal SNR. Used for the video modality.
+            audio_guidance_scale (`float`, *optional* defaults to `4.0`):
+                Audio guidance scale for CFG with respect to the negative prompt. The CFG update rule is the same for
+                video and audio, but they can use different values for the guidance scale. The LTX-2.X authors suggest
+                that the `audio_guidance_scale` should be higher relative to the video `guidance_scale` (e.g. for
+                LTX-2.3 they suggest 3.0 for video and 7.0 for audio).
+            audio_stg_scale (`float`, *optional*, defaults to `0.0`):
+                Audio guidance scale for STG. As with CFG, the STG update rule is otherwise the same for video and
+                audio. For LTX-2.3, a value of 1.0 is suggested for both video and audio.
+            audio_modality_scale (`float`, *optional*, defaults to `1.0`):
+                Audio guidance scale for LTX-2.X modality isolation guidance. As with CFG, the modality guidance rule
+                is otherwise the same for video and audio. For LTX-2.3, a value of 3.0 is suggested for both video and
+                audio.
+            audio_guidance_rescale (`float`, *optional*, defaults to `0.0`):
+                A separate guidance rescale factor for the audio modality.
+            spatio_temporal_guidance_blocks (`list[int]`, *optional*, defaults to `None`):
+                The zero-indexed transformer block indices at which to apply STG. Must be supplied if STG is used
+                (`stg_scale` or `audio_stg_scale` is greater than `0`). A value of `[29]` is recommended for LTX-2.0
+                and `[28]` is recommended for LTX-2.3.
             noise_scale (`float`, *optional*, defaults to `0.0`):
                 The interpolation factor between random noise and denoised latents at each timestep. Applying noise to
                 the `latents` and `audio_latents` before continue denoising.
@@ -844,10 +925,21 @@ def __call__(
             negative_prompt_embeds=negative_prompt_embeds,
             prompt_attention_mask=prompt_attention_mask,
             negative_prompt_attention_mask=negative_prompt_attention_mask,
+            spatio_temporal_guidance_blocks=spatio_temporal_guidance_blocks,
+            stg_scale=stg_scale,
+            audio_stg_scale=audio_stg_scale,
         )
 
+        # Per-modality guidance scales (video, audio)
         self._guidance_scale = guidance_scale
+        self._stg_scale = stg_scale
+        self._modality_scale = modality_scale
         self._guidance_rescale = guidance_rescale
+        self._audio_guidance_scale = audio_guidance_scale
+        self._audio_stg_scale = audio_stg_scale
+        self._audio_modality_scale = audio_modality_scale
+        self._audio_guidance_rescale = audio_guidance_rescale
+
         self._attention_kwargs = attention_kwargs
         self._interrupt = False
         self._current_timestep = None
@@ -995,11 +1087,6 @@ def __call__(
         self._num_timesteps = len(timesteps)
 
         # 6. Prepare micro-conditions
-        rope_interpolation_scale = (
-            self.vae_temporal_compression_ratio / frame_rate,
-            self.vae_spatial_compression_ratio,
-            self.vae_spatial_compression_ratio,
-        )
         # Pre-compute video and audio positional ids as they will be the same at each step of the denoising loop
         video_coords = self.transformer.rope.prepare_video_coords(
             latents.shape[0], latent_num_frames, latent_height, latent_width, latents.device, fps=frame_rate
@@ -1049,7 +1136,9 @@ def __call__(
                         audio_num_frames=audio_num_frames,
                         video_coords=video_coords,
                         audio_coords=audio_coords,
-                        # rope_interpolation_scale=rope_interpolation_scale,
+                        isolate_modalities=False,
+                        spatio_temporal_guidance_blocks=None,
+                        perturbation_mask=None,
                         attention_kwargs=attention_kwargs,
                         return_dict=False,
                     )
@@ -1057,24 +1146,126 @@ def __call__(
                 noise_pred_audio = noise_pred_audio.float()
 
                 if self.do_classifier_free_guidance:
-                    noise_pred_video_uncond, noise_pred_video_text = noise_pred_video.chunk(2)
-                    noise_pred_video = noise_pred_video_uncond + self.guidance_scale * (
-                        noise_pred_video_text - noise_pred_video_uncond
-                    )
+                    noise_pred_video_uncond_text, noise_pred_video = noise_pred_video.chunk(2)
+                    # Use delta formulation as it works more nicely with multiple guidance terms
+                    video_cfg_delta = (self.guidance_scale - 1) * (noise_pred_video - noise_pred_video_uncond_text)
 
-                    noise_pred_audio_uncond, noise_pred_audio_text = noise_pred_audio.chunk(2)
-                    noise_pred_audio = noise_pred_audio_uncond + self.guidance_scale * (
-                        noise_pred_audio_text - noise_pred_audio_uncond
-                    )
+                    noise_pred_audio_uncond_text, noise_pred_audio = noise_pred_audio.chunk(2)
+                    audio_cfg_delta = (self.audio_guidance_scale - 1) * (noise_pred_audio - noise_pred_audio_uncond_text)
 
                     if self.guidance_rescale > 0:
                         # Based on 3.4. in https://huggingface.co/papers/2305.08891
-                        noise_pred_video = rescale_noise_cfg(
-                            noise_pred_video, noise_pred_video_text, guidance_rescale=self.guidance_rescale
+                        video_cfg_delta = rescale_noise_cfg_delta(
+                            noise_cfg=noise_pred_video + video_cfg_delta,
+                            noise_pred_text=noise_pred_video,
+                            guidance_rescale=self.guidance_rescale,
+                        )
+                        audio_cfg_delta = rescale_noise_cfg_delta(
+                            noise_cfg=noise_pred_audio + audio_cfg_delta,
+                            noise_pred_text=noise_pred_audio,
+                            guidance_rescale=self.audio_guidance_rescale,
                         )
-                        noise_pred_audio = rescale_noise_cfg(
-                            noise_pred_audio, noise_pred_audio_text, guidance_rescale=self.guidance_rescale
+                    
+                    # Get positive values from merged CFG inputs in case we need to do other DiT forward passes
+                    if (self.do_spatio_temporal_guidance or self.do_modality_isolation_guidance):
+                        if i == 0:
+                            # Only split values that remain constant throughout the loop once
+                            video_prompt_embeds = connector_prompt_embeds.chunk(2, dim=0)[1]
+                            audio_prompt_embeds = connector_audio_prompt_embeds.chunk(2, dim=0)[1]
+                            prompt_attn_mask = connector_attention_mask.chunk(2, dim=0)[1]
+
+                            video_pos_ids = video_coords.chunk(2, dim=0)[0]
+                            audio_pos_ids = audio_coords.chunk(2, dim=0)[0]
+
+                        # Split values that vary each denoising loop iteration
+                        timestep = timestep.chunk(2, dim=0)[0]
+                else:
+                    video_cfg_delta = audio_cfg_delta = 0
+
+                    video_prompt_embeds = connector_prompt_embeds
+                    audio_prompt_embeds = connector_audio_prompt_embeds
+                    prompt_attn_mask = connector_attention_mask
+
+                    video_pos_ids = video_coords
+                    audio_pos_ids = audio_coords
+                
+                if self.do_spatio_temporal_guidance:
+                    with self.transformer.cache_context("uncond_stg"):
+                        noise_pred_video_uncond_stg, noise_pred_audio_uncond_stg = self.transformer(
+                            hidden_states=latents.to(dtype=prompt_embeds.dtype),
+                            audio_hidden_states=audio_latents.to(dtype=prompt_embeds.dtype),
+                            encoder_hidden_states=video_prompt_embeds,
+                            audio_encoder_hidden_states=audio_prompt_embeds,
+                            timestep=timestep,
+                            sigma=timestep,  # Used by LTX-2.3
+                            encoder_attention_mask=prompt_attn_mask,
+                            audio_encoder_attention_mask=prompt_attn_mask,
+                            self_attention_mask=None,
+                            audio_self_attention_mask=None,
+                            num_frames=latent_num_frames,
+                            height=latent_height,
+                            width=latent_width,
+                            fps=frame_rate,
+                            audio_num_frames=audio_num_frames,
+                            video_coords=video_pos_ids,
+                            audio_coords=audio_pos_ids,
+                            isolate_modalities=False,
+                            # Use STG at given blocks to perturb model
+                            spatio_temporal_guidance_blocks=spatio_temporal_guidance_blocks,
+                            perturbation_mask=None,
+                            attention_kwargs=attention_kwargs,
+                            return_dict=False,
                         )
+                    noise_pred_video_uncond_stg = noise_pred_video_uncond_stg.float()
+                    noise_pred_audio_uncond_stg = noise_pred_audio_uncond_stg.float()
+
+                    video_stg_delta = self.stg_scale * (noise_pred_video - noise_pred_video_uncond_stg)
+                    audio_stg_delta = self.audio_stg_scale * (noise_pred_audio - noise_pred_audio_uncond_stg)
+                else:
+                    video_stg_delta = audio_stg_delta = 0
+                
+                if self.do_modality_isolation_guidance:
+                    with self.transformer.cache_context("uncond_modality"):
+                        noise_pred_video_uncond_modality, noise_pred_audio_uncond_modality = self.transformer(
+                            hidden_states=latents.to(dtype=prompt_embeds.dtype),
+                            audio_hidden_states=audio_latents.to(dtype=prompt_embeds.dtype),
+                            encoder_hidden_states=video_prompt_embeds,
+                            audio_encoder_hidden_states=audio_prompt_embeds,
+                            timestep=timestep,
+                            sigma=timestep,  # Used by LTX-2.3
+                            encoder_attention_mask=prompt_attn_mask,
+                            audio_encoder_attention_mask=prompt_attn_mask,
+                            self_attention_mask=None,
+                            audio_self_attention_mask=None,
+                            num_frames=latent_num_frames,
+                            height=latent_height,
+                            width=latent_width,
+                            fps=frame_rate,
+                            audio_num_frames=audio_num_frames,
+                            video_coords=video_pos_ids,
+                            audio_coords=audio_pos_ids,
+                            # Turn off A2V and V2A cross attn to isolate video and audio modalities
+                            isolate_modalities=True,
+                            spatio_temporal_guidance_blocks=None,
+                            perturbation_mask=None,
+                            attention_kwargs=attention_kwargs,
+                            return_dict=False,
+                        )
+                    noise_pred_video_uncond_modality = noise_pred_video_uncond_modality.float()
+                    noise_pred_audio_uncond_modality = noise_pred_audio_uncond_modality.float()
+
+                    video_modality_delta = (
+                        (self.modality_scale - 1) * (noise_pred_video - noise_pred_video_uncond_modality)
+                    )
+                    audio_modality_delta = (
+                        (self.audio_modality_scale - 1) * (noise_pred_audio - noise_pred_audio_uncond_modality)
+                    )
+                else:
+                    video_modality_delta = audio_modality_delta = 0
+
+                # Now apply all guidance terms
+                noise_pred_video = noise_pred_video + video_cfg_delta + video_stg_delta + video_modality_delta
+                noise_pred_audio = noise_pred_audio + audio_cfg_delta + audio_stg_delta + audio_modality_delta
 
                 # compute the previous noisy sample x_t -> x_t-1
                 latents = self.scheduler.step(noise_pred_video, t, latents, return_dict=False)[0]
diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py
index 058ed27146c5..ed2f87ccdcac 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py
@@ -205,12 +205,13 @@ def retrieve_timesteps(
     return timesteps, num_inference_steps
 
 
-# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.rescale_noise_cfg
-def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
+# Copied from diffusers.pipelines.ltx2.pipeline_ltx2.rescale_noise_cfg_delta
+def rescale_noise_cfg_delta(noise_cfg, noise_pred_text, guidance_rescale=0.0):
     r"""
     Rescales `noise_cfg` tensor based on `guidance_rescale` to improve image quality and fix overexposure. Based on
     Section 3.4 from [Common Diffusion Noise Schedules and Sample Steps are
-    Flawed](https://huggingface.co/papers/2305.08891).
+    Flawed](https://huggingface.co/papers/2305.08891). Returns a delta value with respect to noise_pred_text, which is
+    useful when there are multiple guidance terms.
 
     Args:
         noise_cfg (`torch.Tensor`):
@@ -229,7 +230,9 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
     noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
     # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images
     noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg
-    return noise_cfg
+    # Return as a delta with respect to noise_pred_text
+    rescale_delta = noise_cfg - noise_pred_text
+    return rescale_delta
 
 
 class LTX2ConditionPipeline(DiffusionPipeline, FromSingleFileMixin, LTX2LoraLoaderMixin):
@@ -464,6 +467,9 @@ def check_inputs(
         negative_prompt_attention_mask=None,
         latents=None,
         audio_latents=None,
+        spatio_temporal_guidance_blocks=None,
+        stg_scale=None,
+        audio_stg_scale=None,
     ):
         if height % 32 != 0 or width % 32 != 0:
             raise ValueError(f"`height` and `width` have to be divisible by 32 but are {height} and {width}.")
@@ -520,6 +526,12 @@ def check_inputs(
                 f" using the `_unpack_audio_latents` method)."
             )
 
+        if ((stg_scale > 0.0) or (audio_stg_scale > 0.0)) and not spatio_temporal_guidance_blocks:
+            raise ValueError(
+                "Spatio-Temporal Guidance (STG) is specified but no STG blocks are supplied. Please supply a list of"
+                "block indices at which to apply STG in `spatio_temporal_guidance_blocks`"
+            )
+
     @staticmethod
     # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._pack_latents
     def _pack_latents(latents: torch.Tensor, patch_size: int = 1, patch_size_t: int = 1) -> torch.Tensor:
@@ -915,9 +927,41 @@ def guidance_scale(self):
     def guidance_rescale(self):
         return self._guidance_rescale
 
+    @property
+    def stg_scale(self):
+        return self._stg_scale
+
+    @property
+    def modality_scale(self):
+        return self._modality_scale
+
+    @property
+    def audio_guidance_scale(self):
+        return self._audio_guidance_scale
+
+    @property
+    def audio_guidance_rescale(self):
+        return self._audio_guidance_rescale
+
+    @property
+    def audio_stg_scale(self):
+        return self._audio_stg_scale
+
+    @property
+    def audio_modality_scale(self):
+        return self._audio_modality_scale
+
     @property
     def do_classifier_free_guidance(self):
-        return self._guidance_scale > 1.0
+        return (self._guidance_scale > 1.0) or (self._audio_guidance_scale > 1.0)
+
+    @property
+    def do_spatio_temporal_guidance(self):
+        return (self._stg_scale > 0.0) or (self._audio_stg_scale > 0.0)
+
+    @property
+    def do_modality_isolation_guidance(self):
+        return (self._modality_scale > 1.0) or (self._audio_modality_scale > 1.0)
 
     @property
     def num_timesteps(self):
@@ -950,7 +994,14 @@ def __call__(
         sigmas: list[float] | None = None,
         timesteps: list[float] | None = None,
         guidance_scale: float = 4.0,
+        stg_scale: float = 0.0,
+        modality_scale: float = 1.0,
         guidance_rescale: float = 0.0,
+        audio_guidance_scale: float = 4.0,
+        audio_stg_scale: float = 0.0,
+        audio_modality_scale: float = 1.0,
+        audio_guidance_rescale: float = 0.0,
+        spatio_temporal_guidance_blocks: list[int] | None = None,
         noise_scale: float | None = None,
         num_videos_per_prompt: int | None = 1,
         generator: torch.Generator | list[torch.Generator] | None = None,
@@ -1002,13 +1053,44 @@ def __call__(
                 Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2.
                 of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
                 `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
-                the text `prompt`, usually at the expense of lower image quality.
+                the text `prompt`, usually at the expense of lower image quality. Used for the video modality (there is
+                a separate value `audio_guidance_scale` for the audio modality).
+            stg_scale (`float`, *optional*, defaults to `0.0`):
+                Video guidance scale for Spatio-Temporal Guidance (STG), proposed in [Spatiotemporal Skip Guidance for
+                Enhanced Video Diffusion Sampling](https://arxiv.org/abs/2411.18664). STG uses a CFG-like estimate
+                where we move the sample away from a weak sample from a perturbed version of the denoising model.
+                Enabling STG will result in an additional denoising model forward pass; the default value of `0.0`
+                means that STG is disabled.
+            modality_scale (`float`, *optional*, defaults to `1.0`):
+                Video guidance scale for LTX-2.X modality isolation guidance, where we move the sample away from a
+                weaker sample generated by the denoising model withy cross-modality (audio-to-video and video-to-audio)
+                cross attention disabled using a CFG-like estimate. Enabling modality guidance will result in an
+                additional denoising model forward pass; the default value of `1.0` means that modality guidance is
+                disabled.
             guidance_rescale (`float`, *optional*, defaults to 0.0):
                 Guidance rescale factor proposed by [Common Diffusion Noise Schedules and Sample Steps are
                 Flawed](https://huggingface.co/papers/2305.08891) `guidance_scale` is defined as `φ` in equation 16. of
                 [Common Diffusion Noise Schedules and Sample Steps are
                 Flawed](https://huggingface.co/papers/2305.08891). Guidance rescale factor should fix overexposure when
-                using zero terminal SNR.
+                using zero terminal SNR. Used for the video modality.
+            audio_guidance_scale (`float`, *optional* defaults to `4.0`):
+                Audio guidance scale for CFG with respect to the negative prompt. The CFG update rule is the same for
+                video and audio, but they can use different values for the guidance scale. The LTX-2.X authors suggest
+                that the `audio_guidance_scale` should be higher relative to the video `guidance_scale` (e.g. for
+                LTX-2.3 they suggest 3.0 for video and 7.0 for audio).
+            audio_stg_scale (`float`, *optional*, defaults to `0.0`):
+                Audio guidance scale for STG. As with CFG, the STG update rule is otherwise the same for video and
+                audio. For LTX-2.3, a value of 1.0 is suggested for both video and audio.
+            audio_modality_scale (`float`, *optional*, defaults to `1.0`):
+                Audio guidance scale for LTX-2.X modality isolation guidance. As with CFG, the modality guidance rule
+                is otherwise the same for video and audio. For LTX-2.3, a value of 3.0 is suggested for both video and
+                audio.
+            audio_guidance_rescale (`float`, *optional*, defaults to `0.0`):
+                A separate guidance rescale factor for the audio modality.
+            spatio_temporal_guidance_blocks (`list[int]`, *optional*, defaults to `None`):
+                The zero-indexed transformer block indices at which to apply STG. Must be supplied if STG is used
+                (`stg_scale` or `audio_stg_scale` is greater than `0`). A value of `[29]` is recommended for LTX-2.0
+                and `[28]` is recommended for LTX-2.3.
             noise_scale (`float`, *optional*, defaults to `None`):
                 The interpolation factor between random noise and denoised latents at each timestep. Applying noise to
                 the `latents` and `audio_latents` before continue denoising. If not set, will be inferred from the
@@ -1084,10 +1166,21 @@ def __call__(
             negative_prompt_attention_mask=negative_prompt_attention_mask,
             latents=latents,
             audio_latents=audio_latents,
+            spatio_temporal_guidance_blocks=spatio_temporal_guidance_blocks,
+            stg_scale=stg_scale,
+            audio_stg_scale=audio_stg_scale,
         )
 
+        # Per-modality guidance scales (video, audio)
         self._guidance_scale = guidance_scale
+        self._stg_scale = stg_scale
+        self._modality_scale = modality_scale
         self._guidance_rescale = guidance_rescale
+        self._audio_guidance_scale = audio_guidance_scale
+        self._audio_stg_scale = audio_stg_scale
+        self._audio_modality_scale = audio_modality_scale
+        self._audio_guidance_rescale = audio_guidance_rescale
+
         self._attention_kwargs = attention_kwargs
         self._interrupt = False
         self._current_timestep = None
@@ -1226,11 +1319,6 @@ def __call__(
         self._num_timesteps = len(timesteps)
 
         # 6. Prepare micro-conditions
-        rope_interpolation_scale = (
-            self.vae_temporal_compression_ratio / frame_rate,
-            self.vae_spatial_compression_ratio,
-            self.vae_spatial_compression_ratio,
-        )
         # Pre-compute video and audio positional ids as they will be the same at each step of the denoising loop
         video_coords = self.transformer.rope.prepare_video_coords(
             latents.shape[0], latent_num_frames, latent_height, latent_width, latents.device, fps=frame_rate
@@ -1281,7 +1369,9 @@ def __call__(
                         audio_num_frames=audio_num_frames,
                         video_coords=video_coords,
                         audio_coords=audio_coords,
-                        # rope_interpolation_scale=rope_interpolation_scale,
+                        isolate_modalities=False,
+                        spatio_temporal_guidance_blocks=None,
+                        perturbation_mask=None,
                         attention_kwargs=attention_kwargs,
                         return_dict=False,
                     )
@@ -1289,24 +1379,129 @@ def __call__(
                 noise_pred_audio = noise_pred_audio.float()
 
                 if self.do_classifier_free_guidance:
-                    noise_pred_video_uncond, noise_pred_video_text = noise_pred_video.chunk(2)
-                    noise_pred_video = noise_pred_video_uncond + self.guidance_scale * (
-                        noise_pred_video_text - noise_pred_video_uncond
-                    )
+                    noise_pred_video_uncond_text, noise_pred_video = noise_pred_video.chunk(2)
+                    # Use delta formulation as it works more nicely with multiple guidance terms
+                    video_cfg_delta = (self.guidance_scale - 1) * (noise_pred_video - noise_pred_video_uncond_text)
 
-                    noise_pred_audio_uncond, noise_pred_audio_text = noise_pred_audio.chunk(2)
-                    noise_pred_audio = noise_pred_audio_uncond + self.guidance_scale * (
-                        noise_pred_audio_text - noise_pred_audio_uncond
-                    )
+                    noise_pred_audio_uncond_text, noise_pred_audio = noise_pred_audio.chunk(2)
+                    audio_cfg_delta = (self.audio_guidance_scale - 1) * (noise_pred_audio - noise_pred_audio_uncond_text)
 
                     if self.guidance_rescale > 0:
                         # Based on 3.4. in https://huggingface.co/papers/2305.08891
-                        noise_pred_video = rescale_noise_cfg(
-                            noise_pred_video, noise_pred_video_text, guidance_rescale=self.guidance_rescale
+                        video_cfg_delta = rescale_noise_cfg_delta(
+                            noise_cfg=noise_pred_video + video_cfg_delta,
+                            noise_pred_text=noise_pred_video,
+                            guidance_rescale=self.guidance_rescale,
+                        )
+                        audio_cfg_delta = rescale_noise_cfg_delta(
+                            noise_cfg=noise_pred_audio + audio_cfg_delta,
+                            noise_pred_text=noise_pred_audio,
+                            guidance_rescale=self.audio_guidance_rescale,
                         )
-                        noise_pred_audio = rescale_noise_cfg(
-                            noise_pred_audio, noise_pred_audio_text, guidance_rescale=self.guidance_rescale
+                    
+                    # Get positive values from merged CFG inputs in case we need to do other DiT forward passes
+                    if (self.do_spatio_temporal_guidance or self.do_modality_isolation_guidance):
+                        if i == 0:
+                            # Only split values that remain constant throughout the loop once
+                            video_prompt_embeds = connector_prompt_embeds.chunk(2, dim=0)[1]
+                            audio_prompt_embeds = connector_audio_prompt_embeds.chunk(2, dim=0)[1]
+                            prompt_attn_mask = connector_attention_mask.chunk(2, dim=0)[1]
+
+                            video_pos_ids = video_coords.chunk(2, dim=0)[0]
+                            audio_pos_ids = audio_coords.chunk(2, dim=0)[0]
+
+                        # Split values that vary each denoising loop iteration
+                        timestep = timestep.chunk(2, dim=0)[0]
+                        video_timestep = video_timestep.chunk(2, dim=0)[0]
+                else:
+                    video_cfg_delta = audio_cfg_delta = 0
+
+                    video_prompt_embeds = connector_prompt_embeds
+                    audio_prompt_embeds = connector_audio_prompt_embeds
+                    prompt_attn_mask = connector_attention_mask
+
+                    video_pos_ids = video_coords
+                    audio_pos_ids = audio_coords
+                
+                if self.do_spatio_temporal_guidance:
+                    with self.transformer.cache_context("uncond_stg"):
+                        noise_pred_video_uncond_stg, noise_pred_audio_uncond_stg = self.transformer(
+                            hidden_states=latents.to(dtype=prompt_embeds.dtype),
+                            audio_hidden_states=audio_latents.to(dtype=prompt_embeds.dtype),
+                            encoder_hidden_states=video_prompt_embeds,
+                            audio_encoder_hidden_states=audio_prompt_embeds,
+                            timestep=video_timestep,
+                            audio_timestep=timestep,
+                            sigma=timestep,  # Used by LTX-2.3
+                            encoder_attention_mask=prompt_attn_mask,
+                            audio_encoder_attention_mask=prompt_attn_mask,
+                            self_attention_mask=None,
+                            audio_self_attention_mask=None,
+                            num_frames=latent_num_frames,
+                            height=latent_height,
+                            width=latent_width,
+                            fps=frame_rate,
+                            audio_num_frames=audio_num_frames,
+                            video_coords=video_pos_ids,
+                            audio_coords=audio_pos_ids,
+                            isolate_modalities=False,
+                            # Use STG at given blocks to perturb model
+                            spatio_temporal_guidance_blocks=spatio_temporal_guidance_blocks,
+                            perturbation_mask=None,
+                            attention_kwargs=attention_kwargs,
+                            return_dict=False,
                         )
+                    noise_pred_video_uncond_stg = noise_pred_video_uncond_stg.float()
+                    noise_pred_audio_uncond_stg = noise_pred_audio_uncond_stg.float()
+
+                    video_stg_delta = self.stg_scale * (noise_pred_video - noise_pred_video_uncond_stg)
+                    audio_stg_delta = self.audio_stg_scale * (noise_pred_audio - noise_pred_audio_uncond_stg)
+                else:
+                    video_stg_delta = audio_stg_delta = 0
+                
+                if self.do_modality_isolation_guidance:
+                    with self.transformer.cache_context("uncond_modality"):
+                        noise_pred_video_uncond_modality, noise_pred_audio_uncond_modality = self.transformer(
+                            hidden_states=latents.to(dtype=prompt_embeds.dtype),
+                            audio_hidden_states=audio_latents.to(dtype=prompt_embeds.dtype),
+                            encoder_hidden_states=video_prompt_embeds,
+                            audio_encoder_hidden_states=audio_prompt_embeds,
+                            timestep=video_timestep,
+                            audio_timestep=timestep,
+                            sigma=timestep,  # Used by LTX-2.3
+                            encoder_attention_mask=prompt_attn_mask,
+                            audio_encoder_attention_mask=prompt_attn_mask,
+                            self_attention_mask=None,
+                            audio_self_attention_mask=None,
+                            num_frames=latent_num_frames,
+                            height=latent_height,
+                            width=latent_width,
+                            fps=frame_rate,
+                            audio_num_frames=audio_num_frames,
+                            video_coords=video_pos_ids,
+                            audio_coords=audio_pos_ids,
+                            # Turn off A2V and V2A cross attn to isolate video and audio modalities
+                            isolate_modalities=True,
+                            spatio_temporal_guidance_blocks=None,
+                            perturbation_mask=None,
+                            attention_kwargs=attention_kwargs,
+                            return_dict=False,
+                        )
+                    noise_pred_video_uncond_modality = noise_pred_video_uncond_modality.float()
+                    noise_pred_audio_uncond_modality = noise_pred_audio_uncond_modality.float()
+
+                    video_modality_delta = (
+                        (self.modality_scale - 1) * (noise_pred_video - noise_pred_video_uncond_modality)
+                    )
+                    audio_modality_delta = (
+                        (self.audio_modality_scale - 1) * (noise_pred_audio - noise_pred_audio_uncond_modality)
+                    )
+                else:
+                    video_modality_delta = audio_modality_delta = 0
+
+                # Now apply all guidance terms
+                noise_pred_video = noise_pred_video + video_cfg_delta + video_stg_delta + video_modality_delta
+                noise_pred_audio = noise_pred_audio + audio_cfg_delta + audio_stg_delta + audio_modality_delta
 
                 # NOTE: use only the first chunk of conditioning mask in case it is duplicated for CFG
                 bsz = noise_pred_video.size(0)
diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
index b5371260d3bb..746aab42b668 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
@@ -175,12 +175,13 @@ def retrieve_timesteps(
     return timesteps, num_inference_steps
 
 
-# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.rescale_noise_cfg
-def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
+# Copied from diffusers.pipelines.ltx2.pipeline_ltx2.rescale_noise_cfg_delta
+def rescale_noise_cfg_delta(noise_cfg, noise_pred_text, guidance_rescale=0.0):
     r"""
     Rescales `noise_cfg` tensor based on `guidance_rescale` to improve image quality and fix overexposure. Based on
     Section 3.4 from [Common Diffusion Noise Schedules and Sample Steps are
-    Flawed](https://huggingface.co/papers/2305.08891).
+    Flawed](https://huggingface.co/papers/2305.08891). Returns a delta value with respect to noise_pred_text, which is
+    useful when there are multiple guidance terms.
 
     Args:
         noise_cfg (`torch.Tensor`):
@@ -199,7 +200,9 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
     noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
     # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images
     noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg
-    return noise_cfg
+    # Return as a delta with respect to noise_pred_text
+    rescale_delta = noise_cfg - noise_pred_text
+    return rescale_delta
 
 
 class LTX2ImageToVideoPipeline(DiffusionPipeline, FromSingleFileMixin, LTX2LoraLoaderMixin):
@@ -434,6 +437,9 @@ def check_inputs(
         negative_prompt_embeds=None,
         prompt_attention_mask=None,
         negative_prompt_attention_mask=None,
+        spatio_temporal_guidance_blocks=None,
+        stg_scale=None,
+        audio_stg_scale=None,
     ):
         if height % 32 != 0 or width % 32 != 0:
             raise ValueError(f"`height` and `width` have to be divisible by 32 but are {height} and {width}.")
@@ -477,6 +483,12 @@ def check_inputs(
                     f" {negative_prompt_attention_mask.shape}."
                 )
 
+        if ((stg_scale > 0.0) or (audio_stg_scale > 0.0)) and not spatio_temporal_guidance_blocks:
+            raise ValueError(
+                "Spatio-Temporal Guidance (STG) is specified but no STG blocks are supplied. Please supply a list of"
+                "block indices at which to apply STG in `spatio_temporal_guidance_blocks`"
+            )
+
     @staticmethod
     # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._pack_latents
     def _pack_latents(latents: torch.Tensor, patch_size: int = 1, patch_size_t: int = 1) -> torch.Tensor:
@@ -734,9 +746,41 @@ def guidance_scale(self):
     def guidance_rescale(self):
         return self._guidance_rescale
 
+    @property
+    def stg_scale(self):
+        return self._stg_scale
+
+    @property
+    def modality_scale(self):
+        return self._modality_scale
+
+    @property
+    def audio_guidance_scale(self):
+        return self._audio_guidance_scale
+
+    @property
+    def audio_guidance_rescale(self):
+        return self._audio_guidance_rescale
+
+    @property
+    def audio_stg_scale(self):
+        return self._audio_stg_scale
+
+    @property
+    def audio_modality_scale(self):
+        return self._audio_modality_scale
+
     @property
     def do_classifier_free_guidance(self):
-        return self._guidance_scale > 1.0
+        return (self._guidance_scale > 1.0) or (self._audio_guidance_scale > 1.0)
+
+    @property
+    def do_spatio_temporal_guidance(self):
+        return (self._stg_scale > 0.0) or (self._audio_stg_scale > 0.0)
+
+    @property
+    def do_modality_isolation_guidance(self):
+        return (self._modality_scale > 1.0) or (self._audio_modality_scale > 1.0)
 
     @property
     def num_timesteps(self):
@@ -769,7 +813,14 @@ def __call__(
         sigmas: list[float] | None = None,
         timesteps: list[int] | None = None,
         guidance_scale: float = 4.0,
+        stg_scale: float = 0.0,
+        modality_scale: float = 1.0,
         guidance_rescale: float = 0.0,
+        audio_guidance_scale: float = 4.0,
+        audio_stg_scale: float = 0.0,
+        audio_modality_scale: float = 1.0,
+        audio_guidance_rescale: float = 0.0,
+        spatio_temporal_guidance_blocks: list[int] | None = None,
         noise_scale: float = 0.0,
         num_videos_per_prompt: int = 1,
         generator: torch.Generator | list[torch.Generator] | None = None,
@@ -821,13 +872,44 @@ def __call__(
                 Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2.
                 of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
                 `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
-                the text `prompt`, usually at the expense of lower image quality.
+                the text `prompt`, usually at the expense of lower image quality. Used for the video modality (there is
+                a separate value `audio_guidance_scale` for the audio modality).
+            stg_scale (`float`, *optional*, defaults to `0.0`):
+                Video guidance scale for Spatio-Temporal Guidance (STG), proposed in [Spatiotemporal Skip Guidance for
+                Enhanced Video Diffusion Sampling](https://arxiv.org/abs/2411.18664). STG uses a CFG-like estimate
+                where we move the sample away from a weak sample from a perturbed version of the denoising model.
+                Enabling STG will result in an additional denoising model forward pass; the default value of `0.0`
+                means that STG is disabled.
+            modality_scale (`float`, *optional*, defaults to `1.0`):
+                Video guidance scale for LTX-2.X modality isolation guidance, where we move the sample away from a
+                weaker sample generated by the denoising model withy cross-modality (audio-to-video and video-to-audio)
+                cross attention disabled using a CFG-like estimate. Enabling modality guidance will result in an
+                additional denoising model forward pass; the default value of `1.0` means that modality guidance is
+                disabled.
             guidance_rescale (`float`, *optional*, defaults to 0.0):
                 Guidance rescale factor proposed by [Common Diffusion Noise Schedules and Sample Steps are
                 Flawed](https://huggingface.co/papers/2305.08891) `guidance_scale` is defined as `φ` in equation 16. of
                 [Common Diffusion Noise Schedules and Sample Steps are
                 Flawed](https://huggingface.co/papers/2305.08891). Guidance rescale factor should fix overexposure when
-                using zero terminal SNR.
+                using zero terminal SNR. Used for the video modality.
+            audio_guidance_scale (`float`, *optional* defaults to `4.0`):
+                Audio guidance scale for CFG with respect to the negative prompt. The CFG update rule is the same for
+                video and audio, but they can use different values for the guidance scale. The LTX-2.X authors suggest
+                that the `audio_guidance_scale` should be higher relative to the video `guidance_scale` (e.g. for
+                LTX-2.3 they suggest 3.0 for video and 7.0 for audio).
+            audio_stg_scale (`float`, *optional*, defaults to `0.0`):
+                Audio guidance scale for STG. As with CFG, the STG update rule is otherwise the same for video and
+                audio. For LTX-2.3, a value of 1.0 is suggested for both video and audio.
+            audio_modality_scale (`float`, *optional*, defaults to `1.0`):
+                Audio guidance scale for LTX-2.X modality isolation guidance. As with CFG, the modality guidance rule
+                is otherwise the same for video and audio. For LTX-2.3, a value of 3.0 is suggested for both video and
+                audio.
+            audio_guidance_rescale (`float`, *optional*, defaults to `0.0`):
+                A separate guidance rescale factor for the audio modality.
+            spatio_temporal_guidance_blocks (`list[int]`, *optional*, defaults to `None`):
+                The zero-indexed transformer block indices at which to apply STG. Must be supplied if STG is used
+                (`stg_scale` or `audio_stg_scale` is greater than `0`). A value of `[29]` is recommended for LTX-2.0
+                and `[28]` is recommended for LTX-2.3.
             noise_scale (`float`, *optional*, defaults to `0.0`):
                 The interpolation factor between random noise and denoised latents at each timestep. Applying noise to
                 the `latents` and `audio_latents` before continue denoising.
@@ -900,10 +982,21 @@ def __call__(
             negative_prompt_embeds=negative_prompt_embeds,
             prompt_attention_mask=prompt_attention_mask,
             negative_prompt_attention_mask=negative_prompt_attention_mask,
+            spatio_temporal_guidance_blocks=spatio_temporal_guidance_blocks,
+            stg_scale=stg_scale,
+            audio_stg_scale=audio_stg_scale,
         )
 
+        # Per-modality guidance scales (video, audio)
         self._guidance_scale = guidance_scale
+        self._stg_scale = stg_scale
+        self._modality_scale = modality_scale
         self._guidance_rescale = guidance_rescale
+        self._audio_guidance_scale = audio_guidance_scale
+        self._audio_stg_scale = audio_stg_scale
+        self._audio_modality_scale = audio_modality_scale
+        self._audio_guidance_rescale = audio_guidance_rescale
+
         self._attention_kwargs = attention_kwargs
         self._interrupt = False
         self._current_timestep = None
@@ -1059,11 +1152,6 @@ def __call__(
         self._num_timesteps = len(timesteps)
 
         # 6. Prepare micro-conditions
-        rope_interpolation_scale = (
-            self.vae_temporal_compression_ratio / frame_rate,
-            self.vae_spatial_compression_ratio,
-            self.vae_spatial_compression_ratio,
-        )
         # Pre-compute video and audio positional ids as they will be the same at each step of the denoising loop
         video_coords = self.transformer.rope.prepare_video_coords(
             latents.shape[0], latent_num_frames, latent_height, latent_width, latents.device, fps=frame_rate
@@ -1114,7 +1202,9 @@ def __call__(
                         audio_num_frames=audio_num_frames,
                         video_coords=video_coords,
                         audio_coords=audio_coords,
-                        # rope_interpolation_scale=rope_interpolation_scale,
+                        isolate_modalities=False,
+                        spatio_temporal_guidance_blocks=None,
+                        perturbation_mask=None,
                         attention_kwargs=attention_kwargs,
                         return_dict=False,
                     )
@@ -1122,24 +1212,129 @@ def __call__(
                 noise_pred_audio = noise_pred_audio.float()
 
                 if self.do_classifier_free_guidance:
-                    noise_pred_video_uncond, noise_pred_video_text = noise_pred_video.chunk(2)
-                    noise_pred_video = noise_pred_video_uncond + self.guidance_scale * (
-                        noise_pred_video_text - noise_pred_video_uncond
-                    )
+                    noise_pred_video_uncond_text, noise_pred_video = noise_pred_video.chunk(2)
+                    # Use delta formulation as it works more nicely with multiple guidance terms
+                    video_cfg_delta = (self.guidance_scale - 1) * (noise_pred_video - noise_pred_video_uncond_text)
 
-                    noise_pred_audio_uncond, noise_pred_audio_text = noise_pred_audio.chunk(2)
-                    noise_pred_audio = noise_pred_audio_uncond + self.guidance_scale * (
-                        noise_pred_audio_text - noise_pred_audio_uncond
-                    )
+                    noise_pred_audio_uncond_text, noise_pred_audio = noise_pred_audio.chunk(2)
+                    audio_cfg_delta = (self.audio_guidance_scale - 1) * (noise_pred_audio - noise_pred_audio_uncond_text)
 
                     if self.guidance_rescale > 0:
                         # Based on 3.4. in https://huggingface.co/papers/2305.08891
-                        noise_pred_video = rescale_noise_cfg(
-                            noise_pred_video, noise_pred_video_text, guidance_rescale=self.guidance_rescale
+                        video_cfg_delta = rescale_noise_cfg_delta(
+                            noise_cfg=noise_pred_video + video_cfg_delta,
+                            noise_pred_text=noise_pred_video,
+                            guidance_rescale=self.guidance_rescale,
+                        )
+                        audio_cfg_delta = rescale_noise_cfg_delta(
+                            noise_cfg=noise_pred_audio + audio_cfg_delta,
+                            noise_pred_text=noise_pred_audio,
+                            guidance_rescale=self.audio_guidance_rescale,
                         )
-                        noise_pred_audio = rescale_noise_cfg(
-                            noise_pred_audio, noise_pred_audio_text, guidance_rescale=self.guidance_rescale
+                    
+                    # Get positive values from merged CFG inputs in case we need to do other DiT forward passes
+                    if (self.do_spatio_temporal_guidance or self.do_modality_isolation_guidance):
+                        if i == 0:
+                            # Only split values that remain constant throughout the loop once
+                            video_prompt_embeds = connector_prompt_embeds.chunk(2, dim=0)[1]
+                            audio_prompt_embeds = connector_audio_prompt_embeds.chunk(2, dim=0)[1]
+                            prompt_attn_mask = connector_attention_mask.chunk(2, dim=0)[1]
+
+                            video_pos_ids = video_coords.chunk(2, dim=0)[0]
+                            audio_pos_ids = audio_coords.chunk(2, dim=0)[0]
+
+                        # Split values that vary each denoising loop iteration
+                        timestep = timestep.chunk(2, dim=0)[0]
+                        video_timestep = video_timestep.chunk(2, dim=0)[0]
+                else:
+                    video_cfg_delta = audio_cfg_delta = 0
+
+                    video_prompt_embeds = connector_prompt_embeds
+                    audio_prompt_embeds = connector_audio_prompt_embeds
+                    prompt_attn_mask = connector_attention_mask
+
+                    video_pos_ids = video_coords
+                    audio_pos_ids = audio_coords
+                
+                if self.do_spatio_temporal_guidance:
+                    with self.transformer.cache_context("uncond_stg"):
+                        noise_pred_video_uncond_stg, noise_pred_audio_uncond_stg = self.transformer(
+                            hidden_states=latents.to(dtype=prompt_embeds.dtype),
+                            audio_hidden_states=audio_latents.to(dtype=prompt_embeds.dtype),
+                            encoder_hidden_states=video_prompt_embeds,
+                            audio_encoder_hidden_states=audio_prompt_embeds,
+                            timestep=video_timestep,
+                            audio_timestep=timestep,
+                            sigma=timestep,  # Used by LTX-2.3
+                            encoder_attention_mask=prompt_attn_mask,
+                            audio_encoder_attention_mask=prompt_attn_mask,
+                            self_attention_mask=None,
+                            audio_self_attention_mask=None,
+                            num_frames=latent_num_frames,
+                            height=latent_height,
+                            width=latent_width,
+                            fps=frame_rate,
+                            audio_num_frames=audio_num_frames,
+                            video_coords=video_pos_ids,
+                            audio_coords=audio_pos_ids,
+                            isolate_modalities=False,
+                            # Use STG at given blocks to perturb model
+                            spatio_temporal_guidance_blocks=spatio_temporal_guidance_blocks,
+                            perturbation_mask=None,
+                            attention_kwargs=attention_kwargs,
+                            return_dict=False,
                         )
+                    noise_pred_video_uncond_stg = noise_pred_video_uncond_stg.float()
+                    noise_pred_audio_uncond_stg = noise_pred_audio_uncond_stg.float()
+
+                    video_stg_delta = self.stg_scale * (noise_pred_video - noise_pred_video_uncond_stg)
+                    audio_stg_delta = self.audio_stg_scale * (noise_pred_audio - noise_pred_audio_uncond_stg)
+                else:
+                    video_stg_delta = audio_stg_delta = 0
+                
+                if self.do_modality_isolation_guidance:
+                    with self.transformer.cache_context("uncond_modality"):
+                        noise_pred_video_uncond_modality, noise_pred_audio_uncond_modality = self.transformer(
+                            hidden_states=latents.to(dtype=prompt_embeds.dtype),
+                            audio_hidden_states=audio_latents.to(dtype=prompt_embeds.dtype),
+                            encoder_hidden_states=video_prompt_embeds,
+                            audio_encoder_hidden_states=audio_prompt_embeds,
+                            timestep=video_timestep,
+                            audio_timestep=timestep,
+                            sigma=timestep,  # Used by LTX-2.3
+                            encoder_attention_mask=prompt_attn_mask,
+                            audio_encoder_attention_mask=prompt_attn_mask,
+                            self_attention_mask=None,
+                            audio_self_attention_mask=None,
+                            num_frames=latent_num_frames,
+                            height=latent_height,
+                            width=latent_width,
+                            fps=frame_rate,
+                            audio_num_frames=audio_num_frames,
+                            video_coords=video_pos_ids,
+                            audio_coords=audio_pos_ids,
+                            # Turn off A2V and V2A cross attn to isolate video and audio modalities
+                            isolate_modalities=True,
+                            spatio_temporal_guidance_blocks=None,
+                            perturbation_mask=None,
+                            attention_kwargs=attention_kwargs,
+                            return_dict=False,
+                        )
+                    noise_pred_video_uncond_modality = noise_pred_video_uncond_modality.float()
+                    noise_pred_audio_uncond_modality = noise_pred_audio_uncond_modality.float()
+
+                    video_modality_delta = (
+                        (self.modality_scale - 1) * (noise_pred_video - noise_pred_video_uncond_modality)
+                    )
+                    audio_modality_delta = (
+                        (self.audio_modality_scale - 1) * (noise_pred_audio - noise_pred_audio_uncond_modality)
+                    )
+                else:
+                    video_modality_delta = audio_modality_delta = 0
+
+                # Now apply all guidance terms
+                noise_pred_video = noise_pred_video + video_cfg_delta + video_stg_delta + video_modality_delta
+                noise_pred_audio = noise_pred_audio + audio_cfg_delta + audio_stg_delta + audio_modality_delta
 
                 # compute the previous noisy sample x_t -> x_t-1
                 noise_pred_video = self._unpack_latents(

From 652d363ded7bdbe75afeb4145688424fe215cc58 Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Wed, 11 Mar 2026 06:44:36 +0100
Subject: [PATCH 22/41] make style and make quality

---
 scripts/convert_ltx2_to_diffusers.py          |  4 +--
 .../models/transformers/transformer_ltx2.py   | 33 +++++++++++--------
 src/diffusers/pipelines/ltx2/connectors.py    |  2 +-
 src/diffusers/pipelines/ltx2/pipeline_ltx2.py | 20 ++++++-----
 .../pipelines/ltx2/pipeline_ltx2_condition.py | 20 ++++++-----
 .../ltx2/pipeline_ltx2_image2video.py         | 20 ++++++-----
 src/diffusers/pipelines/ltx2/vocoder.py       | 10 +++---
 7 files changed, 62 insertions(+), 47 deletions(-)

diff --git a/scripts/convert_ltx2_to_diffusers.py b/scripts/convert_ltx2_to_diffusers.py
index b5fca2d33998..7ef322e3039a 100644
--- a/scripts/convert_ltx2_to_diffusers.py
+++ b/scripts/convert_ltx2_to_diffusers.py
@@ -97,7 +97,7 @@
     "conv_post": "conv_out",
 }
 
-LTX_2_3_VOCODER_RENAME_DICT= {
+LTX_2_3_VOCODER_RENAME_DICT = {
     # Handle upsamplers ("ups" --> "upsamplers") due to name clash
     "resblocks": "resnets",
     "conv_pre": "conv_in",
@@ -749,7 +749,7 @@ def get_ltx2_vocoder_config(version: str) -> tuple[dict[str, Any], dict[str, Any
                 "leaky_relu_negative_slope": 0.1,
                 "antialias": False,
                 "final_act_fn": "tanh",
-                "final_bias": True, 
+                "final_bias": True,
                 "output_sampling_rate": 24000,
             },
         }
diff --git a/src/diffusers/models/transformers/transformer_ltx2.py b/src/diffusers/models/transformers/transformer_ltx2.py
index 0aed1124fb4a..e734e4a46b0d 100644
--- a/src/diffusers/models/transformers/transformer_ltx2.py
+++ b/src/diffusers/models/transformers/transformer_ltx2.py
@@ -586,9 +586,8 @@ def get_mod_params(
         scale_shift_table: torch.Tensor, temb: torch.Tensor, batch_size: int
     ) -> tuple[torch.Tensor, ...]:
         num_ada_params = scale_shift_table.shape[0]
-        ada_values = (
-            scale_shift_table[None, None].to(temb.device)
-            + temb.reshape(batch_size, temb.shape[1], num_ada_params, -1)
+        ada_values = scale_shift_table[None, None].to(temb.device) + temb.reshape(
+            batch_size, temb.shape[1], num_ada_params, -1
         )
         ada_params = ada_values.unbind(dim=2)
         return ada_params
@@ -676,7 +675,9 @@ def forward(
             video_prompt_ada_params = self.get_mod_params(self.prompt_scale_shift_table, temb_prompt, batch_size)
             shift_text_kv, scale_text_kv = video_prompt_ada_params
 
-            audio_prompt_ada_params = self.get_mod_params(self.audio_prompt_scale_shift_table, temb_prompt_audio, batch_size)
+            audio_prompt_ada_params = self.get_mod_params(
+                self.audio_prompt_scale_shift_table, temb_prompt_audio, batch_size
+            )
             audio_shift_text_kv, audio_scale_text_kv = audio_prompt_ada_params
 
         # 2.1. Video-Text Cross-Attention (Q: Video; K,V: Test)
@@ -733,7 +734,9 @@ def forward(
             audio_per_layer_ca_scale_shift = self.audio_a2v_cross_attn_scale_shift_table[:4, :]
             audio_per_layer_ca_gate = self.audio_a2v_cross_attn_scale_shift_table[4:, :]
 
-            audio_ca_ada_params = self.get_mod_params(audio_per_layer_ca_scale_shift, temb_ca_audio_scale_shift, batch_size)
+            audio_ca_ada_params = self.get_mod_params(
+                audio_per_layer_ca_scale_shift, temb_ca_audio_scale_shift, batch_size
+            )
             audio_ca_gate_param = self.get_mod_params(audio_per_layer_ca_gate, temb_ca_audio_gate, batch_size)
 
             audio_a2v_ca_scale, audio_a2v_ca_shift, audio_v2a_ca_scale, audio_v2a_ca_shift = audio_ca_ada_params
@@ -741,9 +744,9 @@ def forward(
 
             # 3.2. Audio-to-Video Cross Attention: Q: Video; K,V: Audio
             if use_a2v_cross_attention:
-                mod_norm_hidden_states = norm_hidden_states * (1 + video_a2v_ca_scale.squeeze(2)) + video_a2v_ca_shift.squeeze(
-                    2
-                )
+                mod_norm_hidden_states = norm_hidden_states * (
+                    1 + video_a2v_ca_scale.squeeze(2)
+                ) + video_a2v_ca_shift.squeeze(2)
                 mod_norm_audio_hidden_states = norm_audio_hidden_states * (
                     1 + audio_a2v_ca_scale.squeeze(2)
                 ) + audio_a2v_ca_shift.squeeze(2)
@@ -760,9 +763,9 @@ def forward(
 
             # 3.3. Video-to-Audio Cross Attention: Q: Audio; K,V: Video
             if use_v2a_cross_attention:
-                mod_norm_hidden_states = norm_hidden_states * (1 + video_v2a_ca_scale.squeeze(2)) + video_v2a_ca_shift.squeeze(
-                    2
-                )
+                mod_norm_hidden_states = norm_hidden_states * (
+                    1 + video_v2a_ca_scale.squeeze(2)
+                ) + video_v2a_ca_shift.squeeze(2)
                 mod_norm_audio_hidden_states = norm_audio_hidden_states * (
                     1 + audio_v2a_ca_scale.squeeze(2)
                 ) + audio_v2a_ca_shift.squeeze(2)
@@ -1442,7 +1445,9 @@ def forward(
             # Convert to additive attention mask in log-space where 0 (masked) values get mapped to a large negative
             # number and positive values are mapped to their logarithm.
             dtype_finfo = torch.finfo(hidden_states.dtype)
-            additive_self_attn_mask = torch.full_like(audio_self_attention_mask, dtype_finfo.min, dtype=hidden_states.dtype)
+            additive_self_attn_mask = torch.full_like(
+                audio_self_attention_mask, dtype_finfo.min, dtype=hidden_states.dtype
+            )
             unmasked_entries = audio_self_attention_mask > 0
             if torch.any(unmasked_entries):
                 additive_self_attn_mask[unmasked_entries] = torch.log(
@@ -1548,7 +1553,9 @@ def forward(
             encoder_hidden_states = encoder_hidden_states.view(batch_size, -1, hidden_states.size(-1))
 
             audio_encoder_hidden_states = self.audio_caption_projection(audio_encoder_hidden_states)
-            audio_encoder_hidden_states = audio_encoder_hidden_states.view(batch_size, -1, audio_hidden_states.size(-1))
+            audio_encoder_hidden_states = audio_encoder_hidden_states.view(
+                batch_size, -1, audio_hidden_states.size(-1)
+            )
 
         # 5. Run transformer blocks
         spatio_temporal_guidance_blocks = spatio_temporal_guidance_blocks or []
diff --git a/src/diffusers/pipelines/ltx2/connectors.py b/src/diffusers/pipelines/ltx2/connectors.py
index f945200962d2..3f721a2cfbf9 100644
--- a/src/diffusers/pipelines/ltx2/connectors.py
+++ b/src/diffusers/pipelines/ltx2/connectors.py
@@ -79,7 +79,7 @@ def per_layer_masked_mean_norm(
 
 
 def per_token_rms_norm(text_encoder_hidden_states: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
-    variance = torch.mean(text_encoder_hidden_states ** 2, dim=2, keepdim=True)
+    variance = torch.mean(text_encoder_hidden_states**2, dim=2, keepdim=True)
     norm_text_encoder_hidden_states = text_encoder_hidden_states + torch.rsqrt(variance + eps)
     return norm_text_encoder_hidden_states
 
diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
index 5f3a0fd00bb3..07c0327b8e38 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
@@ -1151,7 +1151,9 @@ def __call__(
                     video_cfg_delta = (self.guidance_scale - 1) * (noise_pred_video - noise_pred_video_uncond_text)
 
                     noise_pred_audio_uncond_text, noise_pred_audio = noise_pred_audio.chunk(2)
-                    audio_cfg_delta = (self.audio_guidance_scale - 1) * (noise_pred_audio - noise_pred_audio_uncond_text)
+                    audio_cfg_delta = (self.audio_guidance_scale - 1) * (
+                        noise_pred_audio - noise_pred_audio_uncond_text
+                    )
 
                     if self.guidance_rescale > 0:
                         # Based on 3.4. in https://huggingface.co/papers/2305.08891
@@ -1165,9 +1167,9 @@ def __call__(
                             noise_pred_text=noise_pred_audio,
                             guidance_rescale=self.audio_guidance_rescale,
                         )
-                    
+
                     # Get positive values from merged CFG inputs in case we need to do other DiT forward passes
-                    if (self.do_spatio_temporal_guidance or self.do_modality_isolation_guidance):
+                    if self.do_spatio_temporal_guidance or self.do_modality_isolation_guidance:
                         if i == 0:
                             # Only split values that remain constant throughout the loop once
                             video_prompt_embeds = connector_prompt_embeds.chunk(2, dim=0)[1]
@@ -1188,7 +1190,7 @@ def __call__(
 
                     video_pos_ids = video_coords
                     audio_pos_ids = audio_coords
-                
+
                 if self.do_spatio_temporal_guidance:
                     with self.transformer.cache_context("uncond_stg"):
                         noise_pred_video_uncond_stg, noise_pred_audio_uncond_stg = self.transformer(
@@ -1223,7 +1225,7 @@ def __call__(
                     audio_stg_delta = self.audio_stg_scale * (noise_pred_audio - noise_pred_audio_uncond_stg)
                 else:
                     video_stg_delta = audio_stg_delta = 0
-                
+
                 if self.do_modality_isolation_guidance:
                     with self.transformer.cache_context("uncond_modality"):
                         noise_pred_video_uncond_modality, noise_pred_audio_uncond_modality = self.transformer(
@@ -1254,11 +1256,11 @@ def __call__(
                     noise_pred_video_uncond_modality = noise_pred_video_uncond_modality.float()
                     noise_pred_audio_uncond_modality = noise_pred_audio_uncond_modality.float()
 
-                    video_modality_delta = (
-                        (self.modality_scale - 1) * (noise_pred_video - noise_pred_video_uncond_modality)
+                    video_modality_delta = (self.modality_scale - 1) * (
+                        noise_pred_video - noise_pred_video_uncond_modality
                     )
-                    audio_modality_delta = (
-                        (self.audio_modality_scale - 1) * (noise_pred_audio - noise_pred_audio_uncond_modality)
+                    audio_modality_delta = (self.audio_modality_scale - 1) * (
+                        noise_pred_audio - noise_pred_audio_uncond_modality
                     )
                 else:
                     video_modality_delta = audio_modality_delta = 0
diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py
index ed2f87ccdcac..9a2e8c70cf32 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py
@@ -1384,7 +1384,9 @@ def __call__(
                     video_cfg_delta = (self.guidance_scale - 1) * (noise_pred_video - noise_pred_video_uncond_text)
 
                     noise_pred_audio_uncond_text, noise_pred_audio = noise_pred_audio.chunk(2)
-                    audio_cfg_delta = (self.audio_guidance_scale - 1) * (noise_pred_audio - noise_pred_audio_uncond_text)
+                    audio_cfg_delta = (self.audio_guidance_scale - 1) * (
+                        noise_pred_audio - noise_pred_audio_uncond_text
+                    )
 
                     if self.guidance_rescale > 0:
                         # Based on 3.4. in https://huggingface.co/papers/2305.08891
@@ -1398,9 +1400,9 @@ def __call__(
                             noise_pred_text=noise_pred_audio,
                             guidance_rescale=self.audio_guidance_rescale,
                         )
-                    
+
                     # Get positive values from merged CFG inputs in case we need to do other DiT forward passes
-                    if (self.do_spatio_temporal_guidance or self.do_modality_isolation_guidance):
+                    if self.do_spatio_temporal_guidance or self.do_modality_isolation_guidance:
                         if i == 0:
                             # Only split values that remain constant throughout the loop once
                             video_prompt_embeds = connector_prompt_embeds.chunk(2, dim=0)[1]
@@ -1422,7 +1424,7 @@ def __call__(
 
                     video_pos_ids = video_coords
                     audio_pos_ids = audio_coords
-                
+
                 if self.do_spatio_temporal_guidance:
                     with self.transformer.cache_context("uncond_stg"):
                         noise_pred_video_uncond_stg, noise_pred_audio_uncond_stg = self.transformer(
@@ -1458,7 +1460,7 @@ def __call__(
                     audio_stg_delta = self.audio_stg_scale * (noise_pred_audio - noise_pred_audio_uncond_stg)
                 else:
                     video_stg_delta = audio_stg_delta = 0
-                
+
                 if self.do_modality_isolation_guidance:
                     with self.transformer.cache_context("uncond_modality"):
                         noise_pred_video_uncond_modality, noise_pred_audio_uncond_modality = self.transformer(
@@ -1490,11 +1492,11 @@ def __call__(
                     noise_pred_video_uncond_modality = noise_pred_video_uncond_modality.float()
                     noise_pred_audio_uncond_modality = noise_pred_audio_uncond_modality.float()
 
-                    video_modality_delta = (
-                        (self.modality_scale - 1) * (noise_pred_video - noise_pred_video_uncond_modality)
+                    video_modality_delta = (self.modality_scale - 1) * (
+                        noise_pred_video - noise_pred_video_uncond_modality
                     )
-                    audio_modality_delta = (
-                        (self.audio_modality_scale - 1) * (noise_pred_audio - noise_pred_audio_uncond_modality)
+                    audio_modality_delta = (self.audio_modality_scale - 1) * (
+                        noise_pred_audio - noise_pred_audio_uncond_modality
                     )
                 else:
                     video_modality_delta = audio_modality_delta = 0
diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
index 746aab42b668..053721c095c5 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
@@ -1217,7 +1217,9 @@ def __call__(
                     video_cfg_delta = (self.guidance_scale - 1) * (noise_pred_video - noise_pred_video_uncond_text)
 
                     noise_pred_audio_uncond_text, noise_pred_audio = noise_pred_audio.chunk(2)
-                    audio_cfg_delta = (self.audio_guidance_scale - 1) * (noise_pred_audio - noise_pred_audio_uncond_text)
+                    audio_cfg_delta = (self.audio_guidance_scale - 1) * (
+                        noise_pred_audio - noise_pred_audio_uncond_text
+                    )
 
                     if self.guidance_rescale > 0:
                         # Based on 3.4. in https://huggingface.co/papers/2305.08891
@@ -1231,9 +1233,9 @@ def __call__(
                             noise_pred_text=noise_pred_audio,
                             guidance_rescale=self.audio_guidance_rescale,
                         )
-                    
+
                     # Get positive values from merged CFG inputs in case we need to do other DiT forward passes
-                    if (self.do_spatio_temporal_guidance or self.do_modality_isolation_guidance):
+                    if self.do_spatio_temporal_guidance or self.do_modality_isolation_guidance:
                         if i == 0:
                             # Only split values that remain constant throughout the loop once
                             video_prompt_embeds = connector_prompt_embeds.chunk(2, dim=0)[1]
@@ -1255,7 +1257,7 @@ def __call__(
 
                     video_pos_ids = video_coords
                     audio_pos_ids = audio_coords
-                
+
                 if self.do_spatio_temporal_guidance:
                     with self.transformer.cache_context("uncond_stg"):
                         noise_pred_video_uncond_stg, noise_pred_audio_uncond_stg = self.transformer(
@@ -1291,7 +1293,7 @@ def __call__(
                     audio_stg_delta = self.audio_stg_scale * (noise_pred_audio - noise_pred_audio_uncond_stg)
                 else:
                     video_stg_delta = audio_stg_delta = 0
-                
+
                 if self.do_modality_isolation_guidance:
                     with self.transformer.cache_context("uncond_modality"):
                         noise_pred_video_uncond_modality, noise_pred_audio_uncond_modality = self.transformer(
@@ -1323,11 +1325,11 @@ def __call__(
                     noise_pred_video_uncond_modality = noise_pred_video_uncond_modality.float()
                     noise_pred_audio_uncond_modality = noise_pred_audio_uncond_modality.float()
 
-                    video_modality_delta = (
-                        (self.modality_scale - 1) * (noise_pred_video - noise_pred_video_uncond_modality)
+                    video_modality_delta = (self.modality_scale - 1) * (
+                        noise_pred_video - noise_pred_video_uncond_modality
                     )
-                    audio_modality_delta = (
-                        (self.audio_modality_scale - 1) * (noise_pred_audio - noise_pred_audio_uncond_modality)
+                    audio_modality_delta = (self.audio_modality_scale - 1) * (
+                        noise_pred_audio - noise_pred_audio_uncond_modality
                     )
                 else:
                     video_modality_delta = audio_modality_delta = 0
diff --git a/src/diffusers/pipelines/ltx2/vocoder.py b/src/diffusers/pipelines/ltx2/vocoder.py
index c17bfb786c9d..f0004f2ec02d 100644
--- a/src/diffusers/pipelines/ltx2/vocoder.py
+++ b/src/diffusers/pipelines/ltx2/vocoder.py
@@ -91,7 +91,7 @@ class UpSample1d(nn.Module):
     def __init__(
         self,
         ratio: int = 2,
-        kernel_size: int | None  = None,
+        kernel_size: int | None = None,
         window_type: str = "kaiser",
         padding_mode: str = "replicate",
         persistent: bool = True,
@@ -134,7 +134,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = F.pad(x, (self.pad, self.pad), mode=self.padding_mode)
         low_pass_filter = self.filter.to(dtype=x.dtype, device=x.device).expand(num_channels, -1, -1)
         x = self.ratio * F.conv_transpose1d(x, low_pass_filter, stride=self.ratio, groups=num_channels)
-        return x[..., self.pad_left:-self.pad_right]
+        return x[..., self.pad_left : -self.pad_right]
 
 
 class AntiAliasAct1d(nn.Module):
@@ -142,6 +142,7 @@ class AntiAliasAct1d(nn.Module):
     Antialiasing activation for a 1D signal: upsamples, applies an activation (usually snakebeta), and then downsamples
     to avoid aliasing.
     """
+
     def __init__(
         self,
         act_fn: str | nn.Module,
@@ -172,6 +173,7 @@ class SnakeBeta(nn.Module):
     """
     Implements the Snake and SnakeBeta activations, which help with learning periodic patterns.
     """
+
     def __init__(
         self,
         channels: int,
@@ -419,8 +421,8 @@ def forward(self, hidden_states: torch.Tensor, time_last: bool = False) -> torch
 class CausalSTFT(nn.Module):
     """
     Performs a causal short-time Fourier transform (STFT) using causal Hann windows on a waveform. The DFT bases
-    multiplied by the Hann windows are pre-calculated and stored as buffers. For exact parity with training, the
-    exact buffers should be loaded from the checkpoint in bfloat16.
+    multiplied by the Hann windows are pre-calculated and stored as buffers. For exact parity with training, the exact
+    buffers should be loaded from the checkpoint in bfloat16.
     """
 
     def __init__(self, filter_length: int = 512, hop_length: int = 80, window_length: int = 512):

From d018534de1fa86bac650ce1777d0946e8eae2354 Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Wed, 11 Mar 2026 07:31:12 +0100
Subject: [PATCH 23/41] Make audio guidance values default to video values by
 default

---
 src/diffusers/pipelines/ltx2/pipeline_ltx2.py | 32 ++++++++++++-------
 .../pipelines/ltx2/pipeline_ltx2_condition.py | 32 ++++++++++++-------
 .../ltx2/pipeline_ltx2_image2video.py         | 32 ++++++++++++-------
 3 files changed, 60 insertions(+), 36 deletions(-)

diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
index 07c0327b8e38..35928a2b1e6e 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
@@ -761,10 +761,10 @@ def __call__(
         stg_scale: float = 0.0,
         modality_scale: float = 1.0,
         guidance_rescale: float = 0.0,
-        audio_guidance_scale: float = 4.0,
-        audio_stg_scale: float = 0.0,
-        audio_modality_scale: float = 1.0,
-        audio_guidance_rescale: float = 0.0,
+        audio_guidance_scale: float | None = None,
+        audio_stg_scale: float | None = None,
+        audio_modality_scale: float | None = None,
+        audio_guidance_rescale: float | None = None,
         spatio_temporal_guidance_blocks: list[int] | None = None,
         noise_scale: float = 0.0,
         num_videos_per_prompt: int = 1,
@@ -835,20 +835,23 @@ def __call__(
                 [Common Diffusion Noise Schedules and Sample Steps are
                 Flawed](https://huggingface.co/papers/2305.08891). Guidance rescale factor should fix overexposure when
                 using zero terminal SNR. Used for the video modality.
-            audio_guidance_scale (`float`, *optional* defaults to `4.0`):
+            audio_guidance_scale (`float`, *optional* defaults to `None`):
                 Audio guidance scale for CFG with respect to the negative prompt. The CFG update rule is the same for
                 video and audio, but they can use different values for the guidance scale. The LTX-2.X authors suggest
                 that the `audio_guidance_scale` should be higher relative to the video `guidance_scale` (e.g. for
-                LTX-2.3 they suggest 3.0 for video and 7.0 for audio).
-            audio_stg_scale (`float`, *optional*, defaults to `0.0`):
+                LTX-2.3 they suggest 3.0 for video and 7.0 for audio). If `None`, defaults to the video value
+                `guidance_scale`.
+            audio_stg_scale (`float`, *optional*, defaults to `None`):
                 Audio guidance scale for STG. As with CFG, the STG update rule is otherwise the same for video and
-                audio. For LTX-2.3, a value of 1.0 is suggested for both video and audio.
-            audio_modality_scale (`float`, *optional*, defaults to `1.0`):
+                audio. For LTX-2.3, a value of 1.0 is suggested for both video and audio. If `None`, defaults to the
+                video value `stg_scale`.
+            audio_modality_scale (`float`, *optional*, defaults to `None`):
                 Audio guidance scale for LTX-2.X modality isolation guidance. As with CFG, the modality guidance rule
                 is otherwise the same for video and audio. For LTX-2.3, a value of 3.0 is suggested for both video and
-                audio.
-            audio_guidance_rescale (`float`, *optional*, defaults to `0.0`):
-                A separate guidance rescale factor for the audio modality.
+                audio. If `None`, defaults to the video value `modality_scale`.
+            audio_guidance_rescale (`float`, *optional*, defaults to `None`):
+                A separate guidance rescale factor for the audio modality. If `None`, defaults to the video value
+                `guidance_rescale`.
             spatio_temporal_guidance_blocks (`list[int]`, *optional*, defaults to `None`):
                 The zero-indexed transformer block indices at which to apply STG. Must be supplied if STG is used
                 (`stg_scale` or `audio_stg_scale` is greater than `0`). A value of `[29]` is recommended for LTX-2.0
@@ -915,6 +918,11 @@ def __call__(
         if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
             callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
 
+        audio_guidance_scale = audio_guidance_scale or guidance_scale
+        audio_stg_scale = audio_stg_scale or stg_scale
+        audio_modality_scale = audio_modality_scale or modality_scale
+        audio_guidance_rescale = audio_guidance_rescale or guidance_rescale
+
         # 1. Check inputs. Raise error if not correct
         self.check_inputs(
             prompt=prompt,
diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py
index 9a2e8c70cf32..6c8d6930e00d 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py
@@ -997,10 +997,10 @@ def __call__(
         stg_scale: float = 0.0,
         modality_scale: float = 1.0,
         guidance_rescale: float = 0.0,
-        audio_guidance_scale: float = 4.0,
-        audio_stg_scale: float = 0.0,
-        audio_modality_scale: float = 1.0,
-        audio_guidance_rescale: float = 0.0,
+        audio_guidance_scale: float | None = None,
+        audio_stg_scale: float | None = None,
+        audio_modality_scale: float | None = None,
+        audio_guidance_rescale: float | None = None,
         spatio_temporal_guidance_blocks: list[int] | None = None,
         noise_scale: float | None = None,
         num_videos_per_prompt: int | None = 1,
@@ -1073,20 +1073,23 @@ def __call__(
                 [Common Diffusion Noise Schedules and Sample Steps are
                 Flawed](https://huggingface.co/papers/2305.08891). Guidance rescale factor should fix overexposure when
                 using zero terminal SNR. Used for the video modality.
-            audio_guidance_scale (`float`, *optional* defaults to `4.0`):
+            audio_guidance_scale (`float`, *optional* defaults to `None`):
                 Audio guidance scale for CFG with respect to the negative prompt. The CFG update rule is the same for
                 video and audio, but they can use different values for the guidance scale. The LTX-2.X authors suggest
                 that the `audio_guidance_scale` should be higher relative to the video `guidance_scale` (e.g. for
-                LTX-2.3 they suggest 3.0 for video and 7.0 for audio).
-            audio_stg_scale (`float`, *optional*, defaults to `0.0`):
+                LTX-2.3 they suggest 3.0 for video and 7.0 for audio). If `None`, defaults to the video value
+                `guidance_scale`.
+            audio_stg_scale (`float`, *optional*, defaults to `None`):
                 Audio guidance scale for STG. As with CFG, the STG update rule is otherwise the same for video and
-                audio. For LTX-2.3, a value of 1.0 is suggested for both video and audio.
-            audio_modality_scale (`float`, *optional*, defaults to `1.0`):
+                audio. For LTX-2.3, a value of 1.0 is suggested for both video and audio. If `None`, defaults to the
+                video value `stg_scale`.
+            audio_modality_scale (`float`, *optional*, defaults to `None`):
                 Audio guidance scale for LTX-2.X modality isolation guidance. As with CFG, the modality guidance rule
                 is otherwise the same for video and audio. For LTX-2.3, a value of 3.0 is suggested for both video and
-                audio.
-            audio_guidance_rescale (`float`, *optional*, defaults to `0.0`):
-                A separate guidance rescale factor for the audio modality.
+                audio. If `None`, defaults to the video value `modality_scale`.
+            audio_guidance_rescale (`float`, *optional*, defaults to `None`):
+                A separate guidance rescale factor for the audio modality. If `None`, defaults to the video value
+                `guidance_rescale`.
             spatio_temporal_guidance_blocks (`list[int]`, *optional*, defaults to `None`):
                 The zero-indexed transformer block indices at which to apply STG. Must be supplied if STG is used
                 (`stg_scale` or `audio_stg_scale` is greater than `0`). A value of `[29]` is recommended for LTX-2.0
@@ -1154,6 +1157,11 @@ def __call__(
         if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
             callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
 
+        audio_guidance_scale = audio_guidance_scale or guidance_scale
+        audio_stg_scale = audio_stg_scale or stg_scale
+        audio_modality_scale = audio_modality_scale or modality_scale
+        audio_guidance_rescale = audio_guidance_rescale or guidance_rescale
+
         # 1. Check inputs. Raise error if not correct
         self.check_inputs(
             prompt=prompt,
diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
index 053721c095c5..37a55da3c6e4 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
@@ -816,10 +816,10 @@ def __call__(
         stg_scale: float = 0.0,
         modality_scale: float = 1.0,
         guidance_rescale: float = 0.0,
-        audio_guidance_scale: float = 4.0,
-        audio_stg_scale: float = 0.0,
-        audio_modality_scale: float = 1.0,
-        audio_guidance_rescale: float = 0.0,
+        audio_guidance_scale: float | None = None,
+        audio_stg_scale: float | None = None,
+        audio_modality_scale: float | None = None,
+        audio_guidance_rescale: float | None = None,
         spatio_temporal_guidance_blocks: list[int] | None = None,
         noise_scale: float = 0.0,
         num_videos_per_prompt: int = 1,
@@ -892,20 +892,23 @@ def __call__(
                 [Common Diffusion Noise Schedules and Sample Steps are
                 Flawed](https://huggingface.co/papers/2305.08891). Guidance rescale factor should fix overexposure when
                 using zero terminal SNR. Used for the video modality.
-            audio_guidance_scale (`float`, *optional* defaults to `4.0`):
+            audio_guidance_scale (`float`, *optional* defaults to `None`):
                 Audio guidance scale for CFG with respect to the negative prompt. The CFG update rule is the same for
                 video and audio, but they can use different values for the guidance scale. The LTX-2.X authors suggest
                 that the `audio_guidance_scale` should be higher relative to the video `guidance_scale` (e.g. for
-                LTX-2.3 they suggest 3.0 for video and 7.0 for audio).
-            audio_stg_scale (`float`, *optional*, defaults to `0.0`):
+                LTX-2.3 they suggest 3.0 for video and 7.0 for audio). If `None`, defaults to the video value
+                `guidance_scale`.
+            audio_stg_scale (`float`, *optional*, defaults to `None`):
                 Audio guidance scale for STG. As with CFG, the STG update rule is otherwise the same for video and
-                audio. For LTX-2.3, a value of 1.0 is suggested for both video and audio.
-            audio_modality_scale (`float`, *optional*, defaults to `1.0`):
+                audio. For LTX-2.3, a value of 1.0 is suggested for both video and audio. If `None`, defaults to the
+                video value `stg_scale`.
+            audio_modality_scale (`float`, *optional*, defaults to `None`):
                 Audio guidance scale for LTX-2.X modality isolation guidance. As with CFG, the modality guidance rule
                 is otherwise the same for video and audio. For LTX-2.3, a value of 3.0 is suggested for both video and
-                audio.
-            audio_guidance_rescale (`float`, *optional*, defaults to `0.0`):
-                A separate guidance rescale factor for the audio modality.
+                audio. If `None`, defaults to the video value `modality_scale`.
+            audio_guidance_rescale (`float`, *optional*, defaults to `None`):
+                A separate guidance rescale factor for the audio modality. If `None`, defaults to the video value
+                `guidance_rescale`.
             spatio_temporal_guidance_blocks (`list[int]`, *optional*, defaults to `None`):
                 The zero-indexed transformer block indices at which to apply STG. Must be supplied if STG is used
                 (`stg_scale` or `audio_stg_scale` is greater than `0`). A value of `[29]` is recommended for LTX-2.0
@@ -972,6 +975,11 @@ def __call__(
         if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
             callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
 
+        audio_guidance_scale = audio_guidance_scale or guidance_scale
+        audio_stg_scale = audio_stg_scale or stg_scale
+        audio_modality_scale = audio_modality_scale or modality_scale
+        audio_guidance_rescale = audio_guidance_rescale or guidance_rescale
+
         # 1. Check inputs. Raise error if not correct
         self.check_inputs(
             prompt=prompt,

From c0bb2ef21ff6db3a77881c3e4bf6cbf98ab81945 Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Wed, 11 Mar 2026 09:25:30 +0100
Subject: [PATCH 24/41] Update to LTX-2.3 style guidance rescaling

---
 src/diffusers/pipelines/ltx2/pipeline_ltx2.py | 46 +++++++++---------
 .../pipelines/ltx2/pipeline_ltx2_condition.py | 47 ++++++++++---------
 .../ltx2/pipeline_ltx2_image2video.py         | 47 ++++++++++---------
 3 files changed, 75 insertions(+), 65 deletions(-)

diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
index 35928a2b1e6e..c7c02f5ae622 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
@@ -155,12 +155,12 @@ def retrieve_timesteps(
     return timesteps, num_inference_steps
 
 
-def rescale_noise_cfg_delta(noise_cfg, noise_pred_text, guidance_rescale=0.0):
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.rescale_noise_cfg
+def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
     r"""
     Rescales `noise_cfg` tensor based on `guidance_rescale` to improve image quality and fix overexposure. Based on
     Section 3.4 from [Common Diffusion Noise Schedules and Sample Steps are
-    Flawed](https://huggingface.co/papers/2305.08891). Returns a delta value with respect to noise_pred_text, which is
-    useful when there are multiple guidance terms.
+    Flawed](https://huggingface.co/papers/2305.08891).
 
     Args:
         noise_cfg (`torch.Tensor`):
@@ -179,9 +179,7 @@ def rescale_noise_cfg_delta(noise_cfg, noise_pred_text, guidance_rescale=0.0):
     noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
     # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images
     noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg
-    # Return as a delta with respect to noise_pred_text
-    rescale_delta = noise_cfg - noise_pred_text
-    return rescale_delta
+    return noise_cfg
 
 
 class LTX2Pipeline(DiffusionPipeline, FromSingleFileMixin, LTX2LoraLoaderMixin):
@@ -1163,19 +1161,6 @@ def __call__(
                         noise_pred_audio - noise_pred_audio_uncond_text
                     )
 
-                    if self.guidance_rescale > 0:
-                        # Based on 3.4. in https://huggingface.co/papers/2305.08891
-                        video_cfg_delta = rescale_noise_cfg_delta(
-                            noise_cfg=noise_pred_video + video_cfg_delta,
-                            noise_pred_text=noise_pred_video,
-                            guidance_rescale=self.guidance_rescale,
-                        )
-                        audio_cfg_delta = rescale_noise_cfg_delta(
-                            noise_cfg=noise_pred_audio + audio_cfg_delta,
-                            noise_pred_text=noise_pred_audio,
-                            guidance_rescale=self.audio_guidance_rescale,
-                        )
-
                     # Get positive values from merged CFG inputs in case we need to do other DiT forward passes
                     if self.do_spatio_temporal_guidance or self.do_modality_isolation_guidance:
                         if i == 0:
@@ -1274,8 +1259,27 @@ def __call__(
                     video_modality_delta = audio_modality_delta = 0
 
                 # Now apply all guidance terms
-                noise_pred_video = noise_pred_video + video_cfg_delta + video_stg_delta + video_modality_delta
-                noise_pred_audio = noise_pred_audio + audio_cfg_delta + audio_stg_delta + audio_modality_delta
+                noise_pred_video_g = noise_pred_video + video_cfg_delta + video_stg_delta + video_modality_delta
+                noise_pred_audio_g = noise_pred_audio + audio_cfg_delta + audio_stg_delta + audio_modality_delta
+
+                # Apply LTX-2.X guidance rescaling
+                if self.guidance_rescale > 0:
+                    video_rescale = self.guidance_rescale
+                    cond_std = noise_pred_video.std(dim=list(range(1, noise_pred_video.ndim)), keepdim=True)
+                    guided_std = noise_pred_video_g.std(dim=list(range(1, noise_pred_video_g.ndim)), keepdim=True)
+                    rescale_factor = video_rescale * (cond_std / guided_std) + (1 - video_rescale)
+                    noise_pred_video = noise_pred_video_g * rescale_factor
+                else:
+                    noise_pred_video = noise_pred_video_g
+
+                if self.audio_guidance_rescale > 0:
+                    audio_rescale = self.audio_guidance_rescale
+                    cond_std = noise_pred_audio.std(dim=list(range(1, noise_pred_audio.ndim)), keepdim=True)
+                    guided_std = noise_pred_audio_g.std(dim=list(range(1, noise_pred_audio_g.ndim)), keepdim=True)
+                    rescale_factor = audio_rescale * (cond_std / guided_std) + (1 - audio_rescale)
+                    noise_pred_audio = noise_pred_audio_g * rescale_factor
+                else:
+                    noise_pred_audio = noise_pred_audio_g
 
                 # compute the previous noisy sample x_t -> x_t-1
                 latents = self.scheduler.step(noise_pred_video, t, latents, return_dict=False)[0]
diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py
index 6c8d6930e00d..e7875aa5426e 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py
@@ -205,13 +205,12 @@ def retrieve_timesteps(
     return timesteps, num_inference_steps
 
 
-# Copied from diffusers.pipelines.ltx2.pipeline_ltx2.rescale_noise_cfg_delta
-def rescale_noise_cfg_delta(noise_cfg, noise_pred_text, guidance_rescale=0.0):
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.rescale_noise_cfg
+def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
     r"""
     Rescales `noise_cfg` tensor based on `guidance_rescale` to improve image quality and fix overexposure. Based on
     Section 3.4 from [Common Diffusion Noise Schedules and Sample Steps are
-    Flawed](https://huggingface.co/papers/2305.08891). Returns a delta value with respect to noise_pred_text, which is
-    useful when there are multiple guidance terms.
+    Flawed](https://huggingface.co/papers/2305.08891).
 
     Args:
         noise_cfg (`torch.Tensor`):
@@ -230,9 +229,7 @@ def rescale_noise_cfg_delta(noise_cfg, noise_pred_text, guidance_rescale=0.0):
     noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
     # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images
     noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg
-    # Return as a delta with respect to noise_pred_text
-    rescale_delta = noise_cfg - noise_pred_text
-    return rescale_delta
+    return noise_cfg
 
 
 class LTX2ConditionPipeline(DiffusionPipeline, FromSingleFileMixin, LTX2LoraLoaderMixin):
@@ -1396,19 +1393,6 @@ def __call__(
                         noise_pred_audio - noise_pred_audio_uncond_text
                     )
 
-                    if self.guidance_rescale > 0:
-                        # Based on 3.4. in https://huggingface.co/papers/2305.08891
-                        video_cfg_delta = rescale_noise_cfg_delta(
-                            noise_cfg=noise_pred_video + video_cfg_delta,
-                            noise_pred_text=noise_pred_video,
-                            guidance_rescale=self.guidance_rescale,
-                        )
-                        audio_cfg_delta = rescale_noise_cfg_delta(
-                            noise_cfg=noise_pred_audio + audio_cfg_delta,
-                            noise_pred_text=noise_pred_audio,
-                            guidance_rescale=self.audio_guidance_rescale,
-                        )
-
                     # Get positive values from merged CFG inputs in case we need to do other DiT forward passes
                     if self.do_spatio_temporal_guidance or self.do_modality_isolation_guidance:
                         if i == 0:
@@ -1510,8 +1494,27 @@ def __call__(
                     video_modality_delta = audio_modality_delta = 0
 
                 # Now apply all guidance terms
-                noise_pred_video = noise_pred_video + video_cfg_delta + video_stg_delta + video_modality_delta
-                noise_pred_audio = noise_pred_audio + audio_cfg_delta + audio_stg_delta + audio_modality_delta
+                noise_pred_video_g = noise_pred_video + video_cfg_delta + video_stg_delta + video_modality_delta
+                noise_pred_audio_g = noise_pred_audio + audio_cfg_delta + audio_stg_delta + audio_modality_delta
+
+                # Apply LTX-2.X guidance rescaling
+                if self.guidance_rescale > 0:
+                    video_rescale = self.guidance_rescale
+                    cond_std = noise_pred_video.std(dim=list(range(1, noise_pred_video.ndim)), keepdim=True)
+                    guided_std = noise_pred_video_g.std(dim=list(range(1, noise_pred_video_g.ndim)), keepdim=True)
+                    rescale_factor = video_rescale * (cond_std / guided_std) + (1 - video_rescale)
+                    noise_pred_video = noise_pred_video_g * rescale_factor
+                else:
+                    noise_pred_video = noise_pred_video_g
+
+                if self.audio_guidance_rescale > 0:
+                    audio_rescale = self.audio_guidance_rescale
+                    cond_std = noise_pred_audio.std(dim=list(range(1, noise_pred_audio.ndim)), keepdim=True)
+                    guided_std = noise_pred_audio_g.std(dim=list(range(1, noise_pred_audio_g.ndim)), keepdim=True)
+                    rescale_factor = audio_rescale * (cond_std / guided_std) + (1 - audio_rescale)
+                    noise_pred_audio = noise_pred_audio_g * rescale_factor
+                else:
+                    noise_pred_audio = noise_pred_audio_g
 
                 # NOTE: use only the first chunk of conditioning mask in case it is duplicated for CFG
                 bsz = noise_pred_video.size(0)
diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
index 37a55da3c6e4..5885f98e43ab 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
@@ -175,13 +175,12 @@ def retrieve_timesteps(
     return timesteps, num_inference_steps
 
 
-# Copied from diffusers.pipelines.ltx2.pipeline_ltx2.rescale_noise_cfg_delta
-def rescale_noise_cfg_delta(noise_cfg, noise_pred_text, guidance_rescale=0.0):
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.rescale_noise_cfg
+def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
     r"""
     Rescales `noise_cfg` tensor based on `guidance_rescale` to improve image quality and fix overexposure. Based on
     Section 3.4 from [Common Diffusion Noise Schedules and Sample Steps are
-    Flawed](https://huggingface.co/papers/2305.08891). Returns a delta value with respect to noise_pred_text, which is
-    useful when there are multiple guidance terms.
+    Flawed](https://huggingface.co/papers/2305.08891).
 
     Args:
         noise_cfg (`torch.Tensor`):
@@ -200,9 +199,7 @@ def rescale_noise_cfg_delta(noise_cfg, noise_pred_text, guidance_rescale=0.0):
     noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
     # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images
     noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg
-    # Return as a delta with respect to noise_pred_text
-    rescale_delta = noise_cfg - noise_pred_text
-    return rescale_delta
+    return noise_cfg
 
 
 class LTX2ImageToVideoPipeline(DiffusionPipeline, FromSingleFileMixin, LTX2LoraLoaderMixin):
@@ -1229,19 +1226,6 @@ def __call__(
                         noise_pred_audio - noise_pred_audio_uncond_text
                     )
 
-                    if self.guidance_rescale > 0:
-                        # Based on 3.4. in https://huggingface.co/papers/2305.08891
-                        video_cfg_delta = rescale_noise_cfg_delta(
-                            noise_cfg=noise_pred_video + video_cfg_delta,
-                            noise_pred_text=noise_pred_video,
-                            guidance_rescale=self.guidance_rescale,
-                        )
-                        audio_cfg_delta = rescale_noise_cfg_delta(
-                            noise_cfg=noise_pred_audio + audio_cfg_delta,
-                            noise_pred_text=noise_pred_audio,
-                            guidance_rescale=self.audio_guidance_rescale,
-                        )
-
                     # Get positive values from merged CFG inputs in case we need to do other DiT forward passes
                     if self.do_spatio_temporal_guidance or self.do_modality_isolation_guidance:
                         if i == 0:
@@ -1343,8 +1327,27 @@ def __call__(
                     video_modality_delta = audio_modality_delta = 0
 
                 # Now apply all guidance terms
-                noise_pred_video = noise_pred_video + video_cfg_delta + video_stg_delta + video_modality_delta
-                noise_pred_audio = noise_pred_audio + audio_cfg_delta + audio_stg_delta + audio_modality_delta
+                noise_pred_video_g = noise_pred_video + video_cfg_delta + video_stg_delta + video_modality_delta
+                noise_pred_audio_g = noise_pred_audio + audio_cfg_delta + audio_stg_delta + audio_modality_delta
+
+                # Apply LTX-2.X guidance rescaling
+                if self.guidance_rescale > 0:
+                    video_rescale = self.guidance_rescale
+                    cond_std = noise_pred_video.std(dim=list(range(1, noise_pred_video.ndim)), keepdim=True)
+                    guided_std = noise_pred_video_g.std(dim=list(range(1, noise_pred_video_g.ndim)), keepdim=True)
+                    rescale_factor = video_rescale * (cond_std / guided_std) + (1 - video_rescale)
+                    noise_pred_video = noise_pred_video_g * rescale_factor
+                else:
+                    noise_pred_video = noise_pred_video_g
+
+                if self.audio_guidance_rescale > 0:
+                    audio_rescale = self.audio_guidance_rescale
+                    cond_std = noise_pred_audio.std(dim=list(range(1, noise_pred_audio.ndim)), keepdim=True)
+                    guided_std = noise_pred_audio_g.std(dim=list(range(1, noise_pred_audio_g.ndim)), keepdim=True)
+                    rescale_factor = audio_rescale * (cond_std / guided_std) + (1 - audio_rescale)
+                    noise_pred_audio = noise_pred_audio_g * rescale_factor
+                else:
+                    noise_pred_audio = noise_pred_audio_g
 
                 # compute the previous noisy sample x_t -> x_t-1
                 noise_pred_video = self._unpack_latents(

From ab0e5b5cbbfbde6ec7fad6b402b6f0329550d7f6 Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Thu, 12 Mar 2026 10:17:44 +0100
Subject: [PATCH 25/41] Support cross timesteps for LTX-2.3 cross attention
 modulation

---
 .../models/transformers/transformer_ltx2.py   | 21 ++++++++++++-------
 src/diffusers/pipelines/ltx2/pipeline_ltx2.py |  8 +++++++
 .../pipelines/ltx2/pipeline_ltx2_condition.py |  8 +++++++
 .../ltx2/pipeline_ltx2_image2video.py         |  8 +++++++
 4 files changed, 38 insertions(+), 7 deletions(-)

diff --git a/src/diffusers/models/transformers/transformer_ltx2.py b/src/diffusers/models/transformers/transformer_ltx2.py
index e734e4a46b0d..a47fb97bdfc9 100644
--- a/src/diffusers/models/transformers/transformer_ltx2.py
+++ b/src/diffusers/models/transformers/transformer_ltx2.py
@@ -1343,6 +1343,7 @@ def forward(
         isolate_modalities: bool = False,
         spatio_temporal_guidance_blocks: list[int] | None = None,
         perturbation_mask: torch.Tensor | None = None,
+        use_cross_timestep: bool = False,
         attention_kwargs: dict[str, Any] | None = None,
         return_dict: bool = True,
     ) -> torch.Tensor:
@@ -1405,6 +1406,10 @@ def forward(
                 Perturbation mask for STG of shape `(batch_size,)` or `(batch_size, 1, 1)`. Should be 0 at batch
                 elements where STG should be applied and 1 elsewhere. If STG is being used but `peturbation_mask` is
                 not supplied, will default to applying STG (perturbing) all batch elements.
+            use_cross_timestep (`bool` *optional*, defaults to `False`):
+                Whether to use the cross modality (audio is the cross modality of video, and vice versa) sigma when
+                calculating the cross attention modulation parameters. `True` is the newer (e.g. LTX-2.3) behavior;
+                `False` is the legacy LTX-2.0 behavior.
             attention_kwargs (`dict[str, Any]`, *optional*):
                 Optional dict of keyword args to be passed to the attention processor.
             return_dict (`bool`, *optional*, defaults to `True`):
@@ -1444,15 +1449,15 @@ def forward(
         if audio_self_attention_mask is not None and audio_self_attention_mask.ndim == 3:
             # Convert to additive attention mask in log-space where 0 (masked) values get mapped to a large negative
             # number and positive values are mapped to their logarithm.
-            dtype_finfo = torch.finfo(hidden_states.dtype)
+            dtype_finfo = torch.finfo(audio_hidden_states.dtype)
             additive_self_attn_mask = torch.full_like(
-                audio_self_attention_mask, dtype_finfo.min, dtype=hidden_states.dtype
+                audio_self_attention_mask, dtype_finfo.min, dtype=audio_hidden_states.dtype
             )
             unmasked_entries = audio_self_attention_mask > 0
             if torch.any(unmasked_entries):
                 additive_self_attn_mask[unmasked_entries] = torch.log(
                     audio_self_attention_mask[unmasked_entries].clamp(min=dtype_finfo.tiny)
-                ).to(hidden_states.dtype)
+                ).to(audio_hidden_states.dtype)
             audio_self_attention_mask = additive_self_attn_mask.unsqueeze(1)  # [batch_size, 1, seq_len, seq_len]
 
         batch_size = hidden_states.size(0)
@@ -1517,13 +1522,14 @@ def forward(
             temb_prompt = temb_prompt_audio = None
 
         # 3.2. Prepare global modality cross attention modulation parameters
+        video_ca_timestep = audio_sigma.flatten() if use_cross_timestep else timestep.flatten()
         video_cross_attn_scale_shift, _ = self.av_cross_attn_video_scale_shift(
-            timestep.flatten(),
+            video_ca_timestep,
             batch_size=batch_size,
             hidden_dtype=hidden_states.dtype,
         )
         video_cross_attn_a2v_gate, _ = self.av_cross_attn_video_a2v_gate(
-            timestep.flatten() * timestep_cross_attn_gate_scale_factor,
+            video_ca_timestep * timestep_cross_attn_gate_scale_factor,
             batch_size=batch_size,
             hidden_dtype=hidden_states.dtype,
         )
@@ -1532,13 +1538,14 @@ def forward(
         )
         video_cross_attn_a2v_gate = video_cross_attn_a2v_gate.view(batch_size, -1, video_cross_attn_a2v_gate.shape[-1])
 
+        audio_ca_timestep = sigma.flatten() if use_cross_timestep else audio_timestep.flatten()
         audio_cross_attn_scale_shift, _ = self.av_cross_attn_audio_scale_shift(
-            audio_timestep.flatten(),
+            audio_ca_timestep,
             batch_size=batch_size,
             hidden_dtype=audio_hidden_states.dtype,
         )
         audio_cross_attn_v2a_gate, _ = self.av_cross_attn_audio_v2a_gate(
-            audio_timestep.flatten() * timestep_cross_attn_gate_scale_factor,
+            audio_ca_timestep * timestep_cross_attn_gate_scale_factor,
             batch_size=batch_size,
             hidden_dtype=audio_hidden_states.dtype,
         )
diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
index c7c02f5ae622..07321518be66 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
@@ -775,6 +775,7 @@ def __call__(
         negative_prompt_attention_mask: torch.Tensor | None = None,
         decode_timestep: float | list[float] = 0.0,
         decode_noise_scale: float | list[float] | None = None,
+        use_cross_timestep: bool = False,
         output_type: str = "pil",
         return_dict: bool = True,
         attention_kwargs: dict[str, Any] | None = None,
@@ -884,6 +885,10 @@ def __call__(
                 The timestep at which generated video is decoded.
             decode_noise_scale (`float`, defaults to `None`):
                 The interpolation factor between random noise and denoised latents at the decode timestep.
+            use_cross_timestep (`bool` *optional*, defaults to `False`):
+                Whether to use the cross modality (audio is the cross modality of video, and vice versa) sigma when
+                calculating the cross attention modulation parameters. `True` is the newer (e.g. LTX-2.3) behavior;
+                `False` is the legacy LTX-2.0 behavior.
             output_type (`str`, *optional*, defaults to `"pil"`):
                 The output format of the generate image. Choose between
                 [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
@@ -1145,6 +1150,7 @@ def __call__(
                         isolate_modalities=False,
                         spatio_temporal_guidance_blocks=None,
                         perturbation_mask=None,
+                        use_cross_timestep=use_cross_timestep,
                         attention_kwargs=attention_kwargs,
                         return_dict=False,
                     )
@@ -1208,6 +1214,7 @@ def __call__(
                             # Use STG at given blocks to perturb model
                             spatio_temporal_guidance_blocks=spatio_temporal_guidance_blocks,
                             perturbation_mask=None,
+                            use_cross_timestep=use_cross_timestep,
                             attention_kwargs=attention_kwargs,
                             return_dict=False,
                         )
@@ -1243,6 +1250,7 @@ def __call__(
                             isolate_modalities=True,
                             spatio_temporal_guidance_blocks=None,
                             perturbation_mask=None,
+                            use_cross_timestep=use_cross_timestep,
                             attention_kwargs=attention_kwargs,
                             return_dict=False,
                         )
diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py
index e7875aa5426e..dd1b45f8de66 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py
@@ -1010,6 +1010,7 @@ def __call__(
         negative_prompt_attention_mask: torch.Tensor | None = None,
         decode_timestep: float | list[float] = 0.0,
         decode_noise_scale: float | list[float] | None = None,
+        use_cross_timestep: bool = False,
         output_type: str = "pil",
         return_dict: bool = True,
         attention_kwargs: dict[str, Any] | None = None,
@@ -1122,6 +1123,10 @@ def __call__(
                 The timestep at which generated video is decoded.
             decode_noise_scale (`float`, defaults to `None`):
                 The interpolation factor between random noise and denoised latents at the decode timestep.
+            use_cross_timestep (`bool` *optional*, defaults to `False`):
+                Whether to use the cross modality (audio is the cross modality of video, and vice versa) sigma when
+                calculating the cross attention modulation parameters. `True` is the newer (e.g. LTX-2.3) behavior;
+                `False` is the legacy LTX-2.0 behavior.
             output_type (`str`, *optional*, defaults to `"pil"`):
                 The output format of the generate image. Choose between
                 [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
@@ -1377,6 +1382,7 @@ def __call__(
                         isolate_modalities=False,
                         spatio_temporal_guidance_blocks=None,
                         perturbation_mask=None,
+                        use_cross_timestep=use_cross_timestep,
                         attention_kwargs=attention_kwargs,
                         return_dict=False,
                     )
@@ -1442,6 +1448,7 @@ def __call__(
                             # Use STG at given blocks to perturb model
                             spatio_temporal_guidance_blocks=spatio_temporal_guidance_blocks,
                             perturbation_mask=None,
+                            use_cross_timestep=use_cross_timestep,
                             attention_kwargs=attention_kwargs,
                             return_dict=False,
                         )
@@ -1478,6 +1485,7 @@ def __call__(
                             isolate_modalities=True,
                             spatio_temporal_guidance_blocks=None,
                             perturbation_mask=None,
+                            use_cross_timestep=use_cross_timestep,
                             attention_kwargs=attention_kwargs,
                             return_dict=False,
                         )
diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
index 5885f98e43ab..df233978e209 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
@@ -829,6 +829,7 @@ def __call__(
         negative_prompt_attention_mask: torch.Tensor | None = None,
         decode_timestep: float | list[float] = 0.0,
         decode_noise_scale: float | list[float] | None = None,
+        use_cross_timestep: bool = False,
         output_type: str = "pil",
         return_dict: bool = True,
         attention_kwargs: dict[str, Any] | None = None,
@@ -940,6 +941,10 @@ def __call__(
                 The timestep at which generated video is decoded.
             decode_noise_scale (`float`, defaults to `None`):
                 The interpolation factor between random noise and denoised latents at the decode timestep.
+            use_cross_timestep (`bool` *optional*, defaults to `False`):
+                Whether to use the cross modality (audio is the cross modality of video, and vice versa) sigma when
+                calculating the cross attention modulation parameters. `True` is the newer (e.g. LTX-2.3) behavior;
+                `False` is the legacy LTX-2.0 behavior.
             output_type (`str`, *optional*, defaults to `"pil"`):
                 The output format of the generate image. Choose between
                 [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
@@ -1210,6 +1215,7 @@ def __call__(
                         isolate_modalities=False,
                         spatio_temporal_guidance_blocks=None,
                         perturbation_mask=None,
+                        use_cross_timestep=use_cross_timestep,
                         attention_kwargs=attention_kwargs,
                         return_dict=False,
                     )
@@ -1275,6 +1281,7 @@ def __call__(
                             # Use STG at given blocks to perturb model
                             spatio_temporal_guidance_blocks=spatio_temporal_guidance_blocks,
                             perturbation_mask=None,
+                            use_cross_timestep=use_cross_timestep,
                             attention_kwargs=attention_kwargs,
                             return_dict=False,
                         )
@@ -1311,6 +1318,7 @@ def __call__(
                             isolate_modalities=True,
                             spatio_temporal_guidance_blocks=None,
                             perturbation_mask=None,
+                            use_cross_timestep=use_cross_timestep,
                             attention_kwargs=attention_kwargs,
                             return_dict=False,
                         )

From f78c3dae5a342a6c06361e5125c0e87931a17ddb Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Thu, 12 Mar 2026 10:46:05 +0100
Subject: [PATCH 26/41] Fix RMS norm bug for LTX-2.3 text connectors

---
 src/diffusers/pipelines/ltx2/connectors.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/diffusers/pipelines/ltx2/connectors.py b/src/diffusers/pipelines/ltx2/connectors.py
index 3f721a2cfbf9..a49de4083342 100644
--- a/src/diffusers/pipelines/ltx2/connectors.py
+++ b/src/diffusers/pipelines/ltx2/connectors.py
@@ -80,7 +80,7 @@ def per_layer_masked_mean_norm(
 
 def per_token_rms_norm(text_encoder_hidden_states: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
     variance = torch.mean(text_encoder_hidden_states**2, dim=2, keepdim=True)
-    norm_text_encoder_hidden_states = text_encoder_hidden_states + torch.rsqrt(variance + eps)
+    norm_text_encoder_hidden_states = text_encoder_hidden_states * torch.rsqrt(variance + eps)
     return norm_text_encoder_hidden_states
 
 

From 63b3c9f223e63c4ee84683878885f1edd7f38966 Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Thu, 12 Mar 2026 12:06:19 +0100
Subject: [PATCH 27/41] Perform guidance rescale in sample (x0) space following
 original code

---
 src/diffusers/pipelines/ltx2/pipeline_ltx2.py | 28 +++++++++++++++----
 .../pipelines/ltx2/pipeline_ltx2_condition.py | 28 +++++++++++++++----
 .../ltx2/pipeline_ltx2_image2video.py         | 28 +++++++++++++++----
 3 files changed, 66 insertions(+), 18 deletions(-)

diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
index 07321518be66..25d578b89884 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
@@ -1272,20 +1272,36 @@ def __call__(
 
                 # Apply LTX-2.X guidance rescaling
                 if self.guidance_rescale > 0:
+                    # Convert from velocity to sample (x0) prediction
+                    video_guided_x0 = latents - noise_pred_video_g * self.scheduler.sigmas[i]
+                    video_cond_x0 = latents - noise_pred_video * self.scheduler.sigmas[i]
+
+                    # Apply guidance rescaling in sample (x0) space, following original code
                     video_rescale = self.guidance_rescale
-                    cond_std = noise_pred_video.std(dim=list(range(1, noise_pred_video.ndim)), keepdim=True)
-                    guided_std = noise_pred_video_g.std(dim=list(range(1, noise_pred_video_g.ndim)), keepdim=True)
+                    cond_std = video_cond_x0.std(dim=list(range(1, video_cond_x0.ndim)), keepdim=True)
+                    guided_std = video_guided_x0.std(dim=list(range(1, video_guided_x0.ndim)), keepdim=True)
                     rescale_factor = video_rescale * (cond_std / guided_std) + (1 - video_rescale)
-                    noise_pred_video = noise_pred_video_g * rescale_factor
+                    video_guided_x0 = video_guided_x0 * rescale_factor
+
+                    # Convert back to velocity space for scheduler
+                    noise_pred_video = (latents - video_guided_x0) / self.scheduler.sigmas[i]
                 else:
                     noise_pred_video = noise_pred_video_g
 
                 if self.audio_guidance_rescale > 0:
+                    # Convert from velocity to sample (x0) prediction
+                    audio_guided_x0 = audio_latents - noise_pred_audio_g * audio_scheduler.sigmas[i]
+                    audio_cond_x0 = audio_latents - noise_pred_audio * audio_scheduler.sigmas[i]
+
+                    # Apply guidance rescaling in sample (x0) space, following original code
                     audio_rescale = self.audio_guidance_rescale
-                    cond_std = noise_pred_audio.std(dim=list(range(1, noise_pred_audio.ndim)), keepdim=True)
-                    guided_std = noise_pred_audio_g.std(dim=list(range(1, noise_pred_audio_g.ndim)), keepdim=True)
+                    cond_std = audio_cond_x0.std(dim=list(range(1, audio_cond_x0.ndim)), keepdim=True)
+                    guided_std = audio_guided_x0.std(dim=list(range(1, audio_guided_x0.ndim)), keepdim=True)
                     rescale_factor = audio_rescale * (cond_std / guided_std) + (1 - audio_rescale)
-                    noise_pred_audio = noise_pred_audio_g * rescale_factor
+                    audio_guided_x0 = audio_guided_x0 * rescale_factor
+
+                    # Convert back to velocity space for scheduler
+                    noise_pred_audio = (audio_latents - audio_guided_x0) / audio_scheduler.sigmas[i]
                 else:
                     noise_pred_audio = noise_pred_audio_g
 
diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py
index dd1b45f8de66..a8aaefd66f21 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py
@@ -1507,20 +1507,36 @@ def __call__(
 
                 # Apply LTX-2.X guidance rescaling
                 if self.guidance_rescale > 0:
+                    # Convert from velocity to sample (x0) prediction
+                    video_guided_x0 = latents - noise_pred_video_g * self.scheduler.sigmas[i]
+                    video_cond_x0 = latents - noise_pred_video * self.scheduler.sigmas[i]
+
+                    # Apply guidance rescaling in sample (x0) space, following original code
                     video_rescale = self.guidance_rescale
-                    cond_std = noise_pred_video.std(dim=list(range(1, noise_pred_video.ndim)), keepdim=True)
-                    guided_std = noise_pred_video_g.std(dim=list(range(1, noise_pred_video_g.ndim)), keepdim=True)
+                    cond_std = video_cond_x0.std(dim=list(range(1, video_cond_x0.ndim)), keepdim=True)
+                    guided_std = video_guided_x0.std(dim=list(range(1, video_guided_x0.ndim)), keepdim=True)
                     rescale_factor = video_rescale * (cond_std / guided_std) + (1 - video_rescale)
-                    noise_pred_video = noise_pred_video_g * rescale_factor
+                    video_guided_x0 = video_guided_x0 * rescale_factor
+
+                    # Convert back to velocity space for scheduler
+                    noise_pred_video = (latents - video_guided_x0) / self.scheduler.sigmas[i]
                 else:
                     noise_pred_video = noise_pred_video_g
 
                 if self.audio_guidance_rescale > 0:
+                    # Convert from velocity to sample (x0) prediction
+                    audio_guided_x0 = audio_latents - noise_pred_audio_g * audio_scheduler.sigmas[i]
+                    audio_cond_x0 = audio_latents - noise_pred_audio * audio_scheduler.sigmas[i]
+
+                    # Apply guidance rescaling in sample (x0) space, following original code
                     audio_rescale = self.audio_guidance_rescale
-                    cond_std = noise_pred_audio.std(dim=list(range(1, noise_pred_audio.ndim)), keepdim=True)
-                    guided_std = noise_pred_audio_g.std(dim=list(range(1, noise_pred_audio_g.ndim)), keepdim=True)
+                    cond_std = audio_cond_x0.std(dim=list(range(1, audio_cond_x0.ndim)), keepdim=True)
+                    guided_std = audio_guided_x0.std(dim=list(range(1, audio_guided_x0.ndim)), keepdim=True)
                     rescale_factor = audio_rescale * (cond_std / guided_std) + (1 - audio_rescale)
-                    noise_pred_audio = noise_pred_audio_g * rescale_factor
+                    audio_guided_x0 = audio_guided_x0 * rescale_factor
+
+                    # Convert back to velocity space for scheduler
+                    noise_pred_audio = (audio_latents - audio_guided_x0) / audio_scheduler.sigmas[i]
                 else:
                     noise_pred_audio = noise_pred_audio_g
 
diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
index df233978e209..f3c734e8a267 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
@@ -1340,20 +1340,36 @@ def __call__(
 
                 # Apply LTX-2.X guidance rescaling
                 if self.guidance_rescale > 0:
+                    # Convert from velocity to sample (x0) prediction
+                    video_guided_x0 = latents - noise_pred_video_g * self.scheduler.sigmas[i]
+                    video_cond_x0 = latents - noise_pred_video * self.scheduler.sigmas[i]
+
+                    # Apply guidance rescaling in sample (x0) space, following original code
                     video_rescale = self.guidance_rescale
-                    cond_std = noise_pred_video.std(dim=list(range(1, noise_pred_video.ndim)), keepdim=True)
-                    guided_std = noise_pred_video_g.std(dim=list(range(1, noise_pred_video_g.ndim)), keepdim=True)
+                    cond_std = video_cond_x0.std(dim=list(range(1, video_cond_x0.ndim)), keepdim=True)
+                    guided_std = video_guided_x0.std(dim=list(range(1, video_guided_x0.ndim)), keepdim=True)
                     rescale_factor = video_rescale * (cond_std / guided_std) + (1 - video_rescale)
-                    noise_pred_video = noise_pred_video_g * rescale_factor
+                    video_guided_x0 = video_guided_x0 * rescale_factor
+
+                    # Convert back to velocity space for scheduler
+                    noise_pred_video = (latents - video_guided_x0) / self.scheduler.sigmas[i]
                 else:
                     noise_pred_video = noise_pred_video_g
 
                 if self.audio_guidance_rescale > 0:
+                    # Convert from velocity to sample (x0) prediction
+                    audio_guided_x0 = audio_latents - noise_pred_audio_g * audio_scheduler.sigmas[i]
+                    audio_cond_x0 = audio_latents - noise_pred_audio * audio_scheduler.sigmas[i]
+
+                    # Apply guidance rescaling in sample (x0) space, following original code
                     audio_rescale = self.audio_guidance_rescale
-                    cond_std = noise_pred_audio.std(dim=list(range(1, noise_pred_audio.ndim)), keepdim=True)
-                    guided_std = noise_pred_audio_g.std(dim=list(range(1, noise_pred_audio_g.ndim)), keepdim=True)
+                    cond_std = audio_cond_x0.std(dim=list(range(1, audio_cond_x0.ndim)), keepdim=True)
+                    guided_std = audio_guided_x0.std(dim=list(range(1, audio_guided_x0.ndim)), keepdim=True)
                     rescale_factor = audio_rescale * (cond_std / guided_std) + (1 - audio_rescale)
-                    noise_pred_audio = noise_pred_audio_g * rescale_factor
+                    audio_guided_x0 = audio_guided_x0 * rescale_factor
+
+                    # Convert back to velocity space for scheduler
+                    noise_pred_audio = (audio_latents - audio_guided_x0) / audio_scheduler.sigmas[i]
                 else:
                     noise_pred_audio = noise_pred_audio_g
 

From 6188af22150e6327229808576c24facca5cd0712 Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Fri, 13 Mar 2026 04:44:02 +0100
Subject: [PATCH 28/41] Support LTX-2.3 Latent Spatial Upsampler model

---
 scripts/convert_ltx2_to_diffusers.py          | 28 ++++++++++++++++---
 .../pipelines/ltx2/latent_upsampler.py        |  5 ++--
 2 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/scripts/convert_ltx2_to_diffusers.py b/scripts/convert_ltx2_to_diffusers.py
index 7ef322e3039a..9a97ad3079dc 100644
--- a/scripts/convert_ltx2_to_diffusers.py
+++ b/scripts/convert_ltx2_to_diffusers.py
@@ -840,6 +840,18 @@ def get_ltx2_spatial_latent_upsampler_config(version: str):
             "spatial_upsample": True,
             "temporal_upsample": False,
             "rational_spatial_scale": 2.0,
+            "use_rational_resampler": True,
+        }
+    elif version == "2.3":
+        config = {
+            "in_channels": 128,
+            "mid_channels": 1024,
+            "num_blocks_per_stage": 4,
+            "dims": 3,
+            "spatial_upsample": True,
+            "temporal_upsample": False,
+            "rational_spatial_scale": 2.0,
+            "use_rational_resampler": False,
         }
     else:
         raise ValueError(f"Unsupported version: {version}")
@@ -1006,6 +1018,12 @@ def none_or_str(value: str):
     parser.add_argument("--text_encoder_dtype", type=str, default="bf16", choices=["fp32", "fp16", "bf16"])
 
     parser.add_argument("--output_path", type=str, required=True, help="Path where converted model should be saved")
+    parser.add_argument(
+        "--upsample_output_path",
+        type=str,
+        default=None,
+        help="Path where converted upsampling pipeline should be saved",
+    )
 
     return parser.parse_args()
 
@@ -1141,10 +1159,12 @@ def main(args):
     if args.upsample_pipeline:
         pipe = LTX2LatentUpsamplePipeline(vae=vae, latent_upsampler=latent_upsampler)
 
-        # Put latent upsampling pipeline in its own subdirectory so it doesn't mess with the full pipeline
-        pipe.save_pretrained(
-            os.path.join(args.output_path, "upsample_pipeline"), safe_serialization=True, max_shard_size="5GB"
-        )
+        # As two diffusers pipelines cannot be in the same directory, save the upsampling pipeline to its own directory
+        if args.upsample_output_path:
+            upsample_output_path = args.upsample_output_path
+        else:
+            upsample_output_path = args.output_path
+        pipe.save_pretrained(upsample_output_path, safe_serialization=True, max_shard_size="5GB")
 
 
 if __name__ == "__main__":
diff --git a/src/diffusers/pipelines/ltx2/latent_upsampler.py b/src/diffusers/pipelines/ltx2/latent_upsampler.py
index f6c589a70ab6..329ced36d45b 100644
--- a/src/diffusers/pipelines/ltx2/latent_upsampler.py
+++ b/src/diffusers/pipelines/ltx2/latent_upsampler.py
@@ -195,7 +195,8 @@ def __init__(
         dims: int = 3,
         spatial_upsample: bool = True,
         temporal_upsample: bool = False,
-        rational_spatial_scale: float | None = 2.0,
+        rational_spatial_scale: float = 2.0,
+        use_rational_resampler: bool = True,
     ):
         super().__init__()
 
@@ -220,7 +221,7 @@ def __init__(
                 PixelShuffleND(3),
             )
         elif spatial_upsample:
-            if rational_spatial_scale is not None:
+            if use_rational_resampler:
                 self.upsampler = SpatialRationalResampler(mid_channels=mid_channels, scale=rational_spatial_scale)
             else:
                 self.upsampler = torch.nn.Sequential(

From 89f8cc43841795e0cde944d2c07acc6842e07acb Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Fri, 13 Mar 2026 05:42:49 +0100
Subject: [PATCH 29/41] Support LTX-2.3 distilled LoRA

---
 src/diffusers/loaders/lora_conversion_utils.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/diffusers/loaders/lora_conversion_utils.py b/src/diffusers/loaders/lora_conversion_utils.py
index 0895d5223e13..67d2e6fefaa9 100644
--- a/src/diffusers/loaders/lora_conversion_utils.py
+++ b/src/diffusers/loaders/lora_conversion_utils.py
@@ -2156,6 +2156,9 @@ def _convert_non_diffusers_ltx2_lora_to_diffusers(state_dict, non_diffusers_pref
             "scale_shift_table_a2v_ca_audio": "audio_a2v_cross_attn_scale_shift_table",
             "q_norm": "norm_q",
             "k_norm": "norm_k",
+            # LTX-2.3
+            "audio_prompt_adaln_single": "audio_prompt_adaln",
+            "prompt_adaln_single": "prompt_adaln",
         }
     else:
         rename_dict = {"aggregate_embed": "text_proj_in"}

From f1a812aa6a31add855e7f3732d58ea2e51ac44eb Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Sat, 14 Mar 2026 02:02:20 +0100
Subject: [PATCH 30/41] Support LTX-2.3 Distilled checkpoint

---
 scripts/convert_ltx2_to_diffusers.py | 28 ++++++++++++++++++++--------
 1 file changed, 20 insertions(+), 8 deletions(-)

diff --git a/scripts/convert_ltx2_to_diffusers.py b/scripts/convert_ltx2_to_diffusers.py
index 9a97ad3079dc..3e7ae4af2ddc 100644
--- a/scripts/convert_ltx2_to_diffusers.py
+++ b/scripts/convert_ltx2_to_diffusers.py
@@ -1134,14 +1134,26 @@ def main(args):
             latent_upsampler.save_pretrained(os.path.join(args.output_path, "latent_upsampler"))
 
     if args.full_pipeline:
-        scheduler = FlowMatchEulerDiscreteScheduler(
-            use_dynamic_shifting=True,
-            base_shift=0.95,
-            max_shift=2.05,
-            base_image_seq_len=1024,
-            max_image_seq_len=4096,
-            shift_terminal=0.1,
-        )
+        is_distilled_ckpt = "distilled" in args.combined_filename
+        if is_distilled_ckpt:
+            # Disable dynamic shifting and terminal shift so that distilled sigmas are used as-is
+            scheduler = FlowMatchEulerDiscreteScheduler(
+                use_dynamic_shifting=False,
+                base_shift=0.95,
+                max_shift=2.05,
+                base_image_seq_len=1024,
+                max_image_seq_len=4096,
+                shift_terminal=None,
+            )
+        else:
+            scheduler = FlowMatchEulerDiscreteScheduler(
+                use_dynamic_shifting=True,
+                base_shift=0.95,
+                max_shift=2.05,
+                base_image_seq_len=1024,
+                max_image_seq_len=4096,
+                shift_terminal=0.1,
+            )
 
         pipe = LTX2Pipeline(
             scheduler=scheduler,

From 145e8e48c8ba9ccbb1339200264b40346774cc89 Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Sat, 14 Mar 2026 09:31:14 +0100
Subject: [PATCH 31/41] Support LTX-2.3 prompt enhancement

---
 scripts/convert_ltx2_to_diffusers.py          | 12 ++-
 src/diffusers/pipelines/ltx2/pipeline_ltx2.py | 74 ++++++++++++++++-
 .../ltx2/pipeline_ltx2_image2video.py         | 82 ++++++++++++++++++-
 3 files changed, 163 insertions(+), 5 deletions(-)

diff --git a/scripts/convert_ltx2_to_diffusers.py b/scripts/convert_ltx2_to_diffusers.py
index 3e7ae4af2ddc..276a011f53ca 100644
--- a/scripts/convert_ltx2_to_diffusers.py
+++ b/scripts/convert_ltx2_to_diffusers.py
@@ -7,7 +7,7 @@
 import torch
 from accelerate import init_empty_weights
 from huggingface_hub import hf_hub_download
-from transformers import AutoTokenizer, Gemma3ForConditionalGeneration
+from transformers import AutoTokenizer, Gemma3ForConditionalGeneration, Gemma3Processor
 
 from diffusers import (
     AutoencoderKLLTX2Audio,
@@ -1010,6 +1010,11 @@ def none_or_str(value: str):
         action="store_true",
         help="Whether to save a latent upsampling pipeline",
     )
+    parser.add_argument(
+        "--add_processor",
+        action="store_true",
+        help="Whether to add a Gemma3Processor to the pipeline for prompt enhancement.",
+    )
 
     parser.add_argument("--vae_dtype", type=str, default="bf16", choices=["fp32", "fp16", "bf16"])
     parser.add_argument("--audio_vae_dtype", type=str, default="bf16", choices=["fp32", "fp16", "bf16"])
@@ -1120,6 +1125,11 @@ def main(args):
         if not args.full_pipeline:
             tokenizer.save_pretrained(os.path.join(args.output_path, "tokenizer"))
 
+        if args.add_processor:
+            processor = Gemma3Processor.from_pretrained(args.text_encoder_model_id)
+            if not args.full_pipeline:
+                processor.save_pretrained(os.path.join(args.output_path, "processor"))
+
     if args.latent_upsampler or args.upsample_pipeline:
         original_latent_upsampler_ckpt = load_hub_or_local_checkpoint(
             repo_id=args.original_state_dict_repo_id, filename=args.latent_upsampler_filename
diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
index 25d578b89884..ab8cfb5abc59 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
@@ -18,7 +18,7 @@
 
 import numpy as np
 import torch
-from transformers import Gemma3ForConditionalGeneration, GemmaTokenizer, GemmaTokenizerFast
+from transformers import Gemma3ForConditionalGeneration, Gemma3Processor, GemmaTokenizer, GemmaTokenizerFast
 
 from ...callbacks import MultiPipelineCallbacks, PipelineCallback
 from ...loaders import FromSingleFileMixin, LTX2LoraLoaderMixin
@@ -209,7 +209,7 @@ class LTX2Pipeline(DiffusionPipeline, FromSingleFileMixin, LTX2LoraLoaderMixin):
     """
 
     model_cpu_offload_seq = "text_encoder->connectors->transformer->vae->audio_vae->vocoder"
-    _optional_components = []
+    _optional_components = ["processor"]
     _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"]
 
     def __init__(
@@ -219,6 +219,7 @@ def __init__(
         audio_vae: AutoencoderKLLTX2Audio,
         text_encoder: Gemma3ForConditionalGeneration,
         tokenizer: GemmaTokenizer | GemmaTokenizerFast,
+        processor: Gemma3Processor,
         connectors: LTX2TextConnectors,
         transformer: LTX2VideoTransformer3DModel,
         vocoder: LTX2Vocoder | LTX2VocoderWithBWE,
@@ -230,6 +231,7 @@ def __init__(
             audio_vae=audio_vae,
             text_encoder=text_encoder,
             tokenizer=tokenizer,
+            processor=processor,
             connectors=connectors,
             transformer=transformer,
             vocoder=vocoder,
@@ -418,6 +420,46 @@ def encode_prompt(
 
         return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask
 
+    @torch.no_grad()
+    def enhance_prompt(
+        self,
+        prompt: str,
+        system_prompt: str,
+        max_new_tokens: int = 512,
+        seed: int = 10,
+        generation_kwargs: dict[str, Any] | None = None,
+        device: str | torch.device | None = None,
+    ):
+        """
+        Enhances the supplied `prompt` by generating a new prompt using the current text encoder (default is a
+        `transformers.Gemma3ForConditionalGeneration` model) from it and a system prompt.
+        """
+        device = device or self._execution_device
+        if generation_kwargs is None:
+            # Set to default generation kwargs
+            generation_kwargs = {"do_sample": True, "temperature": 0.7}
+
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": f"user prompt: {prompt}"},
+        ]
+        template = self.processor.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        model_inputs = self.processor(text=template, images=None, return_tensors="pt").to(device)
+        self.text_encoder.to(device)
+
+        # `transformers.GenerationMixin.generate` does not support using a `torch.Generator` to control randomness,
+        # so manually apply a seed for reproducible generation.
+        torch.manual_seed(seed)
+        generated_sequences = self.text_encoder.generate(
+            **model_inputs,
+            max_new_tokens=max_new_tokens,
+            **generation_kwargs,
+        )  # tensor of shape [batch_size, seq_len]
+
+        generated_ids = [seq[len(model_inputs.input_ids[i]) :] for i, seq in enumerate(generated_sequences)]
+        enhanced_prompt = self.processor.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
+        return enhanced_prompt
+
     def check_inputs(
         self,
         prompt,
@@ -776,6 +818,10 @@ def __call__(
         decode_timestep: float | list[float] = 0.0,
         decode_noise_scale: float | list[float] | None = None,
         use_cross_timestep: bool = False,
+        system_prompt: str | None = None,
+        prompt_max_new_tokens: int = 512,
+        prompt_enhancement_kwargs: dict[str, Any] | None = None,
+        prompt_enhancement_seed: int = 10,
         output_type: str = "pil",
         return_dict: bool = True,
         attention_kwargs: dict[str, Any] | None = None,
@@ -889,6 +935,20 @@ def __call__(
                 Whether to use the cross modality (audio is the cross modality of video, and vice versa) sigma when
                 calculating the cross attention modulation parameters. `True` is the newer (e.g. LTX-2.3) behavior;
                 `False` is the legacy LTX-2.0 behavior.
+            system_prompt (`str`, *optional*, defaults to `None`):
+                Optional system prompt to use for prompt enhancement. The system prompt will be used by the current
+                text encoder (by default, a `Gemma3ForConditionalGeneration` model) to generate an enhanced prompt from
+                the original `prompt` to condition generation. If not supplied, prompt enhancement will not be
+                performed.
+            prompt_max_new_tokens (`int`, *optional*, defaults to `512`):
+                The maximum number of new tokens to generate when performing prompt enhancement.
+            prompt_enhancement_kwargs (`dict[str, Any]`, *optional*, defaults to `None`):
+                Keyword arguments for `self.text_encoder.generate`. If not supplied, default arguments of
+                `do_sample=True` and `temperature=0.7` will be used. See
+                https://huggingface.co/docs/transformers/main/en/main_classes/text_generation#transformers.GenerationMixin.generate
+                for more details.
+            prompt_enhancement_seed (`int`, *optional*, default to `10`):
+                Random seed for any random operations during prompt enhancement.
             output_type (`str`, *optional*, defaults to `"pil"`):
                 The output format of the generate image. Choose between
                 [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
@@ -966,6 +1026,16 @@ def __call__(
         device = self._execution_device
 
         # 3. Prepare text embeddings
+        if system_prompt is not None and prompt is not None:
+            prompt = self.enhance_prompt(
+                prompt=prompt,
+                system_prompt=system_prompt,
+                max_new_tokens=prompt_max_new_tokens,
+                seed=prompt_enhancement_seed,
+                generation_kwargs=prompt_enhancement_kwargs,
+                device=device,
+            )
+
         (
             prompt_embeds,
             prompt_attention_mask,
diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
index f3c734e8a267..09a33b878ab5 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
@@ -18,7 +18,7 @@
 
 import numpy as np
 import torch
-from transformers import Gemma3ForConditionalGeneration, GemmaTokenizer, GemmaTokenizerFast
+from transformers import Gemma3ForConditionalGeneration, Gemma3Processor, GemmaTokenizer, GemmaTokenizerFast
 
 from ...callbacks import MultiPipelineCallbacks, PipelineCallback
 from ...image_processor import PipelineImageInput
@@ -212,7 +212,7 @@ class LTX2ImageToVideoPipeline(DiffusionPipeline, FromSingleFileMixin, LTX2LoraL
     """
 
     model_cpu_offload_seq = "text_encoder->connectors->transformer->vae->audio_vae->vocoder"
-    _optional_components = []
+    _optional_components = ["processor"]
     _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"]
 
     def __init__(
@@ -222,6 +222,7 @@ def __init__(
         audio_vae: AutoencoderKLLTX2Audio,
         text_encoder: Gemma3ForConditionalGeneration,
         tokenizer: GemmaTokenizer | GemmaTokenizerFast,
+        processor: Gemma3Processor,
         connectors: LTX2TextConnectors,
         transformer: LTX2VideoTransformer3DModel,
         vocoder: LTX2Vocoder | LTX2VocoderWithBWE,
@@ -233,6 +234,7 @@ def __init__(
             audio_vae=audio_vae,
             text_encoder=text_encoder,
             tokenizer=tokenizer,
+            processor=processor,
             connectors=connectors,
             transformer=transformer,
             vocoder=vocoder,
@@ -423,6 +425,53 @@ def encode_prompt(
 
         return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask
 
+    @torch.no_grad()
+    def enhance_prompt(
+        self,
+        image: PipelineImageInput,
+        prompt: str,
+        system_prompt: str,
+        max_new_tokens: int = 512,
+        seed: int = 10,
+        generation_kwargs: dict[str, Any] | None = None,
+        device: str | torch.device | None = None,
+    ):
+        """
+        Enhances the supplied `prompt` by generating a new prompt using the current text encoder (default is a
+        `transformers.Gemma3ForConditionalGeneration` model) from it and a system prompt.
+        """
+        device = device or self._execution_device
+        if generation_kwargs is None:
+            # Set to default generation kwargs
+            generation_kwargs = {"do_sample": True, "temperature": 0.7}
+
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {
+                "role": "user",
+                "content": [
+                    {"type": "image"},
+                    {"type": "text", "text": f"User Raw Input Prompt: {prompt}."},
+                ],
+            },
+        ]
+        template = self.processor.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        model_inputs = self.processor(text=template, images=image, return_tensors="pt").to(device)
+        self.text_encoder.to(device)
+
+        # `transformers.GenerationMixin.generate` does not support using a `torch.Generator` to control randomness,
+        # so manually apply a seed for reproducible generation.
+        torch.manual_seed(seed)
+        generated_sequences = self.text_encoder.generate(
+            **model_inputs,
+            max_new_tokens=max_new_tokens,
+            **generation_kwargs,
+        )  # tensor of shape [batch_size, seq_len]
+
+        generated_ids = [seq[len(model_inputs.input_ids[i]) :] for i, seq in enumerate(generated_sequences)]
+        enhanced_prompt = self.processor.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
+        return enhanced_prompt
+
     # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline.check_inputs
     def check_inputs(
         self,
@@ -830,6 +879,10 @@ def __call__(
         decode_timestep: float | list[float] = 0.0,
         decode_noise_scale: float | list[float] | None = None,
         use_cross_timestep: bool = False,
+        system_prompt: str | None = None,
+        prompt_max_new_tokens: int = 512,
+        prompt_enhancement_kwargs: dict[str, Any] | None = None,
+        prompt_enhancement_seed: int = 10,
         output_type: str = "pil",
         return_dict: bool = True,
         attention_kwargs: dict[str, Any] | None = None,
@@ -945,6 +998,20 @@ def __call__(
                 Whether to use the cross modality (audio is the cross modality of video, and vice versa) sigma when
                 calculating the cross attention modulation parameters. `True` is the newer (e.g. LTX-2.3) behavior;
                 `False` is the legacy LTX-2.0 behavior.
+            system_prompt (`str`, *optional*, defaults to `None`):
+                Optional system prompt to use for prompt enhancement. The system prompt will be used by the current
+                text encoder (by default, a `Gemma3ForConditionalGeneration` model) to generate an enhanced prompt from
+                the original `prompt` to condition generation. If not supplied, prompt enhancement will not be
+                performed.
+            prompt_max_new_tokens (`int`, *optional*, defaults to `512`):
+                The maximum number of new tokens to generate when performing prompt enhancement.
+            prompt_enhancement_kwargs (`dict[str, Any]`, *optional*, defaults to `None`):
+                Keyword arguments for `self.text_encoder.generate`. If not supplied, default arguments of
+                `do_sample=True` and `temperature=0.7` will be used. See
+                https://huggingface.co/docs/transformers/main/en/main_classes/text_generation#transformers.GenerationMixin.generate
+                for more details.
+            prompt_enhancement_seed (`int`, *optional*, default to `10`):
+                Random seed for any random operations during prompt enhancement.
             output_type (`str`, *optional*, defaults to `"pil"`):
                 The output format of the generate image. Choose between
                 [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
@@ -1022,6 +1089,17 @@ def __call__(
         device = self._execution_device
 
         # 3. Prepare text embeddings
+        if system_prompt is not None and prompt is not None:
+            prompt = self.enhance_prompt(
+                image=image,
+                prompt=prompt,
+                system_prompt=system_prompt,
+                max_new_tokens=prompt_max_new_tokens,
+                seed=prompt_enhancement_seed,
+                generation_kwargs=prompt_enhancement_kwargs,
+                device=device,
+            )
+
         (
             prompt_embeds,
             prompt_attention_mask,

From 8a5807349dfdd46ba1ffddc024ede2e639fbb54a Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Sat, 14 Mar 2026 09:50:42 +0100
Subject: [PATCH 32/41] Make LTX-2.X processor non-required so that tests pass

---
 src/diffusers/pipelines/ltx2/pipeline_ltx2.py             | 4 ++--
 src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
index ab8cfb5abc59..c0aba10a4f3c 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
@@ -219,10 +219,10 @@ def __init__(
         audio_vae: AutoencoderKLLTX2Audio,
         text_encoder: Gemma3ForConditionalGeneration,
         tokenizer: GemmaTokenizer | GemmaTokenizerFast,
-        processor: Gemma3Processor,
         connectors: LTX2TextConnectors,
         transformer: LTX2VideoTransformer3DModel,
         vocoder: LTX2Vocoder | LTX2VocoderWithBWE,
+        processor: Gemma3Processor | None = None,
     ):
         super().__init__()
 
@@ -231,11 +231,11 @@ def __init__(
             audio_vae=audio_vae,
             text_encoder=text_encoder,
             tokenizer=tokenizer,
-            processor=processor,
             connectors=connectors,
             transformer=transformer,
             vocoder=vocoder,
             scheduler=scheduler,
+            processor=processor,
         )
 
         self.vae_spatial_compression_ratio = (
diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
index 09a33b878ab5..ac3ea98f18f7 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
@@ -222,10 +222,10 @@ def __init__(
         audio_vae: AutoencoderKLLTX2Audio,
         text_encoder: Gemma3ForConditionalGeneration,
         tokenizer: GemmaTokenizer | GemmaTokenizerFast,
-        processor: Gemma3Processor,
         connectors: LTX2TextConnectors,
         transformer: LTX2VideoTransformer3DModel,
         vocoder: LTX2Vocoder | LTX2VocoderWithBWE,
+        processor: Gemma3Processor | None = None,
     ):
         super().__init__()
 
@@ -234,11 +234,11 @@ def __init__(
             audio_vae=audio_vae,
             text_encoder=text_encoder,
             tokenizer=tokenizer,
-            processor=processor,
             connectors=connectors,
             transformer=transformer,
             vocoder=vocoder,
             scheduler=scheduler,
+            processor=processor,
         )
 
         self.vae_spatial_compression_ratio = (

From 93247a0a361068671268c4435cbaab94e9bfa6e4 Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Sat, 14 Mar 2026 10:02:50 +0100
Subject: [PATCH 33/41] Fix test_components_function tests for LTX2 T2V and I2V

---
 tests/pipelines/ltx2/test_ltx2.py             | 1 +
 tests/pipelines/ltx2/test_ltx2_image2video.py | 1 +
 2 files changed, 2 insertions(+)

diff --git a/tests/pipelines/ltx2/test_ltx2.py b/tests/pipelines/ltx2/test_ltx2.py
index 7d1a3bfc9987..0941ae550989 100644
--- a/tests/pipelines/ltx2/test_ltx2.py
+++ b/tests/pipelines/ltx2/test_ltx2.py
@@ -171,6 +171,7 @@ def get_dummy_components(self):
             "tokenizer": tokenizer,
             "connectors": connectors,
             "vocoder": vocoder,
+            "processor": None,
         }
 
         return components
diff --git a/tests/pipelines/ltx2/test_ltx2_image2video.py b/tests/pipelines/ltx2/test_ltx2_image2video.py
index 92c000c7bf7c..a0e4cb803084 100644
--- a/tests/pipelines/ltx2/test_ltx2_image2video.py
+++ b/tests/pipelines/ltx2/test_ltx2_image2video.py
@@ -171,6 +171,7 @@ def get_dummy_components(self):
             "tokenizer": tokenizer,
             "connectors": connectors,
             "vocoder": vocoder,
+            "processor": None,
         }
 
         return components

From 17b53f08661732caca6a546295950fc4b1696ad7 Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Sat, 14 Mar 2026 10:37:10 +0100
Subject: [PATCH 34/41] Fix LTX-2.3 Video VAE configuration bug causing pixel
 jitter

---
 scripts/convert_ltx2_to_diffusers.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/scripts/convert_ltx2_to_diffusers.py b/scripts/convert_ltx2_to_diffusers.py
index 276a011f53ca..f1556557889f 100644
--- a/scripts/convert_ltx2_to_diffusers.py
+++ b/scripts/convert_ltx2_to_diffusers.py
@@ -606,7 +606,7 @@ def get_ltx2_video_vae_config(
                 "decoder_inject_noise": (False, False, False, False, False),
                 "downsample_type": ("spatial", "temporal", "spatiotemporal", "spatiotemporal"),
                 "upsample_type": ("spatiotemporal", "spatiotemporal", "temporal", "spatial"),
-                "upsample_residual": (True, True, True, True),
+                "upsample_residual": (False, False, False, False),
                 "upsample_factor": (2, 2, 1, 2),
                 "timestep_conditioning": timestep_conditioning,
                 "patch_size": 4,

From c016ce56079e50920a5d26e6c63bf2b39bbcb60d Mon Sep 17 00:00:00 2001
From: dg845 <58458699+dg845@users.noreply.github.com>
Date: Mon, 16 Mar 2026 18:37:24 -0700
Subject: [PATCH 35/41] Apply suggestions from code review

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
---
 src/diffusers/models/transformers/transformer_ltx2.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/diffusers/models/transformers/transformer_ltx2.py b/src/diffusers/models/transformers/transformer_ltx2.py
index a47fb97bdfc9..59391b26c99c 100644
--- a/src/diffusers/models/transformers/transformer_ltx2.py
+++ b/src/diffusers/models/transformers/transformer_ltx2.py
@@ -680,7 +680,7 @@ def forward(
             )
             audio_shift_text_kv, audio_scale_text_kv = audio_prompt_ada_params
 
-        # 2.1. Video-Text Cross-Attention (Q: Video; K,V: Test)
+        # 2.1. Video-Text Cross-Attention (Q: Video; K,V: Text)
         norm_hidden_states = self.norm2(hidden_states)
         if self.video_cross_attn_adaln:
             norm_hidden_states = norm_hidden_states * (1 + scale_text_q) + shift_text_q

From 2feb460924b713b44363f31feca1a664bfee08b6 Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Tue, 17 Mar 2026 03:03:19 +0100
Subject: [PATCH 36/41] Refactor LTX-2.X Video VAE upsampler block init logic

---
 .../autoencoders/autoencoder_kl_ltx2.py       | 38 ++++++-------------
 1 file changed, 12 insertions(+), 26 deletions(-)

diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_ltx2.py b/src/diffusers/models/autoencoders/autoencoder_kl_ltx2.py
index 36cf73e4bb2a..f4f7d46628c8 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_ltx2.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_ltx2.py
@@ -615,35 +615,21 @@ def __init__(
             self.upsamplers = nn.ModuleList()
 
             if upsample_type == "spatial":
-                self.upsamplers.append(
-                    LTX2VideoUpsampler3d(
-                        in_channels=out_channels * upscale_factor,
-                        stride=(1, 2, 2),
-                        residual=upsample_residual,
-                        upscale_factor=upscale_factor,
-                        spatial_padding_mode=spatial_padding_mode,
-                    )
-                )
+                upsample_stride = (1, 2, 2)
             elif upsample_type == "temporal":
-                self.upsamplers.append(
-                    LTX2VideoUpsampler3d(
-                        in_channels=out_channels * upscale_factor,
-                        stride=(2, 1, 1),
-                        residual=upsample_residual,
-                        upscale_factor=upscale_factor,
-                        spatial_padding_mode=spatial_padding_mode,
-                    )
-                )
+                upsample_stride = (2, 1, 1)
             elif upsample_type == "spatiotemporal":
-                self.upsamplers.append(
-                    LTX2VideoUpsampler3d(
-                        in_channels=out_channels * upscale_factor,
-                        stride=(2, 2, 2),
-                        residual=upsample_residual,
-                        upscale_factor=upscale_factor,
-                        spatial_padding_mode=spatial_padding_mode,
-                    )
+                upsample_stride = (2, 2, 2)
+
+            self.upsamplers.append(
+                LTX2VideoUpsampler3d(
+                    in_channels=out_channels * upscale_factor,
+                    stride=upsample_stride,
+                    residual=upsample_residual,
+                    upscale_factor=upscale_factor,
+                    spatial_padding_mode=spatial_padding_mode,
                 )
+            )
 
         resnets = []
         for _ in range(num_layers):

From 27404098216df3a7466f8aa482a0db2b168a04d4 Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Wed, 18 Mar 2026 04:44:15 +0100
Subject: [PATCH 37/41] Refactor LTX-2.X guidance rescaling to use
 rescale_noise_cfg

---
 src/diffusers/pipelines/ltx2/pipeline_ltx2.py    | 16 ++++++----------
 .../pipelines/ltx2/pipeline_ltx2_condition.py    | 16 ++++++----------
 .../pipelines/ltx2/pipeline_ltx2_image2video.py  | 16 ++++++----------
 3 files changed, 18 insertions(+), 30 deletions(-)

diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
index c0aba10a4f3c..afd7700f38be 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
@@ -1347,11 +1347,9 @@ def __call__(
                     video_cond_x0 = latents - noise_pred_video * self.scheduler.sigmas[i]
 
                     # Apply guidance rescaling in sample (x0) space, following original code
-                    video_rescale = self.guidance_rescale
-                    cond_std = video_cond_x0.std(dim=list(range(1, video_cond_x0.ndim)), keepdim=True)
-                    guided_std = video_guided_x0.std(dim=list(range(1, video_guided_x0.ndim)), keepdim=True)
-                    rescale_factor = video_rescale * (cond_std / guided_std) + (1 - video_rescale)
-                    video_guided_x0 = video_guided_x0 * rescale_factor
+                    video_guided_x0 = rescale_noise_cfg(
+                        video_guided_x0, video_cond_x0, guidance_rescale=self.guidance_rescale
+                    )
 
                     # Convert back to velocity space for scheduler
                     noise_pred_video = (latents - video_guided_x0) / self.scheduler.sigmas[i]
@@ -1364,11 +1362,9 @@ def __call__(
                     audio_cond_x0 = audio_latents - noise_pred_audio * audio_scheduler.sigmas[i]
 
                     # Apply guidance rescaling in sample (x0) space, following original code
-                    audio_rescale = self.audio_guidance_rescale
-                    cond_std = audio_cond_x0.std(dim=list(range(1, audio_cond_x0.ndim)), keepdim=True)
-                    guided_std = audio_guided_x0.std(dim=list(range(1, audio_guided_x0.ndim)), keepdim=True)
-                    rescale_factor = audio_rescale * (cond_std / guided_std) + (1 - audio_rescale)
-                    audio_guided_x0 = audio_guided_x0 * rescale_factor
+                    audio_guided_x0 = rescale_noise_cfg(
+                        audio_guided_x0, audio_cond_x0, guidance_rescale=self.audio_guidance_rescale
+                    )
 
                     # Convert back to velocity space for scheduler
                     noise_pred_audio = (audio_latents - audio_guided_x0) / audio_scheduler.sigmas[i]
diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py
index a8aaefd66f21..eeb070be62a1 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py
@@ -1512,11 +1512,9 @@ def __call__(
                     video_cond_x0 = latents - noise_pred_video * self.scheduler.sigmas[i]
 
                     # Apply guidance rescaling in sample (x0) space, following original code
-                    video_rescale = self.guidance_rescale
-                    cond_std = video_cond_x0.std(dim=list(range(1, video_cond_x0.ndim)), keepdim=True)
-                    guided_std = video_guided_x0.std(dim=list(range(1, video_guided_x0.ndim)), keepdim=True)
-                    rescale_factor = video_rescale * (cond_std / guided_std) + (1 - video_rescale)
-                    video_guided_x0 = video_guided_x0 * rescale_factor
+                    video_guided_x0 = rescale_noise_cfg(
+                        video_guided_x0, video_cond_x0, guidance_rescale=self.guidance_rescale
+                    )
 
                     # Convert back to velocity space for scheduler
                     noise_pred_video = (latents - video_guided_x0) / self.scheduler.sigmas[i]
@@ -1529,11 +1527,9 @@ def __call__(
                     audio_cond_x0 = audio_latents - noise_pred_audio * audio_scheduler.sigmas[i]
 
                     # Apply guidance rescaling in sample (x0) space, following original code
-                    audio_rescale = self.audio_guidance_rescale
-                    cond_std = audio_cond_x0.std(dim=list(range(1, audio_cond_x0.ndim)), keepdim=True)
-                    guided_std = audio_guided_x0.std(dim=list(range(1, audio_guided_x0.ndim)), keepdim=True)
-                    rescale_factor = audio_rescale * (cond_std / guided_std) + (1 - audio_rescale)
-                    audio_guided_x0 = audio_guided_x0 * rescale_factor
+                    audio_guided_x0 = rescale_noise_cfg(
+                        audio_guided_x0, audio_cond_x0, guidance_rescale=self.audio_guidance_rescale
+                    )
 
                     # Convert back to velocity space for scheduler
                     noise_pred_audio = (audio_latents - audio_guided_x0) / audio_scheduler.sigmas[i]
diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
index ac3ea98f18f7..56bd92b05ade 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
@@ -1423,11 +1423,9 @@ def __call__(
                     video_cond_x0 = latents - noise_pred_video * self.scheduler.sigmas[i]
 
                     # Apply guidance rescaling in sample (x0) space, following original code
-                    video_rescale = self.guidance_rescale
-                    cond_std = video_cond_x0.std(dim=list(range(1, video_cond_x0.ndim)), keepdim=True)
-                    guided_std = video_guided_x0.std(dim=list(range(1, video_guided_x0.ndim)), keepdim=True)
-                    rescale_factor = video_rescale * (cond_std / guided_std) + (1 - video_rescale)
-                    video_guided_x0 = video_guided_x0 * rescale_factor
+                    video_guided_x0 = rescale_noise_cfg(
+                        video_guided_x0, video_cond_x0, guidance_rescale=self.guidance_rescale
+                    )
 
                     # Convert back to velocity space for scheduler
                     noise_pred_video = (latents - video_guided_x0) / self.scheduler.sigmas[i]
@@ -1440,11 +1438,9 @@ def __call__(
                     audio_cond_x0 = audio_latents - noise_pred_audio * audio_scheduler.sigmas[i]
 
                     # Apply guidance rescaling in sample (x0) space, following original code
-                    audio_rescale = self.audio_guidance_rescale
-                    cond_std = audio_cond_x0.std(dim=list(range(1, audio_cond_x0.ndim)), keepdim=True)
-                    guided_std = audio_guided_x0.std(dim=list(range(1, audio_guided_x0.ndim)), keepdim=True)
-                    rescale_factor = audio_rescale * (cond_std / guided_std) + (1 - audio_rescale)
-                    audio_guided_x0 = audio_guided_x0 * rescale_factor
+                    audio_guided_x0 = rescale_noise_cfg(
+                        audio_guided_x0, audio_cond_x0, guidance_rescale=self.audio_guidance_rescale
+                    )
 
                     # Convert back to velocity space for scheduler
                     noise_pred_audio = (audio_latents - audio_guided_x0) / audio_scheduler.sigmas[i]

From 5d8b6342310a6add06bd6efd6533d9eadc942887 Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Thu, 19 Mar 2026 01:02:22 +0100
Subject: [PATCH 38/41] Use generator initial seed to control prompt
 enhancement if available

---
 src/diffusers/pipelines/ltx2/pipeline_ltx2.py             | 5 +++++
 src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py | 5 +++++
 2 files changed, 10 insertions(+)

diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
index afd7700f38be..0bf9724c65e1 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
@@ -427,6 +427,7 @@ def enhance_prompt(
         system_prompt: str,
         max_new_tokens: int = 512,
         seed: int = 10,
+        generator: torch.Generator | None = None,
         generation_kwargs: dict[str, Any] | None = None,
         device: str | torch.device | None = None,
     ):
@@ -449,6 +450,9 @@ def enhance_prompt(
 
         # `transformers.GenerationMixin.generate` does not support using a `torch.Generator` to control randomness,
         # so manually apply a seed for reproducible generation.
+        if generator is not None:
+            # Overwrite seed to generator's initial seed
+            seed = generator.initial_seed()
         torch.manual_seed(seed)
         generated_sequences = self.text_encoder.generate(
             **model_inputs,
@@ -1032,6 +1036,7 @@ def __call__(
                 system_prompt=system_prompt,
                 max_new_tokens=prompt_max_new_tokens,
                 seed=prompt_enhancement_seed,
+                generator=generator,
                 generation_kwargs=prompt_enhancement_kwargs,
                 device=device,
             )
diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
index 56bd92b05ade..da0351f0a5fa 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
@@ -433,6 +433,7 @@ def enhance_prompt(
         system_prompt: str,
         max_new_tokens: int = 512,
         seed: int = 10,
+        generator: torch.Generator | None = None,
         generation_kwargs: dict[str, Any] | None = None,
         device: str | torch.device | None = None,
     ):
@@ -461,6 +462,9 @@ def enhance_prompt(
 
         # `transformers.GenerationMixin.generate` does not support using a `torch.Generator` to control randomness,
         # so manually apply a seed for reproducible generation.
+        if generator is not None:
+            # Overwrite seed to generator's initial seed
+            seed = generator.initial_seed()
         torch.manual_seed(seed)
         generated_sequences = self.text_encoder.generate(
             **model_inputs,
@@ -1096,6 +1100,7 @@ def __call__(
                 system_prompt=system_prompt,
                 max_new_tokens=prompt_max_new_tokens,
                 seed=prompt_enhancement_seed,
+                generator=generator,
                 generation_kwargs=prompt_enhancement_kwargs,
                 device=device,
             )

From b0723de8f7420e566794922c1b04b5ca9d14e9b7 Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Thu, 19 Mar 2026 01:43:38 +0100
Subject: [PATCH 39/41] Remove self attention mask logic as it is not used in
 any current pipelines

---
 .../models/transformers/transformer_ltx2.py   | 38 ++-----------------
 src/diffusers/pipelines/ltx2/pipeline_ltx2.py |  6 ---
 .../pipelines/ltx2/pipeline_ltx2_condition.py |  6 ---
 .../ltx2/pipeline_ltx2_image2video.py         |  6 ---
 4 files changed, 4 insertions(+), 52 deletions(-)

diff --git a/src/diffusers/models/transformers/transformer_ltx2.py b/src/diffusers/models/transformers/transformer_ltx2.py
index 59391b26c99c..a4915ccfb96a 100644
--- a/src/diffusers/models/transformers/transformer_ltx2.py
+++ b/src/diffusers/models/transformers/transformer_ltx2.py
@@ -1331,8 +1331,6 @@ def forward(
         audio_sigma: torch.Tensor | None = None,
         encoder_attention_mask: torch.Tensor | None = None,
         audio_encoder_attention_mask: torch.Tensor | None = None,
-        self_attention_mask: torch.Tensor | None = None,
-        audio_self_attention_mask: torch.Tensor | None = None,
         num_frames: int | None = None,
         height: int | None = None,
         width: int | None = None,
@@ -1376,8 +1374,6 @@ def forward(
                 Optional multiplicative text attention mask of shape `(batch_size, text_seq_len)`.
             audio_encoder_attention_mask (`torch.Tensor`, *optional*):
                 Optional multiplicative text attention mask of shape `(batch_size, text_seq_len)` for audio modeling.
-            self_attention_mask (`torch.Tensor`, *optional*):
-                Optional multiplicative self-attention mask of shape `(batch_size, seq_len, seq_len)`.
             num_frames (`int`, *optional*):
                 The number of latent video frames. Used if calculating the video coordinates for RoPE.
             height (`int`, *optional*):
@@ -1434,32 +1430,6 @@ def forward(
             audio_encoder_attention_mask = (1 - audio_encoder_attention_mask.to(audio_hidden_states.dtype)) * -10000.0
             audio_encoder_attention_mask = audio_encoder_attention_mask.unsqueeze(1)
 
-        if self_attention_mask is not None and self_attention_mask.ndim == 3:
-            # Convert to additive attention mask in log-space where 0 (masked) values get mapped to a large negative
-            # number and positive values are mapped to their logarithm.
-            dtype_finfo = torch.finfo(hidden_states.dtype)
-            additive_self_attn_mask = torch.full_like(self_attention_mask, dtype_finfo.min, dtype=hidden_states.dtype)
-            unmasked_entries = self_attention_mask > 0
-            if torch.any(unmasked_entries):
-                additive_self_attn_mask[unmasked_entries] = torch.log(
-                    self_attention_mask[unmasked_entries].clamp(min=dtype_finfo.tiny)
-                ).to(hidden_states.dtype)
-            self_attention_mask = additive_self_attn_mask.unsqueeze(1)  # [batch_size, 1, seq_len, seq_len]
-
-        if audio_self_attention_mask is not None and audio_self_attention_mask.ndim == 3:
-            # Convert to additive attention mask in log-space where 0 (masked) values get mapped to a large negative
-            # number and positive values are mapped to their logarithm.
-            dtype_finfo = torch.finfo(audio_hidden_states.dtype)
-            additive_self_attn_mask = torch.full_like(
-                audio_self_attention_mask, dtype_finfo.min, dtype=audio_hidden_states.dtype
-            )
-            unmasked_entries = audio_self_attention_mask > 0
-            if torch.any(unmasked_entries):
-                additive_self_attn_mask[unmasked_entries] = torch.log(
-                    audio_self_attention_mask[unmasked_entries].clamp(min=dtype_finfo.tiny)
-                ).to(audio_hidden_states.dtype)
-            audio_self_attention_mask = additive_self_attn_mask.unsqueeze(1)  # [batch_size, 1, seq_len, seq_len]
-
         batch_size = hidden_states.size(0)
 
         # 1. Prepare RoPE positional embeddings
@@ -1599,8 +1569,8 @@ def forward(
                     audio_cross_attn_rotary_emb,
                     encoder_attention_mask,
                     audio_encoder_attention_mask,
-                    self_attention_mask,
-                    audio_self_attention_mask,
+                    None,  # self_attention_mask
+                    None,  # audio_self_attention_mask
                     None,  # a2v_cross_attention_mask
                     None,  # v2a_cross_attention_mask
                     not isolate_modalities,  # use_a2v_cross_attention
@@ -1628,8 +1598,8 @@ def forward(
                     ca_audio_rotary_emb=audio_cross_attn_rotary_emb,
                     encoder_attention_mask=encoder_attention_mask,
                     audio_encoder_attention_mask=audio_encoder_attention_mask,
-                    self_attention_mask=self_attention_mask,
-                    audio_self_attention_mask=audio_self_attention_mask,
+                    self_attention_mask=None,
+                    audio_self_attention_mask=None,
                     a2v_cross_attention_mask=None,
                     v2a_cross_attention_mask=None,
                     use_a2v_cross_attention=not isolate_modalities,
diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
index 0bf9724c65e1..9cef6d1fb5c3 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
@@ -1213,8 +1213,6 @@ def __call__(
                         sigma=timestep,  # Used by LTX-2.3
                         encoder_attention_mask=connector_attention_mask,
                         audio_encoder_attention_mask=connector_attention_mask,
-                        self_attention_mask=None,
-                        audio_self_attention_mask=None,
                         num_frames=latent_num_frames,
                         height=latent_height,
                         width=latent_width,
@@ -1276,8 +1274,6 @@ def __call__(
                             sigma=timestep,  # Used by LTX-2.3
                             encoder_attention_mask=prompt_attn_mask,
                             audio_encoder_attention_mask=prompt_attn_mask,
-                            self_attention_mask=None,
-                            audio_self_attention_mask=None,
                             num_frames=latent_num_frames,
                             height=latent_height,
                             width=latent_width,
@@ -1312,8 +1308,6 @@ def __call__(
                             sigma=timestep,  # Used by LTX-2.3
                             encoder_attention_mask=prompt_attn_mask,
                             audio_encoder_attention_mask=prompt_attn_mask,
-                            self_attention_mask=None,
-                            audio_self_attention_mask=None,
                             num_frames=latent_num_frames,
                             height=latent_height,
                             width=latent_width,
diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py
index eeb070be62a1..a889a3481c18 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py
@@ -1370,8 +1370,6 @@ def __call__(
                         sigma=timestep,  # Used by LTX-2.3
                         encoder_attention_mask=connector_attention_mask,
                         audio_encoder_attention_mask=connector_attention_mask,
-                        self_attention_mask=None,
-                        audio_self_attention_mask=None,
                         num_frames=latent_num_frames,
                         height=latent_height,
                         width=latent_width,
@@ -1435,8 +1433,6 @@ def __call__(
                             sigma=timestep,  # Used by LTX-2.3
                             encoder_attention_mask=prompt_attn_mask,
                             audio_encoder_attention_mask=prompt_attn_mask,
-                            self_attention_mask=None,
-                            audio_self_attention_mask=None,
                             num_frames=latent_num_frames,
                             height=latent_height,
                             width=latent_width,
@@ -1472,8 +1468,6 @@ def __call__(
                             sigma=timestep,  # Used by LTX-2.3
                             encoder_attention_mask=prompt_attn_mask,
                             audio_encoder_attention_mask=prompt_attn_mask,
-                            self_attention_mask=None,
-                            audio_self_attention_mask=None,
                             num_frames=latent_num_frames,
                             height=latent_height,
                             width=latent_width,
diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
index da0351f0a5fa..cfd775915e39 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
@@ -1286,8 +1286,6 @@ def __call__(
                         sigma=timestep,  # Used by LTX-2.3
                         encoder_attention_mask=connector_attention_mask,
                         audio_encoder_attention_mask=connector_attention_mask,
-                        self_attention_mask=None,
-                        audio_self_attention_mask=None,
                         num_frames=latent_num_frames,
                         height=latent_height,
                         width=latent_width,
@@ -1351,8 +1349,6 @@ def __call__(
                             sigma=timestep,  # Used by LTX-2.3
                             encoder_attention_mask=prompt_attn_mask,
                             audio_encoder_attention_mask=prompt_attn_mask,
-                            self_attention_mask=None,
-                            audio_self_attention_mask=None,
                             num_frames=latent_num_frames,
                             height=latent_height,
                             width=latent_width,
@@ -1388,8 +1384,6 @@ def __call__(
                             sigma=timestep,  # Used by LTX-2.3
                             encoder_attention_mask=prompt_attn_mask,
                             audio_encoder_attention_mask=prompt_attn_mask,
-                            self_attention_mask=None,
-                            audio_self_attention_mask=None,
                             num_frames=latent_num_frames,
                             height=latent_height,
                             width=latent_width,

From 67a9ce33c5a210b77d20f224a2e8adca9513c7d2 Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Thu, 19 Mar 2026 02:18:29 +0100
Subject: [PATCH 40/41] Commit fixes suggested by claude code (guidance in
 sample (x0) space, denormalize after timestep conditioning)

---
 src/diffusers/pipelines/ltx2/pipeline_ltx2.py | 80 ++++++++++++-----
 .../pipelines/ltx2/pipeline_ltx2_condition.py | 88 +++++++++++++------
 .../ltx2/pipeline_ltx2_image2video.py         | 80 ++++++++++++-----
 3 files changed, 171 insertions(+), 77 deletions(-)

diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
index 9cef6d1fb5c3..bcc530191268 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
@@ -713,7 +713,6 @@ def prepare_audio_latents(
             latents = self._create_noised_state(latents, noise_scale, generator)
             return latents.to(device=device, dtype=dtype)
 
-        # TODO: confirm whether this logic is correct
         latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio
 
         shape = (batch_size, num_channels_latents, audio_latent_length, latent_mel_bins)
@@ -728,6 +727,24 @@ def prepare_audio_latents(
         latents = self._pack_audio_latents(latents)
         return latents
 
+    def convert_velocity_to_x0(
+        self, sample: torch.Tensor, denoised_output: torch.Tensor, step_idx: int, scheduler: Any | None = None
+    ) -> torch.Tensor:
+        if scheduler is None:
+            scheduler = self.scheduler
+
+        sample_x0 = sample - denoised_output * scheduler.sigmas[step_idx]
+        return sample_x0
+
+    def convert_x0_to_velocity(
+        self, sample: torch.Tensor, denoised_output: torch.Tensor, step_idx: int, scheduler: Any | None = None
+    ) -> torch.Tensor:
+        if scheduler is None:
+            scheduler = self.scheduler
+
+        sample_v = (sample - denoised_output) / scheduler.sigmas[step_idx]
+        return sample_v
+
     @property
     def guidance_scale(self):
         return self._guidance_scale
@@ -1232,10 +1249,18 @@ def __call__(
 
                 if self.do_classifier_free_guidance:
                     noise_pred_video_uncond_text, noise_pred_video = noise_pred_video.chunk(2)
+                    noise_pred_video = self.convert_velocity_to_x0(latents, noise_pred_video, i, self.scheduler)
+                    noise_pred_video_uncond_text = self.convert_velocity_to_x0(
+                        latents, noise_pred_video_uncond_text, i, self.scheduler
+                    )
                     # Use delta formulation as it works more nicely with multiple guidance terms
                     video_cfg_delta = (self.guidance_scale - 1) * (noise_pred_video - noise_pred_video_uncond_text)
 
                     noise_pred_audio_uncond_text, noise_pred_audio = noise_pred_audio.chunk(2)
+                    noise_pred_audio = self.convert_velocity_to_x0(audio_latents, noise_pred_audio, i, audio_scheduler)
+                    noise_pred_audio_uncond_text = self.convert_velocity_to_x0(
+                        audio_latents, noise_pred_audio_uncond_text, i, audio_scheduler
+                    )
                     audio_cfg_delta = (self.audio_guidance_scale - 1) * (
                         noise_pred_audio - noise_pred_audio_uncond_text
                     )
@@ -1263,6 +1288,9 @@ def __call__(
                     video_pos_ids = video_coords
                     audio_pos_ids = audio_coords
 
+                    noise_pred_video = self.convert_velocity_to_x0(latents, noise_pred_video, i, self.scheduler)
+                    noise_pred_audio = self.convert_velocity_to_x0(audio_latents, noise_pred_audio, i, audio_scheduler)
+
                 if self.do_spatio_temporal_guidance:
                     with self.transformer.cache_context("uncond_stg"):
                         noise_pred_video_uncond_stg, noise_pred_audio_uncond_stg = self.transformer(
@@ -1291,6 +1319,12 @@ def __call__(
                         )
                     noise_pred_video_uncond_stg = noise_pred_video_uncond_stg.float()
                     noise_pred_audio_uncond_stg = noise_pred_audio_uncond_stg.float()
+                    noise_pred_video_uncond_stg = self.convert_velocity_to_x0(
+                        latents, noise_pred_video_uncond_stg, i, self.scheduler
+                    )
+                    noise_pred_audio_uncond_stg = self.convert_velocity_to_x0(
+                        audio_latents, noise_pred_audio_uncond_stg, i, audio_scheduler
+                    )
 
                     video_stg_delta = self.stg_scale * (noise_pred_video - noise_pred_video_uncond_stg)
                     audio_stg_delta = self.audio_stg_scale * (noise_pred_audio - noise_pred_audio_uncond_stg)
@@ -1325,6 +1359,12 @@ def __call__(
                         )
                     noise_pred_video_uncond_modality = noise_pred_video_uncond_modality.float()
                     noise_pred_audio_uncond_modality = noise_pred_audio_uncond_modality.float()
+                    noise_pred_video_uncond_modality = self.convert_velocity_to_x0(
+                        latents, noise_pred_video_uncond_modality, i, self.scheduler
+                    )
+                    noise_pred_audio_uncond_modality = self.convert_velocity_to_x0(
+                        audio_latents, noise_pred_audio_uncond_modality, i, audio_scheduler
+                    )
 
                     video_modality_delta = (self.modality_scale - 1) * (
                         noise_pred_video - noise_pred_video_uncond_modality
@@ -1341,35 +1381,23 @@ def __call__(
 
                 # Apply LTX-2.X guidance rescaling
                 if self.guidance_rescale > 0:
-                    # Convert from velocity to sample (x0) prediction
-                    video_guided_x0 = latents - noise_pred_video_g * self.scheduler.sigmas[i]
-                    video_cond_x0 = latents - noise_pred_video * self.scheduler.sigmas[i]
-
-                    # Apply guidance rescaling in sample (x0) space, following original code
-                    video_guided_x0 = rescale_noise_cfg(
-                        video_guided_x0, video_cond_x0, guidance_rescale=self.guidance_rescale
+                    noise_pred_video = rescale_noise_cfg(
+                        noise_pred_video_g, noise_pred_video, guidance_rescale=self.guidance_rescale
                     )
-
-                    # Convert back to velocity space for scheduler
-                    noise_pred_video = (latents - video_guided_x0) / self.scheduler.sigmas[i]
                 else:
                     noise_pred_video = noise_pred_video_g
 
                 if self.audio_guidance_rescale > 0:
-                    # Convert from velocity to sample (x0) prediction
-                    audio_guided_x0 = audio_latents - noise_pred_audio_g * audio_scheduler.sigmas[i]
-                    audio_cond_x0 = audio_latents - noise_pred_audio * audio_scheduler.sigmas[i]
-
-                    # Apply guidance rescaling in sample (x0) space, following original code
-                    audio_guided_x0 = rescale_noise_cfg(
-                        audio_guided_x0, audio_cond_x0, guidance_rescale=self.audio_guidance_rescale
+                    noise_pred_audio = rescale_noise_cfg(
+                        noise_pred_audio_g, noise_pred_audio, guidance_rescale=self.audio_guidance_rescale
                     )
-
-                    # Convert back to velocity space for scheduler
-                    noise_pred_audio = (audio_latents - audio_guided_x0) / audio_scheduler.sigmas[i]
                 else:
                     noise_pred_audio = noise_pred_audio_g
 
+                # Convert back to velocity for scheduler
+                noise_pred_video = self.convert_x0_to_velocity(latents, noise_pred_video, i, self.scheduler)
+                noise_pred_audio = self.convert_x0_to_velocity(audio_latents, noise_pred_audio, i, audio_scheduler)
+
                 # compute the previous noisy sample x_t -> x_t-1
                 latents = self.scheduler.step(noise_pred_video, t, latents, return_dict=False)[0]
                 # NOTE: for now duplicate scheduler for audio latents in case self.scheduler sets internal state in
@@ -1400,9 +1428,6 @@ def __call__(
             self.transformer_spatial_patch_size,
             self.transformer_temporal_patch_size,
         )
-        latents = self._denormalize_latents(
-            latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor
-        )
 
         audio_latents = self._denormalize_audio_latents(
             audio_latents, self.audio_vae.latents_mean, self.audio_vae.latents_std
@@ -1410,6 +1435,9 @@ def __call__(
         audio_latents = self._unpack_audio_latents(audio_latents, audio_num_frames, num_mel_bins=latent_mel_bins)
 
         if output_type == "latent":
+            latents = self._denormalize_latents(
+                latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor
+            )
             video = latents
             audio = audio_latents
         else:
@@ -1432,6 +1460,10 @@ def __call__(
                 ]
                 latents = (1 - decode_noise_scale) * latents + decode_noise_scale * noise
 
+            latents = self._denormalize_latents(
+                latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor
+            )
+
             latents = latents.to(self.vae.dtype)
             video = self.vae.decode(latents, timestep, return_dict=False)[0]
             video = self.video_processor.postprocess_video(video, output_type=output_type)
diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py
index a889a3481c18..e6336a2db85f 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py
@@ -916,6 +916,24 @@ def prepare_audio_latents(
         latents = self._pack_audio_latents(latents)
         return latents
 
+    def convert_velocity_to_x0(
+        self, sample: torch.Tensor, denoised_output: torch.Tensor, step_idx: int, scheduler: Any | None = None
+    ) -> torch.Tensor:
+        if scheduler is None:
+            scheduler = self.scheduler
+
+        sample_x0 = sample - denoised_output * scheduler.sigmas[step_idx]
+        return sample_x0
+
+    def convert_x0_to_velocity(
+        self, sample: torch.Tensor, denoised_output: torch.Tensor, step_idx: int, scheduler: Any | None = None
+    ) -> torch.Tensor:
+        if scheduler is None:
+            scheduler = self.scheduler
+
+        sample_v = (sample - denoised_output) / scheduler.sigmas[step_idx]
+        return sample_v
+
     @property
     def guidance_scale(self):
         return self._guidance_scale
@@ -1389,10 +1407,18 @@ def __call__(
 
                 if self.do_classifier_free_guidance:
                     noise_pred_video_uncond_text, noise_pred_video = noise_pred_video.chunk(2)
+                    noise_pred_video = self.convert_velocity_to_x0(latents, noise_pred_video, i, self.scheduler)
+                    noise_pred_video_uncond_text = self.convert_velocity_to_x0(
+                        latents, noise_pred_video_uncond_text, i, self.scheduler
+                    )
                     # Use delta formulation as it works more nicely with multiple guidance terms
                     video_cfg_delta = (self.guidance_scale - 1) * (noise_pred_video - noise_pred_video_uncond_text)
 
                     noise_pred_audio_uncond_text, noise_pred_audio = noise_pred_audio.chunk(2)
+                    noise_pred_audio = self.convert_velocity_to_x0(audio_latents, noise_pred_audio, i, audio_scheduler)
+                    noise_pred_audio_uncond_text = self.convert_velocity_to_x0(
+                        audio_latents, noise_pred_audio_uncond_text, i, audio_scheduler
+                    )
                     audio_cfg_delta = (self.audio_guidance_scale - 1) * (
                         noise_pred_audio - noise_pred_audio_uncond_text
                     )
@@ -1421,6 +1447,9 @@ def __call__(
                     video_pos_ids = video_coords
                     audio_pos_ids = audio_coords
 
+                    noise_pred_video = self.convert_velocity_to_x0(latents, noise_pred_video, i, self.scheduler)
+                    noise_pred_audio = self.convert_velocity_to_x0(audio_latents, noise_pred_audio, i, audio_scheduler)
+
                 if self.do_spatio_temporal_guidance:
                     with self.transformer.cache_context("uncond_stg"):
                         noise_pred_video_uncond_stg, noise_pred_audio_uncond_stg = self.transformer(
@@ -1450,6 +1479,12 @@ def __call__(
                         )
                     noise_pred_video_uncond_stg = noise_pred_video_uncond_stg.float()
                     noise_pred_audio_uncond_stg = noise_pred_audio_uncond_stg.float()
+                    noise_pred_video_uncond_stg = self.convert_velocity_to_x0(
+                        latents, noise_pred_video_uncond_stg, i, self.scheduler
+                    )
+                    noise_pred_audio_uncond_stg = self.convert_velocity_to_x0(
+                        audio_latents, noise_pred_audio_uncond_stg, i, audio_scheduler
+                    )
 
                     video_stg_delta = self.stg_scale * (noise_pred_video - noise_pred_video_uncond_stg)
                     audio_stg_delta = self.audio_stg_scale * (noise_pred_audio - noise_pred_audio_uncond_stg)
@@ -1485,6 +1520,12 @@ def __call__(
                         )
                     noise_pred_video_uncond_modality = noise_pred_video_uncond_modality.float()
                     noise_pred_audio_uncond_modality = noise_pred_audio_uncond_modality.float()
+                    noise_pred_video_uncond_modality = self.convert_velocity_to_x0(
+                        latents, noise_pred_video_uncond_modality, i, self.scheduler
+                    )
+                    noise_pred_audio_uncond_modality = self.convert_velocity_to_x0(
+                        audio_latents, noise_pred_audio_uncond_modality, i, audio_scheduler
+                    )
 
                     video_modality_delta = (self.modality_scale - 1) * (
                         noise_pred_video - noise_pred_video_uncond_modality
@@ -1501,51 +1542,36 @@ def __call__(
 
                 # Apply LTX-2.X guidance rescaling
                 if self.guidance_rescale > 0:
-                    # Convert from velocity to sample (x0) prediction
-                    video_guided_x0 = latents - noise_pred_video_g * self.scheduler.sigmas[i]
-                    video_cond_x0 = latents - noise_pred_video * self.scheduler.sigmas[i]
-
-                    # Apply guidance rescaling in sample (x0) space, following original code
-                    video_guided_x0 = rescale_noise_cfg(
-                        video_guided_x0, video_cond_x0, guidance_rescale=self.guidance_rescale
+                    noise_pred_video = rescale_noise_cfg(
+                        noise_pred_video_g, noise_pred_video, guidance_rescale=self.guidance_rescale
                     )
-
-                    # Convert back to velocity space for scheduler
-                    noise_pred_video = (latents - video_guided_x0) / self.scheduler.sigmas[i]
                 else:
                     noise_pred_video = noise_pred_video_g
 
                 if self.audio_guidance_rescale > 0:
-                    # Convert from velocity to sample (x0) prediction
-                    audio_guided_x0 = audio_latents - noise_pred_audio_g * audio_scheduler.sigmas[i]
-                    audio_cond_x0 = audio_latents - noise_pred_audio * audio_scheduler.sigmas[i]
-
-                    # Apply guidance rescaling in sample (x0) space, following original code
-                    audio_guided_x0 = rescale_noise_cfg(
-                        audio_guided_x0, audio_cond_x0, guidance_rescale=self.audio_guidance_rescale
+                    noise_pred_audio = rescale_noise_cfg(
+                        noise_pred_audio_g, noise_pred_audio, guidance_rescale=self.audio_guidance_rescale
                     )
-
-                    # Convert back to velocity space for scheduler
-                    noise_pred_audio = (audio_latents - audio_guided_x0) / audio_scheduler.sigmas[i]
                 else:
                     noise_pred_audio = noise_pred_audio_g
 
                 # NOTE: use only the first chunk of conditioning mask in case it is duplicated for CFG
                 bsz = noise_pred_video.size(0)
-                sigma = self.scheduler.sigmas[i]
-                # Convert the noise_pred_video velocity model prediction into a sample (x0) prediction
-                denoised_sample = latents - noise_pred_video * sigma
                 # Apply the (packed) conditioning mask to the denoised (x0) sample and clean conditioning. The
                 # conditioning mask contains conditioning strengths from 0 (always use denoised sample) to 1 (always
                 # use conditions), with intermediate values specifying how strongly to follow the conditions.
+                # NOTE: this operation should be applied in sample (x0) space and not velocity space (which is the
+                # space the denoising model outputs are in)
                 denoised_sample_cond = (
-                    denoised_sample * (1 - conditioning_mask[:bsz]) + clean_latents.float() * conditioning_mask[:bsz]
+                    noise_pred_video * (1 - conditioning_mask[:bsz]) + clean_latents.float() * conditioning_mask[:bsz]
                 ).to(noise_pred_video.dtype)
+
                 # Convert the denoised (x0) sample back to a velocity for the scheduler
-                denoised_latents_cond = ((latents - denoised_sample_cond) / sigma).to(noise_pred_video.dtype)
+                noise_pred_video = self.convert_x0_to_velocity(latents, denoised_sample_cond, i, self.scheduler)
+                noise_pred_audio = self.convert_x0_to_velocity(audio_latents, noise_pred_audio, i, audio_scheduler)
 
                 # Compute the previous noisy sample x_t -> x_t-1
-                latents = self.scheduler.step(denoised_latents_cond, t, latents, return_dict=False)[0]
+                latents = self.scheduler.step(noise_pred_video, t, latents, return_dict=False)[0]
 
                 # NOTE: for now duplicate scheduler for audio latents in case self.scheduler sets internal state in
                 # the step method (such as _step_index)
@@ -1575,9 +1601,6 @@ def __call__(
             self.transformer_spatial_patch_size,
             self.transformer_temporal_patch_size,
         )
-        latents = self._denormalize_latents(
-            latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor
-        )
 
         audio_latents = self._denormalize_audio_latents(
             audio_latents, self.audio_vae.latents_mean, self.audio_vae.latents_std
@@ -1585,6 +1608,9 @@ def __call__(
         audio_latents = self._unpack_audio_latents(audio_latents, audio_num_frames, num_mel_bins=latent_mel_bins)
 
         if output_type == "latent":
+            latents = self._denormalize_latents(
+                latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor
+            )
             video = latents
             audio = audio_latents
         else:
@@ -1607,6 +1633,10 @@ def __call__(
                 ]
                 latents = (1 - decode_noise_scale) * latents + decode_noise_scale * noise
 
+            latents = self._denormalize_latents(
+                latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor
+            )
+
             latents = latents.to(self.vae.dtype)
             video = self.vae.decode(latents, timestep, return_dict=False)[0]
             video = self.video_processor.postprocess_video(video, output_type=output_type)
diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
index cfd775915e39..aaea14f4952c 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
@@ -773,7 +773,6 @@ def prepare_audio_latents(
             latents = self._create_noised_state(latents, noise_scale, generator)
             return latents.to(device=device, dtype=dtype)
 
-        # TODO: confirm whether this logic is correct
         latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio
 
         shape = (batch_size, num_channels_latents, audio_latent_length, latent_mel_bins)
@@ -788,6 +787,24 @@ def prepare_audio_latents(
         latents = self._pack_audio_latents(latents)
         return latents
 
+    def convert_velocity_to_x0(
+        self, sample: torch.Tensor, denoised_output: torch.Tensor, step_idx: int, scheduler: Any | None = None
+    ) -> torch.Tensor:
+        if scheduler is None:
+            scheduler = self.scheduler
+
+        sample_x0 = sample - denoised_output * scheduler.sigmas[step_idx]
+        return sample_x0
+
+    def convert_x0_to_velocity(
+        self, sample: torch.Tensor, denoised_output: torch.Tensor, step_idx: int, scheduler: Any | None = None
+    ) -> torch.Tensor:
+        if scheduler is None:
+            scheduler = self.scheduler
+
+        sample_v = (sample - denoised_output) / scheduler.sigmas[step_idx]
+        return sample_v
+
     @property
     def guidance_scale(self):
         return self._guidance_scale
@@ -1305,10 +1322,18 @@ def __call__(
 
                 if self.do_classifier_free_guidance:
                     noise_pred_video_uncond_text, noise_pred_video = noise_pred_video.chunk(2)
+                    noise_pred_video = self.convert_velocity_to_x0(latents, noise_pred_video, i, self.scheduler)
+                    noise_pred_video_uncond_text = self.convert_velocity_to_x0(
+                        latents, noise_pred_video_uncond_text, i, self.scheduler
+                    )
                     # Use delta formulation as it works more nicely with multiple guidance terms
                     video_cfg_delta = (self.guidance_scale - 1) * (noise_pred_video - noise_pred_video_uncond_text)
 
                     noise_pred_audio_uncond_text, noise_pred_audio = noise_pred_audio.chunk(2)
+                    noise_pred_audio = self.convert_velocity_to_x0(audio_latents, noise_pred_audio, i, audio_scheduler)
+                    noise_pred_audio_uncond_text = self.convert_velocity_to_x0(
+                        audio_latents, noise_pred_audio_uncond_text, i, audio_scheduler
+                    )
                     audio_cfg_delta = (self.audio_guidance_scale - 1) * (
                         noise_pred_audio - noise_pred_audio_uncond_text
                     )
@@ -1337,6 +1362,9 @@ def __call__(
                     video_pos_ids = video_coords
                     audio_pos_ids = audio_coords
 
+                    noise_pred_video = self.convert_velocity_to_x0(latents, noise_pred_video, i, self.scheduler)
+                    noise_pred_audio = self.convert_velocity_to_x0(audio_latents, noise_pred_audio, i, audio_scheduler)
+
                 if self.do_spatio_temporal_guidance:
                     with self.transformer.cache_context("uncond_stg"):
                         noise_pred_video_uncond_stg, noise_pred_audio_uncond_stg = self.transformer(
@@ -1366,6 +1394,12 @@ def __call__(
                         )
                     noise_pred_video_uncond_stg = noise_pred_video_uncond_stg.float()
                     noise_pred_audio_uncond_stg = noise_pred_audio_uncond_stg.float()
+                    noise_pred_video_uncond_stg = self.convert_velocity_to_x0(
+                        latents, noise_pred_video_uncond_stg, i, self.scheduler
+                    )
+                    noise_pred_audio_uncond_stg = self.convert_velocity_to_x0(
+                        audio_latents, noise_pred_audio_uncond_stg, i, audio_scheduler
+                    )
 
                     video_stg_delta = self.stg_scale * (noise_pred_video - noise_pred_video_uncond_stg)
                     audio_stg_delta = self.audio_stg_scale * (noise_pred_audio - noise_pred_audio_uncond_stg)
@@ -1401,6 +1435,12 @@ def __call__(
                         )
                     noise_pred_video_uncond_modality = noise_pred_video_uncond_modality.float()
                     noise_pred_audio_uncond_modality = noise_pred_audio_uncond_modality.float()
+                    noise_pred_video_uncond_modality = self.convert_velocity_to_x0(
+                        latents, noise_pred_video_uncond_modality, i, self.scheduler
+                    )
+                    noise_pred_audio_uncond_modality = self.convert_velocity_to_x0(
+                        audio_latents, noise_pred_audio_uncond_modality, i, audio_scheduler
+                    )
 
                     video_modality_delta = (self.modality_scale - 1) * (
                         noise_pred_video - noise_pred_video_uncond_modality
@@ -1417,32 +1457,16 @@ def __call__(
 
                 # Apply LTX-2.X guidance rescaling
                 if self.guidance_rescale > 0:
-                    # Convert from velocity to sample (x0) prediction
-                    video_guided_x0 = latents - noise_pred_video_g * self.scheduler.sigmas[i]
-                    video_cond_x0 = latents - noise_pred_video * self.scheduler.sigmas[i]
-
-                    # Apply guidance rescaling in sample (x0) space, following original code
-                    video_guided_x0 = rescale_noise_cfg(
-                        video_guided_x0, video_cond_x0, guidance_rescale=self.guidance_rescale
+                    noise_pred_video = rescale_noise_cfg(
+                        noise_pred_video_g, noise_pred_video, guidance_rescale=self.guidance_rescale
                     )
-
-                    # Convert back to velocity space for scheduler
-                    noise_pred_video = (latents - video_guided_x0) / self.scheduler.sigmas[i]
                 else:
                     noise_pred_video = noise_pred_video_g
 
                 if self.audio_guidance_rescale > 0:
-                    # Convert from velocity to sample (x0) prediction
-                    audio_guided_x0 = audio_latents - noise_pred_audio_g * audio_scheduler.sigmas[i]
-                    audio_cond_x0 = audio_latents - noise_pred_audio * audio_scheduler.sigmas[i]
-
-                    # Apply guidance rescaling in sample (x0) space, following original code
-                    audio_guided_x0 = rescale_noise_cfg(
-                        audio_guided_x0, audio_cond_x0, guidance_rescale=self.audio_guidance_rescale
+                    noise_pred_audio = rescale_noise_cfg(
+                        noise_pred_audio_g, noise_pred_audio, guidance_rescale=self.audio_guidance_rescale
                     )
-
-                    # Convert back to velocity space for scheduler
-                    noise_pred_audio = (audio_latents - audio_guided_x0) / audio_scheduler.sigmas[i]
                 else:
                     noise_pred_audio = noise_pred_audio_g
 
@@ -1464,6 +1488,10 @@ def __call__(
                     self.transformer_temporal_patch_size,
                 )
 
+                # Convert back to velocity for scheduler
+                noise_pred_video = self.convert_x0_to_velocity(latents, noise_pred_video, i, self.scheduler)
+                noise_pred_audio = self.convert_x0_to_velocity(audio_latents, noise_pred_audio, i, audio_scheduler)
+
                 noise_pred_video = noise_pred_video[:, :, 1:]
                 noise_latents = latents[:, :, 1:]
                 pred_latents = self.scheduler.step(noise_pred_video, t, noise_latents, return_dict=False)[0]
@@ -1501,9 +1529,6 @@ def __call__(
             self.transformer_spatial_patch_size,
             self.transformer_temporal_patch_size,
         )
-        latents = self._denormalize_latents(
-            latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor
-        )
 
         audio_latents = self._denormalize_audio_latents(
             audio_latents, self.audio_vae.latents_mean, self.audio_vae.latents_std
@@ -1511,6 +1536,9 @@ def __call__(
         audio_latents = self._unpack_audio_latents(audio_latents, audio_num_frames, num_mel_bins=latent_mel_bins)
 
         if output_type == "latent":
+            latents = self._denormalize_latents(
+                latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor
+            )
             video = latents
             audio = audio_latents
         else:
@@ -1533,6 +1561,10 @@ def __call__(
                 ]
                 latents = (1 - decode_noise_scale) * latents + decode_noise_scale * noise
 
+            latents = self._denormalize_latents(
+                latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor
+            )
+
             latents = latents.to(self.vae.dtype)
             video = self.vae.decode(latents, timestep, return_dict=False)[0]
             video = self.video_processor.postprocess_video(video, output_type=output_type)

From 4cbedd7c14aa85b0f36fba44fd604932ec4d8c55 Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Thu, 19 Mar 2026 02:32:03 +0100
Subject: [PATCH 41/41] Use constant shift following original code

---
 src/diffusers/pipelines/ltx2/pipeline_ltx2.py             | 4 ++--
 src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py   | 4 ++--
 src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py | 4 ++--
 3 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
index bcc530191268..73ebac0f173c 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2.py
@@ -1105,7 +1105,7 @@ def __call__(
                 raise ValueError(
                     f"Provided `latents` tensor has shape {latents.shape}, but the expected shape is either [batch_size, seq_len, num_features] or [batch_size, latent_dim, latent_frames, latent_height, latent_width]."
                 )
-        video_sequence_length = latent_num_frames * latent_height * latent_width
+        # video_sequence_length = latent_num_frames * latent_height * latent_width
 
         num_channels_latents = self.transformer.config.in_channels
         latents = self.prepare_latents(
@@ -1162,7 +1162,7 @@ def __call__(
         # 5. Prepare timesteps
         sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
         mu = calculate_shift(
-            video_sequence_length,
+            self.scheduler.config.get("max_image_seq_len", 4096),
             self.scheduler.config.get("base_image_seq_len", 1024),
             self.scheduler.config.get("max_image_seq_len", 4096),
             self.scheduler.config.get("base_shift", 0.95),
diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py
index e6336a2db85f..a80d011015cf 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py
@@ -1268,7 +1268,7 @@ def __call__(
                 "Got latents of shape [batch_size, latent_dim, latent_frames, latent_height, latent_width], `latent_num_frames`, `latent_height`, `latent_width` will be inferred."
             )
             _, _, latent_num_frames, latent_height, latent_width = latents.shape  # [B, C, F, H, W]
-        video_sequence_length = latent_num_frames * latent_height * latent_width
+        # video_sequence_length = latent_num_frames * latent_height * latent_width
 
         num_channels_latents = self.transformer.config.in_channels
         latents, conditioning_mask, clean_latents = self.prepare_latents(
@@ -1318,7 +1318,7 @@ def __call__(
         # 5. Prepare timesteps
         sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
         mu = calculate_shift(
-            video_sequence_length,
+            self.scheduler.config.get("max_image_seq_len", 4096),
             self.scheduler.config.get("base_image_seq_len", 1024),
             self.scheduler.config.get("max_image_seq_len", 4096),
             self.scheduler.config.get("base_shift", 0.95),
diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
index aaea14f4952c..997bfd9fc9dc 100644
--- a/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
+++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py
@@ -1169,7 +1169,7 @@ def __call__(
                 raise ValueError(
                     f"Provided `latents` tensor has shape {latents.shape}, but the expected shape is either [batch_size, seq_len, num_features] or [batch_size, latent_dim, latent_frames, latent_height, latent_width]."
                 )
-        video_sequence_length = latent_num_frames * latent_height * latent_width
+        # video_sequence_length = latent_num_frames * latent_height * latent_width
 
         if latents is None:
             image = self.video_processor.preprocess(image, height=height, width=width)
@@ -1233,7 +1233,7 @@ def __call__(
         # 5. Prepare timesteps
         sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
         mu = calculate_shift(
-            video_sequence_length,
+            self.scheduler.config.get("max_image_seq_len", 4096),
             self.scheduler.config.get("base_image_seq_len", 1024),
             self.scheduler.config.get("max_image_seq_len", 4096),
             self.scheduler.config.get("base_shift", 0.95),