reformatted

Author: prishajain1

src/maxdiffusion/generate_ltx2.py

@@ -132,7 +132,7 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None):
       weights = lora_config.get("weight_name", [None] * len(paths))
       scales = lora_config.get("scale", [1.0] * len(paths))
       ranks = lora_config.get("rank", [64] * len(paths))
-      
+
       for i in range(len(paths)):
         pipeline = lora_loader.load_lora_weights(
             pipeline,

src/maxdiffusion/loaders/lora_conversion_utils.py

@@ -716,41 +716,34 @@ def translate_ltx2_nnx_path_to_diffusers_lora(nnx_path_str, scan_layers=False):
       "attn1.to_k": "attn1.to_k",
       "attn1.to_v": "attn1.to_v",
       "attn1.to_out": "attn1.to_out.0",
-      
       # Audio Self Attention (audio_attn1)
       "audio_attn1.to_q": "audio_attn1.to_q",
       "audio_attn1.to_k": "audio_attn1.to_k",
       "audio_attn1.to_v": "audio_attn1.to_v",
       "audio_attn1.to_out": "audio_attn1.to_out.0",
-      
       # Audio Cross Attention (audio_attn2)
       "audio_attn2.to_q": "audio_attn2.to_q",
       "audio_attn2.to_k": "audio_attn2.to_k",
       "audio_attn2.to_v": "audio_attn2.to_v",
       "audio_attn2.to_out": "audio_attn2.to_out.0",
-      
       # Cross Attention (attn2)
       "attn2.to_q": "attn2.to_q",
       "attn2.to_k": "attn2.to_k",
       "attn2.to_v": "attn2.to_v",
       "attn2.to_out": "attn2.to_out.0",
-      
       # Audio to Video Cross Attention
       "audio_to_video_attn.to_q": "audio_to_video_attn.to_q",
       "audio_to_video_attn.to_k": "audio_to_video_attn.to_k",
       "audio_to_video_attn.to_v": "audio_to_video_attn.to_v",
       "audio_to_video_attn.to_out": "audio_to_video_attn.to_out.0",
-      
       # Video to Audio Cross Attention
       "video_to_audio_attn.to_q": "video_to_audio_attn.to_q",
       "video_to_audio_attn.to_k": "video_to_audio_attn.to_k",
       "video_to_audio_attn.to_v": "video_to_audio_attn.to_v",
       "video_to_audio_attn.to_out": "video_to_audio_attn.to_out.0",
-      
       # Feed Forward
       "ff.net_0": "ff.net.0.proj",
       "ff.net_2": "ff.net.2",
-      
       # Audio Feed Forward
       "audio_ff.net_0": "audio_ff.net.0.proj",
       "audio_ff.net_2": "audio_ff.net.2",
@@ -768,7 +761,6 @@ def translate_ltx2_nnx_path_to_diffusers_lora(nnx_path_str, scan_layers=False):
       "av_cross_attn_audio_v2a_gate.linear": "diffusion_model.av_ca_v2a_gate_adaln_single.linear",
       "av_cross_attn_audio_scale_shift.linear": "diffusion_model.av_ca_audio_scale_shift_adaln_single.linear",
       "av_cross_attn_video_scale_shift.linear": "diffusion_model.av_ca_video_scale_shift_adaln_single.linear",
-      
       # Nested conditioning layers
       "time_embed.emb.timestep_embedder.linear_1": "diffusion_model.adaln_single.emb.timestep_embedder.linear_1",
       "time_embed.emb.timestep_embedder.linear_2": "diffusion_model.adaln_single.emb.timestep_embedder.linear_2",
@@ -786,11 +778,10 @@ def translate_ltx2_nnx_path_to_diffusers_lora(nnx_path_str, scan_layers=False):
       "caption_projection.linear_2": "diffusion_model.caption_projection.linear_2",
       "audio_caption_projection.linear_1": "diffusion_model.audio_caption_projection.linear_1",
       "audio_caption_projection.linear_2": "diffusion_model.audio_caption_projection.linear_2",
-      
       # Connectors
       "feature_extractor.linear": "text_embedding_projection.aggregate_embed",
   }
-  
+
   if nnx_path_str in global_map:
     return global_map[nnx_path_str]
 
