LTX2.3 changes (except vocoder)

Author: prisha0510

src/maxdiffusion/configs/ltx2_3_video.yml

@@ -24,15 +24,18 @@ save_config_to_gcs: False
 max_sequence_length: 1024
 sampler: "from_checkpoint"
 
-# Generation parameters
+# Generation parameters (aligned with Diffusers LTX-2.3 docs: use_cross_timestep, modality + audio CFG)
 global_batch_size_to_train_on: 1
-num_inference_steps: 40
-guidance_scale: 4.0
-audio_guidance_scale: 4.0
+num_inference_steps: 30
+guidance_scale: 3.0
+guidance_rescale: 0.7
+audio_guidance_scale: 7.0
+audio_guidance_rescale: 0.7
 stg_scale: 1.0
 audio_stg_scale: 1.0
-modality_scale: 1.0
-audio_modality_scale: 1.0
+modality_scale: 3.0
+audio_modality_scale: 3.0
+use_cross_timestep: true
 spatio_temporal_guidance_blocks: [28]
 fps: 24
 pipeline_type: multi-scale
@@ -42,9 +45,9 @@ height: 512
 width: 768
 decode_timestep: 0.05
 decode_noise_scale: 0.025
+noise_scale: 0.0
 num_frames: 121
 quantization: "int8"
-seed: 10
 #parallelism
 mesh_axes: ['data', 'fsdp', 'context', 'tensor']
 logical_axis_rules: [
@@ -109,5 +112,5 @@ bwd_quantization_calibration_method: "absmax"
 qwix_module_path: ".*"
 jit_initializers: True 
 enable_single_replica_ckpt_restoring: False
-seed: 0
+seed: 10
 audio_format: "s16"

src/maxdiffusion/generate_ltx2.py

@@ -93,16 +93,20 @@ def call_pipeline(config, pipeline, prompt, negative_prompt):
       num_frames=config.num_frames,
       num_inference_steps=config.num_inference_steps,
       guidance_scale=guidance_scale,
+      guidance_rescale=getattr(config, "guidance_rescale", 0.0),
       generator=generator,
       frame_rate=getattr(config, "fps", 24.0),
       decode_timestep=getattr(config, "decode_timestep", 0.0),
       decode_noise_scale=getattr(config, "decode_noise_scale", None),
       max_sequence_length=getattr(config, "max_sequence_length", 1024),
       audio_guidance_scale=getattr(config, "audio_guidance_scale", None),
+      audio_guidance_rescale=getattr(config, "audio_guidance_rescale", None),
       stg_scale=getattr(config, "stg_scale", 0.0),
       audio_stg_scale=getattr(config, "audio_stg_scale", None),
       modality_scale=getattr(config, "modality_scale", 1.0),
       audio_modality_scale=getattr(config, "audio_modality_scale", None),
+      use_cross_timestep=getattr(config, "use_cross_timestep", None),
+      noise_scale=getattr(config, "noise_scale", 1.0),
       dtype=jnp.bfloat16 if getattr(config, "activations_dtype", "bfloat16") == "bfloat16" else jnp.float32,
   )
   return out

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -112,11 +112,13 @@ def __init__(
       v2a_attention_kernel: str = "dot_product",
       flash_block_sizes: BlockSizes = None,
       flash_min_seq_length: int = 4096,
+      perturbed_attn: bool = False,
   ):
     self.dim = dim
     self.norm_eps = norm_eps
     self.norm_elementwise_affine = norm_elementwise_affine
     self.attention_kernel = attention_kernel
+    self.perturbed_attn = perturbed_attn
 
     # 1. Self-Attention (video and audio)
     self.norm1 = nnx.RMSNorm(
@@ -370,11 +372,11 @@ def __call__(
       audio_rotary_emb: Optional[Tuple[jax.Array, jax.Array]] = None,
       ca_video_rotary_emb: Optional[Tuple[jax.Array, jax.Array]] = None,
       ca_audio_rotary_emb: Optional[Tuple[jax.Array, jax.Array]] = None,
-      attention_mask: Optional[jax.Array] = None,
       encoder_attention_mask: Optional[jax.Array] = None,
       audio_encoder_attention_mask: Optional[jax.Array] = None,
       a2v_cross_attention_mask: Optional[jax.Array] = None,
       v2a_cross_attention_mask: Optional[jax.Array] = None,
+      perturbation_mask: Optional[jax.Array] = None,
   ) -> Tuple[jax.Array, jax.Array]:
     batch_size = hidden_states.shape[0]
 
@@ -419,6 +421,7 @@ def __call__(
           hidden_states=norm_hidden_states,
           encoder_hidden_states=None,
           rotary_emb=video_rotary_emb,
+          perturbation_mask=perturbation_mask if self.perturbed_attn else None,
       )
     hidden_states = hidden_states + attn_hidden_states * gate_msa
 
@@ -449,6 +452,7 @@ def __call__(
           hidden_states=norm_audio_hidden_states,
           encoder_hidden_states=None,
           rotary_emb=audio_rotary_emb,
+          perturbation_mask=perturbation_mask if self.perturbed_attn else None,
       )
     audio_hidden_states = audio_hidden_states + attn_audio_hidden_states * audio_gate_msa
 
@@ -648,6 +652,7 @@ def __init__(
       gated_attn: bool = False,
       cross_attn_mod: bool = False,
       use_prompt_embeddings: bool = True,
+      perturbed_attn: bool = False,
       spatio_temporal_guidance_blocks: Tuple[int, ...] = (),
       **kwargs,
   ):
@@ -700,6 +705,7 @@ def __init__(
     self.attention_kernel = attention_kernel
     self.gated_attn = gated_attn
     self.cross_attn_mod = cross_attn_mod
+    self.perturbed_attn = perturbed_attn
     self.a2v_attention_kernel = a2v_attention_kernel
     self.v2a_attention_kernel = v2a_attention_kernel
     self.flash_min_seq_length = flash_min_seq_length
@@ -943,6 +949,7 @@ def init_block(rngs):
           v2a_attention_kernel=self.v2a_attention_kernel,
           flash_block_sizes=flash_block_sizes,
           flash_min_seq_length=self.flash_min_seq_length,
+          perturbed_attn=self.perturbed_attn,
       )
 
     if self.scan_layers:
@@ -980,6 +987,7 @@ def init_block(rngs):
             v2a_attention_kernel=self.v2a_attention_kernel,
             flash_block_sizes=flash_block_sizes,
             flash_min_seq_length=self.flash_min_seq_length,
+            perturbed_attn=self.perturbed_attn,
         )
         blocks.append(block)
       self.transformer_blocks = nnx.List(blocks)
@@ -1181,7 +1189,7 @@ def scan_fn(carry, block_and_mask):
             ca_audio_rotary_emb=audio_cross_attn_rotary_emb,
             a2v_cross_attention_mask=encoder_attention_mask,
             v2a_cross_attention_mask=audio_encoder_attention_mask,
