resolved comments

Author: prishajain1

src/maxdiffusion/configs/ltx2_3_video.yml

@@ -140,6 +140,11 @@ lora_config: {
 # LTX-2 Latent Upsampler
 run_latent_upsampler: False
 upsampler_model_path: "Lightricks/LTX-2.3"
+# Following upsampler files are supported:
+# ltx-2.3-spatial-upscaler-x2-1.0.safetensors
+# ltx-2.3-spatial-upscaler-x2-1.1.safetensors
+# ltx-2.3-spatial-upscaler-x1.5-1.0.safetensors
+# ltx-2.3-temporal-upscaler-x2-1.0.safetensors
 upsampler_filename: "ltx-2.3-spatial-upscaler-x2-1.0.safetensors"
 upsampler_spatial_patch_size: 1
 upsampler_temporal_patch_size: 1

src/maxdiffusion/models/ltx2/ltx2_utils.py

@@ -26,6 +26,28 @@
 from flax.traverse_util import unflatten_dict, flatten_dict
 from ..modeling_flax_pytorch_utils import (rename_key, rename_key_and_reshape_tensor, torch2jax, validate_flax_state_dict)
 
+KNOWN_UPSAMPLER_CONFIGS = {
+    "ltx-2.3-spatial-upscaler-x2-1.0.safetensors": {
+        "spatial_upsample": True,
+        "temporal_upsample": False,
+        "rational_spatial_scale": None,
+    },
+    "ltx-2.3-spatial-upscaler-x2-1.1.safetensors": {
+        "spatial_upsample": True,
+        "temporal_upsample": False,
+        "rational_spatial_scale": None,
+    },
+    "ltx-2.3-spatial-upscaler-x1.5-1.0.safetensors": {
+        "spatial_upsample": True,
+        "temporal_upsample": False,
+        "rational_spatial_scale": 1.5,
+    },
+    "ltx-2.3-temporal-upscaler-x2-1.0.safetensors": {
+        "spatial_upsample": False,
+        "temporal_upsample": True,
+        "rational_spatial_scale": None,
+    },
+}
 
 def _tuple_str_to_int(in_tuple):
   out_list = []

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -47,6 +47,7 @@
     load_upsampler_weights,
     adain_filter_latent,
     tone_map_latents,
+    KNOWN_UPSAMPLER_CONFIGS,
 )
 from ...models.ltx2.text_encoders.text_encoders_ltx2 import LTX2AudioVideoGemmaTextEncoder
 from ...video_processor import VideoProcessor
@@ -572,17 +573,26 @@ def load_upsampler(cls, devices_array: np.array, mesh: Mesh, rngs: nnx.Rngs, con
     mid_channels = upsampler_config.get("mid_channels", 1024)
 
     if filename is not None:
-      if "temporal" in filename:
-        temporal_upsample = True
-        spatial_upsample = False
-      elif "spatial" in filename:
-        spatial_upsample = True
-        temporal_upsample = False
-
-      if "x1.5" in filename:
-        rational_spatial_scale = 1.5
-      elif "x2" in filename:
-        rational_spatial_scale = None  # Force fallback for x2
+      if filename in KNOWN_UPSAMPLER_CONFIGS:
+        mapped_config = KNOWN_UPSAMPLER_CONFIGS[filename]
+        spatial_upsample = mapped_config.get("spatial_upsample", spatial_upsample)
+        temporal_upsample = mapped_config.get("temporal_upsample", temporal_upsample)
+        rational_spatial_scale = mapped_config.get("rational_spatial_scale", rational_spatial_scale)
+      else:
+        max_logging.log(
+            f"Warning: Filename '{filename}' not in KNOWN_UPSAMPLER_CONFIGS, falling back to heuristic parsing."
+        )
+        if "temporal" in filename:
+          temporal_upsample = True
+          spatial_upsample = False
+        elif "spatial" in filename:
+          spatial_upsample = True
+          temporal_upsample = False
+
+        if "x1.5" in filename:
+          rational_spatial_scale = 1.5
+        elif "x2" in filename:
+          rational_spatial_scale = None  # Force fallback for x2
 
     max_logging.log(
         f"Upsampler config inferred: spatial={spatial_upsample}, temporal={temporal_upsample}, scale={rational_spatial_scale}, mid_channels={mid_channels}"
@@ -1315,7 +1325,7 @@ def __call__(
     audio_channels = (
         self.audio_vae.config.latent_channels
         if hasattr(self.audio_vae, "config") and hasattr(self.audio_vae.config, "latent_channels")
-        else 128
+        else 8
     )
 
     duration_s = num_frames / frame_rate