CFG gating

Author: prishajain1

src/maxdiffusion/configs/ltx2_3_video.yml

@@ -25,6 +25,8 @@ max_sequence_length: 1024
 sampler: "from_checkpoint"
 
 # Generation parameters (aligned with Diffusers LTX-2.3 docs: use_cross_timestep, modality + audio CFG)
+# CFG negative-prompt encoding runs when guidance_scale>1 OR audio_guidance_scale>1 (Diffusers parity).
+# Modality isolation stacks with CFG even when stg_scale is 0 (pipeline stack_kind cfg_mod).
 global_batch_size_to_train_on: 1
 num_inference_steps: 30
 guidance_scale: 3.0

src/maxdiffusion/generate_ltx2.py

@@ -106,7 +106,7 @@ def call_pipeline(config, pipeline, prompt, negative_prompt):
       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),
+      noise_scale=getattr(config, "noise_scale", 0.0),
       dtype=jnp.bfloat16 if getattr(config, "activations_dtype", "bfloat16") == "bfloat16" else jnp.float32,
   )
   return out

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -1236,6 +1236,7 @@ def scan_fn(carry, block_and_mask):
               a2v_cross_attention_mask=encoder_attention_mask,
               v2a_cross_attention_mask=audio_encoder_attention_mask,
               perturbation_mask=mask,
+              modality_mask=modality_mask,
           )
 
     # 6. Output layers