remove debug+ gated attn=true

Author: prishajain1

src/maxdiffusion/models/ltx2/ltx2_3_utils.py

@@ -186,18 +186,12 @@ def load_transformer_weights_2_3(
       renamed_pt_key = rename_key(pt_key)
       renamed_pt_key = rename_for_ltx2_3_transformer(renamed_pt_key)
 
-      if "to_gate_logits" in pt_key:
-        print(f"pt_key: {pt_key} -> renamed_pt_key: {renamed_pt_key}")
-
       pt_tuple_key = tuple(renamed_pt_key.split("."))
 
       flax_key, flax_tensor = get_key_and_value(
           pt_tuple_key, tensor, flax_state_dict, random_flax_state_dict, scan_layers, num_layers
       )
 
-      if "to_gate_logits" in pt_key:
-        print(f"flax_key: {flax_key}")
-
       flax_state_dict[flax_key] = jax.device_put(jnp.asarray(flax_tensor), device=cpu)
 
     validate_flax_state_dict(eval_shapes, flax_state_dict)

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -159,7 +159,7 @@ def create_model(rngs: nnx.Rngs, ltx2_config: dict):
   ltx2_config["names_which_can_be_offloaded"] = config.names_which_can_be_offloaded
 
   if getattr(config, "model_name", "") == "ltx2.3":
-    ltx2_config["gated_attn"] = False
+    ltx2_config["gated_attn"] = True
     ltx2_config["cross_attn_mod"] = True
     ltx2_config["perturbed_attn"] = True
     ltx2_config["use_prompt_embeddings"] = True