ltx2.3 config file and params addition

Author: prishajain1

src/maxdiffusion/common_types.py

@@ -52,6 +52,7 @@
 WAN2_1 = "wan2.1"
 WAN2_2 = "wan2.2"
 LTX2_VIDEO = "ltx2_video"
+LTX2_3 = "ltx2.3"
 
 WAN_MODEL = WAN2_1
 

src/maxdiffusion/configs/ltx2_3_video.yml

@@ -0,0 +1,116 @@
+#hardware
+hardware: 'tpu'
+skip_jax_distributed_system: False
+attention: 'flash'
+a2v_attention_kernel: 'dot_product'
+v2a_attention_kernel: 'dot_product'
+attention_sharding_uniform: True 
+precision: 'bf16'
+scan_layers: True
+names_which_can_be_saved: []
+names_which_can_be_offloaded: []
+remat_policy: "NONE"
+
+jax_cache_dir: ''
+weights_dtype: 'bfloat16'
+activations_dtype: 'bfloat16'
+
+run_name: 'ltx2_inference'
+output_dir: ''
+config_path: ''
+save_config_to_gcs: False
+
+#Checkpoints
+max_sequence_length: 1024
+sampler: "from_checkpoint"
+
+# 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: 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
+use_cross_timestep: true
+spatio_temporal_guidance_blocks: [28]
+fps: 24
+pipeline_type: multi-scale
+prompt: "A man in a brightly lit room talks on a vintage telephone. In a low, heavy voice, he says, 'I understand. I won't call again. Goodbye.' He hangs up the receiver and looks down with a sad expression. He holds the black rotary phone to his right ear with his right hand, his left hand holding a rocks glass with amber liquid. He wears a brown suit jacket over a white shirt, and a gold ring on his left ring finger. His short hair is neatly combed, and he has light skin with visible wrinkles around his eyes. The camera remains stationary, focused on his face and upper body. The room is brightly lit by a warm light source off-screen to the left, casting shadows on the wall behind him. The scene appears to be from a dramatic movie."
+negative_prompt: "shaky, glitchy, low quality, worst quality, deformed, distorted, disfigured, motion smear, motion artifacts, fused fingers, bad anatomy, weird hand, ugly, transition, static."
+height: 512
+width: 768
+decode_timestep: 0.05
+decode_noise_scale: 0.025
+noise_scale: 0.0
+num_frames: 121
+quantization: "int8"
+#parallelism
+mesh_axes: ['data', 'fsdp', 'context', 'tensor']
+logical_axis_rules: [
+                      ['batch', ['data', 'fsdp']],
+                      ['activation_batch', ['data', 'fsdp']],
+                      ['activation_self_attn_heads', ['context', 'tensor']],
+                      ['activation_cross_attn_q_length', ['context', 'tensor']],
+                      ['activation_length', 'context'],
+                      ['activation_heads', 'tensor'],
+                      ['mlp','tensor'],
+                      ['embed', ['context', 'fsdp']],
+                      ['heads', 'tensor'],
+                      ['norm', 'tensor'],
+                      ['conv_batch', ['data', 'context', 'fsdp']],
+                      ['out_channels', 'tensor'],
+                      ['conv_out', 'context'],
+                    ]
+data_sharding: ['data', 'fsdp', 'context', 'tensor']
+
+dcn_data_parallelism: 1  # recommended DCN axis to be auto-sharded
+dcn_fsdp_parallelism: -1
+
+flash_block_sizes: {
+  block_q: 2048,
+  block_kv: 2048,
+  block_kv_compute: 1024,
+  block_q_dkv: 2048,
+  block_kv_dkv: 2048,
+  block_kv_dkv_compute: 2048,
+  use_fused_bwd_kernel: True,
+}
+flash_min_seq_length: 4096
+dcn_context_parallelism: 1
+dcn_tensor_parallelism: 1
+ici_data_parallelism: 1
+ici_fsdp_parallelism: 1  
+ici_context_parallelism: -1 # recommended ICI axis to be auto-sharded
+ici_tensor_parallelism: 1
+enable_profiler: False
+
+replicate_vae: False
+
+allow_split_physical_axes: False
+learning_rate_schedule_steps: -1
+max_train_steps: 500
+pretrained_model_name_or_path: 'dg845/LTX-2.3-Diffusers'
+model_name: "ltx2.3"
+model_type: "T2V"
+unet_checkpoint: ''
+checkpoint_dir: ""
+dataset_name: ''
+train_split: 'train'
+dataset_type: 'tfrecord'
+cache_latents_text_encoder_outputs: True
+per_device_batch_size: 1.0
+compile_topology_num_slices: -1 
+quantization_local_shard_count: -1
+use_qwix_quantization: False
+weight_quantization_calibration_method: "absmax"
+act_quantization_calibration_method: "absmax"
+bwd_quantization_calibration_method: "absmax"
+qwix_module_path: ".*"
+jit_initializers: True 
+enable_single_replica_ckpt_restoring: False
+seed: 10
+audio_format: "s16"

src/maxdiffusion/generate_ltx2.py

@@ -98,6 +98,15 @@ def call_pipeline(config, pipeline, prompt, negative_prompt):
       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),
+      guidance_rescale=getattr(config, "guidance_rescale", 0.0),
+      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,
       output_type=getattr(config, "upsampler_output_type", "pil"),
   )

src/maxdiffusion/pyconfig.py

@@ -34,12 +34,13 @@
     WAN2_1,
     WAN2_2,
     LTX2_VIDEO,
+    LTX2_3,
     RING_ATTENTION_AXIS_RULES,
     SEQUENCE_PARALLEL_AXIS_RULES,
     ULYSSES_ATTENTION_AXIS_RULES,
 )
 
-_ALLOWED_MODEL_NAMES = {WAN2_1, WAN2_2, LTX2_VIDEO}
+_ALLOWED_MODEL_NAMES = {WAN2_1, WAN2_2, LTX2_VIDEO, LTX2_3}
 _ALLOWED_TRAINING_MODEL_NAMES = {WAN2_1}