sharding fix

Author: prishajain1

src/maxdiffusion/configs/ltx2_3_video.yml

@@ -31,8 +31,8 @@ guidance_scale: 3.0
 guidance_rescale: 0.7
 audio_guidance_scale: 7.0
 audio_guidance_rescale: 0.7
-stg_scale: 0.0
-audio_stg_scale: 0.0
+stg_scale: 1.0
+audio_stg_scale: 1.0
 modality_scale: 1.0
 audio_modality_scale: 1.0
 use_cross_timestep: true

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -388,7 +388,7 @@ def __call__(
     if encoder_hidden_states is not None:
       encoder_hidden_states = jax.lax.with_sharding_constraint(encoder_hidden_states, axis_names)
     if audio_encoder_hidden_states is not None:
-      audio_encoder_hidden_states = jax.lax.with_sharding_constraint(audio_encoder_hidden_states, axis_names)
+      audio_encoder_hidden_states = jax.lax.with_sharding_constraint(audio_encoder_hidden_states, axis_names_audio)
 
     # 1. Video and Audio Self-Attention
     norm_hidden_states = self.norm1(hidden_states)