-            attention_mask=mask,
+            perturbation_mask=mask,
             modality_mask=modality_mask,
         )
       return (
@@ -1225,6 +1233,9 @@ def scan_fn(carry, block_and_mask):
               ca_audio_rotary_emb=audio_cross_attn_rotary_emb,
               encoder_attention_mask=encoder_attention_mask,
               audio_encoder_attention_mask=audio_encoder_attention_mask,
+              a2v_cross_attention_mask=encoder_attention_mask,
+              v2a_cross_attention_mask=audio_encoder_attention_mask,
+              perturbation_mask=mask,
           )
 
     # 6. Output layers

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -1257,6 +1257,7 @@ def __call__(
       stg_scale: float = 0.0,
       modality_scale: float = 1.0,
       audio_guidance_scale: Optional[float] = None,
+      audio_guidance_rescale: Optional[float] = None,
       audio_stg_scale: Optional[float] = None,
       audio_modality_scale: Optional[float] = None,
       noise_scale: float = 1.0,
@@ -1274,13 +1275,20 @@ def __call__(
       dtype: Optional[jnp.dtype] = None,
       output_type: str = "pil",
       return_dict: bool = True,
-      use_cross_timestep: bool = False,
+      use_cross_timestep: Optional[bool] = None,
   ):
     # 1. Check inputs
     self.check_inputs(
         prompt, height, width, prompt_embeds, negative_prompt_embeds, prompt_attention_mask, negative_prompt_attention_mask
     )
 
+    if use_cross_timestep is None:
+      use_cross_timestep = getattr(self.config, "model_name", "") == "ltx2.3"
+
+    audio_guidance_rescale = (
+        audio_guidance_rescale if audio_guidance_rescale is not None else guidance_rescale
+    )
+
     # 2. Encode inputs (Text)
     prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask = self.encode_prompt(
         prompt,
@@ -1343,18 +1351,22 @@ def __call__(
         latents=audio_latents,
     )
 
-    # 5. Prepare Timesteps
+    # 5. Prepare Timesteps (match diffusers LTX2: shift uses scheduler config bounds, not latent token count)
     sigmas = jnp.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
 
-    video_sequence_length = (num_frames - 1) // self.vae_temporal_compression_ratio + 1
-    video_sequence_length *= (height // self.vae_spatial_compression_ratio) * (width // self.vae_spatial_compression_ratio)
+    sched_cfg = self.scheduler.config
+
+    def _sched_cfg_get(key: str, default):
+      if hasattr(sched_cfg, "get"):
+        return sched_cfg.get(key, default)
+      return getattr(sched_cfg, key, default)
 
     mu = calculate_shift(
-        video_sequence_length,
-        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),
-        self.scheduler.config.get("max_shift", 2.05),
+        _sched_cfg_get("max_image_seq_len", 4096),
+        _sched_cfg_get("base_image_seq_len", 1024),
+        _sched_cfg_get("max_image_seq_len", 4096),
+        _sched_cfg_get("base_shift", 0.95),
+        _sched_cfg_get("max_shift", 2.05),
     )
 
     scheduler_state = retrieve_timesteps(
@@ -1373,7 +1385,7 @@ def __call__(
     prompt_attention_mask_jax = prompt_attention_mask
 
     do_cfg = guidance_scale > 1.0
-    do_stg = getattr(self.config, "stg_scale", 0.0) > 0.0
+    do_stg = stg_scale > 0.0
 
     if do_cfg and do_stg:
       negative_prompt_embeds_jax = negative_prompt_embeds
@@ -1561,7 +1573,12 @@ def convert_to_vel(lat, x0):
           audio_modality_delta = (audio_modality_scale - 1 if audio_modality_scale is not None else modality_scale - 1) * (x0_audio_text - x0_audio_isolated)
 
           x0_audio_combined = x0_audio_text + cfg_audio_delta + stg_audio_delta + audio_modality_delta
-          
+
+          if audio_guidance_rescale > 0:
+            x0_audio_combined = rescale_noise_cfg(
+                x0_audio_combined, x0_audio_text, guidance_rescale=audio_guidance_rescale
+            )
+
           noise_pred_audio = convert_to_vel(audio_latents_step, x0_audio_combined)
 
         elif do_cfg:
@@ -1586,7 +1603,12 @@ def convert_to_vel(lat, x0):
 
           cfg_audio_delta = (audio_guidance_scale - 1 if audio_guidance_scale is not None else guidance_scale - 1) * (x0_audio_text - x0_audio_uncond)
           x0_audio_combined = x0_audio_text + cfg_audio_delta
-          
+
+          if audio_guidance_rescale > 0:
+            x0_audio_combined = rescale_noise_cfg(
+                x0_audio_combined, x0_audio_text, guidance_rescale=audio_guidance_rescale
+            )
+
           noise_pred_audio = convert_to_vel(audio_latents_step, x0_audio_combined)
 
         elif do_stg:
@@ -1791,8 +1813,14 @@ def transformer_forward_pass(
   else:
     audio_sigma = jnp.expand_dims(audio_sigma, 0).repeat(latents.shape[0])
 
-  N = latents.shape[0] // 4
-  modality_mask = jnp.concatenate([jnp.ones((3 * N, 1, 1), dtype=latents.dtype), jnp.zeros((N, 1, 1), dtype=latents.dtype)], axis=0)
+  b = latents.shape[0]
+  if b % 4 == 0 and b > 0:
+    n = b // 4
+    modality_mask = jnp.concatenate(
+        [jnp.ones((3 * n, 1, 1), dtype=latents.dtype), jnp.zeros((n, 1, 1), dtype=latents.dtype)], axis=0
+    )
+  else:
+    modality_mask = jnp.ones((b, 1, 1), dtype=latents.dtype)
 
   noise_pred, noise_pred_audio = transformer(
       hidden_states=latents,