refactor: improve cast exclusion and guidance conditional logic

Author: Perseus14

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -179,7 +179,6 @@ def _apply_ltx2_guidance(
     audio_guidance_rescale,
     do_cfg,
     do_stg,
-    use_lax_cond=False,
 ):
   """Applies the shared LTX2 CFG/STG delta formulation to video and audio."""
   if not do_cfg:
@@ -199,7 +198,7 @@ def convert_to_vel(latents_to_convert, x0):
     return (latents_to_convert - x0) / sigma_t
 
   def maybe_rescale(noise, text_noise, scale):
-    if use_lax_cond:
+    if isinstance(scale, jax.core.Tracer):
       return jax.lax.cond(
           scale > 0,
           lambda: rescale_noise_cfg(noise, text_noise, guidance_rescale=scale),
@@ -254,15 +253,13 @@ def maybe_rescale(noise, text_noise, scale):
 
 logger = logging.get_logger(__name__)
 
-_CAST_EXCLUSION_KEYWORDS = (
-    "norm",  # All LayerNorm/GroupNorm parameters.
-    "condition_embedder",  # The entire time/text conditioning module.
-    "scale_shift_table",  # The final and AdaLN scale/shift tables.
-)
+import re
+
+_CAST_EXCLUSION_PATTERN = re.compile(r"(^|\.)(norm[0-9]*|.*_norm[0-9]*|norm_out|condition_embedder|scale_shift_table)(\.|$)")
 
 
 def _is_cast_excluded(path_str: str) -> bool:
-  return any(keyword in path_str.lower() for keyword in _CAST_EXCLUSION_KEYWORDS)
+  return bool(_CAST_EXCLUSION_PATTERN.search(path_str.lower()))
 
 
 def _final_param_dtype(flax_key: tuple, dtype_to_cast) -> np.dtype:
@@ -2164,7 +2161,6 @@ def __call__(
               audio_guidance_rescale,
               do_cfg,
               do_stg,
-              use_lax_cond=False,
           )
 
           latents_step, _ = self.scheduler.step(scheduler_state, noise_pred, t, latents_step, return_dict=False)
@@ -2659,7 +2655,6 @@ def scan_body(carry, inputs):
         audio_guidance_rescale,
         do_cfg,
         do_stg,
-        use_lax_cond=True,
     )
 
     # Step scheduler