@@ -807,5 +798,3 @@ def translate_ltx2_nnx_path_to_diffusers_lora(nnx_path_str, scan_layers=False):
         return f"diffusion_model.transformer_blocks.{idx}.{suffix_map[inner_suffix]}"
 
   return None
-
-

src/maxdiffusion/loaders/ltx2_lora_nnx_loader.py

@@ -64,7 +64,7 @@ def translate_fn(nnx_path_str):
       max_logging.log(f"Merging LoRA into connectors with rank={rank}")
       if h_state_dict is None and transformer_weight_name:
         h_state_dict, _ = lora_loader.lora_state_dict(lora_model_path, weight_name=transformer_weight_name, **kwargs)
-      
+
       if h_state_dict is not None:
         # Filter state dict for connector keys to avoid confusing warnings
         connector_state_dict = {k: v for k, v in h_state_dict.items() if k.startswith("text_embedding_projection")}

src/maxdiffusion/models/ltx2/attention_ltx2.py

@@ -455,7 +455,7 @@ def __call__(
       query = self.to_q(hidden_states)
       key = self.to_k(context)
       value = self.to_v(context)
-    
+
     with jax.named_scope("QKV Norm"):
       query = self.norm_q(query)
       key = self.norm_k(key)

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -1026,24 +1026,24 @@ def scan_fn(carry, block):
         )(carry, self.transformer_blocks)
       else:
         for block in self.transformer_blocks:
-            hidden_states, audio_hidden_states = block(
-                hidden_states=hidden_states,
-                audio_hidden_states=audio_hidden_states,
-                encoder_hidden_states=encoder_hidden_states,
-                audio_encoder_hidden_states=audio_encoder_hidden_states,
-                temb=temb,
-                temb_audio=temb_audio,
-                temb_ca_scale_shift=video_cross_attn_scale_shift,
-                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,
-                video_rotary_emb=video_rotary_emb,
-                audio_rotary_emb=audio_rotary_emb,
-                ca_video_rotary_emb=video_cross_attn_rotary_emb,
-                ca_audio_rotary_emb=audio_cross_attn_rotary_emb,
-                encoder_attention_mask=encoder_attention_mask,
-                audio_encoder_attention_mask=audio_encoder_attention_mask,
-            )
+          hidden_states, audio_hidden_states = block(
+              hidden_states=hidden_states,
+              audio_hidden_states=audio_hidden_states,
+              encoder_hidden_states=encoder_hidden_states,
+              audio_encoder_hidden_states=audio_encoder_hidden_states,
+              temb=temb,
+              temb_audio=temb_audio,
+              temb_ca_scale_shift=video_cross_attn_scale_shift,
+              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,
+              video_rotary_emb=video_rotary_emb,
+              audio_rotary_emb=audio_rotary_emb,
+              ca_video_rotary_emb=video_cross_attn_rotary_emb,
+              ca_audio_rotary_emb=audio_cross_attn_rotary_emb,
+              encoder_attention_mask=encoder_attention_mask,
+              audio_encoder_attention_mask=audio_encoder_attention_mask,
+          )
 
     # 6. Output layers
     with jax.named_scope("Output Projection & Norm"):

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -1233,10 +1233,11 @@ def run_connectors(graphdef, state, hidden_states, attention_mask):
       )
 
       import time
+
       timesteps_jax = jnp.array(timesteps, dtype=jnp.float32)
       for i, t_val in enumerate(timesteps):
         t = timesteps_jax[i]
-        
+
         # Isolate input sharding to scan_layers=False to avoid affecting the standard path
         latents_jax_sharded = latents_jax
         audio_latents_jax_sharded = audio_latents_jax
@@ -1340,12 +1341,11 @@ def run_connectors(graphdef, state, hidden_states, attention_mask):
       mesh = latents.sharding.mesh
       replicated_sharding = NamedSharding(mesh, P())
       latents = jax.lax.with_sharding_constraint(latents, replicated_sharding)
-      
+
       # Replicate VAE weights
       graphdef, state = nnx.split(self.vae)
       state = jax.tree_util.tree_map(
-          lambda x: jax.lax.with_sharding_constraint(x, replicated_sharding) if isinstance(x, jax.Array) else x, 
-          state
+          lambda x: jax.lax.with_sharding_constraint(x, replicated_sharding) if isinstance(x, jax.Array) else x, state
       )
       self.vae = nnx.merge(graphdef, state)
     except Exception as e: