From 5c6f635d73d270722a807bfb3dede964329b98e4 Mon Sep 17 00:00:00 2001
From: Shao Duan <shaoxiongduan@gmail.com>
Date: Tue, 20 Jan 2026 05:31:50 +0000
Subject: [PATCH] Use LTX2 config defaults and HF model ID

---
 examples/inference/basic/basic_ltx2.py   |  5 ----
 fastvideo/configs/sample/registry.py     |  9 +------
 fastvideo/models/audio/ltx2_audio_vae.py | 34 ++++++++++++++++++------
 3 files changed, 27 insertions(+), 21 deletions(-)

diff --git a/examples/inference/basic/basic_ltx2.py b/examples/inference/basic/basic_ltx2.py
index f36457403..c76d859a2 100644
--- a/examples/inference/basic/basic_ltx2.py
+++ b/examples/inference/basic/basic_ltx2.py
@@ -35,11 +35,6 @@ def main() -> None:
         vae_cpu_offload=True,
         text_encoder_cpu_offload=True,
         pin_cpu_memory=False,
-        ltx2_vae_tiling=True,
-        ltx2_vae_spatial_tile_size_in_pixels=512,
-        ltx2_vae_spatial_tile_overlap_in_pixels=64,
-        ltx2_vae_temporal_tile_size_in_frames=64,
-        ltx2_vae_temporal_tile_overlap_in_frames=24,
     )
 
     output_path = "outputs_video/ltx2_basic/backyard_drama.mp4"
diff --git a/fastvideo/configs/sample/registry.py b/fastvideo/configs/sample/registry.py
index bda54ff05..6b26133d2 100644
--- a/fastvideo/configs/sample/registry.py
+++ b/fastvideo/configs/sample/registry.py
@@ -105,7 +105,7 @@ SAMPLING_PARAM_REGISTRY: dict[str, Any] = {
 
     # LTX-2 models
     "Lightricks/LTX-2": LTX2SamplingParam,
-    "converted/ltx2_diffusers": LTX2SamplingParam,
+    "FastVideo/LTX2-Distilled-Diffusers": LTX2SamplingParam,
 
     # Add other specific weight variants
 }
@@ -164,13 +164,6 @@ SAMPLING_FALLBACK_PARAM: dict[str, Any] = {
 def get_sampling_param_cls_for_name(pipeline_name_or_path: str) -> Any | None:
     """Get the appropriate sampling param for specific pretrained weights."""
 
-    # Temporary local mapping for LTX-2 diffusers-style exports.
-    if os.path.exists(pipeline_name_or_path):
-        norm_path = os.path.normpath(pipeline_name_or_path)
-        path_parts = norm_path.split(os.sep)
-        if "ltx2_diffusers" in path_parts:
-            return LTX2SamplingParam
-
     # First try exact match for specific weights
     if pipeline_name_or_path in SAMPLING_PARAM_REGISTRY:
         return SAMPLING_PARAM_REGISTRY[pipeline_name_or_path]
diff --git a/fastvideo/models/audio/ltx2_audio_vae.py b/fastvideo/models/audio/ltx2_audio_vae.py
index 0b13780f1..10b16c47c 100644
--- a/fastvideo/models/audio/ltx2_audio_vae.py
+++ b/fastvideo/models/audio/ltx2_audio_vae.py
@@ -12,12 +12,13 @@ import torch.nn as nn
 
 
 def _require_ltx2():
-    repo_root = Path(__file__).resolve().parents[3]
-    local_core = repo_root / "LTX-2" / "packages" / "ltx-core" / "src"
-    if local_core.exists() and str(local_core) not in sys.path:
-        sys.path.insert(0, str(local_core))
+    """Import LTX-2 audio VAE components from ltx_core.
 
+    Attempts to import from an installed ltx_core package first.
+    Falls back to local LTX-2 submodule if the package is not installed.
+    """
     try:
+        # Try importing from installed package first
         from ltx_core.model.audio_vae import (  # type: ignore
             AudioDecoder,
             AudioDecoderConfigurator,
@@ -26,10 +27,27 @@ def _require_ltx2():
             Vocoder,
             VocoderConfigurator,
         )
-    except ImportError as exc:
-        raise ImportError(
-            "LTX-2 sources are required. Ensure FastVideo/LTX-2 is present."
-        ) from exc
+    except ImportError:
+        # Fall back to local LTX-2 submodule
+        repo_root = Path(__file__).resolve().parents[3]
+        local_core = repo_root / "LTX-2" / "packages" / "ltx-core" / "src"
+        if local_core.exists() and str(local_core) not in sys.path:
+            sys.path.insert(0, str(local_core))
+
+        try:
+            from ltx_core.model.audio_vae import (  # type: ignore
+                AudioDecoder,
+                AudioDecoderConfigurator,
+                AudioEncoder,
+                AudioEncoderConfigurator,
+                Vocoder,
+                VocoderConfigurator,
+            )
+        except ImportError as exc:
+            raise ImportError(
+                "ltx_core is required for LTX-2 audio components. "
+                "Install it via pip or ensure FastVideo/LTX-2 submodule is present."
+            ) from exc
 
     return (
         AudioDecoder,