LTX2.3 improvements and bug fixes

Author: Perseus14

src/maxdiffusion/configs/ltx2_3_video.yml

@@ -1,7 +1,12 @@
 #hardware
 hardware: 'tpu'
 skip_jax_distributed_system: False
-attention: 'flash'
+attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring 
+
+# For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this.
+ulysses_shards: -1
+# Splits Ulysses all-to-all into head-group chunks. The last chunk carries any remainder.
+ulysses_attention_chunks: 1
 a2v_attention_kernel: 'flash'
 v2a_attention_kernel: 'dot_product'
 attention_sharding_uniform: True 

src/maxdiffusion/configs/ltx2_video.yml

@@ -1,7 +1,12 @@
 #hardware
 hardware: 'tpu'
 skip_jax_distributed_system: False
-attention: 'flash'
+attention: 'flash'  # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring 
+
+# For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this.
+ulysses_shards: -1
+# Splits Ulysses all-to-all into head-group chunks. The last chunk carries any remainder.
+ulysses_attention_chunks: 1
 a2v_attention_kernel: 'dot_product'
 v2a_attention_kernel: 'dot_product'
 attention_sharding_uniform: True 

src/maxdiffusion/generate_ltx2.py

@@ -19,7 +19,7 @@
 import os
 import subprocess
 from maxdiffusion.checkpointing.ltx2_checkpointer import LTX2Checkpointer
-from maxdiffusion import pyconfig, max_logging, max_utils
+from maxdiffusion import aot_cache, pyconfig, max_logging, max_utils
 from absl import app
 from google.cloud import storage
 from google.api_core.exceptions import GoogleAPIError
@@ -189,14 +189,37 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None):
   original_enable_mld = config.get_keys().get("enable_ml_diagnostics", False)
   original_num_steps = config.get_keys().get("num_inference_steps", 40)
 
+  # Per-shape AOT executable cache
+  aot_cache.install(
+      getattr(config, "aot_cache_dir", ""),
+      meta={
+          "model": config.pretrained_model_name_or_path,
+          "attention": getattr(config, "attention", ""),
+          "flash_block_sizes": str(getattr(config, "flash_block_sizes", "")),
+          "mesh_shape": str(pipeline.mesh.shape) if pipeline and hasattr(pipeline, "mesh") and pipeline.mesh else "",
+          "weights_dtype": str(getattr(config, "weights_dtype", "bfloat16")),
+          "activations_dtype": str(getattr(config, "activations_dtype", "bfloat16")),
+          "scan_layers": str(getattr(config, "scan_layers", True)),
+          "jax": jax.__version__,
+      },
+      mesh=pipeline.mesh if pipeline else None,
+  )
+  aot_cache.wait_for_loads()
+
   # ---------------------------------------------------------
   # Run 1: Warmup Compilation (Original steps, NO profiling)
   # ---------------------------------------------------------
   config.get_keys()["enable_profiler"] = False
   config.get_keys()["enable_ml_diagnostics"] = False
+  warmup_steps = min(2, original_num_steps)
+  config.get_keys()["num_inference_steps"] = warmup_steps
+
+  max_logging.log(f"🚀 Starting warmup compilation pass ({warmup_steps} steps)...")
+  with aot_cache.warmup_mode():
+    _ = call_pipeline(config, pipeline, prompt, negative_prompt)
 
-  max_logging.log(f"🚀 Starting warmup compilation pass ({original_num_steps} steps)...")
-  _ = call_pipeline(config, pipeline, prompt, negative_prompt)
+  aot_cache.save_pending()
+  config.get_keys()["num_inference_steps"] = original_num_steps
 
   compile_time = time.perf_counter() - s0
   max_logging.log(f"compile_time: {compile_time}")
@@ -237,7 +260,9 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None):
 
   # Export videos
   for i in range(len(videos)):
-    video_path = f"{filename_prefix}ltx2_output_{getattr(config, 'seed', 0)}_{i}.mp4"
+    model_name = getattr(config, "model_name", "ltx2") or "ltx2"
+    model_name_prefix = model_name.replace(".", "_")
+    video_path = f"{filename_prefix}{model_name_prefix}_output_{getattr(config, 'seed', 0)}_{i}.mp4"
     audio_i = audios[i] if audios is not None else None
 
     audio_format = getattr(config, "audio_format", "s16")

src/maxdiffusion/models/attention_flax.py

@@ -1914,7 +1914,7 @@ def __init__(
       self,
       mesh: Mesh,
       attention_kernel: str,
-      scale: int,
+      scale: float,
       heads: int,
       dim_head: int,
       use_memory_efficient_attention: bool = False,
@@ -2009,7 +2009,7 @@ def apply_attention(self, query: Array, key: Array, value: Array, attention_mask
 class AttentionOp(nn.Module):
   mesh: Mesh
   attention_kernel: str
-  scale: int
+  scale: float
   heads: int
   dim_head: int
   use_memory_efficient_attention: bool = False

src/maxdiffusion/models/ltx2/attention_ltx2.py

@@ -88,13 +88,8 @@ def apply_split_rotary_emb(x: Array, freqs: Tuple[Array, Array]) -> Array:
   first_x = split_x[..., 0, :]
   second_x = split_x[..., 1, :]
 
-  cos_u = jnp.expand_dims(cos, axis=-2)
-  sin_u = jnp.expand_dims(sin, axis=-2)
-
-  out = split_x * cos_u
-
-  out_first = out[..., 0, :] - second_x * sin_u.squeeze(-2)
-  out_second = out[..., 1, :] + first_x * sin_u.squeeze(-2)
+  out_first = first_x * cos - second_x * sin
+  out_second = second_x * cos + first_x * sin
 
   out = jnp.stack([out_first, out_second], axis=-2)
   out = out.reshape(*out.shape[:-2], last_dim)
@@ -176,12 +171,6 @@ def prepare_video_coords(
     patch_ends = grid + patch_size_delta
 
     # Combine start and end coordinates
-    latent_coords = jnp.stack([grid, patch_ends], axis=-1)  # [3, N_F, N_H, N_W, 2]
-    latent_coords = latent_coords.transpose(1, 2, 3, 0, 4)  # [N_F, N_H, N_W, 3, 2]
-    latent_coords = latent_coords.reshape(-1, 3, 2)  # [num_patches, 3, 2]
-    latent_coords = jnp.expand_dims(latent_coords, 0)  # [1, num_patches, 3, 2]
-    latent_coords = jnp.tile(latent_coords, (batch_size, 1, 1, 1))  # [B, num_patches, 3, 2]
-
     latent_coords = jnp.stack([grid, patch_ends], axis=-1)  # [3, N_F, N_H, N_W, 2]
     latent_coords = latent_coords.reshape(3, -1, 2)  # [3, num_patches, 2]
     latent_coords = jnp.expand_dims(latent_coords, 0)  # [1, 3, num_patches, 2]
@@ -352,6 +341,8 @@ def __init__(
       flash_min_seq_length: int = 4096,
       sharding_specs: Optional[LTX2DiTShardingSpecs] = None,
       gated_attn: bool = False,
+      ulysses_shards: int = -1,
+      ulysses_attention_chunks: int = 1,
   ):
     self.heads = heads
     self.rope_type = rope_type
@@ -445,17 +436,37 @@ def __init__(
           dtype=dtype,
       )
 
+    is_self_attention = context_dim is None
+    if not is_self_attention:
+      if attention_kernel in ("tokamax_ring", "tokamax_ring_custom", "ulysses_ring"):
+        attention_kernel = "tokamax_flash"  # do not use ring attention for cross attention
+      if attention_kernel in ("ulysses_ring_custom", "ulysses_ring_custom_bidir"):
+        attention_kernel = "ulysses_custom"  # plain ulysses (no ring) for cross attention
+
+      axis_names_q = (common_types.BATCH, common_types.CROSS_ATTN_HEAD, common_types.CROSS_ATTN_Q_LENGTH, common_types.D_KV)
+      axis_names_kv = (
+          common_types.BATCH,
+          common_types.CROSS_ATTN_HEAD,
+          common_types.CROSS_ATTN_KV_LENGTH,
+          common_types.D_KV,
+      )
+    else:
+      axis_names_q = (common_types.BATCH, common_types.SELF_ATTN_HEAD, common_types.SELF_ATTN_Q_LENGTH, common_types.D_KV)
+      axis_names_kv = (common_types.BATCH, common_types.SELF_ATTN_HEAD, common_types.SELF_ATTN_KV_LENGTH, common_types.D_KV)
+
     self.attention_op = NNXAttentionOp(
         mesh=mesh,
         attention_kernel=attention_kernel,
         scale=dim_head**-0.5,
         heads=heads,
         dim_head=dim_head,
         dtype=dtype,
-        axis_names_q=(common_types.BATCH, common_types.SELF_ATTN_HEAD, common_types.SELF_ATTN_Q_LENGTH, common_types.D_KV),
-        axis_names_kv=(common_types.BATCH, common_types.SELF_ATTN_HEAD, common_types.SELF_ATTN_KV_LENGTH, common_types.D_KV),
+        axis_names_q=axis_names_q,
+        axis_names_kv=axis_names_kv,
         flash_block_sizes=flash_block_sizes,
         flash_min_seq_length=flash_min_seq_length,
+        ulysses_shards=ulysses_shards,
+        ulysses_attention_chunks=ulysses_attention_chunks,
     )
 
   def __call__(
@@ -485,7 +496,7 @@ def __call__(
     # 3. Apply RoPE
     with jax.named_scope("Apply RoPE"):
       if rotary_emb is not None:
-        if hasattr(self, "rope_type") and self.rope_type == "split":
+        if self.rope_type == "split":
           # Split RoPE: passing full freqs [B, H, S, D//2]
           # apply_split_rotary_emb handles reshaping query/key
 

src/maxdiffusion/models/ltx2/ltx2_utils.py

@@ -17,6 +17,7 @@
 import json
 from typing import Optional
 import torch
+import numpy as np
 import jax
 import jax.numpy as jnp
 from maxdiffusion import max_logging
@@ -184,6 +185,14 @@ def load_sharded_checkpoint(pretrained_model_name_or_path, subfolder, device, fi
   return tensors
 
 
+def _torch_tensor_to_numpy(tensor: torch.Tensor) -> np.ndarray:
+  import ml_dtypes
+
+  if tensor.dtype == torch.bfloat16:
+    return tensor.view(torch.uint16).numpy().view(ml_dtypes.bfloat16)
+  return tensor.numpy()
+
+
 def load_transformer_weights(
     pretrained_model_name_or_path: str,
     eval_shapes: dict,
@@ -193,38 +202,110 @@ def load_transformer_weights(
     scan_layers: bool = True,
     subfolder: str = "transformer",
 ):
+  import threading
+  import concurrent.futures
+  import time
+
   device = jax.local_devices(backend=device)[0]
   max_logging.log(f"Load and port {pretrained_model_name_or_path} {subfolder} on {device}")
 
-  with jax.default_device(device):
-    # Support sharded loading
-    tensors = load_sharded_checkpoint(pretrained_model_name_or_path, subfolder, device)
+  index_file = "diffusion_pytorch_model.safetensors.index.json"
+  try:
+    index_path = hf_hub_download(pretrained_model_name_or_path, subfolder=subfolder, filename=index_file)
+    with open(index_path, "r") as f:
+      index_data = json.load(f)
+    weight_map = index_data["weight_map"]
+    shards = sorted(set(weight_map.values()))
 
-    flax_state_dict = {}
-    cpu = jax.local_devices(backend="cpu")[0]
-    flattened_dict = flatten_dict(eval_shapes)
+    def resolve_shard_path(model_file):
+      return hf_hub_download(pretrained_model_name_or_path, subfolder=subfolder, filename=model_file)
 
-    random_flax_state_dict = {}
-    for key in flattened_dict:
-      random_flax_state_dict[tuple(str(item) for item in key)] = flattened_dict[key]
+  except EntryNotFoundError:
+    shards = ["diffusion_pytorch_model.safetensors"]
 
-    for pt_key, tensor in tensors.items():
-      renamed_pt_key = rename_key(pt_key)
-      renamed_pt_key = rename_for_ltx2_transformer(renamed_pt_key)
+    def resolve_shard_path(model_file):
+      try:
+        return hf_hub_download(pretrained_model_name_or_path, subfolder=subfolder, filename=model_file)
+      except EntryNotFoundError:
+        return hf_hub_download(pretrained_model_name_or_path, subfolder=subfolder, filename="diffusion_pytorch_model.bin")
 
-      pt_tuple_key = tuple(renamed_pt_key.split("."))
+  t_start = time.perf_counter()
 
-      flax_key, flax_tensor = get_key_and_value(
-          pt_tuple_key, tensor, flax_state_dict, random_flax_state_dict, scan_layers, num_layers
-      )
+  flattened_dict = flatten_dict(eval_shapes)
+  random_flax_state_dict = {}
+  for key in flattened_dict:
+    random_flax_state_dict[tuple(str(item) for item in key)] = flattened_dict[key]
 
-      flax_state_dict[flax_key] = jax.device_put(jnp.asarray(flax_tensor), device=cpu)
+  flax_state_dict = {}
+  dict_lock = threading.Lock()
+
+  def convert_chunk(ckpt_shard_path, chunk_keys):
+    if ckpt_shard_path.endswith(".safetensors"):
+      with safe_open(ckpt_shard_path, framework="pt") as f:
+        for pt_key in chunk_keys:
+          tensor = _torch_tensor_to_numpy(f.get_tensor(pt_key))
+          process_tensor(pt_key, tensor)
+    else:
+      loaded_state_dict = torch.load(ckpt_shard_path, map_location="cpu")
+      for pt_key in chunk_keys:
+        tensor = _torch_tensor_to_numpy(loaded_state_dict[pt_key])
+        process_tensor(pt_key, tensor)
+
+  def process_tensor(pt_key, tensor):
+    renamed_pt_key = rename_key(pt_key)
+    renamed_pt_key = rename_for_ltx2_transformer(renamed_pt_key)
+    pt_tuple_key = tuple(renamed_pt_key.split("."))
+
+    block_index = None
+    if scan_layers and len(pt_tuple_key) > 0 and "transformer_blocks_" in pt_tuple_key[0]:
+      import re
+
+      m = re.match(r"transformer_blocks_(\d+)", pt_tuple_key[0])
+      if m:
+        block_index = int(m.group(1))
+        pt_tuple_key = ("transformer_blocks",) + pt_tuple_key[1:]
 
-    validate_flax_state_dict(eval_shapes, flax_state_dict)
-    flax_state_dict = unflatten_dict(flax_state_dict)
-    del tensors
-    jax.clear_caches()
-    return flax_state_dict
+    flax_key, flax_tensor = rename_key_and_reshape_tensor(pt_tuple_key, tensor, random_flax_state_dict, scan_layers)
+    flax_key_str = [str(k) for k in flax_key]
+    if "scale_shift_table" in flax_key_str and flax_key_str[-1] in ["kernel", "weight"]:
+      flax_key_str.pop()
+    flax_key = tuple(flax_key_str)
+    flax_key = _tuple_str_to_int(flax_key)
+
+    if block_index is not None:
+      with dict_lock:
+        stacked = flax_state_dict.get(flax_key)
+        if stacked is None:
+          stacked = np.empty((num_layers,) + flax_tensor.shape, dtype=flax_tensor.dtype)
+          flax_state_dict[flax_key] = stacked
+      stacked[block_index] = flax_tensor
+    else:
+      value = np.array(flax_tensor, dtype=flax_tensor.dtype, copy=True, order="C")
+      with dict_lock:
+        flax_state_dict[flax_key] = value
+
+  chunk_size = 32
+  tasks = []
+  for model_file in shards:
+    ckpt_shard_path = resolve_shard_path(model_file)
+    if ckpt_shard_path.endswith(".safetensors"):
+      with safe_open(ckpt_shard_path, framework="pt") as f:
+        shard_keys = list(f.keys())
+    else:
+      loaded_state_dict = torch.load(ckpt_shard_path, map_location="cpu")
+      shard_keys = list(loaded_state_dict.keys())
+    for i in range(0, len(shard_keys), chunk_size):
+      tasks.append((ckpt_shard_path, shard_keys[i : i + chunk_size]))
+
+  with concurrent.futures.ThreadPoolExecutor(max_workers=32) as executor:
+    futures = [executor.submit(convert_chunk, path, keys) for path, keys in tasks]
+    for future in concurrent.futures.as_completed(futures):
+      future.result()
+
+  validate_flax_state_dict(eval_shapes, flax_state_dict)
+  flax_state_dict = unflatten_dict(flax_state_dict)
+  max_logging.log(f"Converted weights in {time.perf_counter() - t_start:.1f}s")
+  return flax_state_dict
 
 
 def load_vae_weights(

src/maxdiffusion/models/ltx2/text_encoders/torchax_text_encoder.py

@@ -20,11 +20,9 @@
 import jax
 from torchax import interop, default_env
 
-# --- Monkeypatch transformers masking_utils to avoid torchax integer tracing bug ---
+import contextlib
 import transformers.masking_utils
 
-_orig_sliding_window_overlay = transformers.masking_utils.sliding_window_overlay
-
 
 def _patched_sliding_window_overlay(sliding_window: int):
   # pylint: disable=unused-argument
@@ -44,6 +42,16 @@ def inner_mask(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool:
   return inner_mask
 
 
+@contextlib.contextmanager
+def patch_sliding_window_overlay():
+  orig = transformers.masking_utils.sliding_window_overlay
+  transformers.masking_utils.sliding_window_overlay = _patched_sliding_window_overlay
+  try:
+    yield
+  finally:
+    transformers.masking_utils.sliding_window_overlay = orig
+
+
 class TorchaxGemma3TextEncoder(interop.JittableModule):
   """
   A jittable Torchax module for wrapping the HuggingFace PyTorch
@@ -57,8 +65,7 @@ def __call__(
       self, input_ids: jax.Array, attention_mask: jax.Array, output_hidden_states: bool = True
   ) -> Tuple[jax.Array, ...]:
     # Dynamically patch transformers.masking_utils only during the duration of this call
-    transformers.masking_utils.sliding_window_overlay = _patched_sliding_window_overlay
-    try:
+    with patch_sliding_window_overlay():
       with default_env():
         input_ids = interop.torch_view(input_ids)
         attention_mask = interop.torch_view(attention_mask)
@@ -72,9 +79,6 @@ def __call__(
             output_hidden_states=output_hidden_states,
         )
       return interop.jax_view(output)
-    finally:
-      # Restore original behavior to prevent side effects on other potential models in same env
-      transformers.masking_utils.sliding_window_overlay = _orig_sliding_window_overlay
 
   @staticmethod
   def _forward_inner(model, input_ids, attention_mask, output_hidden_states=True):