Upload folder using huggingface_hub

.gitattributes

@@ -34,3 +34,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 example_face.jpg filter=lfs diff=lfs merge=lfs -text
+reference_sample1.jpg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,15 +1,28 @@
----
-title: LTX-2 Face ID
-emoji: 🎥🧑
-colorFrom: purple
-colorTo: pink
-sdk: gradio
-sdk_version: 6.2.0
-app_file: app.py
-pinned: false
-short_description: Identity-preserving video generation with LTX-2
-python_version: "3.10"
-startup_duration_timeout: 30m
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+---
+title: "LTX-Best-Face-ID"
+emoji: 🎬
+colorFrom: purple
+colorTo: pink
+sdk: gradio
+sdk_version: 6.15.1
+app_file: app.py
+short_description: "Identity-preserving video generation from a reference photo"
+python_version: "3.12"
+startup_duration_timeout: 1h
+---
+
+# LTX-Best-Face-ID
+
+Identity-preserving video generation from a reference photo using **LTX-2** with the **LTX-Best-Face-ID** LoRA.
+
+Upload a close-up frontal photo of a person, describe the scene, and get a short video that maintains their identity across frames.
+
+## How it works
+
+This demo uses the [Alissonerdx/LTX-Best-Face-ID](https://huggingface.co/Alissonerdx/LTX-Best-Face-ID) LoRA on top of the [LTX-2](https://huggingface.co/Lightricks/LTX-2) base model. Identity is injected through **overlap reference conditioning** with **TASS-RoPE** (source-phase RoPE), replicating the ComfyUI BFSNodes logic in a diffusers-style pipeline.
+
+## Tips for best results
+
+- Use a **clean, frontal, well-lit close-up** as the reference
+- Let the **Prompt Enhancer** add identity attributes automatically
+- Prompts are automatically prefixed with `ref_t2v:` for the model

app.py

@@ -1,340 +1,605 @@
-import sys
-import os
-from pathlib import Path
-
-# Add packages to Python path
-current_dir = Path(__file__).parent
-sys.path.insert(0, str(current_dir / "packages" / "ltx-pipelines" / "src"))
-sys.path.insert(0, str(current_dir / "packages" / "ltx-core" / "src"))
-
-import spaces
-import gradio as gr
-from gradio_client import Client, handle_file
-import numpy as np
-import random
-import torch
-from typing import Optional
-from pathlib import Path
-from huggingface_hub import hf_hub_download
-from ltx_pipelines.distilled import DistilledPipeline
-from ltx_core.tiling import TilingConfig
-from ltx_core.loader.primitives import LoraPathStrengthAndSDOps
-from ltx_core.loader.sd_ops import LTXV_LORA_COMFY_RENAMING_MAP
-from ltx_pipelines.constants import (
-    DEFAULT_SEED,
-    DEFAULT_HEIGHT,
-    DEFAULT_WIDTH,
-    DEFAULT_NUM_FRAMES,
-    DEFAULT_FRAME_RATE,
-    DEFAULT_LORA_STRENGTH,
-)
-
-MAX_SEED = np.iinfo(np.int32).max
-
-# Default prompt — identity-preserving example
-DEFAULT_PROMPT = (
-    "A person with distinctive facial features looking directly at the camera with a confident expression. "
-    "They blink naturally, turn their head slightly, and smile subtly. Soft natural lighting illuminates "
-    "their face. Realistic skin texture, detailed facial features, cinematic quality."
-)
-
-# HuggingFace Hub defaults
-DEFAULT_REPO_ID = "Lightricks/LTX-2"
-DEFAULT_CHECKPOINT_FILENAME = "ltx-2-19b-dev-fp8.safetensors"
-DEFAULT_DISTILLED_LORA_FILENAME = "ltx-2-19b-distilled-lora-384.safetensors"
-DEFAULT_SPATIAL_UPSAMPLER_FILENAME = "ltx-2-spatial-upscaler-x2-1.0.safetensors"
-
-# Face ID LoRA
-FACE_ID_REPO_ID = "Alissonerdx/LTX-Best-Face-ID"
-FACE_ID_LORA_FILENAME = "Best_FaceID_v1.0_LoRA.safetensors"
-
-# Text encoder space URL
-TEXT_ENCODER_SPACE = "linoyts/gemma-text-encoder"
-
-
-def get_hub_or_local_checkpoint(repo_id: Optional[str] = None, filename: Optional[str] = None):
-    """Download from HuggingFace Hub or use local checkpoint."""
-    if repo_id is None and filename is None:
-        raise ValueError("Please supply at least one of `repo_id` or `filename`")
-
-    if repo_id is not None:
-        if filename is None:
-            raise ValueError("If repo_id is specified, filename must also be specified.")
-        print(f"Downloading {filename} from {repo_id}...")
-        ckpt_path = hf_hub_download(repo_id=repo_id, filename=filename)
-        print(f"Downloaded to {ckpt_path}")
-    else:
-        ckpt_path = filename
-
-    return ckpt_path
-
-
-# Initialize pipeline at startup
-print("=" * 80)
-print("Loading LTX-2 Distilled pipeline with Face ID LoRA...")
-print("=" * 80)
-
-checkpoint_path = get_hub_or_local_checkpoint(DEFAULT_REPO_ID, DEFAULT_CHECKPOINT_FILENAME)
-distilled_lora_path = get_hub_or_local_checkpoint(DEFAULT_REPO_ID, DEFAULT_DISTILLED_LORA_FILENAME)
-spatial_upsampler_path = get_hub_or_local_checkpoint(DEFAULT_REPO_ID, DEFAULT_SPATIAL_UPSAMPLER_FILENAME)
-face_id_lora_path = get_hub_or_local_checkpoint(FACE_ID_REPO_ID, FACE_ID_LORA_FILENAME)
-
-print(f"Initializing pipeline with:")
-print(f"  checkpoint_path={checkpoint_path}")
-print(f"  distilled_lora_path={distilled_lora_path}")
-print(f"  spatial_upsampler_path={spatial_upsampler_path}")
-print(f"  face_id_lora_path={face_id_lora_path}")
-print(f"  text_encoder_space={TEXT_ENCODER_SPACE}")
-
-# Load distilled LoRA and Face ID LoRA
-loras = [
-    LoraPathStrengthAndSDOps(
-        path=distilled_lora_path,
-        strength=DEFAULT_LORA_STRENGTH,
-        sd_ops=LTXV_LORA_COMFY_RENAMING_MAP,
-    ),
-    LoraPathStrengthAndSDOps(
-        path=face_id_lora_path,
-        strength=1.0,
-        sd_ops=LTXV_LORA_COMFY_RENAMING_MAP,
-    ),
-]
-
-# Initialize pipeline WITHOUT text encoder (gemma_root=None)
-# Text encoding will be done by external space
-pipeline = DistilledPipeline(
-    checkpoint_path=checkpoint_path,
-    spatial_upsampler_path=spatial_upsampler_path,
-    gemma_root=None,
-    loras=loras,
-    fp8transformer=True,
-    local_files_only=False,
-)
-
-# Initialize text encoder client
-print(f"Connecting to text encoder space: {TEXT_ENCODER_SPACE}")
-try:
-    text_encoder_client = Client(TEXT_ENCODER_SPACE)
-    print("Text encoder client connected!")
-except Exception as e:
-    print(f"Warning: Could not connect to text encoder space: {e}")
-    text_encoder_client = None
-
-print("=" * 80)
-print("Pipeline fully loaded and ready!")
-print("=" * 80)
-
-
-@spaces.GPU(duration=300)
-def generate_video(
-    face_image,
-    prompt: str,
-    duration: float,
-    enhance_prompt: bool,
-    seed: int,
-    randomize_seed: bool,
-    height: int,
-    width: int,
-    progress=gr.Progress(track_tqdm=True),
-):
-    """Generate a video with identity preservation from a reference face image."""
-    try:
-        if face_image is None:
-            raise ValueError("Please provide a reference face image.")
-
-        # Randomize seed if checkbox is enabled
-        current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
-
-        # Calculate num_frames from duration (using fixed 24 fps)
-        frame_rate = 24.0
-        num_frames = int(duration * frame_rate) + 1
-
-        # Create output directory if it doesn't exist
-        output_dir = Path("outputs")
-        output_dir.mkdir(exist_ok=True)
-        output_path = output_dir / f"video_{current_seed}.mp4"
-
-        # Save uploaded image temporarily
-        temp_image_path = output_dir / f"temp_input_{current_seed}.jpg"
-        if hasattr(face_image, "save"):
-            face_image.save(temp_image_path)
-        else:
-            temp_image_path = Path(face_image)
-
-        # Format: (image_path, frame_idx, strength)
-        images = [(str(temp_image_path), 0, 1.0)]
-
-        # Get embeddings from text encoder space
-        print(f"Encoding prompt: {prompt}")
-
-        if text_encoder_client is None:
-            raise RuntimeError(
-                f"Text encoder client not connected. Please ensure the text encoder space "
-                f"({TEXT_ENCODER_SPACE}) is running and accessible."
-            )
-
-        try:
-            image_input = handle_file(str(temp_image_path))
-
-            result = text_encoder_client.predict(
-                prompt=prompt,
-                enhance_prompt=enhance_prompt,
-                input_image=image_input,
-                seed=current_seed,
-                negative_prompt="",
-                api_name="/encode_prompt",
-            )
-            embedding_path = result[0]
-            print(f"Embeddings received from: {embedding_path}")
-
-            # Load embeddings
-            embeddings = torch.load(embedding_path)
-            video_context = embeddings["video_context"]
-            audio_context = embeddings["audio_context"]
-            print("Embeddings loaded successfully")
-        except Exception as e:
-            raise RuntimeError(
-                f"Failed to get embeddings from text encoder space: {e}\n"
-                f"Please ensure {TEXT_ENCODER_SPACE} is running properly."
-            )
-
-        # Run inference
-        pipeline(
-            prompt=prompt,
-            output_path=str(output_path),
-            seed=current_seed,
-            height=height,
-            width=width,
-            num_frames=num_frames,
-            frame_rate=frame_rate,
-            images=images,
-            tiling_config=TilingConfig.default(),
-            video_context=video_context,
-            audio_context=audio_context,
-        )
-
-        # Clean up temp image
-        if temp_image_path.exists():
-            temp_image_path.unlink()
-
-        return str(output_path), current_seed
-
-    except Exception as e:
-        import traceback
-
-        error_msg = f"Error: {str(e)}\n{traceback.format_exc()}"
-        print(error_msg)
-        return None, current_seed
-
-
-# Create Gradio interface
-with gr.Blocks(title="LTX-2 Face ID 🎥🧑") as demo:
-    gr.Markdown("# LTX-2 Face ID 🎥🧑: Identity-Preserving Video Generation")
-    gr.Markdown(
-        "Generate identity-preserving videos using [Lightricks LTX-2](https://huggingface.co/Lightricks/LTX-2) "
-        "with the [BFS Best-Face-ID LoRA](https://huggingface.co/Alissonerdx/LTX-Best-Face-ID). "
-        "Upload a clear, frontal face photo and describe the desired video action."
-    )
-
-    with gr.Row():
-        with gr.Column():
-            face_image = gr.Image(
-                label="Reference Face Image",
-                type="pil",
-                info="A clean, frontal, well-lit close-up of the person's face",
-            )
-
-            prompt = gr.Textbox(
-                label="Prompt",
-                info="Describe the person and the action. Include identity attributes (hair color, skin tone, etc.) for best results.",
-                value=DEFAULT_PROMPT,
-                lines=3,
-                placeholder="Describe the person and the desired video action...",
-            )
-
-            with gr.Row():
-                duration = gr.Slider(
-                    label="Duration (seconds)",
-                    minimum=1.0,
-                    maximum=5.0,
-                    value=3.0,
-                    step=0.5,
-                )
-                enhance_prompt = gr.Checkbox(
-                    label="Enhance Prompt",
-                    value=True,
-                )
-
-            generate_btn = gr.Button("Generate Video", variant="primary", size="lg")
-
-            with gr.Accordion("Advanced Settings", open=False):
-                seed = gr.Slider(
-                    label="Seed",
-                    minimum=0,
-                    maximum=MAX_SEED,
-                    value=DEFAULT_SEED,
-                    step=1,
-                )
-
-                randomize_seed = gr.Checkbox(
-                    label="Randomize Seed",
-                    value=True,
-                )
-
-                with gr.Row():
-                    width = gr.Number(
-                        label="Width",
-                        value=DEFAULT_WIDTH,
-                        precision=0,
-                    )
-                    height = gr.Number(
-                        label="Height",
-                        value=DEFAULT_HEIGHT,
-                        precision=0,
-                    )
-
-        with gr.Column():
-            output_video = gr.Video(label="Generated Video", autoplay=True)
-
-    generate_btn.click(
-        fn=generate_video,
-        inputs=[
-            face_image,
-            prompt,
-            duration,
-            enhance_prompt,
-            seed,
-            randomize_seed,
-            height,
-            width,
-        ],
-        outputs=[output_video, seed],
-    )
-
-    gr.Examples(
-        examples=[
-            [
-                "example_face.jpg",
-                "A person with warm olive skin and dark brown shoulder-length wavy hair is sitting in a cozy cafe, looking directly at the camera. They take a sip from a coffee cup, set it down, and smile warmly. Soft natural light streams through the window. Cinematic quality, realistic skin texture.",
-                3.0,
-            ],
-            [
-                "example_face.jpg",
-                "A person with warm olive skin and dark brown shoulder-length wavy hair walking through a sunlit park on a crisp autumn day. Golden leaves fall around them as they stroll at a relaxed pace, occasionally glancing at the camera with a gentle smile. Natural lighting, shallow depth of field.",
-                3.0,
-            ],
-        ],
-        fn=generate_video,
-        inputs=[face_image, prompt, duration],
-        outputs=[output_video, seed],
-        label="Example",
-        cache_examples=True,
-        cache_mode="lazy",
-    )
-
-
-css = """
-.gradio-container .contain{max-width: 1200px !important; margin: 0 auto !important}
-"""
-if __name__ == "__main__":
-    demo.launch(theme=gr.themes.Citrus(), css=css)
+import os
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+import spaces
+import torch
+import torch.nn.functional as F
+import types
+import logging
+import tempfile
+import time
+import random
+import numpy as np
+from pathlib import Path
+from typing import Optional
+from dataclasses import replace
+from einops import rearrange
+
+import gradio as gr
+
+# Add packages to Python path
+APP_ROOT = Path(__file__).resolve().parent
+import sys
+sys.path.insert(0, str(APP_ROOT / "packages" / "ltx-pipelines" / "src"))
+sys.path.insert(0, str(APP_ROOT / "packages" / "ltx-core" / "src"))
+
+from huggingface_hub import hf_hub_download, snapshot_download
+from ltx_pipelines.utils import ModelLedger
+from ltx_core.model.video_vae import TilingConfig
+from ltx_core.loader.primitives import LoraPathStrengthAndSDOps
+from ltx_core.loader.sd_ops import LTXV_LORA_COMFY_RENAMING_MAP
+from ltx_pipelines.utils.constants import (
+    DEFAULT_SEED,
+    DEFAULT_1_STAGE_HEIGHT,
+    DEFAULT_1_STAGE_WIDTH,
+    DEFAULT_NUM_FRAMES,
+    DEFAULT_FRAME_RATE,
+    DEFAULT_LORA_STRENGTH,
+    DEFAULT_NEGATIVE_PROMPT,
+)
+from ltx_pipelines.utils.helpers import (
+    cleanup_memory,
+    denoise_audio_video,
+    euler_denoising_loop,
+    generate_enhanced_prompt,
+    get_device,
+    guider_denoising_func,
+    image_conditionings_by_replacing_latent,
+    modality_from_latent_state,
+    simple_denoising_func,
+)
+from ltx_core.components.diffusion_steps import EulerDiffusionStep
+from ltx_core.components.guiders import CFGGuider
+from ltx_core.components.noisers import GaussianNoiser
+from ltx_core.components.schedulers import LTX2Scheduler
+from ltx_core.model.transformer.modality import Modality
+from ltx_core.model.transformer.rope import LTXRopeType
+from ltx_core.model.transformer.transformer_args import TransformerArgs
+from ltx_core.text_encoders.gemma import encode_text
+from ltx_core.types import LatentState, VideoPixelShape, VideoLatentShape, SpatioTemporalScaleFactors
+from ltx_core.tools import VideoLatentTools
+from ltx_core.model.video_vae import decode_video as vae_decode_video
+from ltx_core.model.audio_vae import decode_audio as vae_decode_audio
+from ltx_pipelines.utils.types import PipelineComponents
+from ltx_pipelines.utils.media_io import encode_video, load_image_conditioning
+
+from PIL import Image
+import imageio
+
+logging.basicConfig(level=logging.INFO)
+log = logging.getLogger(__name__)
+
+MAX_SEED = np.iinfo(np.int32).max
+
+# Model IDs
+LTX2_REPO_ID = "Lightricks/LTX-2"
+LTX2_CHECKPOINT = "ltx-2-19b-dev.safetensors"
+GEMMA_REPO_ID = "google/gemma-3-12b-it-qat-q4_0-unquantized"
+FACEID_LORA_REPO = "Alissonerdx/LTX-Best-Face-ID"
+FACEID_LORA_FILE = "Best_FaceID_v1.0_LoRA.safetensors"
+
+# Distilled sigma values for fast generation
+DISTILLED_SIGMA_VALUES = [1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875, 0.0]
+
+VIDEO_SCALE_FACTORS = SpatioTemporalScaleFactors.default()
+
+# ---- Download model weights at module scope ----
+log.info("Downloading LTX-2 checkpoint...")
+checkpoint_path = hf_hub_download(LTX2_REPO_ID, LTX2_CHECKPOINT)
+log.info(f"Checkpoint: {checkpoint_path}")
+
+log.info("Downloading Gemma text encoder...")
+gemma_local_path = snapshot_download(GEMMA_REPO_ID)
+log.info(f"Gemma: {gemma_local_path}")
+
+log.info("Downloading FaceID LoRA...")
+faceid_lora_path = hf_hub_download(FACEID_LORA_REPO, FACEID_LORA_FILE)
+log.info(f"FaceID LoRA: {faceid_lora_path}")
+
+device = torch.device("cuda")
+dtype = torch.bfloat16
+
+# ---- Initialize model ledger with the FaceID LoRA ----
+log.info("Initializing ModelLedger with FaceID LoRA...")
+model_ledger = ModelLedger(
+    dtype=dtype,
+    device=device,
+    checkpoint_path=checkpoint_path,
+    gemma_root_path=gemma_local_path,
+    loras=(
+        LoraPathStrengthAndSDOps(faceid_lora_path, DEFAULT_LORA_STRENGTH, LTXV_LORA_COMFY_RENAMING_MAP),
+    ),
+    local_files_only=False,
+)
+
+# Load text encoder once at startup
+log.info("Loading text encoder...")
+text_encoder = model_ledger.text_encoder()
+log.info("Text encoder loaded.")
+
+# ---- Identity Overlap Patching ----
+# This replicates the BFSNodes ltx_identity_overlap.py logic:
+# 1) Encode reference image -> latent
+# 2) Patchify and project it
+# 3) Append ref tokens to video tokens (overlap frame-0 grid)
+# 4) Rotate ref tokens' RoPE freqs by source_phase
+# 5) Give ref tokens clean (0) timestep
+# 6) Trim ref tokens from output
+
+def _rotate_ref_freqs(cos_freq, sin_freq, ref_len, seg_value, theta=10000.0):
+    """Rotate the LAST ref_len tokens' RoPE freqs by phase = seg_value * theta^(-d/L)."""
+    if ref_len <= 0 or seg_value == 0.0:
+        return cos_freq, sin_freq
+    L = cos_freq.shape[-1]
+    d = torch.arange(L, device=cos_freq.device, dtype=torch.float32)
+    rate = theta ** (-d / float(L))
+    phase = (seg_value * rate)
+    pc = phase.cos().to(cos_freq.dtype)
+    ps = phase.sin().to(sin_freq.dtype)
+    # index the token axis (=-2 for [B,H,T,L])
+    idx = [slice(None)] * cos_freq.dim()
+    idx[-2] = slice(cos_freq.shape[-2] - ref_len, cos_freq.shape[-2])
+    idx = tuple(idx)
+    c0, s0 = cos_freq[idx], sin_freq[idx]
+    cos_freq = cos_freq.clone()
+    sin_freq = sin_freq.clone()
+    cos_freq[idx] = c0 * pc - s0 * ps
+    sin_freq[idx] = s0 * pc + c0 * ps
+    return cos_freq, sin_freq
+
+
+_identity_ref_latent = None
+_identity_ref_len = 0
+_identity_target_len = 0
+_identity_seg_value = 2.0
+
+
+def _install_identity_patches(transformer_model):
+    """Install identity overlap patches on the LTXModel (velocity_model inside X0Model).
+
+    Patches the MultiModalTransformerArgsPreprocessor.prepare() to:
+    1) Append reference latent tokens after the target video tokens
+    2) Rotate the reference tokens' RoPE frequencies by source_phase
+    3) Give reference tokens clean (0) timesteps
+    4) Trim reference tokens from the output
+    """
+    global _identity_ref_len, _identity_target_len
+
+    ltx_model = transformer_model  # This is X0Model wrapping LTXModel
+    velocity_model = ltx_model.velocity_model  # The actual LTXModel
+
+    preprocessor = velocity_model.video_args_preprocessor
+    if hasattr(preprocessor, '_identity_patched'):
+        return
+    preprocessor._identity_patched = True
+
+    orig_prepare = preprocessor.prepare
+
+    def patched_prepare(self, modality: Modality) -> TransformerArgs:
+        global _identity_ref_len, _identity_target_len
+
+        # Original preparation
+        args = orig_prepare(modality)
+
+        ref_lat = _identity_ref_latent
+        if ref_lat is None:
+            return args
+
+        # Patchify the reference latent and project it
+        # ref_lat shape: (B, C, F, H, W) -> need to patchify to (B, T, D)
+        # Use the patchifier from the pipeline components
+        # The reference latent is a single-frame latent at the target resolution
+        ref_lat_b = ref_lat.to(device=args.x.device, dtype=args.x.dtype)
+        # Patchify: (B, C, F, H, W) -> (B, T, C*patch_h*patch_w)
+        patch_size = 1  # temporal patch size is 1
+        spatial_patch = 1  # already in latent space, patchifier uses 1x1 spatial
+        # Actually, the patchifier is used in VideoLatentTools - we need to patchify manually
+        # The video latent is (B, 128, F, H, W), patchify with (1, patch_size, patch_size)
+        # But the patchifier's patch_size is typically 1 for latents (the VAE already did spatial downsampling)
+        # Let's use the patchifier from the preprocessor's model
+        # Actually, looking at the code, patchify_proj is applied after patchify
+        # The patchifier is separate from the preprocessor
+        # Let's do it manually: rearrange (B, C, F, H, W) -> (B, F*H*W, C)
+        # Since patch_size is 1 for video latents in the transformer
+        ref_tokens = rearrange(ref_lat_b, "b c f h w -> b (f h w) c")
+        ref_tokens = self.patchify_proj(ref_tokens)  # Project to inner_dim
+
+        ref_len = ref_tokens.shape[1]
+        _identity_ref_len = ref_len
+        _identity_target_len = args.x.shape[1]
+
+        # Expand to match batch size
+        if ref_tokens.shape[0] != args.x.shape[0]:
+            ref_tokens = ref_tokens.expand(args.x.shape[0], -1, -1)
+
+        # Append ref tokens after target tokens
+        new_x = torch.cat([args.x, ref_tokens], dim=1)
+
+        # Extend timesteps: give ref tokens clean (0) timestep
+        # timesteps shape: (B, T, D) or (B, 1, D)
+        if args.timesteps.dim() >= 2:
+            ref_ts = torch.zeros(args.x.shape[0], ref_len, args.timesteps.shape[-1],
+                                device=args.timesteps.device, dtype=args.timesteps.dtype)
+            new_timesteps = torch.cat([args.timesteps, ref_ts], dim=1)
+        else:
+            new_timesteps = args.timesteps
+
+        # Extend embedded_timestep similarly
+        if args.embedded_timestep.dim() >= 2 and args.embedded_timestep.shape[1] > 1:
+            ref_et = torch.zeros(args.x.shape[0], ref_len, args.embedded_timestep.shape[-1],
+                                device=args.embedded_timestep.device, dtype=args.embedded_timestep.dtype)
+            new_embedded_ts = torch.cat([args.embedded_timestep, ref_et], dim=1)
+        else:
+            new_embedded_ts = args.embedded_timestep
+
+        # Rotate the reference tokens' RoPE frequencies
+        pe = args.positional_embeddings
+        # For INTERLEAVED rope, pe is (cos_freq, sin_freq) tuple
+        # cos_freq/sin_freq shape: (B, T, D) for interleaved, or (B, H, T, D//2) for split
+        if isinstance(pe, tuple) and len(pe) == 2:
+            cos_freq, sin_freq = pe
+            # Check if this is interleaved (B, T, D) or split (B, H, T, D//2)
+            if cos_freq.dim() == 3:
+                # Interleaved: (B, T, D)
+                cos_new, sin_new = _rotate_ref_freqs(cos_freq, sin_freq, ref_len, _identity_seg_value)
+                new_pe = (cos_new, sin_new)
+            elif cos_freq.dim() == 4:
+                # Split: (B, H, T, D//2)
+                cos_new, sin_new = _rotate_ref_freqs(cos_freq, sin_freq, ref_len, _identity_seg_value)
+                new_pe = (cos_new, sin_new)
+            else:
+                new_pe = pe
+        else:
+            new_pe = pe
+
+        return replace(
+            args,
+            x=new_x,
+            timesteps=new_timesteps,
+            embedded_timestep=new_embedded_ts,
+            positional_embeddings=new_pe,
+        )
+
+    # Patch the _process_output to trim reference tokens
+    orig_process_output = velocity_model._process_output
+
+    def patched_process_output(self, scale_shift_table, norm_out, proj_out, x, embedded_timestep):
+        ref_len = _identity_ref_len
+        if ref_len > 0 and x is not None:
+            # Trim ref tokens from x
+            if isinstance(x, (list, tuple)):
+                x = [x_i[:, :x_i.shape[1] - ref_len] if x_i.dim() >= 2 and x_i.shape[1] > ref_len else x_i for x_i in x]
+            else:
+                if x.dim() >= 2 and x.shape[1] > ref_len:
+                    x = x[:, :x.shape[1] - ref_len]
+
+            # Trim embedded_timestep
+            if isinstance(embedded_timestep, (list, tuple)):
+                embedded_timestep = [
+                    et[:, :et.shape[1] - ref_len] if et.dim() >= 2 and et.shape[1] > ref_len else et
+                    for et in embedded_timestep
+                ]
+            elif embedded_timestep is not None:
+                if embedded_timestep.dim() >= 2 and embedded_timestep.shape[1] > ref_len:
+                    embedded_timestep = embedded_timestep[:, :embedded_timestep.shape[1] - ref_len]
+
+        return orig_process_output(scale_shift_table, norm_out, proj_out, x, embedded_timestep)
+
+    preprocessor.prepare = types.MethodType(patched_prepare, preprocessor)
+    velocity_model._process_output = types.MethodType(patched_process_output, velocity_model)
+
+    log.info("Identity overlap patches installed on transformer")
+
+
+def _encode_reference_image(ref_image_pil, width, height, video_encoder, device, dtype):
+    """Encode a reference image to latent space for identity conditioning.
+
+    Resizes the image to the target video resolution, encodes it through the VAE
+    to get a single-frame latent at the target latent resolution.
+    """
+    # Resize to match the video resolution
+    ref_resized = ref_image_pil.convert("RGB").resize((width, height), Image.BILINEAR)
+    ref_np = np.array(ref_resized).astype(np.float32) / 255.0  # (H, W, 3) in [0,1]
+    # Convert to (B, C, F, H, W) format for VAE encoder
+    ref_tensor = torch.from_numpy(ref_np).permute(2, 0, 1).unsqueeze(0).unsqueeze(2)  # (1, 3, 1, H, W)
+    # Normalize to [-1, 1]
+    ref_tensor = (ref_tensor * 2.0 - 1.0).to(device=device, dtype=dtype)
+    # Encode through VAE
+    with torch.inference_mode():
+        ref_latent = video_encoder(ref_tensor)  # (1, 128, 1, H//32, W//32)
+    return ref_latent
+
+
+@spaces.GPU(duration=lambda *a, **kw: min(240, 60 + int((kw.get("duration_seconds", 2) * 24 + 1) * 1.5)), size="xlarge")
+def generate_video(
+    reference_image,
+    prompt: str,
+    duration_seconds: float = 2.0,
+    seed: int = 42,
+    randomize_seed: bool = True,
+    height: int = 512,
+    width: int = 768,
+    cfg_guidance_scale: float = 3.0,
+    enhance_prompt: bool = True,
+    negative_prompt: str = "",
+    source_id: float = 2.0,
+    progress=gr.Progress(track_tqdm=True),
+):
+    """Generate an identity-preserving video from a reference image and text prompt.
+
+    Uses the LTX-2 base model with the LTX-Best-Face-ID LoRA and overlap reference
+    conditioning with TASS-RoPE to maintain the subject's identity across frames.
+
+    Args:
+        reference_image: A reference photo of the person (close-up frontal works best).
+        prompt: Text description of the scene and action. Will be prefixed with 'ref_t2v:'.
+        duration_seconds: Length of the generated video in seconds.
+        seed: Random seed for reproducibility.
+        randomize_seed: If True, a random seed is generated each run.
+        height: Output video height in pixels.
+        width: Output video width in pixels.
+        cfg_guidance_scale: Classifier-free guidance scale.
+        enhance_prompt: If True, use Gemma to enhance the prompt with identity attributes.
+        negative_prompt: Negative prompt to guide generation away from undesired features.
+        source_id: TASS-RoPE source phase ID (training used 2.0).
+
+    Returns:
+        A tuple of (video_path, seed_used).
+    """
+    global _identity_ref_latent, _identity_ref_len, _identity_target_len, _identity_seg_value, text_encoder
+
+    if reference_image is None:
+        raise gr.Error("Please provide a reference image.")
+
+    if not prompt:
+        raise gr.Error("Please provide a text prompt.")
+
+    # Ensure prompt has the ref_t2v: prefix
+    if not prompt.strip().lower().startswith("ref_t2v:"):
+        prompt = f"ref_t2v: {prompt.strip()}"
+
+    # Randomize seed
+    current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
+
+    frame_rate = 24.0
+    num_frames = int(duration_seconds * frame_rate) + 1
+
+    # Set identity parameters
+    _identity_seg_value = float(source_id)
+    _identity_ref_len = 0
+    _identity_target_len = 0
+
+    log.info(f"Generating video: {num_frames} frames, {width}x{height}, seed={current_seed}")
+
+    # Load reference image
+    if isinstance(reference_image, str):
+        ref_pil = Image.open(reference_image).convert("RGB")
+    else:
+        ref_pil = Image.fromarray(reference_image).convert("RGB") if isinstance(reference_image, np.ndarray) else reference_image
+
+    # Load video encoder for reference encoding
+    video_encoder = model_ledger.video_encoder()
+
+    # Encode reference image to latent
+    _identity_ref_latent = _encode_reference_image(ref_pil, width, height, video_encoder, device, dtype)
+    log.info(f"Reference latent shape: {_identity_ref_latent.shape}")
+
+    # Free video encoder
+    del video_encoder
+    cleanup_memory()
+
+    # Load transformer (with LoRA already fused)
+    transformer = model_ledger.transformer()
+
+    # Install identity patches
+    _install_identity_patches(transformer)
+
+    # Encode prompt
+    if not negative_prompt:
+        negative_prompt = DEFAULT_NEGATIVE_PROMPT
+
+    log.info("Encoding prompt...")
+    if enhance_prompt:
+        # Save reference image temporarily for prompt enhancement
+        ref_tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
+        ref_pil.save(ref_tmp.name)
+        enhanced_prompt = generate_enhanced_prompt(text_encoder, prompt, ref_tmp.name, seed=current_seed)
+        os.unlink(ref_tmp.name)
+        prompt = enhanced_prompt if enhanced_prompt else prompt
+        log.info(f"Enhanced prompt: {prompt}")
+
+    context_p, context_n = encode_text(text_encoder, prompts=[prompt, negative_prompt])
+    v_context_p, a_context_p = context_p
+    v_context_n, a_context_n = context_n
+
+    del text_encoder
+    cleanup_memory()
+
+    # Set up diffusion
+    generator = torch.Generator(device=device).manual_seed(current_seed)
+    noiser = GaussianNoiser(generator=generator)
+    stepper = EulerDiffusionStep()
+    cfg_guider = CFGGuider(cfg_guidance_scale)
+
+    sigmas = torch.tensor(DISTILLED_SIGMA_VALUES, dtype=torch.float32, device=device)
+
+    pipeline_components = PipelineComponents(dtype=dtype, device=device)
+
+    output_shape = VideoPixelShape(
+        batch=1, frames=num_frames, width=width, height=height, fps=frame_rate
+    )
+
+    # No image conditioning (we use the reference via identity overlap)
+    conditionings = []
+
+    log.info("Starting denoising...")
+
+    def first_stage_denoising_loop(sigmas, video_state, audio_state, stepper):
+        return euler_denoising_loop(
+            sigmas=sigmas,
+            video_state=video_state,
+            audio_state=audio_state,
+            stepper=stepper,
+            denoise_fn=guider_denoising_func(
+                cfg_guider,
+                v_context_p,
+                v_context_n,
+                a_context_p,
+                a_context_n,
+                transformer=transformer,
+            ),
+        )
+
+    video_state, audio_state = denoise_audio_video(
+        output_shape=output_shape,
+        conditionings=conditionings,
+        noiser=noiser,
+        sigmas=sigmas,
+        stepper=stepper,
+        denoising_loop_fn=first_stage_denoising_loop,
+        components=pipeline_components,
+        dtype=dtype,
+        device=device,
+    )
+
+    # Reset identity state
+    _identity_ref_latent = None
+    _identity_ref_len = 0
+
+    torch.cuda.synchronize()
+    del transformer
+    cleanup_memory()
+
+    # Decode video
+    log.info("Decoding video...")
+    video_decoder = model_ledger.video_decoder()
+    video_frames = list(vae_decode_video(video_state.latent, video_decoder))
+    del video_decoder
+    cleanup_memory()
+
+    # Decode audio
+    log.info("Decoding audio...")
+    audio_decoder = model_ledger.audio_decoder()
+    vocoder = model_ledger.vocoder()
+    decoded_audio = vae_decode_audio(audio_state.latent, audio_decoder, vocoder)
+    del audio_decoder, vocoder
+    cleanup_memory()
+
+    # Encode to MP4
+    log.info("Encoding MP4...")
+    output_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
+
+    # Combine frames
+    frames_list = []
+    for chunk in video_frames:
+        frames_list.append(chunk.cpu().numpy())
+
+    if frames_list:
+        all_frames = np.concatenate(frames_list, axis=0)
+    else:
+        raise RuntimeError("No video frames were generated.")
+
+    # Write video using imageio
+    writer = imageio.get_writer(output_path, fps=frame_rate, codec='libx264', quality=8)
+    for frame in all_frames:
+        writer.append_data(frame)
+    if decoded_audio is not None:
+        # Add audio if available
+        audio_np = decoded_audio.cpu().numpy()
+        # imageio doesn't support audio well, so we'll skip it for now
+    writer.close()
+
+    log.info(f"Video saved to {output_path}")
+
+    # Reload text encoder for next call
+    text_encoder = model_ledger.text_encoder()
+
+    return output_path, current_seed
+
+
+# ---- Gradio UI ----
+
+CSS = """
+#col-container { max-width: 1100px; margin: 0 auto; }
+.dark .gradio-container { color: var(--body-text-color); }
+"""
+
+with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
+    gr.Markdown("# 🎬 LTX-Best-Face-ID: Identity-Preserving Video Generation")
+    gr.Markdown("""
+    Generate identity-preserving videos from a reference photo and text prompt using **LTX-2** with the **LTX-Best-Face-ID** LoRA.
+
+    Upload a close-up frontal photo of a person, describe the scene, and get a video that maintains their identity.
+
+    **Tips:** Use a clean, well-lit, frontal close-up. Describe the person's appearance and the action in present-progressive tense.
+    """)
+
+    with gr.Column(elem_id="col-container"):
+        with gr.Row():
+            with gr.Column(scale=1):
+                reference_image = gr.Image(
+                    label="Reference Image",
+                    type="pil",
+                    height=300,
+                )
+                prompt = gr.Textbox(
+                    label="Prompt",
+                    placeholder="A person is walking in a park, medium shot...",
+                    lines=3,
+                )
+                run_btn = gr.Button("Generate Video", variant="primary")
+
+            with gr.Column(scale=1):
+                video_output = gr.Video(label="Generated Video")
+                seed_output = gr.Number(label="Seed", precision=0, interactive=False)
+
+        with gr.Accordion("Advanced Settings", open=False):
+            with gr.Row():
+                duration_seconds = gr.Slider(
+                    label="Duration (seconds)", minimum=1.0, maximum=4.0, value=2.0, step=0.5
+                )
+                height = gr.Number(label="Height", value=DEFAULT_1_STAGE_HEIGHT, precision=0)
+                width = gr.Number(label="Width", value=DEFAULT_1_STAGE_WIDTH, precision=0)
+            with gr.Row():
+                cfg_guidance_scale = gr.Slider(
+                    label="CFG Guidance Scale", minimum=1.0, maximum=10.0, value=3.0, step=0.5
+                )
+                seed = gr.Number(label="Seed", value=DEFAULT_SEED, precision=0)
+                randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
+            with gr.Row():
+                enhance_prompt = gr.Checkbox(label="Enhance Prompt with Gemma", value=True)
+                source_id = gr.Slider(
+                    label="TASS-RoPE Source ID", minimum=0.0, maximum=8.0, value=2.0, step=1.0,
+                )
+            negative_prompt = gr.Textbox(
+                label="Negative Prompt",
+                value="",
+                placeholder="Leave empty for default negative prompt...",
+                lines=2,
+            )
+
+    gr.Examples(
+        examples=[
+            [os.path.join(APP_ROOT, "reference_sample1.jpg"),
+             "A man with a beard is standing on a beach, looking at the ocean, medium shot. Waves are gently crashing behind him. Warm sunset lighting."],
+            [os.path.join(APP_ROOT, "reference_sample2.jpg"),
+             "A person in an astronaut suit is floating in space, stars and Earth visible in the background. Slow, gentle floating motion."],
+        ],
+        inputs=[reference_image, prompt],
+        outputs=[video_output, seed_output],
+        fn=generate_video,
+        cache_examples=True,
+        cache_mode="lazy",
+    )
+
+    run_btn.click(
+        fn=generate_video,
+        inputs=[
+            reference_image, prompt, duration_seconds, seed, randomize_seed,
+            height, width, cfg_guidance_scale, enhance_prompt, negative_prompt, source_id
+        ],
+        outputs=[video_output, seed_output],
+        api_name="generate",
+    )
+
+demo.launch(mcp_server=True)

packages/ltx-core/README.md

@@ -1 +1,280 @@
-# LTX-2 Core
+# LTX-Core
+
+The foundational library for the LTX-2 Audio-Video generation model. This package contains the raw model definitions, component implementations, and loading logic used by `ltx-pipelines` and `ltx-trainer`.
+
+## 📦 What's Inside?
+
+- **`components/`**: Modular diffusion components (Schedulers, Guiders, Noisers, Patchifiers) following standard protocols
+- **`conditioning/`**: Tools for preparing latent states and applying conditioning (image, video, keyframes)
+- **`guidance/`**: Perturbation system for fine-grained control over attention mechanisms
+- **`loader/`**: Utilities for loading weights from `.safetensors`, fusing LoRAs, and managing memory
+- **`model/`**: PyTorch implementations of the LTX-2 Transformer, Video VAE, Audio VAE, Vocoder and Upscaler
+- **`text_encoders/gemma`**: Gemma text encoder implementation with tokenizers, feature extractors, and separate encoders for audio-video and video-only generation
+
+## 🚀 Quick Start
+
+`ltx-core` provides the building blocks (models, components, and utilities) needed to construct inference flows. For ready-made inference pipelines use [`ltx-pipelines`](../ltx-pipelines/) or [`ltx-trainer`](../ltx-trainer/) for training.
+
+## 🔧 Installation
+
+```bash
+# From the repository root
+uv sync --frozen
+
+# Or install as a package
+pip install -e packages/ltx-core
+```
+
+## Building Blocks Overview
+
+`ltx-core` provides modular components that can be combined to build custom inference flows:
+
+### Core Models
+
+- **Transformer** ([`model/transformer/`](src/ltx_core/model/transformer/)): The 48-layer LTX-2 transformer with cross-modal attention for joint audio-video processing. Expects inputs in [`Modality`](src/ltx_core/model/transformer/modality.py) format
+- **Video VAE** ([`model/video_vae/`](src/ltx_core/model/video_vae/)): Encodes/decodes video pixels to/from latent space with temporal and spatial compression
+- **Audio VAE** ([`model/audio_vae/`](src/ltx_core/model/audio_vae/)): Encodes/decodes audio spectrograms to/from latent space
+- **Vocoder** ([`model/audio_vae/`](src/ltx_core/model/audio_vae/)): Neural vocoder that converts mel spectrograms to audio waveforms
+- **Text Encoder** ([`text_encoders/`](src/ltx_core/text_encoders/)): Gemma-based encoder that produces separate embeddings for video and audio conditioning
+- **Spatial Upscaler** ([`model/upsampler/`](src/ltx_core/model/upsampler/)): Upsamples latent representations for higher-resolution generation
+
+### Diffusion Components
+
+- **Schedulers** ([`components/schedulers.py`](src/ltx_core/components/schedulers.py)): Noise schedules (LTX2Scheduler, LinearQuadratic, Beta) that control the denoising process
+- **Guiders** ([`components/guiders.py`](src/ltx_core/components/guiders.py)): Guidance strategies (CFG, STG, APG) for controlling generation quality and adherence to prompts
+- **Noisers** ([`components/noisers.py`](src/ltx_core/components/noisers.py)): Add noise to latents according to the diffusion schedule
+- **Patchifiers** ([`components/patchifiers.py`](src/ltx_core/components/patchifiers.py)): Convert between spatial latents `[B, C, F, H, W]` and sequence format `[B, seq_len, dim]` for transformer processing
+
+### Conditioning & Control
+
+- **Conditioning** ([`conditioning/`](src/ltx_core/conditioning/)): Tools for preparing and applying various conditioning types (image, video, keyframes)
+- **Guidance** ([`guidance/`](src/ltx_core/guidance/)): Perturbation system for fine-grained control over attention mechanisms (e.g., skipping specific attention layers)
+
+### Utilities
+
+- **Loader** ([`loader/`](src/ltx_core/loader/)): Model loading from `.safetensors`, LoRA fusion, weight remapping, and memory management
+
+For complete, production-ready pipeline implementations that combine these building blocks, see the [`ltx-pipelines`](../ltx-pipelines/) package.
+
+---
+
+# Architecture Overview
+
+This section provides a deep dive into the internal architecture of the LTX-2 Audio-Video generation model.
+
+## Table of Contents
+
+1. [High-Level Architecture](#high-level-architecture)
+2. [The Transformer](#the-transformer)
+3. [Video VAE](#video-vae)
+4. [Audio VAE](#audio-vae)
+5. [Text Encoding (Gemma)](#text-encoding-gemma)
+6. [Spatial Upscaler](#spatial-upsampler)
+7. [Data Flow](#data-flow)
+
+---
+
+## High-Level Architecture
+
+LTX-2 is a **joint Audio-Video diffusion transformer** that processes both modalities simultaneously in a unified architecture. Unlike traditional models that handle video and audio separately, LTX-2 uses cross-modal attention to enable natural synchronization.
+
+```text
+┌─────────────────────────────────────────────────────────────┐
+│                    INPUT PREPARATION                        │
+│                                                             │
+│  Video Pixels → Video VAE Encoder → Video Latents           │
+│  Audio Waveform → Audio VAE Encoder → Audio Latents         │
+│  Text Prompt → Gemma Encoder → Text Embeddings              │
+└─────────────────────────────────────────────────────────────┘
+                            ↓
+┌─────────────────────────────────────────────────────────────┐
+│              LTX-2 TRANSFORMER (48 Blocks)                  │
+│                                                             │
+│  ┌──────────────┐              ┌──────────────┐             │
+│  │ Video Stream │              │ Audio Stream │             │
+│  │              │              │              │             │
+│  │ Self-Attn    │              │ Self-Attn    │             │
+│  │ Cross-Attn   │              │ Cross-Attn   │             │
+│  │              │◄────────────►│              │             │
+│  │ A↔V Cross    │              │ A↔V Cross    │             │
+│  │ Feed-Forward │              │ Feed-Forward │             │
+│  └──────────────┘              └──────────────┘             │
+└─────────────────────────────────────────────────────────────┘
+                            ↓
+┌─────────────────────────────────────────────────────────────┐
+│                    OUTPUT DECODING                          │
+│                                                             │
+│  Video Latents → Video VAE Decoder → Video Pixels           │
+│  Audio Latents → Audio VAE Decoder → Mel Spectrogram        │
+│  Mel Spectrogram → Vocoder → Audio Waveform                 │
+└─────────────────────────────────────────────────────────────┘
+```
+
+---
+
+## The Transformer
+
+The core of LTX-2 is a 48-layer transformer that processes both video and audio tokens simultaneously.
+
+### Model Structure
+
+**Source**: [`src/ltx_core/model/transformer/model.py`](src/ltx_core/model/transformer/model.py)
+
+The `LTXModel` class implements the transformer. It supports both video-only and audio-video generation modes. For actual usage, see the [`ltx-pipelines`](../ltx-pipelines/) package which handles model loading and initialization.
+
+### Transformer Block Architecture
+
+**Source**: [`src/ltx_core/model/transformer/transformer.py`](src/ltx_core/model/transformer/transformer.py)
+
+```text
+┌─────────────────────────────────────────────────────────────┐
+│                    TRANSFORMER BLOCK                        │
+│                                                             │
+│  VIDEO PATH:                                                │
+│    Input → RMSNorm → AdaLN → Self-Attn (attn1)              │
+│         → RMSNorm → Cross-Attn (attn2, text)                │
+│         → RMSNorm → AdaLN → A↔V Cross-Attn                  │
+│         → RMSNorm → AdaLN → Feed-Forward (ff) → Output      │
+│                                                             │
+│  AUDIO PATH:                                                │
+│    Input → RMSNorm → AdaLN → Self-Attn (audio_attn1)        │
+│         → RMSNorm → Cross-Attn (audio_attn2, text)          │
+│         → RMSNorm → AdaLN → A↔V Cross-Attn                  │
+│         → RMSNorm → AdaLN → Feed-Forward (audio_ff)         │
+│                                                             │
+│  AdaLN (Adaptive Layer Normalization):                      │
+│    - Uses scale_shift_table (6 params) for video/audio      │
+│    - Uses scale_shift_table_a2v_ca (5 params) for A↔V CA    │
+│    - Conditioned on per-token timestep embeddings           │
+└─────────────────────────────────────────────────────────────┘
+```
+
+### Perturbations
+
+The transformer supports [**perturbations**](src/ltx_core/guidance/perturbations.py) that selectively skip attention operations.
+
+Perturbations allow you to disable specific attention mechanisms during inference, which is useful for guidance techniques like STG (Spatio-Temporal Guidance).
+
+**Supported Perturbation Types**:
+
+- `SKIP_VIDEO_SELF_ATTN`: Skip video self-attention
+- `SKIP_AUDIO_SELF_ATTN`: Skip audio self-attention
+- `SKIP_A2V_CROSS_ATTN`: Skip audio-to-video cross-attention
+- `SKIP_V2A_CROSS_ATTN`: Skip video-to-audio cross-attention
+
+Perturbations are used internally by guidance mechanisms like STG (Spatio-Temporal Guidance). For usage examples, see the [`ltx-pipelines`](../ltx-pipelines/) package.
+
+---
+
+## Video VAE
+
+The Video VAE ([`src/ltx_core/model/video_vae/`](src/ltx_core/model/video_vae/)) encodes video pixels into latent representations and decodes them back.
+
+### Architecture
+
+- **Encoder**: Compresses `[B, 3, F, H, W]` pixels → `[B, 128, F', H/32, W/32]` latents
+  - Where `F' = 1 + (F-1)/8` (frame count must satisfy `(F-1) % 8 == 0`)
+  - Example: `[B, 3, 33, 512, 512]` → `[B, 128, 5, 16, 16]`
+- **Decoder**: Expands `[B, 128, F, H, W]` latents → `[B, 3, F', H*32, W*32]` pixels
+  - Where `F' = 1 + (F-1)*8`
+  - Example: `[B, 128, 5, 16, 16]` → `[B, 3, 33, 512, 512]`
+
+The Video VAE is used internally by pipelines for encoding video pixels to latents and decoding latents back to pixels. For usage examples, see the [`ltx-pipelines`](../ltx-pipelines/) package.
+
+---
+
+## Audio VAE
+
+The Audio VAE ([`src/ltx_core/model/audio_vae/`](src/ltx_core/model/audio_vae/)) processes audio spectrograms.
+
+### Audio VAE Architecture
+
+- **Encoder**: Compresses mel spectrogram `[B, mel_bins, T]` → `[B, 8, T/4, 16]` latents
+  - Temporal downsampling: 4× (`LATENT_DOWNSAMPLE_FACTOR = 4`)
+  - Frequency bins: Fixed 16 mel bins in latent space
+  - Latent channels: 8
+- **Decoder**: Expands `[B, 8, T, 16]` latents → mel spectrogram `[B, mel_bins, T*4]`
+- **Vocoder**: Converts mel spectrogram → audio waveform
+
+**Downsampling**:
+
+- Temporal: 4× (time steps)
+- Frequency: Variable (input mel_bins → fixed 16 in latent space)
+
+The Audio VAE is used internally by pipelines for encoding mel spectrograms to latents and decoding latents back to mel spectrograms. The vocoder converts mel spectrograms to audio waveforms. For usage examples, see the [`ltx-pipelines`](../ltx-pipelines/) package.
+
+---
+
+## Text Encoding (Gemma)
+
+LTX-2 uses **Gemma** (Google's open LLM) as the text encoder, located in [`src/ltx_core/text_encoders/gemma/`](src/ltx_core/text_encoders/gemma/).
+
+### Text Encoder Architecture
+
+- **Tokenizer**: Converts text → token IDs
+- **Gemma Model**: Processes tokens → embeddings
+- **Text Projection**: Uses `PixArtAlphaTextProjection` to project caption embeddings
+  - Two-layer MLP with GELU (tanh approximation) or SiLU activation
+  - Projects from caption channels (3840) to model dimensions
+- **Feature Extractor**: Extracts video/audio-specific embeddings
+- **Separate Encoders**:
+  - `AVEncoder`: For audio-video generation (outputs separate video and audio contexts)
+  - `VideoOnlyEncoder`: For video-only generation
+
+### System Prompts
+
+System prompts are also used to enhance user's prompts.
+
+- **Text-to-Video**: [`gemma_t2v_system_prompt.txt`](src/ltx_core/text_encoders/gemma/encoders/prompts/gemma_t2v_system_prompt.txt)
+- **Image-to-Video**: [`gemma_i2v_system_prompt.txt`](src/ltx_core/text_encoders/gemma/encoders/prompts/gemma_i2v_system_prompt.txt)
+
+**Important**: Video and audio receive **different** context embeddings, even from the same prompt. This allows better modality-specific conditioning.
+
+**Output Format**:
+
+- Video context: `[B, seq_len, 4096]` - Video-specific text embeddings
+- Audio context: `[B, seq_len, 2048]` - Audio-specific text embeddings
+
+The text encoder is used internally by pipelines. For usage examples, see the [`ltx-pipelines`](../ltx-pipelines/) package.
+
+---
+
+## Upscaler
+
+The Upscaler ([`src/ltx_core/model/upsampler/`](src/ltx_core/model/upsampler/)) upsamples latent representations for higher-resolution output.
+
+The spatial upsampler is used internally by two-stage pipelines (e.g., [`TI2VidTwoStagesPipeline`](../ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages.py), [`ICLoraPipeline`](../ltx-pipelines/src/ltx_pipelines/ic_lora.py)) to upsample low-resolution latents before final VAE decoding. For usage examples, see the [`ltx-pipelines`](../ltx-pipelines/) package.
+
+---
+
+## Data Flow
+
+### Complete Generation Pipeline
+
+Here's how all the components work together conceptually ([`src/ltx_core/components/`](src/ltx_core/components/)):
+
+**Pipeline Steps**:
+
+1. **Text Encoding**: Text prompt → Gemma encoder → separate video/audio embeddings
+2. **Latent Initialization**: Initialize noise latents in spatial format `[B, C, F, H, W]`
+3. **Patchification**: Convert spatial latents to sequence format `[B, seq_len, dim]` for transformer
+4. **Sigma Schedule**: Generate noise schedule (adapts to token count)
+5. **Denoising Loop**: Iteratively denoise using transformer predictions
+   - Create Modality inputs with per-token timesteps and RoPE positions
+   - Forward pass through transformer (conditional and unconditional for CFG)
+   - Apply guidance (CFG, STG, etc.)
+   - Update latents using diffusion step (Euler, etc.)
+6. **Unpatchification**: Convert sequence back to spatial format
+7. **VAE Decoding**: Decode latents to pixel space (with optional upsampling for two-stage)
+
+- [`TI2VidTwoStagesPipeline`](../ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages.py) - Two-stage text-to-video (recommended)
+- [`ICLoraPipeline`](../ltx-pipelines/src/ltx_pipelines/ic_lora.py) - Video-to-video with IC-LoRA control
+- [`DistilledPipeline`](../ltx-pipelines/src/ltx_pipelines/distilled.py) - Fast inference with distilled model
+- [`KeyframeInterpolationPipeline`](../ltx-pipelines/src/ltx_pipelines/keyframe_interpolation.py) - Keyframe-based interpolation
+
+See the [ltx-pipelines README](../ltx-pipelines/README.md) for usage examples.
+
+## 🔗 Related Projects
+
+- **[ltx-pipelines](../ltx-pipelines/)** - High-level pipeline implementations for text-to-video, image-to-video, and video-to-video
+- **[ltx-trainer](../ltx-trainer/)** - Training and fine-tuning tools

packages/ltx-core/pyproject.toml

@@ -1,38 +1,37 @@
-[project]
-name = "ltx-core"
-version = "0.1.0"
-description = "Core implementation of Lightricks' LTX-2 model"
-readme = "README.md"
-requires-python = ">=3.12"
-dependencies = [
-    "torch~=2.7",
-    "torchaudio",
-    "einops",
-    "numpy",
-    "transformers",
-    "safetensors",
-    "accelerate",
-    "scipy>=1.14",
-]
-
-[project.optional-dependencies]
-flashpack = ["flashpack==0.1.2"]
-xformers = ["xformers"]
-
-
-[tool.uv.sources]
-xformers = { index = "pytorch" }
-
-[[tool.uv.index]]
-name = "pytorch"
-url = "https://download.pytorch.org/whl/cu129"
-explicit = true
-
-[build-system]
-requires = ["uv_build>=0.9.8,<0.10.0"]
-build-backend = "uv_build"
-
-[dependency-groups]
-dev = [
-    "scikit-image>=0.25.2",
-]
+[project]
+name = "ltx-core"
+version = "1.0.0"
+description = "Core implementation of Lightricks' LTX-2 model"
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = [
+    "torch~=2.7",
+    "torchaudio",
+    "einops",
+    "numpy",
+    "transformers",
+    "safetensors",
+    "accelerate",
+    "scipy>=1.14",
+]
+
+[project.optional-dependencies]
+xformers = ["xformers"]
+
+
+[tool.uv.sources]
+xformers = { index = "pytorch" }
+
+[[tool.uv.index]]
+name = "pytorch"
+url = "https://download.pytorch.org/whl/cu129"
+explicit = true
+
+[build-system]
+requires = ["uv_build>=0.9.8,<0.10.0"]
+build-backend = "uv_build"
+
+[dependency-groups]
+dev = [
+    "scikit-image>=0.25.2",
+]

packages/ltx-core/src/ltx_core/components/__init__.py

@@ -0,0 +1,10 @@
+"""
+Diffusion pipeline components.
+Submodules:
+    diffusion_steps - Diffusion stepping algorithms (EulerDiffusionStep)
+    guiders         - Guidance strategies (CFGGuider, STGGuider, APG variants)
+    noisers         - Noise samplers (GaussianNoiser)
+    patchifiers     - Latent patchification (VideoLatentPatchifier, AudioPatchifier)
+    protocols       - Protocol definitions (Patchifier, etc.)
+    schedulers      - Sigma schedulers (LTX2Scheduler, LinearQuadraticScheduler)
+"""

packages/ltx-core/src/ltx_core/components/diffusion_steps.py

@@ -0,0 +1,22 @@
+import torch
+
+from ltx_core.components.protocols import DiffusionStepProtocol
+from ltx_core.utils import to_velocity
+
+
+class EulerDiffusionStep(DiffusionStepProtocol):
+    """
+    First-order Euler method for diffusion sampling.
+    Takes a single step from the current noise level (sigma) to the next by
+    computing velocity from the denoised prediction and applying: sample + velocity * dt.
+    """
+
+    def step(
+        self, sample: torch.Tensor, denoised_sample: torch.Tensor, sigmas: torch.Tensor, step_index: int
+    ) -> torch.Tensor:
+        sigma = sigmas[step_index]
+        sigma_next = sigmas[step_index + 1]
+        dt = sigma_next - sigma
+        velocity = to_velocity(sample, sigma, denoised_sample)
+
+        return (sample.to(torch.float32) + velocity.to(torch.float32) * dt).to(sample.dtype)

packages/ltx-core/src/ltx_core/components/guiders.py

@@ -0,0 +1,198 @@
+from dataclasses import dataclass
+
+import torch
+
+from ltx_core.components.protocols import GuiderProtocol
+
+
+@dataclass(frozen=True)
+class CFGGuider(GuiderProtocol):
+    """
+    Classifier-free guidance (CFG) guider.
+    Computes the guidance delta as (scale - 1) * (cond - uncond), steering the
+    denoising process toward the conditioned prediction.
+    Attributes:
+        scale: Guidance strength. 1.0 means no guidance, higher values increase
+            adherence to the conditioning.
+    """
+
+    scale: float
+
+    def delta(self, cond: torch.Tensor, uncond: torch.Tensor) -> torch.Tensor:
+        return (self.scale - 1) * (cond - uncond)
+
+    def enabled(self) -> bool:
+        return self.scale != 1.0
+
+
+@dataclass(frozen=True)
+class CFGStarRescalingGuider(GuiderProtocol):
+    """
+    Calculates the CFG delta between conditioned and unconditioned samples.
+    To minimize offset in the denoising direction and move mostly along the
+    conditioning axis within the distribution, the unconditioned sample is
+    rescaled in accordance with the norm of the conditioned sample.
+    Attributes:
+        scale (float):
+            Global guidance strength. A value of 1.0 corresponds to no extra
+            guidance beyond the base model prediction. Values > 1.0 increase
+            the influence of the conditioned sample relative to the
+            unconditioned one.
+    """
+
+    scale: float
+
+    def delta(self, cond: torch.Tensor, uncond: torch.Tensor) -> torch.Tensor:
+        rescaled_neg = projection_coef(cond, uncond) * uncond
+        return (self.scale - 1) * (cond - rescaled_neg)
+
+    def enabled(self) -> bool:
+        return self.scale != 1.0
+
+
+@dataclass(frozen=True)
+class STGGuider(GuiderProtocol):
+    """
+    Calculates the STG delta between conditioned and perturbed denoised samples.
+    Perturbed samples are the result of the denoising process with perturbations,
+    e.g. attentions acting as passthrough for certain layers and modalities.
+    Attributes:
+        scale (float):
+            Global strength of the STG guidance. A value of 0.0 disables the
+            guidance. Larger values increase the correction applied in the
+            direction of (pos_denoised - perturbed_denoised).
+    """
+
+    scale: float
+
+    def delta(self, pos_denoised: torch.Tensor, perturbed_denoised: torch.Tensor) -> torch.Tensor:
+        return self.scale * (pos_denoised - perturbed_denoised)
+
+    def enabled(self) -> bool:
+        return self.scale != 0.0
+
+
+@dataclass(frozen=True)
+class LtxAPGGuider(GuiderProtocol):
+    """
+    Calculates the APG (adaptive projected guidance) delta between conditioned
+    and unconditioned samples.
+    To minimize offset in the denoising direction and move mostly along the
+    conditioning axis within the distribution, the (cond - uncond) delta is
+    decomposed into components parallel and orthogonal to the conditioned
+    sample. The `eta` parameter weights the parallel component, while `scale`
+    is applied to the orthogonal component. Optionally, a norm threshold can
+    be used to suppress guidance when the magnitude of the correction is small.
+    Attributes:
+        scale (float):
+            Strength applied to the component of the guidance that is orthogonal
+            to the conditioned sample. Controls how aggressively we move in
+            directions that change semantics but stay consistent with the
+            conditioning manifold.
+        eta (float):
+            Weight of the component of the guidance that is parallel to the
+            conditioned sample. A value of 1.0 keeps the full parallel
+            component; values in [0, 1] attenuate it, and values > 1.0 amplify
+            motion along the conditioning direction.
+        norm_threshold (float):
+            Minimum L2 norm of the guidance delta below which the guidance
+            can be reduced or ignored (depending on implementation).
+            This is useful for avoiding noisy or unstable updates when the
+            guidance signal is very small.
+    """
+
+    scale: float
+    eta: float = 1.0
+    norm_threshold: float = 0.0
+
+    def delta(self, cond: torch.Tensor, uncond: torch.Tensor) -> torch.Tensor:
+        guidance = cond - uncond
+        if self.norm_threshold > 0:
+            ones = torch.ones_like(guidance)
+            guidance_norm = guidance.norm(p=2, dim=[-1, -2, -3], keepdim=True)
+            scale_factor = torch.minimum(ones, self.norm_threshold / guidance_norm)
+            guidance = guidance * scale_factor
+        proj_coeff = projection_coef(guidance, cond)
+        g_parallel = proj_coeff * cond
+        g_orth = guidance - g_parallel
+        g_apg = g_parallel * self.eta + g_orth
+
+        return g_apg * (self.scale - 1)
+
+    def enabled(self) -> bool:
+        return self.scale != 1.0
+
+
+@dataclass(frozen=False)
+class LegacyStatefulAPGGuider(GuiderProtocol):
+    """
+    Calculates the APG (adaptive projected guidance) delta between conditioned
+    and unconditioned samples.
+    To minimize offset in the denoising direction and move mostly along the
+    conditioning axis within the distribution, the (cond - uncond) delta is
+    decomposed into components parallel and orthogonal to the conditioned
+    sample. The `eta` parameter weights the parallel component, while `scale`
+    is applied to the orthogonal component. Optionally, a norm threshold can
+    be used to suppress guidance when the magnitude of the correction is small.
+    Attributes:
+        scale (float):
+            Strength applied to the component of the guidance that is orthogonal
+            to the conditioned sample. Controls how aggressively we move in
+            directions that change semantics but stay consistent with the
+            conditioning manifold.
+        eta (float):
+            Weight of the component of the guidance that is parallel to the
+            conditioned sample. A value of 1.0 keeps the full parallel
+            component; values in [0, 1] attenuate it, and values > 1.0 amplify
+            motion along the conditioning direction.
+        norm_threshold (float):
+            Minimum L2 norm of the guidance delta below which the guidance
+            can be reduced or ignored (depending on implementation).
+            This is useful for avoiding noisy or unstable updates when the
+            guidance signal is very small.
+        momentum (float):
+            Exponential moving-average coefficient for accumulating guidance
+            over time. running_avg = momentum * running_avg + guidance
+    """
+
+    scale: float
+    eta: float
+    norm_threshold: float = 5.0
+    momentum: float = 0.0
+    # it is user's responsibility not to use same APGGuider for several denoisings or different modalities
+    # in order not to share accumulated average across different denoisings or modalities
+    running_avg: torch.Tensor | None = None
+
+    def delta(self, cond: torch.Tensor, uncond: torch.Tensor) -> torch.Tensor:
+        guidance = cond - uncond
+        if self.momentum != 0:
+            if self.running_avg is None:
+                self.running_avg = guidance.clone()
+            else:
+                self.running_avg = self.momentum * self.running_avg + guidance
+            guidance = self.running_avg
+
+        if self.norm_threshold > 0:
+            ones = torch.ones_like(guidance)
+            guidance_norm = guidance.norm(p=2, dim=[-1, -2, -3], keepdim=True)
+            scale_factor = torch.minimum(ones, self.norm_threshold / guidance_norm)
+            guidance = guidance * scale_factor
+
+        proj_coeff = projection_coef(guidance, cond)
+        g_parallel = proj_coeff * cond
+        g_orth = guidance - g_parallel
+        g_apg = g_parallel * self.eta + g_orth
+
+        return g_apg * self.scale
+
+    def enabled(self) -> bool:
+        return self.scale != 0.0
+
+
+def projection_coef(to_project: torch.Tensor, project_onto: torch.Tensor) -> torch.Tensor:
+    batch_size = to_project.shape[0]
+    positive_flat = to_project.reshape(batch_size, -1)
+    negative_flat = project_onto.reshape(batch_size, -1)
+    dot_product = torch.sum(positive_flat * negative_flat, dim=1, keepdim=True)
+    squared_norm = torch.sum(negative_flat**2, dim=1, keepdim=True) + 1e-8
+    return dot_product / squared_norm

packages/ltx-core/src/ltx_core/components/noisers.py

@@ -0,0 +1,35 @@
+from dataclasses import replace
+from typing import Protocol
+
+import torch
+
+from ltx_core.types import LatentState
+
+
+class Noiser(Protocol):
+    """Protocol for adding noise to a latent state during diffusion."""
+
+    def __call__(self, latent_state: LatentState, noise_scale: float) -> LatentState: ...
+
+
+class GaussianNoiser(Noiser):
+    """Adds Gaussian noise to a latent state, scaled by the denoise mask."""
+
+    def __init__(self, generator: torch.Generator):
+        super().__init__()
+
+        self.generator = generator
+
+    def __call__(self, latent_state: LatentState, noise_scale: float = 1.0) -> LatentState:
+        noise = torch.randn(
+            *latent_state.latent.shape,
+            device=latent_state.latent.device,
+            dtype=latent_state.latent.dtype,
+            generator=self.generator,
+        )
+        scaled_mask = latent_state.denoise_mask * noise_scale
+        latent = noise * scaled_mask + latent_state.latent * (1 - scaled_mask)
+        return replace(
+            latent_state,
+            latent=latent.to(latent_state.latent.dtype),
+        )

packages/ltx-core/src/ltx_core/components/patchifiers.py

@@ -0,0 +1,348 @@
+import math
+from typing import Optional, Tuple
+
+import einops
+import torch
+
+from ltx_core.components.protocols import Patchifier
+from ltx_core.types import AudioLatentShape, SpatioTemporalScaleFactors, VideoLatentShape
+
+
+class VideoLatentPatchifier(Patchifier):
+    def __init__(self, patch_size: int):
+        # Patch sizes for video latents.
+        self._patch_size = (
+            1,  # temporal dimension
+            patch_size,  # height dimension
+            patch_size,  # width dimension
+        )
+
+    @property
+    def patch_size(self) -> Tuple[int, int, int]:
+        return self._patch_size
+
+    def get_token_count(self, tgt_shape: VideoLatentShape) -> int:
+        return math.prod(tgt_shape.to_torch_shape()[2:]) // math.prod(self._patch_size)
+
+    def patchify(
+        self,
+        latents: torch.Tensor,
+    ) -> torch.Tensor:
+        latents = einops.rearrange(
+            latents,
+            "b c (f p1) (h p2) (w p3) -> b (f h w) (c p1 p2 p3)",
+            p1=self._patch_size[0],
+            p2=self._patch_size[1],
+            p3=self._patch_size[2],
+        )
+
+        return latents
+
+    def unpatchify(
+        self,
+        latents: torch.Tensor,
+        output_shape: VideoLatentShape,
+    ) -> torch.Tensor:
+        assert self._patch_size[0] == 1, "Temporal patch size must be 1 for symmetric patchifier"
+
+        patch_grid_frames = output_shape.frames // self._patch_size[0]
+        patch_grid_height = output_shape.height // self._patch_size[1]
+        patch_grid_width = output_shape.width // self._patch_size[2]
+
+        latents = einops.rearrange(
+            latents,
+            "b (f h w) (c p q) -> b c f (h p) (w q)",
+            f=patch_grid_frames,
+            h=patch_grid_height,
+            w=patch_grid_width,
+            p=self._patch_size[1],
+            q=self._patch_size[2],
+        )
+
+        return latents
+
+    def get_patch_grid_bounds(
+        self,
+        output_shape: AudioLatentShape | VideoLatentShape,
+        device: Optional[torch.device] = None,
+    ) -> torch.Tensor:
+        """
+        Return the per-dimension bounds [inclusive start, exclusive end) for every
+        patch produced by `patchify`. The bounds are expressed in the original
+        video grid coordinates: frame/time, height, and width.
+        The resulting tensor is shaped `[batch_size, 3, num_patches, 2]`, where:
+            - axis 1 (size 3) enumerates (frame/time, height, width) dimensions
+            - axis 3 (size 2) stores `[start, end)` indices within each dimension
+        Args:
+            output_shape: Video grid description containing frames, height, and width.
+            device: Device of the latent tensor.
+        """
+        if not isinstance(output_shape, VideoLatentShape):
+            raise ValueError("VideoLatentPatchifier expects VideoLatentShape when computing coordinates")
+
+        frames = output_shape.frames
+        height = output_shape.height
+        width = output_shape.width
+        batch_size = output_shape.batch
+
+        # Validate inputs to ensure positive dimensions
+        assert frames > 0, f"frames must be positive, got {frames}"
+        assert height > 0, f"height must be positive, got {height}"
+        assert width > 0, f"width must be positive, got {width}"
+        assert batch_size > 0, f"batch_size must be positive, got {batch_size}"
+
+        # Generate grid coordinates for each dimension (frame, height, width)
+        # We use torch.arange to create the starting coordinates for each patch.
+        # indexing='ij' ensures the dimensions are in the order (frame, height, width).
+        grid_coords = torch.meshgrid(
+            torch.arange(start=0, end=frames, step=self._patch_size[0], device=device),
+            torch.arange(start=0, end=height, step=self._patch_size[1], device=device),
+            torch.arange(start=0, end=width, step=self._patch_size[2], device=device),
+            indexing="ij",
+        )
+
+        # Stack the grid coordinates to create the start coordinates tensor.
+        # Shape becomes (3, grid_f, grid_h, grid_w)
+        patch_starts = torch.stack(grid_coords, dim=0)
+
+        # Create a tensor containing the size of a single patch:
+        # (frame_patch_size, height_patch_size, width_patch_size).
+        # Reshape to (3, 1, 1, 1) to enable broadcasting when adding to the start coordinates.
+        patch_size_delta = torch.tensor(
+            self._patch_size,
+            device=patch_starts.device,
+            dtype=patch_starts.dtype,
+        ).view(3, 1, 1, 1)
+
+        # Calculate end coordinates: start + patch_size
+        # Shape becomes (3, grid_f, grid_h, grid_w)
+        patch_ends = patch_starts + patch_size_delta
+
+        # Stack start and end coordinates together along the last dimension
+        # Shape becomes (3, grid_f, grid_h, grid_w, 2), where the last dimension is [start, end]
+        latent_coords = torch.stack((patch_starts, patch_ends), dim=-1)
+
+        # Broadcast to batch size and flatten all spatial/temporal dimensions into one sequence.
+        # Final Shape: (batch_size, 3, num_patches, 2)
+        latent_coords = einops.repeat(
+            latent_coords,
+            "c f h w bounds -> b c (f h w) bounds",
+            b=batch_size,
+            bounds=2,
+        )
+
+        return latent_coords
+
+
+def get_pixel_coords(
+    latent_coords: torch.Tensor,
+    scale_factors: SpatioTemporalScaleFactors,
+    causal_fix: bool = False,
+) -> torch.Tensor:
+    """
+    Map latent-space `[start, end)` coordinates to their pixel-space equivalents by scaling
+    each axis (frame/time, height, width) with the corresponding VAE downsampling factors.
+    Optionally compensate for causal encoding that keeps the first frame at unit temporal scale.
+    Args:
+        latent_coords: Tensor of latent bounds shaped `(batch, 3, num_patches, 2)`.
+        scale_factors: SpatioTemporalScaleFactors tuple `(temporal, height, width)` with integer scale factors applied
+        per axis.
+        causal_fix: When True, rewrites the temporal axis of the first frame so causal VAEs
+            that treat frame zero differently still yield non-negative timestamps.
+    """
+    # Broadcast the VAE scale factors so they align with the `(batch, axis, patch, bound)` layout.
+    broadcast_shape = [1] * latent_coords.ndim
+    broadcast_shape[1] = -1  # axis dimension corresponds to (frame/time, height, width)
+    scale_tensor = torch.tensor(scale_factors, device=latent_coords.device).view(*broadcast_shape)
+
+    # Apply per-axis scaling to convert latent bounds into pixel-space coordinates.
+    pixel_coords = latent_coords * scale_tensor
+
+    if causal_fix:
+        # VAE temporal stride for the very first frame is 1 instead of `scale_factors[0]`.
+        # Shift and clamp to keep the first-frame timestamps causal and non-negative.
+        pixel_coords[:, 0, ...] = (pixel_coords[:, 0, ...] + 1 - scale_factors[0]).clamp(min=0)
+
+    return pixel_coords
+
+
+class AudioPatchifier(Patchifier):
+    def __init__(
+        self,
+        patch_size: int,
+        sample_rate: int = 16000,
+        hop_length: int = 160,
+        audio_latent_downsample_factor: int = 4,
+        is_causal: bool = True,
+        shift: int = 0,
+    ):
+        """
+        Patchifier tailored for spectrogram/audio latents.
+        Args:
+            patch_size: Number of mel bins combined into a single patch. This
+                controls the resolution along the frequency axis.
+            sample_rate: Original waveform sampling rate. Used to map latent
+                indices back to seconds so downstream consumers can align audio
+                and video cues.
+            hop_length: Window hop length used for the spectrogram. Determines
+                how many real-time samples separate two consecutive latent frames.
+            audio_latent_downsample_factor: Ratio between spectrogram frames and
+                latent frames; compensates for additional downsampling inside the
+                VAE encoder.
+            is_causal: When True, timing is shifted to account for causal
+                receptive fields so timestamps do not peek into the future.
+            shift: Integer offset applied to the latent indices. Enables
+                constructing overlapping windows from the same latent sequence.
+        """
+        self.hop_length = hop_length
+        self.sample_rate = sample_rate
+        self.audio_latent_downsample_factor = audio_latent_downsample_factor
+        self.is_causal = is_causal
+        self.shift = shift
+        self._patch_size = (1, patch_size, patch_size)
+
+    @property
+    def patch_size(self) -> Tuple[int, int, int]:
+        return self._patch_size
+
+    def get_token_count(self, tgt_shape: AudioLatentShape) -> int:
+        return tgt_shape.frames
+
+    def _get_audio_latent_time_in_sec(
+        self,
+        start_latent: int,
+        end_latent: int,
+        dtype: torch.dtype,
+        device: Optional[torch.device] = None,
+    ) -> torch.Tensor:
+        """
+        Converts latent indices into real-time seconds while honoring causal
+        offsets and the configured hop length.
+        Args:
+            start_latent: Inclusive start index inside the latent sequence. This
+                sets the first timestamp returned.
+            end_latent: Exclusive end index. Determines how many timestamps get
+                generated.
+            dtype: Floating-point dtype used for the returned tensor, allowing
+                callers to control precision.
+            device: Target device for the timestamp tensor. When omitted the
+                computation occurs on CPU to avoid surprising GPU allocations.
+        """
+        if device is None:
+            device = torch.device("cpu")
+
+        audio_latent_frame = torch.arange(start_latent, end_latent, dtype=dtype, device=device)
+
+        audio_mel_frame = audio_latent_frame * self.audio_latent_downsample_factor
+
+        if self.is_causal:
+            # Frame offset for causal alignment.
+            # The "+1" ensures the timestamp corresponds to the first sample that is fully available.
+            causal_offset = 1
+            audio_mel_frame = (audio_mel_frame + causal_offset - self.audio_latent_downsample_factor).clip(min=0)
+
+        return audio_mel_frame * self.hop_length / self.sample_rate
+
+    def _compute_audio_timings(
+        self,
+        batch_size: int,
+        num_steps: int,
+        device: Optional[torch.device] = None,
+    ) -> torch.Tensor:
+        """
+        Builds a `(B, 1, T, 2)` tensor containing timestamps for each latent frame.
+        This helper method underpins `get_patch_grid_bounds` for the audio patchifier.
+        Args:
+            batch_size: Number of sequences to broadcast the timings over.
+            num_steps: Number of latent frames (time steps) to convert into timestamps.
+            device: Device on which the resulting tensor should reside.
+        """
+        resolved_device = device
+        if resolved_device is None:
+            resolved_device = torch.device("cpu")
+
+        start_timings = self._get_audio_latent_time_in_sec(
+            self.shift,
+            num_steps + self.shift,
+            torch.float32,
+            resolved_device,
+        )
+        start_timings = start_timings.unsqueeze(0).expand(batch_size, -1).unsqueeze(1)
+
+        end_timings = self._get_audio_latent_time_in_sec(
+            self.shift + 1,
+            num_steps + self.shift + 1,
+            torch.float32,
+            resolved_device,
+        )
+        end_timings = end_timings.unsqueeze(0).expand(batch_size, -1).unsqueeze(1)
+
+        return torch.stack([start_timings, end_timings], dim=-1)
+
+    def patchify(
+        self,
+        audio_latents: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Flattens the audio latent tensor along time. Use `get_patch_grid_bounds`
+        to derive timestamps for each latent frame based on the configured hop
+        length and downsampling.
+        Args:
+            audio_latents: Latent tensor to patchify.
+        Returns:
+            Flattened patch tokens tensor. Use `get_patch_grid_bounds` to compute the
+            corresponding timing metadata when needed.
+        """
+        audio_latents = einops.rearrange(
+            audio_latents,
+            "b c t f -> b t (c f)",
+        )
+
+        return audio_latents
+
+    def unpatchify(
+        self,
+        audio_latents: torch.Tensor,
+        output_shape: AudioLatentShape,
+    ) -> torch.Tensor:
+        """
+        Restores the `(B, C, T, F)` spectrogram tensor from flattened patches.
+        Use `get_patch_grid_bounds` to recompute the timestamps that describe each
+        frame's position in real time.
+        Args:
+            audio_latents: Latent tensor to unpatchify.
+            output_shape: Shape of the unpatched output tensor.
+        Returns:
+            Unpatched latent tensor. Use `get_patch_grid_bounds` to compute the timing
+            metadata associated with the restored latents.
+        """
+        # audio_latents shape: (batch, time, freq * channels)
+        audio_latents = einops.rearrange(
+            audio_latents,
+            "b t (c f) -> b c t f",
+            c=output_shape.channels,
+            f=output_shape.mel_bins,
+        )
+
+        return audio_latents
+
+    def get_patch_grid_bounds(
+        self,
+        output_shape: AudioLatentShape | VideoLatentShape,
+        device: Optional[torch.device] = None,
+    ) -> torch.Tensor:
+        """
+        Return the temporal bounds `[inclusive start, exclusive end)` for every
+        patch emitted by `patchify`. For audio this corresponds to timestamps in
+        seconds aligned with the original spectrogram grid.
+        The returned tensor has shape `[batch_size, 1, time_steps, 2]`, where:
+            - axis 1 (size 1) represents the temporal dimension
+            - axis 3 (size 2) stores the `[start, end)` timestamps per patch
+        Args:
+            output_shape: Audio grid specification describing the number of time steps.
+            device: Target device for the returned tensor.
+        """
+        if not isinstance(output_shape, AudioLatentShape):
+            raise ValueError("AudioPatchifier expects AudioLatentShape when computing coordinates")
+
+        return self._compute_audio_timings(output_shape.batch, output_shape.frames, device)

packages/ltx-core/src/ltx_core/components/protocols.py

@@ -0,0 +1,101 @@
+from typing import Protocol, Tuple
+
+import torch
+
+from ltx_core.types import AudioLatentShape, VideoLatentShape
+
+
+class Patchifier(Protocol):
+    """
+    Protocol for patchifiers that convert latent tensors into patches and assemble them back.
+    """
+
+    def patchify(
+        self,
+        latents: torch.Tensor,
+    ) -> torch.Tensor:
+        ...
+        """
+        Convert latent tensors into flattened patch tokens.
+        Args:
+            latents: Latent tensor to patchify.
+        Returns:
+            Flattened patch tokens tensor.
+        """
+
+    def unpatchify(
+        self,
+        latents: torch.Tensor,
+        output_shape: AudioLatentShape | VideoLatentShape,
+    ) -> torch.Tensor:
+        """
+        Converts latent tensors between spatio-temporal formats and flattened sequence representations.
+        Args:
+            latents: Patch tokens that must be rearranged back into the latent grid constructed by `patchify`.
+            output_shape: Shape of the output tensor. Note that output_shape is either AudioLatentShape or
+            VideoLatentShape.
+        Returns:
+            Dense latent tensor restored from the flattened representation.
+        """
+
+    @property
+    def patch_size(self) -> Tuple[int, int, int]:
+        ...
+        """
+        Returns the patch size as a tuple of (temporal, height, width) dimensions
+        """
+
+    def get_patch_grid_bounds(
+        self,
+        output_shape: AudioLatentShape | VideoLatentShape,
+        device: torch.device | None = None,
+    ) -> torch.Tensor:
+        ...
+        """
+        Compute metadata describing where each latent patch resides within the
+        grid specified by `output_shape`.
+        Args:
+            output_shape: Target grid layout for the patches.
+            device: Target device for the returned tensor.
+        Returns:
+            Tensor containing patch coordinate metadata such as spatial or temporal intervals.
+        """
+
+
+class SchedulerProtocol(Protocol):
+    """
+    Protocol for schedulers that provide a sigmas schedule tensor for a
+    given number of steps. Device is cpu.
+    """
+
+    def execute(self, steps: int, **kwargs) -> torch.FloatTensor: ...
+
+
+class GuiderProtocol(Protocol):
+    """
+    Protocol for guiders that compute a delta tensor given conditioning inputs.
+    The returned delta should be added to the conditional output (cond), enabling
+    multiple guiders to be chained together by accumulating their deltas.
+    """
+
+    scale: float
+
+    def delta(self, cond: torch.Tensor, uncond: torch.Tensor) -> torch.Tensor: ...
+
+    def enabled(self) -> bool:
+        """
+        Returns whether the corresponding perturbation is enabled. E.g. for CFG, this should return False if the scale
+        is 1.0.
+        """
+        ...
+
+
+class DiffusionStepProtocol(Protocol):
+    """
+    Protocol for diffusion steps that provide a next sample tensor for a given current sample tensor,
+    current denoised sample tensor, and sigmas tensor.
+    """
+
+    def step(
+        self, sample: torch.Tensor, denoised_sample: torch.Tensor, sigmas: torch.Tensor, step_index: int
+    ) -> torch.Tensor: ...

packages/ltx-core/src/ltx_core/components/schedulers.py

@@ -0,0 +1,129 @@
+import math
+from functools import lru_cache
+
+import numpy
+import scipy
+import torch
+
+from ltx_core.components.protocols import SchedulerProtocol
+
+BASE_SHIFT_ANCHOR = 1024
+MAX_SHIFT_ANCHOR = 4096
+
+
+class LTX2Scheduler(SchedulerProtocol):
+    """
+    Default scheduler for LTX-2 diffusion sampling.
+    Generates a sigma schedule with token-count-dependent shifting and optional
+    stretching to a terminal value.
+    """
+
+    def execute(
+        self,
+        steps: int,
+        latent: torch.Tensor | None = None,
+        max_shift: float = 2.05,
+        base_shift: float = 0.95,
+        stretch: bool = True,
+        terminal: float = 0.1,
+        **_kwargs,
+    ) -> torch.FloatTensor:
+        tokens = math.prod(latent.shape[2:]) if latent is not None else MAX_SHIFT_ANCHOR
+        sigmas = torch.linspace(1.0, 0.0, steps + 1)
+
+        x1 = BASE_SHIFT_ANCHOR
+        x2 = MAX_SHIFT_ANCHOR
+        mm = (max_shift - base_shift) / (x2 - x1)
+        b = base_shift - mm * x1
+        sigma_shift = (tokens) * mm + b
+
+        power = 1
+        sigmas = torch.where(
+            sigmas != 0,
+            math.exp(sigma_shift) / (math.exp(sigma_shift) + (1 / sigmas - 1) ** power),
+            0,
+        )
+
+        # Stretch sigmas so that its final value matches the given terminal value.
+        if stretch:
+            non_zero_mask = sigmas != 0
+            non_zero_sigmas = sigmas[non_zero_mask]
+            one_minus_z = 1.0 - non_zero_sigmas
+            scale_factor = one_minus_z[-1] / (1.0 - terminal)
+            stretched = 1.0 - (one_minus_z / scale_factor)
+            sigmas[non_zero_mask] = stretched
+
+        return sigmas.to(torch.float32)
+
+
+class LinearQuadraticScheduler(SchedulerProtocol):
+    """
+    Scheduler with linear steps followed by quadratic steps.
+    Produces a sigma schedule that transitions linearly up to a threshold,
+    then follows a quadratic curve for the remaining steps.
+    """
+
+    def execute(
+        self, steps: int, threshold_noise: float = 0.025, linear_steps: int | None = None, **_kwargs
+    ) -> torch.FloatTensor:
+        if steps == 1:
+            return torch.FloatTensor([1.0, 0.0])
+
+        if linear_steps is None:
+            linear_steps = steps // 2
+        linear_sigma_schedule = [i * threshold_noise / linear_steps for i in range(linear_steps)]
+        threshold_noise_step_diff = linear_steps - threshold_noise * steps
+        quadratic_steps = steps - linear_steps
+        quadratic_sigma_schedule = []
+        if quadratic_steps > 0:
+            quadratic_coef = threshold_noise_step_diff / (linear_steps * quadratic_steps**2)
+            linear_coef = threshold_noise / linear_steps - 2 * threshold_noise_step_diff / (quadratic_steps**2)
+            const = quadratic_coef * (linear_steps**2)
+            quadratic_sigma_schedule = [
+                quadratic_coef * (i**2) + linear_coef * i + const for i in range(linear_steps, steps)
+            ]
+        sigma_schedule = linear_sigma_schedule + quadratic_sigma_schedule + [1.0]
+        sigma_schedule = [1.0 - x for x in sigma_schedule]
+        return torch.FloatTensor(sigma_schedule)
+
+
+class BetaScheduler(SchedulerProtocol):
+    """
+    Scheduler using a beta distribution to sample timesteps.
+    Based on: https://arxiv.org/abs/2407.12173
+    """
+
+    shift = 2.37
+    timesteps_length = 10000
+
+    def execute(self, steps: int, alpha: float = 0.6, beta: float = 0.6) -> torch.FloatTensor:
+        """
+        Execute the beta scheduler.
+        Args:
+            steps: The number of steps to execute the scheduler for.
+            alpha: The alpha parameter for the beta distribution.
+            beta: The beta parameter for the beta distribution.
+        Warnings:
+            The number of steps within `sigmas` theoretically might be less than `steps+1`,
+            because of the deduplication of the identical timesteps
+        Returns:
+            A tensor of sigmas.
+        """
+        model_sampling_sigmas = _precalculate_model_sampling_sigmas(self.shift, self.timesteps_length)
+        total_timesteps = len(model_sampling_sigmas) - 1
+        ts = 1 - numpy.linspace(0, 1, steps, endpoint=False)
+        ts = numpy.rint(scipy.stats.beta.ppf(ts, alpha, beta) * total_timesteps).tolist()
+        ts = list(dict.fromkeys(ts))
+
+        sigmas = [float(model_sampling_sigmas[int(t)]) for t in ts] + [0.0]
+        return torch.FloatTensor(sigmas)
+
+
+@lru_cache(maxsize=5)
+def _precalculate_model_sampling_sigmas(shift: float, timesteps_length: int) -> torch.Tensor:
+    timesteps = torch.arange(1, timesteps_length + 1, 1) / timesteps_length
+    return torch.Tensor([flux_time_shift(shift, 1.0, t) for t in timesteps])
+
+
+def flux_time_shift(mu: float, sigma: float, t: float) -> float:
+    return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)

packages/ltx-core/src/ltx_core/conditioning/__init__.py

@@ -0,0 +1,12 @@
+"""Conditioning utilities: latent state, tools, and conditioning types."""
+
+from ltx_core.conditioning.exceptions import ConditioningError
+from ltx_core.conditioning.item import ConditioningItem
+from ltx_core.conditioning.types import VideoConditionByKeyframeIndex, VideoConditionByLatentIndex
+
+__all__ = [
+    "ConditioningError",
+    "ConditioningItem",
+    "VideoConditionByKeyframeIndex",
+    "VideoConditionByLatentIndex",
+]

packages/ltx-core/src/ltx_core/conditioning/exceptions.py

@@ -0,0 +1,4 @@
+class ConditioningError(Exception):
+    """
+    Class for conditioning-related errors.
+    """

packages/ltx-core/src/ltx_core/conditioning/item.py

@@ -0,0 +1,20 @@
+from typing import Protocol
+
+from ltx_core.tools import LatentTools
+from ltx_core.types import LatentState
+
+
+class ConditioningItem(Protocol):
+    """Protocol for conditioning items that modify latent state during diffusion."""
+
+    def apply_to(self, latent_state: LatentState, latent_tools: LatentTools) -> LatentState:
+        """
+        Apply the conditioning to the latent state.
+        Args:
+            latent_state: The latent state to apply the conditioning to. This is state always patchified.
+        Returns:
+            The latent state after the conditioning has been applied.
+        IMPORTANT: If the conditioning needs to add extra tokens to the latent, it should add them to the end of the
+        latent.
+        """
+        ...

packages/ltx-core/src/ltx_core/conditioning/types/__init__.py

@@ -0,0 +1,9 @@
+"""Conditioning type implementations."""
+
+from ltx_core.conditioning.types.keyframe_cond import VideoConditionByKeyframeIndex
+from ltx_core.conditioning.types.latent_cond import VideoConditionByLatentIndex
+
+__all__ = [
+    "VideoConditionByKeyframeIndex",
+    "VideoConditionByLatentIndex",
+]

packages/ltx-core/src/ltx_core/conditioning/types/keyframe_cond.py

@@ -0,0 +1,53 @@
+import torch
+
+from ltx_core.components.patchifiers import get_pixel_coords
+from ltx_core.conditioning.item import ConditioningItem
+from ltx_core.tools import VideoLatentTools
+from ltx_core.types import LatentState, VideoLatentShape
+
+
+class VideoConditionByKeyframeIndex(ConditioningItem):
+    """
+    Conditions video generation on keyframe latents at a specific frame index.
+    Appends keyframe tokens to the latent state with positions offset by frame_idx,
+    and sets denoise strength according to the strength parameter.
+    """
+
+    def __init__(self, keyframes: torch.Tensor, frame_idx: int, strength: float):
+        self.keyframes = keyframes
+        self.frame_idx = frame_idx
+        self.strength = strength
+
+    def apply_to(
+        self,
+        latent_state: LatentState,
+        latent_tools: VideoLatentTools,
+    ) -> LatentState:
+        tokens = latent_tools.patchifier.patchify(self.keyframes)
+        latent_coords = latent_tools.patchifier.get_patch_grid_bounds(
+            output_shape=VideoLatentShape.from_torch_shape(self.keyframes.shape),
+            device=self.keyframes.device,
+        )
+        positions = get_pixel_coords(
+            latent_coords=latent_coords,
+            scale_factors=latent_tools.scale_factors,
+            causal_fix=latent_tools.causal_fix if self.frame_idx == 0 else False,
+        )
+
+        positions[:, 0, ...] += self.frame_idx
+        positions = positions.to(dtype=torch.float32)
+        positions[:, 0, ...] /= latent_tools.fps
+
+        denoise_mask = torch.full(
+            size=(*tokens.shape[:2], 1),
+            fill_value=1.0 - self.strength,
+            device=self.keyframes.device,
+            dtype=self.keyframes.dtype,
+        )
+
+        return LatentState(
+            latent=torch.cat([latent_state.latent, tokens], dim=1),
+            denoise_mask=torch.cat([latent_state.denoise_mask, denoise_mask], dim=1),
+            positions=torch.cat([latent_state.positions, positions], dim=2),
+            clean_latent=torch.cat([latent_state.clean_latent, tokens], dim=1),
+        )

packages/ltx-core/src/ltx_core/conditioning/types/latent_cond.py

@@ -0,0 +1,44 @@
+import torch
+
+from ltx_core.conditioning.exceptions import ConditioningError
+from ltx_core.conditioning.item import ConditioningItem
+from ltx_core.tools import LatentTools
+from ltx_core.types import LatentState
+
+
+class VideoConditionByLatentIndex(ConditioningItem):
+    """
+    Conditions video generation by injecting latents at a specific latent frame index.
+    Replaces tokens in the latent state at positions corresponding to latent_idx,
+    and sets denoise strength according to the strength parameter.
+    """
+
+    def __init__(self, latent: torch.Tensor, strength: float, latent_idx: int):
+        self.latent = latent
+        self.strength = strength
+        self.latent_idx = latent_idx
+
+    def apply_to(self, latent_state: LatentState, latent_tools: LatentTools) -> LatentState:
+        cond_batch, cond_channels, _, cond_height, cond_width = self.latent.shape
+        tgt_batch, tgt_channels, tgt_frames, tgt_height, tgt_width = latent_tools.target_shape.to_torch_shape()
+
+        if (cond_batch, cond_channels, cond_height, cond_width) != (tgt_batch, tgt_channels, tgt_height, tgt_width):
+            raise ConditioningError(
+                f"Can't apply image conditioning item to latent with shape {latent_tools.target_shape}, expected "
+                f"shape is ({tgt_batch}, {tgt_channels}, {tgt_frames}, {tgt_height}, {tgt_width}). Make sure "
+                "the image and latent have the same spatial shape."
+            )
+
+        tokens = latent_tools.patchifier.patchify(self.latent)
+        start_token = latent_tools.patchifier.get_token_count(
+            latent_tools.target_shape._replace(frames=self.latent_idx)
+        )
+        stop_token = start_token + tokens.shape[1]
+
+        latent_state = latent_state.clone()
+
+        latent_state.latent[:, start_token:stop_token] = tokens
+        latent_state.clean_latent[:, start_token:stop_token] = tokens
+        latent_state.denoise_mask[:, start_token:stop_token] = 1.0 - self.strength
+
+        return latent_state

packages/ltx-core/src/ltx_core/guidance/__init__.py

@@ -0,0 +1,15 @@
+"""Guidance and perturbation utilities for attention manipulation."""
+
+from ltx_core.guidance.perturbations import (
+    BatchedPerturbationConfig,
+    Perturbation,
+    PerturbationConfig,
+    PerturbationType,
+)
+
+__all__ = [
+    "BatchedPerturbationConfig",
+    "Perturbation",
+    "PerturbationConfig",
+    "PerturbationType",
+]

packages/ltx-core/src/ltx_core/guidance/perturbations.py

@@ -1,74 +1,79 @@
-# Copyright (c) 2025 Lightricks. All rights reserved.
-# Created by Andrew Kvochko
-
-from dataclasses import dataclass
-from enum import Enum
-
-import torch
-from torch._prims_common import DeviceLikeType
-
-
-class PerturbationType(Enum):
-    SKIP_A2V_CROSS_ATTN = "skip_a2v_cross_attn"
-    SKIP_V2A_CROSS_ATTN = "skip_v2a_cross_attn"
-    SKIP_VIDEO_SELF_ATTN = "skip_video_self_attn"
-    SKIP_AUDIO_SELF_ATTN = "skip_audio_self_attn"
-
-
-@dataclass(frozen=True)
-class Perturbation:
-    type: PerturbationType
-    blocks: list[int] | None  # None means all blocks
-
-    def is_perturbed(self, perturbation_type: PerturbationType, block: int) -> bool:
-        if self.type != perturbation_type:
-            return False
-
-        if self.blocks is None:
-            return True
-
-        return block in self.blocks
-
-
-@dataclass(frozen=True)
-class PerturbationConfig:
-    perturbations: list[Perturbation] | None
-
-    def is_perturbed(self, perturbation_type: PerturbationType, block: int) -> bool:
-        if self.perturbations is None:
-            return False
-
-        return any(perturbation.is_perturbed(perturbation_type, block) for perturbation in self.perturbations)
-
-    @staticmethod
-    def empty() -> "PerturbationConfig":
-        return PerturbationConfig([])
-
-
-@dataclass(frozen=True)
-class BatchedPerturbationConfig:
-    perturbations: list[PerturbationConfig]
-
-    def mask(
-        self, perturbation_type: PerturbationType, block: int, device: DeviceLikeType, dtype: torch.dtype
-    ) -> torch.Tensor:
-        mask = torch.ones((len(self.perturbations),), device=device, dtype=dtype)
-        for batch_idx, perturbation in enumerate(self.perturbations):
-            if perturbation.is_perturbed(perturbation_type, block):
-                mask[batch_idx] = 0
-
-        return mask
-
-    def mask_like(self, perturbation_type: PerturbationType, block: int, values: torch.Tensor) -> torch.Tensor:
-        mask = self.mask(perturbation_type, block, values.device, values.dtype)
-        return mask.view(mask.numel(), *([1] * len(values.shape[1:])))
-
-    def any_in_batch(self, perturbation_type: PerturbationType, block: int) -> bool:
-        return any(perturbation.is_perturbed(perturbation_type, block) for perturbation in self.perturbations)
-
-    def all_in_batch(self, perturbation_type: PerturbationType, block: int) -> bool:
-        return all(perturbation.is_perturbed(perturbation_type, block) for perturbation in self.perturbations)
-
-    @staticmethod
-    def empty(batch_size: int) -> "BatchedPerturbationConfig":
-        return BatchedPerturbationConfig([PerturbationConfig.empty() for _ in range(batch_size)])
+from dataclasses import dataclass
+from enum import Enum
+
+import torch
+from torch._prims_common import DeviceLikeType
+
+
+class PerturbationType(Enum):
+    """Types of attention perturbations for STG (Spatio-Temporal Guidance)."""
+
+    SKIP_A2V_CROSS_ATTN = "skip_a2v_cross_attn"
+    SKIP_V2A_CROSS_ATTN = "skip_v2a_cross_attn"
+    SKIP_VIDEO_SELF_ATTN = "skip_video_self_attn"
+    SKIP_AUDIO_SELF_ATTN = "skip_audio_self_attn"
+
+
+@dataclass(frozen=True)
+class Perturbation:
+    """A single perturbation specifying which attention type to skip and in which blocks."""
+
+    type: PerturbationType
+    blocks: list[int] | None  # None means all blocks
+
+    def is_perturbed(self, perturbation_type: PerturbationType, block: int) -> bool:
+        if self.type != perturbation_type:
+            return False
+
+        if self.blocks is None:
+            return True
+
+        return block in self.blocks
+
+
+@dataclass(frozen=True)
+class PerturbationConfig:
+    """Configuration holding a list of perturbations for a single sample."""
+
+    perturbations: list[Perturbation] | None
+
+    def is_perturbed(self, perturbation_type: PerturbationType, block: int) -> bool:
+        if self.perturbations is None:
+            return False
+
+        return any(perturbation.is_perturbed(perturbation_type, block) for perturbation in self.perturbations)
+
+    @staticmethod
+    def empty() -> "PerturbationConfig":
+        return PerturbationConfig([])
+
+
+@dataclass(frozen=True)
+class BatchedPerturbationConfig:
+    """Perturbation configurations for a batch, with utilities for generating attention masks."""
+
+    perturbations: list[PerturbationConfig]
+
+    def mask(
+        self, perturbation_type: PerturbationType, block: int, device: DeviceLikeType, dtype: torch.dtype
+    ) -> torch.Tensor:
+        mask = torch.ones((len(self.perturbations),), device=device, dtype=dtype)
+        for batch_idx, perturbation in enumerate(self.perturbations):
+            if perturbation.is_perturbed(perturbation_type, block):
+                mask[batch_idx] = 0
+
+        return mask
+
+    def mask_like(self, perturbation_type: PerturbationType, block: int, values: torch.Tensor) -> torch.Tensor:
+        mask = self.mask(perturbation_type, block, values.device, values.dtype)
+        return mask.view(mask.numel(), *([1] * len(values.shape[1:])))
+
+    def any_in_batch(self, perturbation_type: PerturbationType, block: int) -> bool:
+        return any(perturbation.is_perturbed(perturbation_type, block) for perturbation in self.perturbations)
+
+    def all_in_batch(self, perturbation_type: PerturbationType, block: int) -> bool:
+        return all(perturbation.is_perturbed(perturbation_type, block) for perturbation in self.perturbations)
+
+    @staticmethod
+    def empty(batch_size: int) -> "BatchedPerturbationConfig":
+        return BatchedPerturbationConfig([PerturbationConfig.empty() for _ in range(batch_size)])

packages/ltx-core/src/ltx_core/loader/__init__.py

@@ -0,0 +1,48 @@
+"""Loader utilities for model weights, LoRAs, and safetensor operations."""
+
+from ltx_core.loader.fuse_loras import apply_loras
+from ltx_core.loader.module_ops import ModuleOps
+from ltx_core.loader.primitives import (
+    LoRAAdaptableProtocol,
+    LoraPathStrengthAndSDOps,
+    LoraStateDictWithStrength,
+    ModelBuilderProtocol,
+    StateDict,
+    StateDictLoader,
+)
+from ltx_core.loader.registry import DummyRegistry, Registry, StateDictRegistry
+from ltx_core.loader.sd_ops import (
+    LTXV_LORA_COMFY_RENAMING_MAP,
+    ContentMatching,
+    ContentReplacement,
+    KeyValueOperation,
+    KeyValueOperationResult,
+    SDKeyValueOperation,
+    SDOps,
+)
+from ltx_core.loader.sft_loader import SafetensorsModelStateDictLoader, SafetensorsStateDictLoader
+from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder
+
+__all__ = [
+    "LTXV_LORA_COMFY_RENAMING_MAP",
+    "ContentMatching",
+    "ContentReplacement",
+    "DummyRegistry",
+    "KeyValueOperation",
+    "KeyValueOperationResult",
+    "LoRAAdaptableProtocol",
+    "LoraPathStrengthAndSDOps",
+    "LoraStateDictWithStrength",
+    "ModelBuilderProtocol",
+    "ModuleOps",
+    "Registry",
+    "SDKeyValueOperation",
+    "SDOps",
+    "SafetensorsModelStateDictLoader",
+    "SafetensorsStateDictLoader",
+    "SingleGPUModelBuilder",
+    "StateDict",
+    "StateDictLoader",
+    "StateDictRegistry",
+    "apply_loras",
+]

packages/ltx-core/src/ltx_core/loader/fuse_loras.py

@@ -1,102 +1,118 @@
-# Copyright (c) 2025 Lightricks. All rights reserved.
-# Created by Alexey Kravtsov
-import torch
-import triton
-
-from ltx_core.loader.kernels import fused_add_round_kernel
-from ltx_core.loader.primitives import LoraStateDictWithStrength, StateDict
-
-BLOCK_SIZE = 1024
-
-
-def fused_add_round_launch(target_weight: torch.Tensor, original_weight: torch.Tensor, seed: int) -> torch.Tensor:
-    if original_weight.dtype == torch.float8_e4m3fn:
-        exponent_bits, mantissa_bits, exponent_bias = 4, 3, 7
-    elif original_weight.dtype == torch.float8_e5m2:
-        exponent_bits, mantissa_bits, exponent_bias = 5, 2, 15  # noqa: F841
-    else:
-        raise ValueError("Unsupported dtype")
-
-    if target_weight.dtype != torch.bfloat16:
-        raise ValueError("target_weight dtype must be bfloat16")
-
-    # Calculate grid and block sizes
-    n_elements = original_weight.numel()
-    grid = (triton.cdiv(n_elements, BLOCK_SIZE),)
-
-    # Launch kernel
-    fused_add_round_kernel[grid](
-        original_weight,
-        target_weight,
-        seed,
-        n_elements,
-        exponent_bias,
-        mantissa_bits,
-        BLOCK_SIZE,
-    )
-    return target_weight
-
-
-def calculate_weight_float8_(target_weights: torch.Tensor, original_weights: torch.Tensor) -> torch.Tensor:
-    result = fused_add_round_launch(target_weights, original_weights, seed=0).to(target_weights.dtype)
-    target_weights.copy_(result, non_blocking=True)
-    return target_weights
-
-
-def _prepare_deltas(
-    lora_sd_and_strengths: list[LoraStateDictWithStrength], key: str, dtype: torch.dtype, device: torch.device
-) -> torch.Tensor | None:
-    deltas = []
-    prefix = key[: -len(".weight")]
-    key_a = f"{prefix}.lora_A.weight"
-    key_b = f"{prefix}.lora_B.weight"
-    for lsd, coef in lora_sd_and_strengths:
-        if key_a not in lsd.sd or key_b not in lsd.sd:
-            continue
-        product = torch.matmul(lsd.sd[key_b] * coef, lsd.sd[key_a])
-        deltas.append(product.to(dtype=dtype, device=device))
-    if len(deltas) == 0:
-        return None
-    elif len(deltas) == 1:
-        return deltas[0]
-    return torch.sum(torch.stack(deltas, dim=0), dim=0)
-
-
-def apply_loras(
-    model_sd: StateDict,
-    lora_sd_and_strengths: list[LoraStateDictWithStrength],
-    dtype: torch.dtype,
-    destination_sd: StateDict | None = None,
-) -> StateDict:
-    sd = {}
-    if destination_sd is not None:
-        sd = destination_sd.sd
-    size = 0
-    device = torch.device("meta")
-    inner_dtypes = set()
-    for key, weight in model_sd.sd.items():
-        if weight is None:
-            continue
-        device = weight.device
-        target_dtype = dtype if dtype is not None else weight.dtype
-        deltas_dtype = target_dtype if target_dtype not in [torch.float8_e4m3fn, torch.float8_e5m2] else torch.bfloat16
-        deltas = _prepare_deltas(lora_sd_and_strengths, key, deltas_dtype, device)
-        if deltas is None:
-            if key in sd:
-                continue
-            deltas = weight.clone().to(dtype=target_dtype, device=device)
-        elif weight.dtype == torch.float8_e4m3fn:
-            if str(device).startswith("cuda"):
-                deltas = calculate_weight_float8_(deltas, weight)
-            else:
-                deltas.add_(weight.to(dtype=deltas.dtype, device=device))
-        elif weight.dtype == torch.bfloat16:
-            deltas.add_(weight)
-        else:
-            raise ValueError(f"Unsupported dtype: {weight.dtype}")
-        sd[key] = deltas.to(dtype=target_dtype)
-        inner_dtypes.add(target_dtype)
-        size += deltas.nbytes
-    if destination_sd is not None:
-        return destination_sd
-    return StateDict(sd, device, size, inner_dtypes)
+import torch
+import triton
+
+from ltx_core.loader.kernels import fused_add_round_kernel
+from ltx_core.loader.primitives import LoraStateDictWithStrength, StateDict
+from typing import Iterable
+
+BLOCK_SIZE = 1024
+
+
+def fused_add_round_launch(target_weight: torch.Tensor, original_weight: torch.Tensor, seed: int) -> torch.Tensor:
+    if original_weight.dtype == torch.float8_e4m3fn:
+        exponent_bits, mantissa_bits, exponent_bias = 4, 3, 7
+    elif original_weight.dtype == torch.float8_e5m2:
+        exponent_bits, mantissa_bits, exponent_bias = 5, 2, 15  # noqa: F841
+    else:
+        raise ValueError("Unsupported dtype")
+
+    if target_weight.dtype != torch.bfloat16:
+        raise ValueError("target_weight dtype must be bfloat16")
+
+    # Calculate grid and block sizes
+    n_elements = original_weight.numel()
+    grid = (triton.cdiv(n_elements, BLOCK_SIZE),)
+
+    # Launch kernel
+    fused_add_round_kernel[grid](
+        original_weight,
+        target_weight,
+        seed,
+        n_elements,
+        exponent_bias,
+        mantissa_bits,
+        BLOCK_SIZE,
+    )
+    return target_weight
+
+
+def calculate_weight_float8_(target_weights: torch.Tensor, original_weights: torch.Tensor) -> torch.Tensor:
+    result = fused_add_round_launch(target_weights, original_weights, seed=0).to(target_weights.dtype)
+    target_weights.copy_(result, non_blocking=True)
+    return target_weights
+
+
+def _prepare_deltas(
+    lora_sd_and_strengths: list[LoraStateDictWithStrength], key: str, dtype: torch.dtype, device: torch.device
+) -> torch.Tensor | None:
+    deltas = []
+    prefix = key[: -len(".weight")]
+    key_a = f"{prefix}.lora_A.weight"
+    key_b = f"{prefix}.lora_B.weight"
+    for lsd, coef in lora_sd_and_strengths:
+        if key_a not in lsd.sd or key_b not in lsd.sd:
+            continue
+        product = torch.matmul(lsd.sd[key_b] * coef, lsd.sd[key_a])
+        deltas.append(product.to(dtype=dtype, device=device))
+    if len(deltas) == 0:
+        return None
+    elif len(deltas) == 1:
+        return deltas[0]
+    return torch.sum(torch.stack(deltas, dim=0), dim=0)
+
+def apply_loras(
+    model_sd: StateDict,
+    lora_sd_and_strengths: list[LoraStateDictWithStrength],
+    dtype: torch.dtype,
+    destination_sd: StateDict | None = None,
+    return_affected: bool = False,
+) -> StateDict | tuple[StateDict, list[str]]:
+    sd = destination_sd.sd if destination_sd is not None else {}
+    size = 0
+    device = torch.device("meta")
+    inner_dtypes = set()
+
+    affected_weight_keys: list[str] = []
+    affected_module_prefixes: set[str] = set()
+
+    for key, weight in model_sd.sd.items():
+        if weight is None:
+            continue
+        if not key.endswith(".weight"):
+            # optional: skip non-weight tensors if your SD has them
+            continue
+
+        device = weight.device
+        target_dtype = dtype if dtype is not None else weight.dtype
+        deltas_dtype = target_dtype  # you said ignore fp8 path
+
+        deltas = _prepare_deltas(lora_sd_and_strengths, key, deltas_dtype, device)
+
+        # Record which weights are actually modified by LoRA
+        if deltas is not None:
+            affected_weight_keys.append(key)
+            affected_module_prefixes.add(key[: -len(".weight")])
+
+        if deltas is None:
+            if key in sd:
+                continue
+            out = weight.clone().to(dtype=target_dtype, device=device)
+        else:
+            # normal add_ path (bf16 etc)
+            out = deltas.to(dtype=target_dtype)
+            # IMPORTANT: add base weight
+            out.add_(weight.to(dtype=out.dtype, device=device))
+
+        sd[key] = out
+        inner_dtypes.add(target_dtype)
+        size += out.nbytes
+
+    result = destination_sd if destination_sd is not None else StateDict(sd, device, size, inner_dtypes)
+
+    if return_affected:
+        # sorted for stable output
+        affected = sorted(affected_module_prefixes)
+        return result, affected
+
+    return result
+

packages/ltx-core/src/ltx_core/loader/kernels.py

@@ -1,74 +1,72 @@
-# ruff: noqa: ANN001, ANN201, ERA001, N803, N806
-# Copyright (c) 2025 Lightricks. All rights reserved.
-# Created by Alexey Kravtsov
-import triton
-import triton.language as tl
-
-
-@triton.jit
-def fused_add_round_kernel(
-    x_ptr,
-    output_ptr,  # contents will be added to the output
-    seed,
-    n_elements,
-    EXPONENT_BIAS,
-    MANTISSA_BITS,
-    BLOCK_SIZE: tl.constexpr,
-):
-    """
-    A kernel to upcast 8bit quantized weights to bfloat16 with stochastic rounding
-    and add them to bfloat16 output weights. Might be used to upcast original model weights
-    and to further add them to precalculated deltas coming from LoRAs.
-    """
-    # Get program ID and compute offsets
-    pid = tl.program_id(axis=0)
-    block_start = pid * BLOCK_SIZE
-    offsets = block_start + tl.arange(0, BLOCK_SIZE)
-    mask = offsets < n_elements
-
-    # Load data
-    x = tl.load(x_ptr + offsets, mask=mask)
-    rand_vals = tl.rand(seed, offsets) - 0.5
-
-    x = tl.cast(x, tl.float16)
-    delta = tl.load(output_ptr + offsets, mask=mask)
-    delta = tl.cast(delta, tl.float16)
-    x = x + delta
-
-    x_bits = tl.cast(x, tl.int16, bitcast=True)
-
-    # Calculate the exponent. Unbiased fp16 exponent is ((x_bits & 0x7C00) >> 10) - 15 for
-    # normal numbers and -14 for subnormals.
-    fp16_exponent_bits = (x_bits & 0x7C00) >> 10
-    fp16_normals = fp16_exponent_bits > 0
-    fp16_exponent = tl.where(fp16_normals, fp16_exponent_bits - 15, -14)
-
-    # Add the target dtype's exponent bias and clamp to the target dtype's exponent range.
-    exponent = fp16_exponent + EXPONENT_BIAS
-    MAX_EXPONENT = 2 * EXPONENT_BIAS + 1
-    exponent = tl.where(exponent > MAX_EXPONENT, MAX_EXPONENT, exponent)
-    exponent = tl.where(exponent < 0, 0, exponent)
-
-    # Normal ULP exponent, expressed as an fp16 exponent field:
-    # (exponent - EXPONENT_BIAS - MANTISSA_BITS) + 15
-    # Simplifies to: fp16_exponent - MANTISSA_BITS + 15
-    # See https://en.wikipedia.org/wiki/Unit_in_the_last_place
-    eps_exp = tl.maximum(0, tl.minimum(31, exponent - EXPONENT_BIAS - MANTISSA_BITS + 15))
-
-    # Calculate epsilon in the target dtype
-    eps_normal = tl.cast(tl.cast(eps_exp << 10, tl.int16), tl.float16, bitcast=True)
-
-    # Subnormal ULP: 2^(1 - EXPONENT_BIAS - MANTISSA_BITS) ->
-    # fp16 exponent bits: (1 - EXPONENT_BIAS - MANTISSA_BITS) + 15 =
-    # 16 - EXPONENT_BIAS - MANTISSA_BITS
-    eps_subnormal = tl.cast(tl.cast((16 - EXPONENT_BIAS - MANTISSA_BITS) << 10, tl.int16), tl.float16, bitcast=True)
-    eps = tl.where(exponent > 0, eps_normal, eps_subnormal)
-
-    # Apply zero mask to epsilon
-    eps = tl.where(x == 0, 0.0, eps)
-
-    # Apply stochastic rounding
-    output = tl.cast(x + rand_vals * eps, tl.bfloat16)
-
-    # Store the result
-    tl.store(output_ptr + offsets, output, mask=mask)
+# ruff: noqa: ANN001, ANN201, ERA001, N803, N806
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def fused_add_round_kernel(
+    x_ptr,
+    output_ptr,  # contents will be added to the output
+    seed,
+    n_elements,
+    EXPONENT_BIAS,
+    MANTISSA_BITS,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """
+    A kernel to upcast 8bit quantized weights to bfloat16 with stochastic rounding
+    and add them to bfloat16 output weights. Might be used to upcast original model weights
+    and to further add them to precalculated deltas coming from LoRAs.
+    """
+    # Get program ID and compute offsets
+    pid = tl.program_id(axis=0)
+    block_start = pid * BLOCK_SIZE
+    offsets = block_start + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_elements
+
+    # Load data
+    x = tl.load(x_ptr + offsets, mask=mask)
+    rand_vals = tl.rand(seed, offsets) - 0.5
+
+    x = tl.cast(x, tl.float16)
+    delta = tl.load(output_ptr + offsets, mask=mask)
+    delta = tl.cast(delta, tl.float16)
+    x = x + delta
+
+    x_bits = tl.cast(x, tl.int16, bitcast=True)
+
+    # Calculate the exponent. Unbiased fp16 exponent is ((x_bits & 0x7C00) >> 10) - 15 for
+    # normal numbers and -14 for subnormals.
+    fp16_exponent_bits = (x_bits & 0x7C00) >> 10
+    fp16_normals = fp16_exponent_bits > 0
+    fp16_exponent = tl.where(fp16_normals, fp16_exponent_bits - 15, -14)
+
+    # Add the target dtype's exponent bias and clamp to the target dtype's exponent range.
+    exponent = fp16_exponent + EXPONENT_BIAS
+    MAX_EXPONENT = 2 * EXPONENT_BIAS + 1
+    exponent = tl.where(exponent > MAX_EXPONENT, MAX_EXPONENT, exponent)
+    exponent = tl.where(exponent < 0, 0, exponent)
+
+    # Normal ULP exponent, expressed as an fp16 exponent field:
+    # (exponent - EXPONENT_BIAS - MANTISSA_BITS) + 15
+    # Simplifies to: fp16_exponent - MANTISSA_BITS + 15
+    # See https://en.wikipedia.org/wiki/Unit_in_the_last_place
+    eps_exp = tl.maximum(0, tl.minimum(31, exponent - EXPONENT_BIAS - MANTISSA_BITS + 15))
+
+    # Calculate epsilon in the target dtype
+    eps_normal = tl.cast(tl.cast(eps_exp << 10, tl.int16), tl.float16, bitcast=True)
+
+    # Subnormal ULP: 2^(1 - EXPONENT_BIAS - MANTISSA_BITS) ->
+    # fp16 exponent bits: (1 - EXPONENT_BIAS - MANTISSA_BITS) + 15 =
+    # 16 - EXPONENT_BIAS - MANTISSA_BITS
+    eps_subnormal = tl.cast(tl.cast((16 - EXPONENT_BIAS - MANTISSA_BITS) << 10, tl.int16), tl.float16, bitcast=True)
+    eps = tl.where(exponent > 0, eps_normal, eps_subnormal)
+
+    # Apply zero mask to epsilon
+    eps = tl.where(x == 0, 0.0, eps)
+
+    # Apply stochastic rounding
+    output = tl.cast(x + rand_vals * eps, tl.bfloat16)
+
+    # Store the result
+    tl.store(output_ptr + offsets, output, mask=mask)

packages/ltx-core/src/ltx_core/loader/module_ops.py

@@ -1,11 +1,14 @@
-# Copyright (c) 2025 Lightricks. All rights reserved.
-# Created by Alexey Kravtsov
-from typing import Callable, NamedTuple
-
-import torch
-
-
-class ModuleOps(NamedTuple):
-    name: str
-    matcher: Callable[[torch.nn.Module], bool]
-    mutator: Callable[[torch.nn.Module], torch.nn.Module]
+from typing import Callable, NamedTuple
+
+import torch
+
+
+class ModuleOps(NamedTuple):
+    """
+    Defines a named operation for matching and mutating PyTorch modules.
+    Used to selectively transform modules in a model (e.g., replacing layers with quantized versions).
+    """
+
+    name: str
+    matcher: Callable[[torch.nn.Module], bool]
+    mutator: Callable[[torch.nn.Module], torch.nn.Module]

packages/ltx-core/src/ltx_core/loader/primitives.py

@@ -1,63 +1,109 @@
-# Copyright (c) 2025 Lightricks. All rights reserved.
-# Created by Alexey Kravtsov
-from dataclasses import dataclass
-from typing import NamedTuple, Protocol
-
-import torch
-
-from ltx_core.loader.module_ops import ModuleOps
-from ltx_core.loader.sd_ops import SDOps
-from ltx_core.model.model_protocol import ModelType
-
-
-@dataclass(frozen=True)
-class StateDict:
-    sd: dict
-    device: torch.device
-    size: int
-    dtype: set[torch.dtype]
-
-    def footprint(self) -> tuple[int, torch.device]:
-        return self.size, self.device
-
-
-class StateDictLoader(Protocol):
-    def metadata(self, path: str) -> dict:
-        """
-        Load metadata from path
-        """
-
-    def load(self, path: str | list[str], sd_ops: SDOps | None = None, device: torch.device | None = None) -> StateDict:
-        """
-        Load state dict from path or paths (for sharded model storage) and apply sd_ops
-        """
-
-
-class ModelBuilderProtocol(Protocol[ModelType]):
-    def meta_model(self, config: dict, module_ops: list[ModuleOps] | None = None) -> ModelType: ...
-
-    def build(self, dtype: torch.dtype | None = None) -> ModelType:
-        """
-        Build the model
-        Args:
-            dtype: Target dtype for the model, if None, uses the dtype of the model_path model
-        Returns:
-            Model instance
-        """
-        ...
-
-
-class LoRAAdaptableProtocol(Protocol):
-    def lora(self, lora_path: str, strength: float) -> "LoRAAdaptableProtocol":
-        pass
-
-
-class LoraPathStrengthAndSDOps(NamedTuple):
-    path: str
-    strength: float
-    sd_ops: SDOps
-
-
-class LoraStateDictWithStrength(NamedTuple):
-    state_dict: StateDict
-    strength: float
+from dataclasses import dataclass
+from typing import NamedTuple, Protocol
+
+import torch
+
+from ltx_core.loader.module_ops import ModuleOps
+from ltx_core.loader.sd_ops import SDOps
+from ltx_core.model.model_protocol import ModelType
+
+
+@dataclass(frozen=True)
+class StateDict:
+    """
+    Immutable container for a PyTorch state dictionary.
+    Contains:
+    - sd: Dictionary of tensors (weights, buffers, etc.)
+    - device: Device where tensors are stored
+    - size: Total memory footprint in bytes
+    - dtype: Set of tensor dtypes present
+    """
+
+    sd: dict
+    device: torch.device
+    size: int
+    dtype: set[torch.dtype]
+
+    def footprint(self) -> tuple[int, torch.device]:
+        return self.size, self.device
+
+
+class StateDictLoader(Protocol):
+    """
+    Protocol for loading state dictionaries from various sources.
+    Implementations must provide:
+    - metadata: Extract model metadata from a single path
+    - load: Load state dict from path(s) and apply SDOps transformations
+    """
+
+    def metadata(self, path: str) -> dict:
+        """
+        Load metadata from path
+        """
+
+    def load(self, path: str | list[str], sd_ops: SDOps | None = None, device: torch.device | None = None) -> StateDict:
+        """
+        Load state dict from path or paths (for sharded model storage) and apply sd_ops
+        """
+
+
+class ModelBuilderProtocol(Protocol[ModelType]):
+    """
+    Protocol for building PyTorch models from configuration dictionaries.
+    Implementations must provide:
+    - meta_model: Create a model from configuration dictionary and apply module operations
+    - build: Create and initialize a model from state dictionary and apply dtype transformations
+    """
+
+    def meta_model(self, config: dict, module_ops: list[ModuleOps] | None = None) -> ModelType:
+        """
+        Create a model on the meta device from a configuration dictionary.
+        This decouples model creation from weight loading, allowing the model
+        architecture to be instantiated without allocating memory for parameters.
+        Args:
+            config: Model configuration dictionary.
+            module_ops: Optional list of module operations to apply (e.g., quantization).
+        Returns:
+            Model instance on meta device (no actual memory allocated for parameters).
+        """
+        ...
+
+    def build(self, dtype: torch.dtype | None = None) -> ModelType:
+        """
+        Build the model
+        Args:
+            dtype: Target dtype for the model, if None, uses the dtype of the model_path model
+        Returns:
+            Model instance
+        """
+        ...
+
+
+class LoRAAdaptableProtocol(Protocol):
+    """
+    Protocol for models that can be adapted with LoRAs.
+    Implementations must provide:
+    - lora: Add a LoRA to the model
+    """
+
+    def lora(self, lora_path: str, strength: float) -> "LoRAAdaptableProtocol":
+        pass
+
+
+class LoraPathStrengthAndSDOps(NamedTuple):
+    """
+    Tuple containing a LoRA path, strength, and SDOps for applying to the LoRA state dict.
+    """
+
+    path: str
+    strength: float
+    sd_ops: SDOps
+
+
+class LoraStateDictWithStrength(NamedTuple):
+    """
+    Tuple containing a LoRA state dict and strength for applying to the model.
+    """
+
+    state_dict: StateDict
+    strength: float

packages/ltx-core/src/ltx_core/loader/registry.py

@@ -1,68 +1,84 @@
-# Copyright (c) 2025 Lightricks. All rights reserved.
-# Created by Alexey Kravtsov
-import hashlib
-import threading
-from dataclasses import dataclass, field
-from pathlib import Path
-from typing import Protocol
-
-from ltx_core.loader.primitives import StateDict
-from ltx_core.loader.sd_ops import SDOps
-
-
-class Registry(Protocol):
-    def add(self, paths: list[str], sd_ops: SDOps | None, state_dict: StateDict) -> None: ...
-
-    def pop(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None: ...
-
-    def get(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None: ...
-
-    def clear(self) -> None: ...
-
-
-class DummyRegistry(Registry):
-    def add(self, paths: list[str], sd_ops: SDOps | None, state_dict: StateDict) -> None:
-        pass
-
-    def pop(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None:
-        pass
-
-    def get(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None:
-        pass
-
-    def clear(self) -> None:
-        pass
-
-
-@dataclass
-class StateDictRegistry(Registry):
-    _state_dicts: dict[str, StateDict] = field(default_factory=dict)
-    _lock: threading.Lock = field(default_factory=threading.Lock)
-
-    def _generate_id(self, paths: list[str], sd_ops: SDOps) -> str:
-        m = hashlib.sha256()
-        parts = [str(Path(p).resolve()) for p in paths]
-        if sd_ops is not None:
-            parts.append(sd_ops.name)
-        m.update("\0".join(parts).encode("utf-8"))
-        return m.hexdigest()
-
-    def add(self, paths: list[str], sd_ops: SDOps | None, state_dict: StateDict) -> str:
-        sd_id = self._generate_id(paths, sd_ops)
-        with self._lock:
-            if sd_id in self._state_dicts:
-                raise ValueError(f"State dict retrieved from {paths} with {sd_ops} already added, check with get first")
-            self._state_dicts[sd_id] = state_dict
-        return sd_id
-
-    def pop(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None:
-        with self._lock:
-            return self._state_dicts.pop(self._generate_id(paths, sd_ops), None)
-
-    def get(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None:
-        with self._lock:
-            return self._state_dicts.get(self._generate_id(paths, sd_ops), None)
-
-    def clear(self) -> None:
-        with self._lock:
-            self._state_dicts.clear()
+import hashlib
+import threading
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Protocol
+
+from ltx_core.loader.primitives import StateDict
+from ltx_core.loader.sd_ops import SDOps
+
+
+class Registry(Protocol):
+    """
+    Protocol for managing state dictionaries in a registry.
+    It is used to store state dictionaries and reuse them later without loading them again.
+    Implementations must provide:
+    - add: Add a state dictionary to the registry
+    - pop: Remove a state dictionary from the registry
+    - get: Retrieve a state dictionary from the registry
+    - clear: Clear all state dictionaries from the registry
+    """
+
+    def add(self, paths: list[str], sd_ops: SDOps | None, state_dict: StateDict) -> None: ...
+
+    def pop(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None: ...
+
+    def get(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None: ...
+
+    def clear(self) -> None: ...
+
+
+class DummyRegistry(Registry):
+    """
+    Dummy registry that does not store state dictionaries.
+    """
+
+    def add(self, paths: list[str], sd_ops: SDOps | None, state_dict: StateDict) -> None:
+        pass
+
+    def pop(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None:
+        pass
+
+    def get(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None:
+        pass
+
+    def clear(self) -> None:
+        pass
+
+
+@dataclass
+class StateDictRegistry(Registry):
+    """
+    Registry that stores state dictionaries in a dictionary.
+    """
+
+    _state_dicts: dict[str, StateDict] = field(default_factory=dict)
+    _lock: threading.Lock = field(default_factory=threading.Lock)
+
+    def _generate_id(self, paths: list[str], sd_ops: SDOps) -> str:
+        m = hashlib.sha256()
+        parts = [str(Path(p).resolve()) for p in paths]
+        if sd_ops is not None:
+            parts.append(sd_ops.name)
+        m.update("\0".join(parts).encode("utf-8"))
+        return m.hexdigest()
+
+    def add(self, paths: list[str], sd_ops: SDOps | None, state_dict: StateDict) -> str:
+        sd_id = self._generate_id(paths, sd_ops)
+        with self._lock:
+            if sd_id in self._state_dicts:
+                raise ValueError(f"State dict retrieved from {paths} with {sd_ops} already added, check with get first")
+            self._state_dicts[sd_id] = state_dict
+        return sd_id
+
+    def pop(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None:
+        with self._lock:
+            return self._state_dicts.pop(self._generate_id(paths, sd_ops), None)
+
+    def get(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None:
+        with self._lock:
+            return self._state_dicts.get(self._generate_id(paths, sd_ops), None)
+
+    def clear(self) -> None:
+        with self._lock:
+            self._state_dicts.clear()

packages/ltx-core/src/ltx_core/loader/sd_ops.py

@@ -1,107 +1,127 @@
-# Copyright (c) 2025 Lightricks. All rights reserved.
-# Created by Alexey Kravtsov
-
-from dataclasses import dataclass, replace
-#from typing import NamedTuple, Protocol, Self
-from typing import NamedTuple, Protocol
-from typing_extensions import Self
-
-import torch
-
-
-@dataclass(frozen=True, slots=True)
-class ContentReplacement:
-    content: str
-    replacement: str
-
-
-@dataclass(frozen=True, slots=True)
-class ContentMatching:
-    prefix: str = ""
-    suffix: str = ""
-
-
-class KeyValueOperationResult(NamedTuple):
-    new_key: str
-    new_value: torch.Tensor
-
-
-class KeyValueOperation(Protocol):
-    def __call__(self, tensor_key: str, tensor_value: torch.Tensor) -> list[KeyValueOperationResult]: ...
-
-
-@dataclass(frozen=True, slots=True)
-class SDKeyValueOperation:
-    key_matcher: ContentMatching
-    kv_operation: KeyValueOperation
-
-
-@dataclass(frozen=True, slots=True)
-class SDOps:
-    """Immutable class representing state dict key operations."""
-
-    name: str
-    mapping: tuple[
-        ContentReplacement | ContentMatching | SDKeyValueOperation, ...
-    ] = ()  # Immutable tuple of (key, value) pairs
-
-    def with_replacement(self, content: str, replacement: str) -> Self:
-        """Create a new SDOps instance with the specified replacement added to the mapping."""
-
-        new_mapping = (*self.mapping, ContentReplacement(content, replacement))
-        return replace(self, mapping=new_mapping)
-
-    def with_matching(self, prefix: str = "", suffix: str = "") -> Self:
-        """Create a new SDOps instance with the specified prefix and suffix matching added to the mapping."""
-
-        new_mapping = (*self.mapping, ContentMatching(prefix, suffix))
-        return replace(self, mapping=new_mapping)
-
-    def with_kv_operation(
-        self,
-        operation: KeyValueOperation,
-        key_prefix: str = "",
-        key_suffix: str = "",
-    ) -> Self:
-        """Create a new SDOps instance with the specified value operation added to the mapping."""
-        key_matcher = ContentMatching(key_prefix, key_suffix)
-        sd_kv_operation = SDKeyValueOperation(key_matcher, operation)
-        new_mapping = (*self.mapping, sd_kv_operation)
-        return replace(self, mapping=new_mapping)
-
-    def apply_to_key(self, key: str) -> str | None:
-        """Apply the mapping to the given name."""
-        matchers = [content for content in self.mapping if isinstance(content, ContentMatching)]
-        valid = any(key.startswith(f.prefix) and key.endswith(f.suffix) for f in matchers)
-        if not valid:
-            return None
-
-        for replacement in self.mapping:
-            if not isinstance(replacement, ContentReplacement):
-                continue
-            if replacement.content in key:
-                key = key.replace(replacement.content, replacement.replacement)
-        return key
-
-    def apply_to_key_value(self, key: str, value: torch.Tensor) -> list[KeyValueOperationResult]:
-        """Apply the value operation to the given name and associated value."""
-        for operation in self.mapping:
-            if not isinstance(operation, SDKeyValueOperation):
-                continue
-            if key.startswith(operation.key_matcher.prefix) and key.endswith(operation.key_matcher.suffix):
-                return operation.kv_operation(key, value)
-        return [KeyValueOperationResult(key, value)]
-
-
-# Predefined SDOps instances
-LTXV_LORA_COMFY_RENAMING_MAP = (
-    SDOps("LTXV_LORA_COMFY_PREFIX_MAP").with_matching().with_replacement("diffusion_model.", "")
-)
-
-LTXV_LORA_COMFY_TARGET_MAP = (
-    SDOps("LTXV_LORA_COMFY_TARGET_MAP")
-    .with_matching()
-    .with_replacement("diffusion_model.", "")
-    .with_replacement(".lora_A.weight", ".weight")
-    .with_replacement(".lora_B.weight", ".weight")
-)
+from dataclasses import dataclass, replace
+from typing import NamedTuple, Protocol
+
+import torch
+
+
+@dataclass(frozen=True, slots=True)
+class ContentReplacement:
+    """
+    Represents a content replacement operation.
+    Used to replace a specific content with a replacement in a state dict key.
+    """
+
+    content: str
+    replacement: str
+
+
+@dataclass(frozen=True, slots=True)
+class ContentMatching:
+    """
+    Represents a content matching operation.
+    Used to match a specific prefix and suffix in a state dict key.
+    """
+
+    prefix: str = ""
+    suffix: str = ""
+
+
+class KeyValueOperationResult(NamedTuple):
+    """
+    Represents the result of a key-value operation.
+    Contains the new key and value after the operation has been applied.
+    """
+
+    new_key: str
+    new_value: torch.Tensor
+
+
+class KeyValueOperation(Protocol):
+    """
+    Protocol for key-value operations.
+    Used to apply operations to a specific key and value in a state dict.
+    """
+
+    def __call__(self, tensor_key: str, tensor_value: torch.Tensor) -> list[KeyValueOperationResult]: ...
+
+
+@dataclass(frozen=True, slots=True)
+class SDKeyValueOperation:
+    """
+    Represents a key-value operation.
+    Used to apply operations to a specific key and value in a state dict.
+    """
+
+    key_matcher: ContentMatching
+    kv_operation: KeyValueOperation
+
+
+@dataclass(frozen=True, slots=True)
+class SDOps:
+    """Immutable class representing state dict key operations."""
+
+    name: str
+    mapping: tuple[
+        ContentReplacement | ContentMatching | SDKeyValueOperation, ...
+    ] = ()  # Immutable tuple of (key, value) pairs
+
+    def with_replacement(self, content: str, replacement: str) -> "SDOps":
+        """Create a new SDOps instance with the specified replacement added to the mapping."""
+
+        new_mapping = (*self.mapping, ContentReplacement(content, replacement))
+        return replace(self, mapping=new_mapping)
+
+    def with_matching(self, prefix: str = "", suffix: str = "") -> "SDOps":
+        """Create a new SDOps instance with the specified prefix and suffix matching added to the mapping."""
+
+        new_mapping = (*self.mapping, ContentMatching(prefix, suffix))
+        return replace(self, mapping=new_mapping)
+
+    def with_kv_operation(
+        self,
+        operation: KeyValueOperation,
+        key_prefix: str = "",
+        key_suffix: str = "",
+    ) -> "SDOps":
+        """Create a new SDOps instance with the specified value operation added to the mapping."""
+        key_matcher = ContentMatching(key_prefix, key_suffix)
+        sd_kv_operation = SDKeyValueOperation(key_matcher, operation)
+        new_mapping = (*self.mapping, sd_kv_operation)
+        return replace(self, mapping=new_mapping)
+
+    def apply_to_key(self, key: str) -> str | None:
+        """Apply the mapping to the given name."""
+        matchers = [content for content in self.mapping if isinstance(content, ContentMatching)]
+        valid = any(key.startswith(f.prefix) and key.endswith(f.suffix) for f in matchers)
+        if not valid:
+            return None
+
+        for replacement in self.mapping:
+            if not isinstance(replacement, ContentReplacement):
+                continue
+            if replacement.content in key:
+                key = key.replace(replacement.content, replacement.replacement)
+        return key
+
+    def apply_to_key_value(self, key: str, value: torch.Tensor) -> list[KeyValueOperationResult]:
+        """Apply the value operation to the given name and associated value."""
+        for operation in self.mapping:
+            if not isinstance(operation, SDKeyValueOperation):
+                continue
+            if key.startswith(operation.key_matcher.prefix) and key.endswith(operation.key_matcher.suffix):
+                return operation.kv_operation(key, value)
+        return [KeyValueOperationResult(key, value)]
+
+
+# Predefined SDOps instances
+LTXV_LORA_COMFY_RENAMING_MAP = (
+    SDOps("LTXV_LORA_COMFY_PREFIX_MAP").with_matching().with_replacement("diffusion_model.", "")
+)
+
+LTXV_LORA_COMFY_TARGET_MAP = (
+    SDOps("LTXV_LORA_COMFY_TARGET_MAP")
+    .with_matching()
+    .with_replacement("diffusion_model.", "")
+    .with_replacement(".lora_A.weight", ".weight")
+    .with_replacement(".lora_B.weight", ".weight")
+)

packages/ltx-core/src/ltx_core/loader/sft_loader.py

@@ -1,53 +1,63 @@
-# Copyright (c) 2025 Lightricks. All rights reserved.
-# Created by Alexey Kravtsov
-import json
-
-import safetensors
-import torch
-
-from ltx_core.loader.primitives import StateDict, StateDictLoader
-from ltx_core.loader.sd_ops import SDOps
-
-
-class SafetensorsStateDictLoader(StateDictLoader):
-    def metadata(self, path: str) -> dict:
-        raise NotImplementedError("Not implemented")
-
-    def load(self, path: str | list[str], sd_ops: SDOps, device: torch.device | None = None) -> StateDict:
-        """
-        Load state dict from path or paths (for sharded model storage) and apply sd_ops
-        """
-        sd = {}
-        size = 0
-        dtype = set()
-        device = device or torch.device("cpu")
-        model_paths = path if isinstance(path, list) else [path]
-        for shard_path in model_paths:
-            with safetensors.safe_open(shard_path, framework="pt", device=str(device)) as f:
-                safetensor_keys = f.keys()
-                for name in safetensor_keys:
-                    expected_name = name if sd_ops is None else sd_ops.apply_to_key(name)
-                    if expected_name is None:
-                        continue
-                    value = f.get_tensor(name).to(device=device, non_blocking=True, copy=False)
-                    key_value_pairs = ((expected_name, value),)
-                    if sd_ops is not None:
-                        key_value_pairs = sd_ops.apply_to_key_value(expected_name, value)
-                    for key, value in key_value_pairs:
-                        size += value.nbytes
-                        dtype.add(value.dtype)
-                        sd[key] = value
-
-        return StateDict(sd=sd, device=device, size=size, dtype=dtype)
-
-
-class SafetensorsModelStateDictLoader(StateDictLoader):
-    def __init__(self, weight_loader: SafetensorsStateDictLoader | None = None):
-        self.weight_loader = weight_loader if weight_loader is not None else SafetensorsStateDictLoader()
-
-    def metadata(self, path: str) -> dict:
-        with safetensors.safe_open(path, framework="pt") as f:
-            return json.loads(f.metadata()["config"])
-
-    def load(self, path: str | list[str], sd_ops: SDOps | None = None, device: torch.device | None = None) -> StateDict:
-        return self.weight_loader.load(path, sd_ops, device)
+import json
+
+import safetensors
+import torch
+
+from ltx_core.loader.primitives import StateDict, StateDictLoader
+from ltx_core.loader.sd_ops import SDOps
+
+
+class SafetensorsStateDictLoader(StateDictLoader):
+    """
+    Loads weights from safetensors files without metadata support.
+    Use this for loading raw weight files. For model files that include
+    configuration metadata, use SafetensorsModelStateDictLoader instead.
+    """
+
+    def metadata(self, path: str) -> dict:
+        raise NotImplementedError("Not implemented")
+
+    def load(self, path: str | list[str], sd_ops: SDOps, device: torch.device | None = None) -> StateDict:
+        """
+        Load state dict from path or paths (for sharded model storage) and apply sd_ops
+        """
+        sd = {}
+        size = 0
+        dtype = set()
+        device = device or torch.device("cpu")
+        model_paths = path if isinstance(path, list) else [path]
+        for shard_path in model_paths:
+            with safetensors.safe_open(shard_path, framework="pt", device=str(device)) as f:
+                safetensor_keys = f.keys()
+                for name in safetensor_keys:
+                    expected_name = name if sd_ops is None else sd_ops.apply_to_key(name)
+                    if expected_name is None:
+                        continue
+                    value = f.get_tensor(name).to(device=device, non_blocking=True, copy=False)
+                    key_value_pairs = ((expected_name, value),)
+                    if sd_ops is not None:
+                        key_value_pairs = sd_ops.apply_to_key_value(expected_name, value)
+                    for key, value in key_value_pairs:
+                        size += value.nbytes
+                        dtype.add(value.dtype)
+                        sd[key] = value
+
+        return StateDict(sd=sd, device=device, size=size, dtype=dtype)
+
+
+class SafetensorsModelStateDictLoader(StateDictLoader):
+    """
+    Loads weights and configuration metadata from safetensors model files.
+    Unlike SafetensorsStateDictLoader, this loader can read model configuration
+    from the safetensors file metadata via the metadata() method.
+    """
+
+    def __init__(self, weight_loader: SafetensorsStateDictLoader | None = None):
+        self.weight_loader = weight_loader if weight_loader is not None else SafetensorsStateDictLoader()
+
+    def metadata(self, path: str) -> dict:
+        with safetensors.safe_open(path, framework="pt") as f:
+            return json.loads(f.metadata()["config"])
+
+    def load(self, path: str | list[str], sd_ops: SDOps | None = None, device: torch.device | None = None) -> StateDict:
+        return self.weight_loader.load(path, sd_ops, device)

packages/ltx-core/src/ltx_core/loader/single_gpu_model_builder.py

@@ -1,99 +1,329 @@
-# Copyright (c) 2025 Lightricks. All rights reserved.
-# Created by Alexey Kravtsov
-import logging
-from dataclasses import dataclass, field, replace
-from typing import Generic
-
-import torch
-
-from ltx_core.loader.fuse_loras import apply_loras
-from ltx_core.loader.module_ops import ModuleOps
-from ltx_core.loader.primitives import (
-    LoRAAdaptableProtocol,
-    LoraPathStrengthAndSDOps,
-    LoraStateDictWithStrength,
-    ModelBuilderProtocol,
-    StateDict,
-    StateDictLoader,
-)
-from ltx_core.loader.registry import DummyRegistry, Registry
-from ltx_core.loader.sd_ops import SDOps
-from ltx_core.loader.sft_loader import SafetensorsModelStateDictLoader
-from ltx_core.model.model_protocol import ModelConfigurator, ModelType
-
-logger: logging.Logger = logging.getLogger(__name__)
-
-
-@dataclass(frozen=True)
-class SingleGPUModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType], LoRAAdaptableProtocol):
-    model_class_configurator: type[ModelConfigurator[ModelType]]
-    model_path: str | tuple[str, ...]
-    model_sd_ops: SDOps | None = None
-    module_ops: tuple[ModuleOps, ...] = field(default_factory=tuple)
-    loras: tuple[LoraPathStrengthAndSDOps, ...] = field(default_factory=tuple)
-    model_loader: StateDictLoader = field(default_factory=SafetensorsModelStateDictLoader)
-    registry: Registry = field(default_factory=DummyRegistry)
-
-    def lora(self, lora_path: str, strength: float = 1.0, sd_ops: SDOps | None = None) -> "SingleGPUModelBuilder":
-        return replace(self, loras=(*self.loras, LoraPathStrengthAndSDOps(lora_path, strength, sd_ops)))
-
-    def model_config(self) -> dict:
-        first_shard_path = self.model_path[0] if isinstance(self.model_path, tuple) else self.model_path
-        return self.model_loader.metadata(first_shard_path)
-
-    def meta_model(self, config: dict, module_ops: tuple[ModuleOps, ...]) -> ModelType:
-        with torch.device("meta"):
-            model = self.model_class_configurator.from_config(config)
-        for module_op in module_ops:
-            if module_op.matcher(model):
-                model = module_op.mutator(model)
-        return model
-
-    def load_sd(
-        self, paths: list[str], registry: Registry, device: torch.device | None, sd_ops: SDOps | None = None
-    ) -> StateDict:
-        state_dict = registry.get(paths, sd_ops)
-        if state_dict is None:
-            state_dict = self.model_loader.load(paths, sd_ops=sd_ops, device=device)
-            registry.add(paths, sd_ops=sd_ops, state_dict=state_dict)
-        return state_dict
-
-    def _return_model(self, meta_model: ModelType, device: torch.device) -> ModelType:
-        uninitialized_params = [name for name, param in meta_model.named_parameters() if str(param.device) == "meta"]
-        uninitialized_buffers = [name for name, buffer in meta_model.named_buffers() if str(buffer.device) == "meta"]
-        if uninitialized_params or uninitialized_buffers:
-            logger.warning(f"Uninitialized parameters or buffers: {uninitialized_params + uninitialized_buffers}")
-            return meta_model
-        retval = meta_model.to(device)
-        return retval
-
-    def build(self, device: torch.device | None = None, dtype: torch.dtype | None = None) -> ModelType:
-        device = torch.device("cuda") if device is None else device
-        config = self.model_config()
-        meta_model = self.meta_model(config, self.module_ops)
-        model_paths = self.model_path if isinstance(self.model_path, tuple) else [self.model_path]
-        model_state_dict = self.load_sd(model_paths, sd_ops=self.model_sd_ops, registry=self.registry, device=device)
-
-        lora_strengths = [lora.strength for lora in self.loras]
-        if not lora_strengths or (min(lora_strengths) == 0 and max(lora_strengths) == 0):
-            sd = model_state_dict.sd
-            if dtype is not None:
-                sd = {key: value.to(dtype=dtype) for key, value in model_state_dict.sd.items()}
-            meta_model.load_state_dict(sd, strict=False, assign=True)
-            return self._return_model(meta_model, device)
-
-        lora_state_dicts = [
-            self.load_sd([lora.path], sd_ops=lora.sd_ops, registry=self.registry, device=device) for lora in self.loras
-        ]
-        lora_sd_and_strengths = [
-            LoraStateDictWithStrength(sd, strength)
-            for sd, strength in zip(lora_state_dicts, lora_strengths, strict=True)
-        ]
-        final_sd = apply_loras(
-            model_sd=model_state_dict,
-            lora_sd_and_strengths=lora_sd_and_strengths,
-            dtype=dtype,
-            destination_sd=model_state_dict if isinstance(self.registry, DummyRegistry) else None,
-        )
-        meta_model.load_state_dict(final_sd.sd, strict=False, assign=True)
-        return self._return_model(meta_model, device)
+import logging
+from dataclasses import dataclass, field, replace
+from typing import Generic
+
+import torch
+import torch.nn as nn
+
+from ltx_core.loader.fuse_loras import apply_loras
+from ltx_core.loader.module_ops import ModuleOps
+from ltx_core.loader.primitives import (
+    LoRAAdaptableProtocol,
+    LoraPathStrengthAndSDOps,
+    LoraStateDictWithStrength,
+    ModelBuilderProtocol,
+    StateDict,
+    StateDictLoader,
+)
+from ltx_core.loader.registry import DummyRegistry, Registry
+from ltx_core.loader.sd_ops import SDOps
+from ltx_core.loader.sft_loader import SafetensorsModelStateDictLoader
+from ltx_core.model.model_protocol import ModelConfigurator, ModelType
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+
+def get_submodule_and_parent(root: nn.Module, path: str):
+    """
+    Returns (parent_module, child_name, child_module)
+    where child_module is reachable at `path` from root.
+    Supports numeric segments for Sequential/ModuleList.
+    """
+    parts = path.split(".")
+    parent = root
+    for p in parts[:-1]:
+        if p.isdigit():
+            parent = parent[int(p)]  # Sequential/ModuleList
+        else:
+            parent = getattr(parent, p)
+    last = parts[-1]
+    if last.isdigit():
+        child = parent[int(last)]
+    else:
+        child = getattr(parent, last)
+    return parent, last, child
+
+def set_submodule(root: nn.Module, path: str, new_module: nn.Module):
+    parent, last, _ = get_submodule_and_parent(root, path)
+    if last.isdigit():
+        parent[int(last)] = new_module
+    else:
+        setattr(parent, last, new_module)
+
+
+class MultiLoraLinear(nn.Module):
+    def __init__(self, base: nn.Linear):
+        super().__init__()
+        self.base = base
+        # Each entry: dict with CPU tensors always, and optional CUDA cache
+        self.adapters: list[dict] = []
+        self.enabled: list[bool] = []
+
+    def add_adapter_cpu(self, A_cpu, B_cpu, scale=1.0, enabled=False, pin_memory=False):
+        if A_cpu is None or B_cpu is None:
+            self.adapters.append({
+                "A_cpu": None,
+                "B_cpu": None,
+                "scale": float(scale),
+                "A_gpu": None,
+                "B_gpu": None,
+                "gpu_dtype": None,
+                "gpu_device": None,
+            })
+            self.enabled.append(bool(enabled))
+            return
+        
+        A_cpu = A_cpu.contiguous()
+        B_cpu = B_cpu.contiguous()
+        self.adapters.append({
+            "A_cpu": A_cpu,
+            "B_cpu": B_cpu,
+            "scale": float(scale),
+            "A_gpu": None,
+            "B_gpu": None,
+            "gpu_dtype": None,
+            "gpu_device": None,
+        })
+        self.enabled.append(bool(enabled))
+
+
+
+    def _materialize_to_gpu(self, idx: int):
+        entry = self.adapters[idx]
+        if entry["A_cpu"] is None or entry["B_cpu"] is None:
+            return
+        """Move adapter idx to the base weight device/dtype if not already there."""
+        entry = self.adapters[idx]
+        dev = self.base.weight.device
+        dt = self.base.weight.dtype
+
+        if (
+            entry["A_gpu"] is not None
+            and entry["B_gpu"] is not None
+            and entry["gpu_device"] == dev
+            and entry["gpu_dtype"] == dt
+        ):
+            return  # already good
+
+        A = entry["A_cpu"].to(device=dev, dtype=dt, non_blocking=True)
+        B = entry["B_cpu"].to(device=dev, dtype=dt, non_blocking=True)
+
+        entry["A_gpu"] = A
+        entry["B_gpu"] = B
+        entry["gpu_device"] = dev
+        entry["gpu_dtype"] = dt
+
+    def _evict_from_gpu(self, idx: int):
+        """Drop CUDA copies (free VRAM). CPU tensors remain."""
+        entry = self.adapters[idx]
+        entry["A_gpu"] = None
+        entry["B_gpu"] = None
+        entry["gpu_device"] = None
+        entry["gpu_dtype"] = None
+
+    def set_enabled(self, idx: int, enabled: bool, offload_when_disabled: bool = True):
+        if not (0 <= idx < len(self.enabled)):
+            return
+        self.enabled[idx] = enabled
+        if not enabled and offload_when_disabled:
+            self._evict_from_gpu(idx)
+
+    def forward(self, x):
+        out = self.base(x)
+        for i, on in enumerate(self.enabled):
+            if not on:
+                continue
+            entry = self.adapters[i]
+            if entry["A_cpu"] is None or entry["B_cpu"] is None:
+                continue
+            if entry["A_gpu"] is None or entry["B_gpu"] is None:
+                self._materialize_to_gpu(i)
+            A = entry["A_gpu"]
+            B = entry["B_gpu"]
+            out = out + ((x @ A.t()) @ B.t()) * entry["scale"]
+        return out
+
+def set_lora_enabled(model: nn.Module, adapter_idx: int, enabled: bool, offload_when_disabled: bool = True):
+    for m in model.modules():
+        if isinstance(m, MultiLoraLinear):
+            m.set_enabled(adapter_idx, enabled, offload_when_disabled=offload_when_disabled)
+
+def enable_only_lora(model: nn.Module, adapter_idx: int | None):
+    # disable all
+    # (assumes all layers have same number of adapters; true if you patch consistently)
+    for m in model.modules():
+        if isinstance(m, MultiLoraLinear):
+            for i in range(len(m.enabled)):
+                m.set_enabled(i, False, offload_when_disabled=True)
+    
+    torch.cuda.empty_cache()
+
+    # enable selected
+    if adapter_idx is not None and adapter_idx >= 0:
+        set_lora_enabled(model, adapter_idx, True)
+
+    torch.cuda.empty_cache()
+
+
+
+def patch_only_affected_linears(
+    model: nn.Module,
+    lora_sd: dict,                 # can be CPU state dict
+    affected_modules: list[str],
+    strength: float,
+    adapter_idx: int,
+    default_enabled: bool = False,
+):
+    for prefix in affected_modules:
+        _, _, mod = get_submodule_and_parent(model, prefix)
+
+        if isinstance(mod, MultiLoraLinear):
+            wrapped = mod
+        else:
+            if not isinstance(mod, nn.Linear):
+                continue
+            wrapped = MultiLoraLinear(mod)
+            set_submodule(model, prefix, wrapped)
+
+        # ensure adapter slots exist up to adapter_idx
+        while len(wrapped.adapters) <= adapter_idx:
+            wrapped.add_adapter_cpu(None, None, scale=0.0, enabled=False)
+
+        key_a = f"{prefix}.lora_A.weight"
+        key_b = f"{prefix}.lora_B.weight"
+        if key_a not in lora_sd or key_b not in lora_sd:
+            # leave the padded empty slot
+            continue
+
+        A_cpu = lora_sd[key_a]
+        B_cpu = lora_sd[key_b]
+
+        # overwrite the placeholder slot
+        wrapped.adapters[adapter_idx] = {
+            "A_cpu": A_cpu.contiguous(),
+            "B_cpu": B_cpu.contiguous(),
+            "scale": float(strength),
+            "A_gpu": None,
+            "B_gpu": None,
+            "gpu_dtype": None,
+            "gpu_device": None,
+        }
+        wrapped.enabled[adapter_idx] = default_enabled
+
+
+@dataclass(frozen=True)
+class SingleGPUModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType], LoRAAdaptableProtocol):
+    """
+    Builder for PyTorch models residing on a single GPU.
+    """
+
+    model_class_configurator: type[ModelConfigurator[ModelType]]
+    model_path: str | tuple[str, ...]
+    model_sd_ops: SDOps | None = None
+    module_ops: tuple[ModuleOps, ...] = field(default_factory=tuple)
+    loras: tuple[LoraPathStrengthAndSDOps, ...] = field(default_factory=tuple)
+    model_loader: StateDictLoader = field(default_factory=SafetensorsModelStateDictLoader)
+    registry: Registry = field(default_factory=DummyRegistry)
+
+    def lora(self, lora_path: str, strength: float = 1.0, sd_ops: SDOps | None = None) -> "SingleGPUModelBuilder":
+        return replace(self, loras=(*self.loras, LoraPathStrengthAndSDOps(lora_path, strength, sd_ops)))
+
+    def model_config(self) -> dict:
+        first_shard_path = self.model_path[0] if isinstance(self.model_path, tuple) else self.model_path
+        return self.model_loader.metadata(first_shard_path)
+
+    def meta_model(self, config: dict, module_ops: tuple[ModuleOps, ...]) -> ModelType:
+        with torch.device("meta"):
+            model = self.model_class_configurator.from_config(config)
+        for module_op in module_ops:
+            if module_op.matcher(model):
+                model = module_op.mutator(model)
+        return model
+
+    def load_sd(
+        self, paths: list[str], registry: Registry, device: torch.device | None, sd_ops: SDOps | None = None
+    ) -> StateDict:
+        state_dict = registry.get(paths, sd_ops)
+        if state_dict is None:
+            state_dict = self.model_loader.load(paths, sd_ops=sd_ops, device=device)
+            registry.add(paths, sd_ops=sd_ops, state_dict=state_dict)
+        return state_dict
+
+    def _return_model(self, meta_model: ModelType, device: torch.device) -> ModelType:
+        uninitialized_params = [name for name, param in meta_model.named_parameters() if str(param.device) == "meta"]
+        uninitialized_buffers = [name for name, buffer in meta_model.named_buffers() if str(buffer.device) == "meta"]
+        if uninitialized_params or uninitialized_buffers:
+            logger.warning(f"Uninitialized parameters or buffers: {uninitialized_params + uninitialized_buffers}")
+            return meta_model
+        retval = meta_model.to(device)
+        return retval
+
+    def build(self, device: torch.device | None = None, dtype: torch.dtype | None = None) -> ModelType:
+        device = torch.device("cuda") if device is None else device
+        config = self.model_config()
+        meta_model = self.meta_model(config, self.module_ops)
+        model_paths = self.model_path if isinstance(self.model_path, tuple) else [self.model_path]
+        model_state_dict = self.load_sd(model_paths, sd_ops=self.model_sd_ops, registry=self.registry, device=device)
+
+        lora_strengths = [lora.strength for lora in self.loras]
+        if not lora_strengths or (min(lora_strengths) == 0 and max(lora_strengths) == 0):
+            sd = model_state_dict.sd
+            if dtype is not None:
+                sd = {key: value.to(dtype=dtype) for key, value in model_state_dict.sd.items()}
+            meta_model.load_state_dict(sd, strict=False, assign=True)
+            return self._return_model(meta_model, device)
+
+        # Load LoRA[0] (fused) on GPU (or CPU—GPU is fine since you fuse immediately)
+        lora0_sd = self.load_sd(
+            [self.loras[0].path],
+            sd_ops=self.loras[0].sd_ops,
+            registry=self.registry,
+            device=device,
+        )
+
+        # Load runtime LoRAs on CPU so they don't sit in VRAM
+        runtime_lora_sds = [
+            self.load_sd(
+                [lora.path],
+                sd_ops=lora.sd_ops,
+                registry=self.registry,
+                device=torch.device("cpu"),
+            )
+            for lora in self.loras[1:]
+        ]
+
+        # Rebuild lists to match your later code expectations
+        lora_state_dicts = [lora0_sd, *runtime_lora_sds]
+
+        lora_sd_and_strengths = [
+            LoraStateDictWithStrength(sd, strength)
+            for sd, strength in zip(lora_state_dicts, lora_strengths, strict=True)
+        ]
+        final_sd = apply_loras(
+            model_sd=model_state_dict,
+            lora_sd_and_strengths=[lora_sd_and_strengths[0]],
+            dtype=dtype,
+            destination_sd=model_state_dict if isinstance(self.registry, DummyRegistry) else None,
+        )
+        meta_model.load_state_dict(final_sd.sd, strict=False, assign=True)
+        model = self._return_model(meta_model, device)
+
+        _, affected_modules = apply_loras(
+            model_sd=model_state_dict,
+            lora_sd_and_strengths=lora_sd_and_strengths[1:],
+            dtype=dtype,
+            destination_sd=None,
+            return_affected=True,
+        )
+
+        for runtime_idx, (lora_sd, strength) in enumerate(zip(lora_state_dicts[1:], lora_strengths[1:], strict=True)):
+            patch_only_affected_linears(
+                model,
+                lora_sd.sd,
+                affected_modules,
+                strength=strength,
+                adapter_idx=runtime_idx,
+                default_enabled=False,  # start off
+            )
+
+        return model

packages/ltx-core/src/ltx_core/model/__init__.py

@@ -0,0 +1,8 @@
+"""Model definitions for LTX-2."""
+
+from ltx_core.model.model_protocol import ModelConfigurator, ModelType
+
+__all__ = [
+    "ModelConfigurator",
+    "ModelType",
+]

packages/ltx-core/src/ltx_core/model/audio_vae/__init__.py

@@ -0,0 +1,28 @@
+"""Audio VAE model components."""
+
+from ltx_core.model.audio_vae.audio_vae import AudioDecoder, AudioEncoder, decode_audio, encode_audio
+from ltx_core.model.audio_vae.model_configurator import (
+    AUDIO_VAE_DECODER_COMFY_KEYS_FILTER,
+    AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER,
+    VOCODER_COMFY_KEYS_FILTER,
+    AudioDecoderConfigurator,
+    AudioEncoderConfigurator,
+    VocoderConfigurator,
+)
+from ltx_core.model.audio_vae.ops import AudioProcessor
+from ltx_core.model.audio_vae.vocoder import Vocoder
+
+__all__ = [
+    "AUDIO_VAE_DECODER_COMFY_KEYS_FILTER",
+    "AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER",
+    "VOCODER_COMFY_KEYS_FILTER",
+    "AudioDecoder",
+    "AudioDecoderConfigurator",
+    "AudioEncoder",
+    "AudioEncoderConfigurator",
+    "AudioProcessor",
+    "Vocoder",
+    "VocoderConfigurator",
+    "decode_audio",
+    "encode_audio",
+]

packages/ltx-core/src/ltx_core/model/audio_vae/attention.py

@@ -1,71 +1,71 @@
-from enum import Enum
-
-import torch
-
-from ltx_core.model.common.normalization import NormType, build_normalization_layer
-
-
-class AttentionType(Enum):
-    """Enum for specifying the attention mechanism type."""
-
-    VANILLA = "vanilla"
-    LINEAR = "linear"
-    NONE = "none"
-
-
-class AttnBlock(torch.nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        norm_type: NormType = NormType.GROUP,
-    ) -> None:
-        super().__init__()
-        self.in_channels = in_channels
-
-        self.norm = build_normalization_layer(in_channels, normtype=norm_type)
-        self.q = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
-        self.k = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
-        self.v = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
-        self.proj_out = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        h_ = x
-        h_ = self.norm(h_)
-        q = self.q(h_)
-        k = self.k(h_)
-        v = self.v(h_)
-
-        # compute attention
-        b, c, h, w = q.shape
-        q = q.reshape(b, c, h * w).contiguous()
-        q = q.permute(0, 2, 1).contiguous()  # b,hw,c
-        k = k.reshape(b, c, h * w).contiguous()  # b,c,hw
-        w_ = torch.bmm(q, k).contiguous()  # b,hw,hw    w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
-        w_ = w_ * (int(c) ** (-0.5))
-        w_ = torch.nn.functional.softmax(w_, dim=2)
-
-        # attend to values
-        v = v.reshape(b, c, h * w).contiguous()
-        w_ = w_.permute(0, 2, 1).contiguous()  # b,hw,hw (first hw of k, second of q)
-        h_ = torch.bmm(v, w_).contiguous()  # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
-        h_ = h_.reshape(b, c, h, w).contiguous()
-
-        h_ = self.proj_out(h_)
-
-        return x + h_
-
-
-def make_attn(
-    in_channels: int,
-    attn_type: AttentionType = AttentionType.VANILLA,
-    norm_type: NormType = NormType.GROUP,
-) -> torch.nn.Module:
-    match attn_type:
-        case AttentionType.VANILLA:
-            return AttnBlock(in_channels, norm_type=norm_type)
-        case AttentionType.NONE:
-            return torch.nn.Identity()
-        case AttentionType.LINEAR:
-            raise NotImplementedError(f"Attention type {attn_type.value} is not supported yet.")
-        case _:
-            raise ValueError(f"Unknown attention type: {attn_type}")
+from enum import Enum
+
+import torch
+
+from ltx_core.model.common.normalization import NormType, build_normalization_layer
+
+
+class AttentionType(Enum):
+    """Enum for specifying the attention mechanism type."""
+
+    VANILLA = "vanilla"
+    LINEAR = "linear"
+    NONE = "none"
+
+
+class AttnBlock(torch.nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        norm_type: NormType = NormType.GROUP,
+    ) -> None:
+        super().__init__()
+        self.in_channels = in_channels
+
+        self.norm = build_normalization_layer(in_channels, normtype=norm_type)
+        self.q = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.k = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.v = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.proj_out = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h_ = x
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        # compute attention
+        b, c, h, w = q.shape
+        q = q.reshape(b, c, h * w).contiguous()
+        q = q.permute(0, 2, 1).contiguous()  # b,hw,c
+        k = k.reshape(b, c, h * w).contiguous()  # b,c,hw
+        w_ = torch.bmm(q, k).contiguous()  # b,hw,hw    w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
+        w_ = w_ * (int(c) ** (-0.5))
+        w_ = torch.nn.functional.softmax(w_, dim=2)
+
+        # attend to values
+        v = v.reshape(b, c, h * w).contiguous()
+        w_ = w_.permute(0, 2, 1).contiguous()  # b,hw,hw (first hw of k, second of q)
+        h_ = torch.bmm(v, w_).contiguous()  # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
+        h_ = h_.reshape(b, c, h, w).contiguous()
+
+        h_ = self.proj_out(h_)
+
+        return x + h_
+
+
+def make_attn(
+    in_channels: int,
+    attn_type: AttentionType = AttentionType.VANILLA,
+    norm_type: NormType = NormType.GROUP,
+) -> torch.nn.Module:
+    match attn_type:
+        case AttentionType.VANILLA:
+            return AttnBlock(in_channels, norm_type=norm_type)
+        case AttentionType.NONE:
+            return torch.nn.Identity()
+        case AttentionType.LINEAR:
+            raise NotImplementedError(f"Attention type {attn_type.value} is not supported yet.")
+        case _:
+            raise ValueError(f"Unknown attention type: {attn_type}")

packages/ltx-core/src/ltx_core/model/audio_vae/audio_vae.py

@@ -1,483 +1,531 @@
-# Copyright (c) 2025 Lightricks. All rights reserved.
-# Created by Ivan Zorin
-
-
-from typing import Set, Tuple
-
-import torch
-import torch.nn.functional as F
-
-from ltx_core.model.audio_vae.attention import AttentionType, make_attn
-from ltx_core.model.audio_vae.causal_conv_2d import make_conv2d
-from ltx_core.model.audio_vae.causality_axis import CausalityAxis
-from ltx_core.model.audio_vae.downsample import build_downsampling_path
-from ltx_core.model.audio_vae.ops import PerChannelStatistics
-from ltx_core.model.audio_vae.resnet import ResnetBlock
-from ltx_core.model.audio_vae.upsample import build_upsampling_path
-from ltx_core.model.common.normalization import NormType, build_normalization_layer
-from ltx_core.pipeline.components.patchifiers import AudioPatchifier
-from ltx_core.pipeline.components.protocols import AudioLatentShape
-
-LATENT_DOWNSAMPLE_FACTOR = 4
-
-
-def build_mid_block(
-    channels: int,
-    temb_channels: int,
-    dropout: float,
-    norm_type: NormType,
-    causality_axis: CausalityAxis,
-    attn_type: AttentionType,
-    add_attention: bool,
-) -> torch.nn.Module:
-    """Build the middle block with two ResNet blocks and optional attention."""
-    mid = torch.nn.Module()
-    mid.block_1 = ResnetBlock(
-        in_channels=channels,
-        out_channels=channels,
-        temb_channels=temb_channels,
-        dropout=dropout,
-        norm_type=norm_type,
-        causality_axis=causality_axis,
-    )
-    mid.attn_1 = make_attn(channels, attn_type=attn_type, norm_type=norm_type) if add_attention else torch.nn.Identity()
-    mid.block_2 = ResnetBlock(
-        in_channels=channels,
-        out_channels=channels,
-        temb_channels=temb_channels,
-        dropout=dropout,
-        norm_type=norm_type,
-        causality_axis=causality_axis,
-    )
-    return mid
-
-
-def run_mid_block(mid: torch.nn.Module, features: torch.Tensor) -> torch.Tensor:
-    """Run features through the middle block."""
-    features = mid.block_1(features, temb=None)
-    features = mid.attn_1(features)
-    return mid.block_2(features, temb=None)
-
-
-class Encoder(torch.nn.Module):
-    """
-    Encoder that compresses audio spectrograms into latent representations.
-
-    The encoder uses a series of downsampling blocks with residual connections,
-    attention mechanisms, and configurable causal convolutions.
-    """
-
-    def __init__(  # noqa: PLR0913
-        self,
-        *,
-        ch: int,
-        ch_mult: Tuple[int, ...] = (1, 2, 4, 8),
-        num_res_blocks: int,
-        attn_resolutions: Set[int],
-        dropout: float = 0.0,
-        resamp_with_conv: bool = True,
-        in_channels: int,
-        resolution: int,
-        z_channels: int,
-        double_z: bool = True,
-        attn_type: AttentionType = AttentionType.VANILLA,
-        mid_block_add_attention: bool = True,
-        norm_type: NormType = NormType.GROUP,
-        causality_axis: CausalityAxis = CausalityAxis.WIDTH,
-        sample_rate: int = 16000,
-        mel_hop_length: int = 160,
-        n_fft: int = 1024,
-        is_causal: bool = True,
-        mel_bins: int = 64,
-        **_ignore_kwargs,
-    ) -> None:
-        """
-        Initialize the Encoder.
-
-        Args:
-            Arguments are configuration parameters, loaded from the audio VAE checkpoint config
-            (audio_vae.model.params.ddconfig):
-
-            ch: Base number of feature channels used in the first convolution layer.
-            ch_mult: Multiplicative factors for the number of channels at each resolution level.
-            num_res_blocks: Number of residual blocks to use at each resolution level.
-            attn_resolutions: Spatial resolutions (e.g., in time/frequency) at which to apply attention.
-            resolution: Input spatial resolution of the spectrogram (height, width).
-            z_channels: Number of channels in the latent representation.
-            norm_type: Normalization layer type to use within the network (e.g., group, batch).
-            causality_axis: Axis along which convolutions should be causal (e.g., time axis).
-            sample_rate: Audio sample rate in Hz for the input signals.
-            mel_hop_length: Hop length used when computing the mel spectrogram.
-            n_fft: FFT size used to compute the spectrogram.
-            mel_bins: Number of mel-frequency bins in the input spectrogram.
-            in_channels: Number of channels in the input spectrogram tensor.
-            double_z: If True, predict both mean and log-variance (doubling latent channels).
-            is_causal: If True, use causal convolutions suitable for streaming setups.
-            dropout: Dropout probability used in residual and mid blocks.
-            attn_type: Type of attention mechanism to use in attention blocks.
-            resamp_with_conv: If True, perform resolution changes using strided convolutions.
-            mid_block_add_attention: If True, add an attention block in the mid-level of the encoder.
-        """
-        super().__init__()
-
-        self.per_channel_statistics = PerChannelStatistics(latent_channels=ch)
-        self.sample_rate = sample_rate
-        self.mel_hop_length = mel_hop_length
-        self.n_fft = n_fft
-        self.is_causal = is_causal
-        self.mel_bins = mel_bins
-
-        self.patchifier = AudioPatchifier(
-            patch_size=1,
-            audio_latent_downsample_factor=LATENT_DOWNSAMPLE_FACTOR,
-            sample_rate=sample_rate,
-            hop_length=mel_hop_length,
-            is_causal=is_causal,
-        )
-
-        self.ch = ch
-        self.temb_ch = 0
-        self.num_resolutions = len(ch_mult)
-        self.num_res_blocks = num_res_blocks
-        self.resolution = resolution
-        self.in_channels = in_channels
-        self.z_channels = z_channels
-        self.double_z = double_z
-        self.norm_type = norm_type
-        self.causality_axis = causality_axis
-        self.attn_type = attn_type
-
-        # downsampling
-        self.conv_in = make_conv2d(
-            in_channels,
-            self.ch,
-            kernel_size=3,
-            stride=1,
-            causality_axis=self.causality_axis,
-        )
-
-        self.non_linearity = torch.nn.SiLU()
-
-        self.down, block_in = build_downsampling_path(
-            ch=ch,
-            ch_mult=ch_mult,
-            num_resolutions=self.num_resolutions,
-            num_res_blocks=num_res_blocks,
-            resolution=resolution,
-            temb_channels=self.temb_ch,
-            dropout=dropout,
-            norm_type=self.norm_type,
-            causality_axis=self.causality_axis,
-            attn_type=self.attn_type,
-            attn_resolutions=attn_resolutions,
-            resamp_with_conv=resamp_with_conv,
-        )
-
-        self.mid = build_mid_block(
-            channels=block_in,
-            temb_channels=self.temb_ch,
-            dropout=dropout,
-            norm_type=self.norm_type,
-            causality_axis=self.causality_axis,
-            attn_type=self.attn_type,
-            add_attention=mid_block_add_attention,
-        )
-
-        self.norm_out = build_normalization_layer(block_in, normtype=self.norm_type)
-        self.conv_out = make_conv2d(
-            block_in,
-            2 * z_channels if double_z else z_channels,
-            kernel_size=3,
-            stride=1,
-            causality_axis=self.causality_axis,
-        )
-
-    def forward(self, spectrogram: torch.Tensor) -> torch.Tensor:
-        """
-        Encode audio spectrogram into latent representations.
-
-        Args:
-            spectrogram: Input spectrogram of shape (batch, channels, time, frequency)
-
-        Returns:
-            Encoded latent representation of shape (batch, channels, frames, mel_bins)
-        """
-        h = self.conv_in(spectrogram)
-        h = self._run_downsampling_path(h)
-        h = run_mid_block(self.mid, h)
-        h = self._finalize_output(h)
-
-        return self._normalize_latents(h)
-
-    def _run_downsampling_path(self, h: torch.Tensor) -> torch.Tensor:
-        for level in range(self.num_resolutions):
-            stage = self.down[level]
-            for block_idx in range(self.num_res_blocks):
-                h = stage.block[block_idx](h, temb=None)
-                if stage.attn:
-                    h = stage.attn[block_idx](h)
-
-            if level != self.num_resolutions - 1:
-                h = stage.downsample(h)
-
-        return h
-
-    def _finalize_output(self, h: torch.Tensor) -> torch.Tensor:
-        h = self.norm_out(h)
-        h = self.non_linearity(h)
-        return self.conv_out(h)
-
-    def _normalize_latents(self, latent_output: torch.Tensor) -> torch.Tensor:
-        """
-        Normalize encoder latents using per-channel statistics.
-
-        When the encoder is configured with ``double_z=True``, the final
-        convolution produces twice the number of latent channels, typically
-        interpreted as two concatenated tensors along the channel dimension
-        (e.g., mean and variance or other auxiliary parameters).
-
-        This method intentionally uses only the first half of the channels
-        (the "mean" component) as input to the patchifier and normalization
-        logic. The remaining channels are left unchanged by this method and
-        are expected to be consumed elsewhere in the VAE pipeline.
-
-        If ``double_z=False``, the encoder output already contains only the
-        mean latents and the chunking operation simply returns that tensor.
-        """
-        means = torch.chunk(latent_output, 2, dim=1)[0]
-        latent_shape = AudioLatentShape(
-            batch=means.shape[0],
-            channels=means.shape[1],
-            frames=means.shape[2],
-            mel_bins=means.shape[3],
-        )
-        latent_patched = self.patchifier.patchify(means)
-        latent_normalized = self.per_channel_statistics.normalize(latent_patched)
-        return self.patchifier.unpatchify(latent_normalized, latent_shape)
-
-
-class Decoder(torch.nn.Module):
-    """
-    Symmetric decoder that reconstructs audio spectrograms from latent features.
-
-    The decoder mirrors the encoder structure with configurable channel multipliers,
-    attention resolutions, and causal convolutions.
-    """
-
-    def __init__(  # noqa: PLR0913
-        self,
-        *,
-        ch: int,
-        out_ch: int,
-        ch_mult: Tuple[int, ...] = (1, 2, 4, 8),
-        num_res_blocks: int,
-        attn_resolutions: Set[int],
-        resolution: int,
-        z_channels: int,
-        norm_type: NormType = NormType.GROUP,
-        causality_axis: CausalityAxis = CausalityAxis.WIDTH,
-        dropout: float = 0.0,
-        mid_block_add_attention: bool = True,
-        sample_rate: int = 16000,
-        mel_hop_length: int = 160,
-        is_causal: bool = True,
-        mel_bins: int | None = None,
-    ) -> None:
-        """
-        Initialize the Decoder.
-
-        Args:
-            Arguments are configuration parameters, loaded from the audio VAE checkpoint config
-            (audio_vae.model.params.ddconfig):
-            - ch, out_ch, ch_mult, num_res_blocks, attn_resolutions
-            - resolution, z_channels
-            - norm_type, causality_axis
-        """
-        super().__init__()
-
-        # Internal behavioural defaults that are not driven by the checkpoint.
-        resamp_with_conv = True
-        attn_type = AttentionType.VANILLA
-
-        # Per-channel statistics for denormalizing latents
-        self.per_channel_statistics = PerChannelStatistics(latent_channels=ch)
-        self.sample_rate = sample_rate
-        self.mel_hop_length = mel_hop_length
-        self.is_causal = is_causal
-        self.mel_bins = mel_bins
-        self.patchifier = AudioPatchifier(
-            patch_size=1,
-            audio_latent_downsample_factor=LATENT_DOWNSAMPLE_FACTOR,
-            sample_rate=sample_rate,
-            hop_length=mel_hop_length,
-            is_causal=is_causal,
-        )
-
-        self.ch = ch
-        self.temb_ch = 0
-        self.num_resolutions = len(ch_mult)
-        self.num_res_blocks = num_res_blocks
-        self.resolution = resolution
-        self.out_ch = out_ch
-        self.give_pre_end = False
-        self.tanh_out = False
-        self.norm_type = norm_type
-        self.z_channels = z_channels
-        self.channel_multipliers = ch_mult
-        self.attn_resolutions = attn_resolutions
-        self.causality_axis = causality_axis
-        self.attn_type = attn_type
-
-        base_block_channels = ch * self.channel_multipliers[-1]
-        base_resolution = resolution // (2 ** (self.num_resolutions - 1))
-        self.z_shape = (1, z_channels, base_resolution, base_resolution)
-
-        self.conv_in = make_conv2d(
-            z_channels, base_block_channels, kernel_size=3, stride=1, causality_axis=self.causality_axis
-        )
-        self.non_linearity = torch.nn.SiLU()
-        self.mid = build_mid_block(
-            channels=base_block_channels,
-            temb_channels=self.temb_ch,
-            dropout=dropout,
-            norm_type=self.norm_type,
-            causality_axis=self.causality_axis,
-            attn_type=self.attn_type,
-            add_attention=mid_block_add_attention,
-        )
-        self.up, final_block_channels = build_upsampling_path(
-            ch=ch,
-            ch_mult=ch_mult,
-            num_resolutions=self.num_resolutions,
-            num_res_blocks=num_res_blocks,
-            resolution=resolution,
-            temb_channels=self.temb_ch,
-            dropout=dropout,
-            norm_type=self.norm_type,
-            causality_axis=self.causality_axis,
-            attn_type=self.attn_type,
-            attn_resolutions=attn_resolutions,
-            resamp_with_conv=resamp_with_conv,
-            initial_block_channels=base_block_channels,
-        )
-
-        self.norm_out = build_normalization_layer(final_block_channels, normtype=self.norm_type)
-        self.conv_out = make_conv2d(
-            final_block_channels, out_ch, kernel_size=3, stride=1, causality_axis=self.causality_axis
-        )
-
-    def forward(self, sample: torch.Tensor) -> torch.Tensor:
-        """
-        Decode latent features back to audio spectrograms.
-
-        Args:
-            sample: Encoded latent representation of shape (batch, channels, frames, mel_bins)
-
-        Returns:
-            Reconstructed audio spectrogram of shape (batch, channels, time, frequency)
-        """
-        sample, target_shape = self._denormalize_latents(sample)
-
-        h = self.conv_in(sample)
-        h = run_mid_block(self.mid, h)
-        h = self._run_upsampling_path(h)
-        h = self._finalize_output(h)
-
-        return self._adjust_output_shape(h, target_shape)
-
-    def _denormalize_latents(self, sample: torch.Tensor) -> tuple[torch.Tensor, AudioLatentShape]:
-        latent_shape = AudioLatentShape(
-            batch=sample.shape[0],
-            channels=sample.shape[1],
-            frames=sample.shape[2],
-            mel_bins=sample.shape[3],
-        )
-
-        sample_patched = self.patchifier.patchify(sample)
-        sample_denormalized = self.per_channel_statistics.un_normalize(sample_patched)
-        sample = self.patchifier.unpatchify(sample_denormalized, latent_shape)
-
-        target_frames = latent_shape.frames * LATENT_DOWNSAMPLE_FACTOR
-        if self.causality_axis != CausalityAxis.NONE:
-            target_frames = max(target_frames - (LATENT_DOWNSAMPLE_FACTOR - 1), 1)
-
-        target_shape = AudioLatentShape(
-            batch=latent_shape.batch,
-            channels=self.out_ch,
-            frames=target_frames,
-            mel_bins=self.mel_bins if self.mel_bins is not None else latent_shape.mel_bins,
-        )
-
-        return sample, target_shape
-
-    def _adjust_output_shape(
-        self,
-        decoded_output: torch.Tensor,
-        target_shape: AudioLatentShape,
-    ) -> torch.Tensor:
-        """
-        Adjust output shape to match target dimensions for variable-length audio.
-
-        This function handles the common case where decoded audio spectrograms need to be
-        resized to match a specific target shape.
-
-        Args:
-            decoded_output: Tensor of shape (batch, channels, time, frequency)
-            target_shape: AudioLatentShape describing (batch, channels, time, mel bins)
-
-        Returns:
-            Tensor adjusted to match target_shape exactly
-        """
-        # Current output shape: (batch, channels, time, frequency)
-        _, _, current_time, current_freq = decoded_output.shape
-        target_channels = target_shape.channels
-        target_time = target_shape.frames
-        target_freq = target_shape.mel_bins
-
-        # Step 1: Crop first to avoid exceeding target dimensions
-        decoded_output = decoded_output[
-            :, :target_channels, : min(current_time, target_time), : min(current_freq, target_freq)
-        ]
-
-        # Step 2: Calculate padding needed for time and frequency dimensions
-        time_padding_needed = target_time - decoded_output.shape[2]
-        freq_padding_needed = target_freq - decoded_output.shape[3]
-
-        # Step 3: Apply padding if needed
-        if time_padding_needed > 0 or freq_padding_needed > 0:
-            # PyTorch padding format: (pad_left, pad_right, pad_top, pad_bottom)
-            # For audio: pad_left/right = frequency, pad_top/bottom = time
-            padding = (
-                0,
-                max(freq_padding_needed, 0),  # frequency padding (left, right)
-                0,
-                max(time_padding_needed, 0),  # time padding (top, bottom)
-            )
-            decoded_output = F.pad(decoded_output, padding)
-
-        # Step 4: Final safety crop to ensure exact target shape
-        decoded_output = decoded_output[:, :target_channels, :target_time, :target_freq]
-
-        return decoded_output
-
-    def _run_upsampling_path(self, h: torch.Tensor) -> torch.Tensor:
-        for level in reversed(range(self.num_resolutions)):
-            stage = self.up[level]
-            for block_idx, block in enumerate(stage.block):
-                h = block(h, temb=None)
-                if stage.attn:
-                    h = stage.attn[block_idx](h)
-
-            if level != 0 and hasattr(stage, "upsample"):
-                h = stage.upsample(h)
-
-        return h
-
-    def _finalize_output(self, h: torch.Tensor) -> torch.Tensor:
-        if self.give_pre_end:
-            return h
-
-        h = self.norm_out(h)
-        h = self.non_linearity(h)
-        h = self.conv_out(h)
-        return torch.tanh(h) if self.tanh_out else h
+from typing import Set, Tuple
+
+import torch
+import torch.nn.functional as F
+
+from ltx_core.components.patchifiers import AudioPatchifier
+from ltx_core.model.audio_vae.attention import AttentionType, make_attn
+from ltx_core.model.audio_vae.causal_conv_2d import make_conv2d
+from ltx_core.model.audio_vae.causality_axis import CausalityAxis
+from ltx_core.model.audio_vae.downsample import build_downsampling_path
+from ltx_core.model.audio_vae.ops import PerChannelStatistics, AudioProcessor
+from ltx_core.model.audio_vae.resnet import ResnetBlock
+from ltx_core.model.audio_vae.upsample import build_upsampling_path
+from ltx_core.model.audio_vae.vocoder import Vocoder
+from ltx_core.model.common.normalization import NormType, build_normalization_layer
+from ltx_core.types import AudioLatentShape
+
+LATENT_DOWNSAMPLE_FACTOR = 4
+
+
+def build_mid_block(
+    channels: int,
+    temb_channels: int,
+    dropout: float,
+    norm_type: NormType,
+    causality_axis: CausalityAxis,
+    attn_type: AttentionType,
+    add_attention: bool,
+) -> torch.nn.Module:
+    """Build the middle block with two ResNet blocks and optional attention."""
+    mid = torch.nn.Module()
+    mid.block_1 = ResnetBlock(
+        in_channels=channels,
+        out_channels=channels,
+        temb_channels=temb_channels,
+        dropout=dropout,
+        norm_type=norm_type,
+        causality_axis=causality_axis,
+    )
+    mid.attn_1 = make_attn(channels, attn_type=attn_type, norm_type=norm_type) if add_attention else torch.nn.Identity()
+    mid.block_2 = ResnetBlock(
+        in_channels=channels,
+        out_channels=channels,
+        temb_channels=temb_channels,
+        dropout=dropout,
+        norm_type=norm_type,
+        causality_axis=causality_axis,
+    )
+    return mid
+
+
+def run_mid_block(mid: torch.nn.Module, features: torch.Tensor) -> torch.Tensor:
+    """Run features through the middle block."""
+    features = mid.block_1(features, temb=None)
+    features = mid.attn_1(features)
+    return mid.block_2(features, temb=None)
+
+
+class AudioEncoder(torch.nn.Module):
+    """
+    Encoder that compresses audio spectrograms into latent representations.
+    The encoder uses a series of downsampling blocks with residual connections,
+    attention mechanisms, and configurable causal convolutions.
+    """
+
+    def __init__(  # noqa: PLR0913
+        self,
+        *,
+        ch: int,
+        ch_mult: Tuple[int, ...] = (1, 2, 4, 8),
+        num_res_blocks: int,
+        attn_resolutions: Set[int],
+        dropout: float = 0.0,
+        resamp_with_conv: bool = True,
+        in_channels: int,
+        resolution: int,
+        z_channels: int,
+        double_z: bool = True,
+        attn_type: AttentionType = AttentionType.VANILLA,
+        mid_block_add_attention: bool = True,
+        norm_type: NormType = NormType.GROUP,
+        causality_axis: CausalityAxis = CausalityAxis.WIDTH,
+        sample_rate: int = 16000,
+        mel_hop_length: int = 160,
+        n_fft: int = 1024,
+        is_causal: bool = True,
+        mel_bins: int = 64,
+        **_ignore_kwargs,
+    ) -> None:
+        """
+        Initialize the Encoder.
+        Args:
+            Arguments are configuration parameters, loaded from the audio VAE checkpoint config
+            (audio_vae.model.params.ddconfig):
+            ch: Base number of feature channels used in the first convolution layer.
+            ch_mult: Multiplicative factors for the number of channels at each resolution level.
+            num_res_blocks: Number of residual blocks to use at each resolution level.
+            attn_resolutions: Spatial resolutions (e.g., in time/frequency) at which to apply attention.
+            resolution: Input spatial resolution of the spectrogram (height, width).
+            z_channels: Number of channels in the latent representation.
+            norm_type: Normalization layer type to use within the network (e.g., group, batch).
+            causality_axis: Axis along which convolutions should be causal (e.g., time axis).
+            sample_rate: Audio sample rate in Hz for the input signals.
+            mel_hop_length: Hop length used when computing the mel spectrogram.
+            n_fft: FFT size used to compute the spectrogram.
+            mel_bins: Number of mel-frequency bins in the input spectrogram.
+            in_channels: Number of channels in the input spectrogram tensor.
+            double_z: If True, predict both mean and log-variance (doubling latent channels).
+            is_causal: If True, use causal convolutions suitable for streaming setups.
+            dropout: Dropout probability used in residual and mid blocks.
+            attn_type: Type of attention mechanism to use in attention blocks.
+            resamp_with_conv: If True, perform resolution changes using strided convolutions.
+            mid_block_add_attention: If True, add an attention block in the mid-level of the encoder.
+        """
+        super().__init__()
+
+        self.per_channel_statistics = PerChannelStatistics(latent_channels=ch)
+        self.sample_rate = sample_rate
+        self.mel_hop_length = mel_hop_length
+        self.n_fft = n_fft
+        self.is_causal = is_causal
+        self.mel_bins = mel_bins
+
+        self.patchifier = AudioPatchifier(
+            patch_size=1,
+            audio_latent_downsample_factor=LATENT_DOWNSAMPLE_FACTOR,
+            sample_rate=sample_rate,
+            hop_length=mel_hop_length,
+            is_causal=is_causal,
+        )
+
+        self.ch = ch
+        self.temb_ch = 0
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.z_channels = z_channels
+        self.double_z = double_z
+        self.norm_type = norm_type
+        self.causality_axis = causality_axis
+        self.attn_type = attn_type
+
+        # downsampling
+        self.conv_in = make_conv2d(
+            in_channels,
+            self.ch,
+            kernel_size=3,
+            stride=1,
+            causality_axis=self.causality_axis,
+        )
+
+        self.non_linearity = torch.nn.SiLU()
+
+        self.down, block_in = build_downsampling_path(
+            ch=ch,
+            ch_mult=ch_mult,
+            num_resolutions=self.num_resolutions,
+            num_res_blocks=num_res_blocks,
+            resolution=resolution,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+            norm_type=self.norm_type,
+            causality_axis=self.causality_axis,
+            attn_type=self.attn_type,
+            attn_resolutions=attn_resolutions,
+            resamp_with_conv=resamp_with_conv,
+        )
+
+        self.mid = build_mid_block(
+            channels=block_in,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+            norm_type=self.norm_type,
+            causality_axis=self.causality_axis,
+            attn_type=self.attn_type,
+            add_attention=mid_block_add_attention,
+        )
+
+        self.norm_out = build_normalization_layer(block_in, normtype=self.norm_type)
+        self.conv_out = make_conv2d(
+            block_in,
+            2 * z_channels if double_z else z_channels,
+            kernel_size=3,
+            stride=1,
+            causality_axis=self.causality_axis,
+        )
+
+    def forward(self, spectrogram: torch.Tensor) -> torch.Tensor:
+        """
+        Encode audio spectrogram into latent representations.
+        Args:
+            spectrogram: Input spectrogram of shape (batch, channels, time, frequency)
+        Returns:
+            Encoded latent representation of shape (batch, channels, frames, mel_bins)
+        """
+        h = self.conv_in(spectrogram)
+        h = self._run_downsampling_path(h)
+        h = run_mid_block(self.mid, h)
+        h = self._finalize_output(h)
+
+        return self._normalize_latents(h)
+
+    def _run_downsampling_path(self, h: torch.Tensor) -> torch.Tensor:
+        for level in range(self.num_resolutions):
+            stage = self.down[level]
+            for block_idx in range(self.num_res_blocks):
+                h = stage.block[block_idx](h, temb=None)
+                if stage.attn:
+                    h = stage.attn[block_idx](h)
+
+            if level != self.num_resolutions - 1:
+                h = stage.downsample(h)
+
+        return h
+
+    def _finalize_output(self, h: torch.Tensor) -> torch.Tensor:
+        h = self.norm_out(h)
+        h = self.non_linearity(h)
+        return self.conv_out(h)
+
+    def _normalize_latents(self, latent_output: torch.Tensor) -> torch.Tensor:
+        """
+        Normalize encoder latents using per-channel statistics.
+        When the encoder is configured with ``double_z=True``, the final
+        convolution produces twice the number of latent channels, typically
+        interpreted as two concatenated tensors along the channel dimension
+        (e.g., mean and variance or other auxiliary parameters).
+        This method intentionally uses only the first half of the channels
+        (the "mean" component) as input to the patchifier and normalization
+        logic. The remaining channels are left unchanged by this method and
+        are expected to be consumed elsewhere in the VAE pipeline.
+        If ``double_z=False``, the encoder output already contains only the
+        mean latents and the chunking operation simply returns that tensor.
+        """
+        means = torch.chunk(latent_output, 2, dim=1)[0]
+        latent_shape = AudioLatentShape(
+            batch=means.shape[0],
+            channels=means.shape[1],
+            frames=means.shape[2],
+            mel_bins=means.shape[3],
+        )
+        latent_patched = self.patchifier.patchify(means)
+        latent_normalized = self.per_channel_statistics.normalize(latent_patched)
+        return self.patchifier.unpatchify(latent_normalized, latent_shape)
+
+
+class AudioDecoder(torch.nn.Module):
+    """
+    Symmetric decoder that reconstructs audio spectrograms from latent features.
+    The decoder mirrors the encoder structure with configurable channel multipliers,
+    attention resolutions, and causal convolutions.
+    """
+
+    def __init__(  # noqa: PLR0913
+        self,
+        *,
+        ch: int,
+        out_ch: int,
+        ch_mult: Tuple[int, ...] = (1, 2, 4, 8),
+        num_res_blocks: int,
+        attn_resolutions: Set[int],
+        resolution: int,
+        z_channels: int,
+        norm_type: NormType = NormType.GROUP,
+        causality_axis: CausalityAxis = CausalityAxis.WIDTH,
+        dropout: float = 0.0,
+        mid_block_add_attention: bool = True,
+        sample_rate: int = 16000,
+        mel_hop_length: int = 160,
+        is_causal: bool = True,
+        mel_bins: int | None = None,
+    ) -> None:
+        """
+        Initialize the Decoder.
+        Args:
+            Arguments are configuration parameters, loaded from the audio VAE checkpoint config
+            (audio_vae.model.params.ddconfig):
+            - ch, out_ch, ch_mult, num_res_blocks, attn_resolutions
+            - resolution, z_channels
+            - norm_type, causality_axis
+        """
+        super().__init__()
+
+        # Internal behavioural defaults that are not driven by the checkpoint.
+        resamp_with_conv = True
+        attn_type = AttentionType.VANILLA
+
+        # Per-channel statistics for denormalizing latents
+        self.per_channel_statistics = PerChannelStatistics(latent_channels=ch)
+        self.sample_rate = sample_rate
+        self.mel_hop_length = mel_hop_length
+        self.is_causal = is_causal
+        self.mel_bins = mel_bins
+        self.patchifier = AudioPatchifier(
+            patch_size=1,
+            audio_latent_downsample_factor=LATENT_DOWNSAMPLE_FACTOR,
+            sample_rate=sample_rate,
+            hop_length=mel_hop_length,
+            is_causal=is_causal,
+        )
+
+        self.ch = ch
+        self.temb_ch = 0
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.out_ch = out_ch
+        self.give_pre_end = False
+        self.tanh_out = False
+        self.norm_type = norm_type
+        self.z_channels = z_channels
+        self.channel_multipliers = ch_mult
+        self.attn_resolutions = attn_resolutions
+        self.causality_axis = causality_axis
+        self.attn_type = attn_type
+
+        base_block_channels = ch * self.channel_multipliers[-1]
+        base_resolution = resolution // (2 ** (self.num_resolutions - 1))
+        self.z_shape = (1, z_channels, base_resolution, base_resolution)
+
+        self.conv_in = make_conv2d(
+            z_channels, base_block_channels, kernel_size=3, stride=1, causality_axis=self.causality_axis
+        )
+        self.non_linearity = torch.nn.SiLU()
+        self.mid = build_mid_block(
+            channels=base_block_channels,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+            norm_type=self.norm_type,
+            causality_axis=self.causality_axis,
+            attn_type=self.attn_type,
+            add_attention=mid_block_add_attention,
+        )
+        self.up, final_block_channels = build_upsampling_path(
+            ch=ch,
+            ch_mult=ch_mult,
+            num_resolutions=self.num_resolutions,
+            num_res_blocks=num_res_blocks,
+            resolution=resolution,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+            norm_type=self.norm_type,
+            causality_axis=self.causality_axis,
+            attn_type=self.attn_type,
+            attn_resolutions=attn_resolutions,
+            resamp_with_conv=resamp_with_conv,
+            initial_block_channels=base_block_channels,
+        )
+
+        self.norm_out = build_normalization_layer(final_block_channels, normtype=self.norm_type)
+        self.conv_out = make_conv2d(
+            final_block_channels, out_ch, kernel_size=3, stride=1, causality_axis=self.causality_axis
+        )
+
+    def forward(self, sample: torch.Tensor) -> torch.Tensor:
+        """
+        Decode latent features back to audio spectrograms.
+        Args:
+            sample: Encoded latent representation of shape (batch, channels, frames, mel_bins)
+        Returns:
+            Reconstructed audio spectrogram of shape (batch, channels, time, frequency)
+        """
+        sample, target_shape = self._denormalize_latents(sample)
+
+        h = self.conv_in(sample)
+        h = run_mid_block(self.mid, h)
+        h = self._run_upsampling_path(h)
+        h = self._finalize_output(h)
+
+        return self._adjust_output_shape(h, target_shape)
+
+    def _denormalize_latents(self, sample: torch.Tensor) -> tuple[torch.Tensor, AudioLatentShape]:
+        latent_shape = AudioLatentShape(
+            batch=sample.shape[0],
+            channels=sample.shape[1],
+            frames=sample.shape[2],
+            mel_bins=sample.shape[3],
+        )
+
+        sample_patched = self.patchifier.patchify(sample)
+        sample_denormalized = self.per_channel_statistics.un_normalize(sample_patched)
+        sample = self.patchifier.unpatchify(sample_denormalized, latent_shape)
+
+        target_frames = latent_shape.frames * LATENT_DOWNSAMPLE_FACTOR
+        if self.causality_axis != CausalityAxis.NONE:
+            target_frames = max(target_frames - (LATENT_DOWNSAMPLE_FACTOR - 1), 1)
+
+        target_shape = AudioLatentShape(
+            batch=latent_shape.batch,
+            channels=self.out_ch,
+            frames=target_frames,
+            mel_bins=self.mel_bins if self.mel_bins is not None else latent_shape.mel_bins,
+        )
+
+        return sample, target_shape
+
+    def _adjust_output_shape(
+        self,
+        decoded_output: torch.Tensor,
+        target_shape: AudioLatentShape,
+    ) -> torch.Tensor:
+        """
+        Adjust output shape to match target dimensions for variable-length audio.
+        This function handles the common case where decoded audio spectrograms need to be
+        resized to match a specific target shape.
+        Args:
+            decoded_output: Tensor of shape (batch, channels, time, frequency)
+            target_shape: AudioLatentShape describing (batch, channels, time, mel bins)
+        Returns:
+            Tensor adjusted to match target_shape exactly
+        """
+        # Current output shape: (batch, channels, time, frequency)
+        _, _, current_time, current_freq = decoded_output.shape
+        target_channels = target_shape.channels
+        target_time = target_shape.frames
+        target_freq = target_shape.mel_bins
+
+        # Step 1: Crop first to avoid exceeding target dimensions
+        decoded_output = decoded_output[
+            :, :target_channels, : min(current_time, target_time), : min(current_freq, target_freq)
+        ]
+
+        # Step 2: Calculate padding needed for time and frequency dimensions
+        time_padding_needed = target_time - decoded_output.shape[2]
+        freq_padding_needed = target_freq - decoded_output.shape[3]
+
+        # Step 3: Apply padding if needed
+        if time_padding_needed > 0 or freq_padding_needed > 0:
+            # PyTorch padding format: (pad_left, pad_right, pad_top, pad_bottom)
+            # For audio: pad_left/right = frequency, pad_top/bottom = time
+            padding = (
+                0,
+                max(freq_padding_needed, 0),  # frequency padding (left, right)
+                0,
+                max(time_padding_needed, 0),  # time padding (top, bottom)
+            )
+            decoded_output = F.pad(decoded_output, padding)
+
+        # Step 4: Final safety crop to ensure exact target shape
+        decoded_output = decoded_output[:, :target_channels, :target_time, :target_freq]
+
+        return decoded_output
+
+    def _run_upsampling_path(self, h: torch.Tensor) -> torch.Tensor:
+        for level in reversed(range(self.num_resolutions)):
+            stage = self.up[level]
+            for block_idx, block in enumerate(stage.block):
+                h = block(h, temb=None)
+                if stage.attn:
+                    h = stage.attn[block_idx](h)
+
+            if level != 0 and hasattr(stage, "upsample"):
+                h = stage.upsample(h)
+
+        return h
+
+    def _finalize_output(self, h: torch.Tensor) -> torch.Tensor:
+        if self.give_pre_end:
+            return h
+
+        h = self.norm_out(h)
+        h = self.non_linearity(h)
+        h = self.conv_out(h)
+        return torch.tanh(h) if self.tanh_out else h
+
+@torch.no_grad()
+def encode_audio(
+    waveform: torch.Tensor,
+    waveform_sample_rate: int,
+    *,
+    audio_encoder: "AudioEncoder",
+    audio_processor: "AudioProcessor",
+    return_mean_only: bool = False,
+) -> torch.Tensor:
+    """
+    Encode a waveform into an audio latent representation.
+
+    Args:
+        waveform: Audio waveform tensor.
+            Expected shapes:
+              - (T,)                -> treated as (1,1,T)
+              - (B,T)               -> treated as (B,1,T)
+              - (B,C,T)             -> used as-is
+        waveform_sample_rate: Sample rate of the provided waveform.
+        audio_encoder: AudioEncoder that consumes (B, C, frames, mel_bins).
+        audio_processor: AudioProcessor from ops.py that produces log-mel features.
+        return_mean_only: If True and encoder outputs double_z, return only the mean half.
+
+    Returns:
+        Latent tensor from AudioEncoder.
+        If return_mean_only=True and double_z=True: returns (B, z_channels, frames, mel_bins).
+        Otherwise returns the raw encoder output (often (B, 2*z_channels, frames, mel_bins)).
+    """
+    # --- normalize waveform shape to (B, C, T) ---
+    if waveform.dim() == 1:
+        waveform = waveform.unsqueeze(0).unsqueeze(0)
+    elif waveform.dim() == 2:
+        waveform = waveform.unsqueeze(1)
+    elif waveform.dim() == 3:
+        pass
+    else:
+        raise ValueError(f"Unexpected waveform shape: {tuple(waveform.shape)}")
+
+    waveform = waveform.float()
+
+    # --- waveform -> log-mel spectrogram (B, C, frames, mel_bins) ---
+    mel = audio_processor.waveform_to_mel(waveform, waveform_sample_rate)
+
+    # --- mel -> latent ---
+    latent = audio_encoder(mel)
+
+    if return_mean_only and getattr(audio_encoder, "double_z", False):
+        latent = torch.chunk(latent, 2, dim=1)[0]
+
+    return latent
+
+
+def decode_audio(latent: torch.Tensor, audio_decoder: "AudioDecoder", vocoder: "Vocoder") -> torch.Tensor:
+    """
+    Decode an audio latent representation using the provided audio decoder and vocoder.
+    Args:
+        latent: Input audio latent tensor.
+        audio_decoder: Model to decode the latent to waveform features.
+        vocoder: Model to convert decoded features to audio waveform.
+    Returns:
+        Decoded audio as a float tensor.
+    """
+    decoded_audio = audio_decoder(latent)
+    decoded_audio = vocoder(decoded_audio).squeeze(0).float()
+    return decoded_audio

packages/ltx-core/src/ltx_core/model/audio_vae/causal_conv_2d.py

@@ -1,113 +1,110 @@
-import torch
-import torch.nn.functional as F
-
-from ltx_core.model.audio_vae.causality_axis import CausalityAxis
-
-
-class CausalConv2d(torch.nn.Module):
-    """
-    A causal 2D convolution.
-
-    This layer ensures that the output at time `t` only depends on inputs
-    at time `t` and earlier. It achieves this by applying asymmetric padding
-    to the time dimension (width) before the convolution.
-    """
-
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        kernel_size: int | tuple[int, int],
-        stride: int = 1,
-        dilation: int | tuple[int, int] = 1,
-        groups: int = 1,
-        bias: bool = True,
-        causality_axis: CausalityAxis = CausalityAxis.HEIGHT,
-    ) -> None:
-        super().__init__()
-
-        self.causality_axis = causality_axis
-
-        # Ensure kernel_size and dilation are tuples
-        kernel_size = torch.nn.modules.utils._pair(kernel_size)
-        dilation = torch.nn.modules.utils._pair(dilation)
-
-        # Calculate padding dimensions
-        pad_h = (kernel_size[0] - 1) * dilation[0]
-        pad_w = (kernel_size[1] - 1) * dilation[1]
-
-        # The padding tuple for F.pad is (pad_left, pad_right, pad_top, pad_bottom)
-        match self.causality_axis:
-            case CausalityAxis.NONE:
-                self.padding = (pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2)
-            case CausalityAxis.WIDTH | CausalityAxis.WIDTH_COMPATIBILITY:
-                self.padding = (pad_w, 0, pad_h // 2, pad_h - pad_h // 2)
-            case CausalityAxis.HEIGHT:
-                self.padding = (pad_w // 2, pad_w - pad_w // 2, pad_h, 0)
-            case _:
-                raise ValueError(f"Invalid causality_axis: {causality_axis}")
-
-        # The internal convolution layer uses no padding, as we handle it manually
-        self.conv = torch.nn.Conv2d(
-            in_channels,
-            out_channels,
-            kernel_size,
-            stride=stride,
-            padding=0,
-            dilation=dilation,
-            groups=groups,
-            bias=bias,
-        )
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        # Apply causal padding before convolution
-        x = F.pad(x, self.padding)
-        return self.conv(x)
-
-
-def make_conv2d(
-    in_channels: int,
-    out_channels: int,
-    kernel_size: int | tuple[int, int],
-    stride: int = 1,
-    padding: tuple[int, int, int, int] | None = None,
-    dilation: int = 1,
-    groups: int = 1,
-    bias: bool = True,
-    causality_axis: CausalityAxis | None = None,
-) -> torch.nn.Module:
-    """
-    Create a 2D convolution layer that can be either causal or non-causal.
-
-    Args:
-        in_channels: Number of input channels
-        out_channels: Number of output channels
-        kernel_size: Size of the convolution kernel
-        stride: Convolution stride
-        padding: Padding (if None, will be calculated based on causal flag)
-        dilation: Dilation rate
-        groups: Number of groups for grouped convolution
-        bias: Whether to use bias
-        causality_axis: Dimension along which to apply causality.
-
-    Returns:
-        Either a regular Conv2d or CausalConv2d layer
-    """
-    if causality_axis is not None:
-        # For causal convolution, padding is handled internally by CausalConv2d
-        return CausalConv2d(in_channels, out_channels, kernel_size, stride, dilation, groups, bias, causality_axis)
-    else:
-        # For non-causal convolution, use symmetric padding if not specified
-        if padding is None:
-            padding = kernel_size // 2 if isinstance(kernel_size, int) else tuple(k // 2 for k in kernel_size)
-
-        return torch.nn.Conv2d(
-            in_channels,
-            out_channels,
-            kernel_size,
-            stride,
-            padding,
-            dilation,
-            groups,
-            bias,
-        )
+import torch
+import torch.nn.functional as F
+
+from ltx_core.model.audio_vae.causality_axis import CausalityAxis
+
+
+class CausalConv2d(torch.nn.Module):
+    """
+    A causal 2D convolution.
+    This layer ensures that the output at time `t` only depends on inputs
+    at time `t` and earlier. It achieves this by applying asymmetric padding
+    to the time dimension (width) before the convolution.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int | tuple[int, int],
+        stride: int = 1,
+        dilation: int | tuple[int, int] = 1,
+        groups: int = 1,
+        bias: bool = True,
+        causality_axis: CausalityAxis = CausalityAxis.HEIGHT,
+    ) -> None:
+        super().__init__()
+
+        self.causality_axis = causality_axis
+
+        # Ensure kernel_size and dilation are tuples
+        kernel_size = torch.nn.modules.utils._pair(kernel_size)
+        dilation = torch.nn.modules.utils._pair(dilation)
+
+        # Calculate padding dimensions
+        pad_h = (kernel_size[0] - 1) * dilation[0]
+        pad_w = (kernel_size[1] - 1) * dilation[1]
+
+        # The padding tuple for F.pad is (pad_left, pad_right, pad_top, pad_bottom)
+        match self.causality_axis:
+            case CausalityAxis.NONE:
+                self.padding = (pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2)
+            case CausalityAxis.WIDTH | CausalityAxis.WIDTH_COMPATIBILITY:
+                self.padding = (pad_w, 0, pad_h // 2, pad_h - pad_h // 2)
+            case CausalityAxis.HEIGHT:
+                self.padding = (pad_w // 2, pad_w - pad_w // 2, pad_h, 0)
+            case _:
+                raise ValueError(f"Invalid causality_axis: {causality_axis}")
+
+        # The internal convolution layer uses no padding, as we handle it manually
+        self.conv = torch.nn.Conv2d(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride=stride,
+            padding=0,
+            dilation=dilation,
+            groups=groups,
+            bias=bias,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Apply causal padding before convolution
+        x = F.pad(x, self.padding)
+        return self.conv(x)
+
+
+def make_conv2d(
+    in_channels: int,
+    out_channels: int,
+    kernel_size: int | tuple[int, int],
+    stride: int = 1,
+    padding: tuple[int, int, int, int] | None = None,
+    dilation: int = 1,
+    groups: int = 1,
+    bias: bool = True,
+    causality_axis: CausalityAxis | None = None,
+) -> torch.nn.Module:
+    """
+    Create a 2D convolution layer that can be either causal or non-causal.
+    Args:
+        in_channels: Number of input channels
+        out_channels: Number of output channels
+        kernel_size: Size of the convolution kernel
+        stride: Convolution stride
+        padding: Padding (if None, will be calculated based on causal flag)
+        dilation: Dilation rate
+        groups: Number of groups for grouped convolution
+        bias: Whether to use bias
+        causality_axis: Dimension along which to apply causality.
+    Returns:
+        Either a regular Conv2d or CausalConv2d layer
+    """
+    if causality_axis is not None:
+        # For causal convolution, padding is handled internally by CausalConv2d
+        return CausalConv2d(in_channels, out_channels, kernel_size, stride, dilation, groups, bias, causality_axis)
+    else:
+        # For non-causal convolution, use symmetric padding if not specified
+        if padding is None:
+            padding = kernel_size // 2 if isinstance(kernel_size, int) else tuple(k // 2 for k in kernel_size)
+
+        return torch.nn.Conv2d(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride,
+            padding,
+            dilation,
+            groups,
+            bias,
+        )

packages/ltx-core/src/ltx_core/model/audio_vae/causality_axis.py

@@ -1,10 +1,10 @@
-from enum import Enum
-
-
-class CausalityAxis(Enum):
-    """Enum for specifying the causality axis in causal convolutions."""
-
-    NONE = None
-    WIDTH = "width"
-    HEIGHT = "height"
-    WIDTH_COMPATIBILITY = "width-compatibility"
+from enum import Enum
+
+
+class CausalityAxis(Enum):
+    """Enum for specifying the causality axis in causal convolutions."""
+
+    NONE = None
+    WIDTH = "width"
+    HEIGHT = "height"
+    WIDTH_COMPATIBILITY = "width-compatibility"

packages/ltx-core/src/ltx_core/model/audio_vae/downsample.py

@@ -1,110 +1,110 @@
-from typing import Set, Tuple
-
-import torch
-
-from ltx_core.model.audio_vae.attention import AttentionType, make_attn
-from ltx_core.model.audio_vae.causality_axis import CausalityAxis
-from ltx_core.model.audio_vae.resnet import ResnetBlock
-from ltx_core.model.common.normalization import NormType
-
-
-class Downsample(torch.nn.Module):
-    """
-    A downsampling layer that can use either a strided convolution
-    or average pooling. Supports standard and causal padding for the
-    convolutional mode.
-    """
-
-    def __init__(
-        self,
-        in_channels: int,
-        with_conv: bool,
-        causality_axis: CausalityAxis = CausalityAxis.WIDTH,
-    ) -> None:
-        super().__init__()
-        self.with_conv = with_conv
-        self.causality_axis = causality_axis
-
-        if self.causality_axis != CausalityAxis.NONE and not self.with_conv:
-            raise ValueError("causality is only supported when `with_conv=True`.")
-
-        if self.with_conv:
-            # Do time downsampling here
-            # no asymmetric padding in torch conv, must do it ourselves
-            self.conv = torch.nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        if self.with_conv:
-            # Padding tuple is in the order: (left, right, top, bottom).
-            match self.causality_axis:
-                case CausalityAxis.NONE:
-                    pad = (0, 1, 0, 1)
-                case CausalityAxis.WIDTH:
-                    pad = (2, 0, 0, 1)
-                case CausalityAxis.HEIGHT:
-                    pad = (0, 1, 2, 0)
-                case CausalityAxis.WIDTH_COMPATIBILITY:
-                    pad = (1, 0, 0, 1)
-                case _:
-                    raise ValueError(f"Invalid causality_axis: {self.causality_axis}")
-
-            x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
-            x = self.conv(x)
-        else:
-            # This branch is only taken if with_conv=False, which implies causality_axis is NONE.
-            x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
-
-        return x
-
-
-def build_downsampling_path(  # noqa: PLR0913
-    *,
-    ch: int,
-    ch_mult: Tuple[int, ...],
-    num_resolutions: int,
-    num_res_blocks: int,
-    resolution: int,
-    temb_channels: int,
-    dropout: float,
-    norm_type: NormType,
-    causality_axis: CausalityAxis,
-    attn_type: AttentionType,
-    attn_resolutions: Set[int],
-    resamp_with_conv: bool,
-) -> tuple[torch.nn.ModuleList, int]:
-    """Build the downsampling path with residual blocks, attention, and downsampling layers."""
-    down_modules = torch.nn.ModuleList()
-    curr_res = resolution
-    in_ch_mult = (1, *tuple(ch_mult))
-    block_in = ch
-
-    for i_level in range(num_resolutions):
-        block = torch.nn.ModuleList()
-        attn = torch.nn.ModuleList()
-        block_in = ch * in_ch_mult[i_level]
-        block_out = ch * ch_mult[i_level]
-
-        for _ in range(num_res_blocks):
-            block.append(
-                ResnetBlock(
-                    in_channels=block_in,
-                    out_channels=block_out,
-                    temb_channels=temb_channels,
-                    dropout=dropout,
-                    norm_type=norm_type,
-                    causality_axis=causality_axis,
-                )
-            )
-            block_in = block_out
-            if curr_res in attn_resolutions:
-                attn.append(make_attn(block_in, attn_type=attn_type, norm_type=norm_type))
-
-        down = torch.nn.Module()
-        down.block = block
-        down.attn = attn
-        if i_level != num_resolutions - 1:
-            down.downsample = Downsample(block_in, resamp_with_conv, causality_axis=causality_axis)
-            curr_res = curr_res // 2
-        down_modules.append(down)
-
-    return down_modules, block_in
+from typing import Set, Tuple
+
+import torch
+
+from ltx_core.model.audio_vae.attention import AttentionType, make_attn
+from ltx_core.model.audio_vae.causality_axis import CausalityAxis
+from ltx_core.model.audio_vae.resnet import ResnetBlock
+from ltx_core.model.common.normalization import NormType
+
+
+class Downsample(torch.nn.Module):
+    """
+    A downsampling layer that can use either a strided convolution
+    or average pooling. Supports standard and causal padding for the
+    convolutional mode.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        with_conv: bool,
+        causality_axis: CausalityAxis = CausalityAxis.WIDTH,
+    ) -> None:
+        super().__init__()
+        self.with_conv = with_conv
+        self.causality_axis = causality_axis
+
+        if self.causality_axis != CausalityAxis.NONE and not self.with_conv:
+            raise ValueError("causality is only supported when `with_conv=True`.")
+
+        if self.with_conv:
+            # Do time downsampling here
+            # no asymmetric padding in torch conv, must do it ourselves
+            self.conv = torch.nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.with_conv:
+            # Padding tuple is in the order: (left, right, top, bottom).
+            match self.causality_axis:
+                case CausalityAxis.NONE:
+                    pad = (0, 1, 0, 1)
+                case CausalityAxis.WIDTH:
+                    pad = (2, 0, 0, 1)
+                case CausalityAxis.HEIGHT:
+                    pad = (0, 1, 2, 0)
+                case CausalityAxis.WIDTH_COMPATIBILITY:
+                    pad = (1, 0, 0, 1)
+                case _:
+                    raise ValueError(f"Invalid causality_axis: {self.causality_axis}")
+
+            x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
+            x = self.conv(x)
+        else:
+            # This branch is only taken if with_conv=False, which implies causality_axis is NONE.
+            x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
+
+        return x
+
+
+def build_downsampling_path(  # noqa: PLR0913
+    *,
+    ch: int,
+    ch_mult: Tuple[int, ...],
+    num_resolutions: int,
+    num_res_blocks: int,
+    resolution: int,
+    temb_channels: int,
+    dropout: float,
+    norm_type: NormType,
+    causality_axis: CausalityAxis,
+    attn_type: AttentionType,
+    attn_resolutions: Set[int],
+    resamp_with_conv: bool,
+) -> tuple[torch.nn.ModuleList, int]:
+    """Build the downsampling path with residual blocks, attention, and downsampling layers."""
+    down_modules = torch.nn.ModuleList()
+    curr_res = resolution
+    in_ch_mult = (1, *tuple(ch_mult))
+    block_in = ch
+
+    for i_level in range(num_resolutions):
+        block = torch.nn.ModuleList()
+        attn = torch.nn.ModuleList()
+        block_in = ch * in_ch_mult[i_level]
+        block_out = ch * ch_mult[i_level]
+
+        for _ in range(num_res_blocks):
+            block.append(
+                ResnetBlock(
+                    in_channels=block_in,
+                    out_channels=block_out,
+                    temb_channels=temb_channels,
+                    dropout=dropout,
+                    norm_type=norm_type,
+                    causality_axis=causality_axis,
+                )
+            )
+            block_in = block_out
+            if curr_res in attn_resolutions:
+                attn.append(make_attn(block_in, attn_type=attn_type, norm_type=norm_type))
+
+        down = torch.nn.Module()
+        down.block = block
+        down.attn = attn
+        if i_level != num_resolutions - 1:
+            down.downsample = Downsample(block_in, resamp_with_conv, causality_axis=causality_axis)
+            curr_res = curr_res // 2
+        down_modules.append(down)
+
+    return down_modules, block_in

packages/ltx-core/src/ltx_core/model/audio_vae/model_configurator.py

@@ -1,123 +1,123 @@
-from ltx_core.loader.sd_ops import SDOps
-from ltx_core.model.audio_vae.attention import AttentionType
-from ltx_core.model.audio_vae.audio_vae import Decoder, Encoder
-from ltx_core.model.audio_vae.causality_axis import CausalityAxis
-from ltx_core.model.audio_vae.vocoder import Vocoder
-from ltx_core.model.common.normalization import NormType
-from ltx_core.model.model_protocol import ModelConfigurator
-
-
-class VocoderConfigurator(ModelConfigurator[Vocoder]):
-    @classmethod
-    def from_config(cls: type[Vocoder], config: dict) -> Vocoder:
-        config = config.get("vocoder", {})
-        return Vocoder(
-            resblock_kernel_sizes=config.get("resblock_kernel_sizes", [3, 7, 11]),
-            upsample_rates=config.get("upsample_rates", [6, 5, 2, 2, 2]),
-            upsample_kernel_sizes=config.get("upsample_kernel_sizes", [16, 15, 8, 4, 4]),
-            resblock_dilation_sizes=config.get("resblock_dilation_sizes", [[1, 3, 5], [1, 3, 5], [1, 3, 5]]),
-            upsample_initial_channel=config.get("upsample_initial_channel", 1024),
-            stereo=config.get("stereo", True),
-            resblock=config.get("resblock", "1"),
-            output_sample_rate=config.get("output_sample_rate", 24000),
-        )
-
-
-VOCODER_COMFY_KEYS_FILTER = (
-    SDOps("VOCODER_COMFY_KEYS_FILTER").with_matching(prefix="vocoder.").with_replacement("vocoder.", "")
-)
-
-
-class VAEDecoderConfigurator(ModelConfigurator[Decoder]):
-    @classmethod
-    def from_config(cls: type[Decoder], config: dict) -> Decoder:
-        audio_vae_cfg = config.get("audio_vae", {})
-        model_cfg = audio_vae_cfg.get("model", {})
-        model_params = model_cfg.get("params", {})
-        ddconfig = model_params.get("ddconfig", {})
-        preprocessing_cfg = audio_vae_cfg.get("preprocessing", {})
-        stft_cfg = preprocessing_cfg.get("stft", {})
-        mel_cfg = preprocessing_cfg.get("mel", {})
-        variables_cfg = audio_vae_cfg.get("variables", {})
-
-        sample_rate = model_params.get("sampling_rate", 16000)
-        mel_hop_length = stft_cfg.get("hop_length", 160)
-        is_causal = stft_cfg.get("causal", True)
-        mel_bins = ddconfig.get("mel_bins") or mel_cfg.get("n_mel_channels") or variables_cfg.get("mel_bins")
-
-        return Decoder(
-            ch=ddconfig.get("ch", 128),
-            out_ch=ddconfig.get("out_ch", 2),
-            ch_mult=tuple(ddconfig.get("ch_mult", (1, 2, 4))),
-            num_res_blocks=ddconfig.get("num_res_blocks", 2),
-            attn_resolutions=ddconfig.get("attn_resolutions", {8, 16, 32}),
-            resolution=ddconfig.get("resolution", 256),
-            z_channels=ddconfig.get("z_channels", 8),
-            norm_type=NormType(ddconfig.get("norm_type", "pixel")),
-            causality_axis=CausalityAxis(ddconfig.get("causality_axis", "height")),
-            dropout=ddconfig.get("dropout", 0.0),
-            mid_block_add_attention=ddconfig.get("mid_block_add_attention", True),
-            sample_rate=sample_rate,
-            mel_hop_length=mel_hop_length,
-            is_causal=is_causal,
-            mel_bins=mel_bins,
-        )
-
-
-class VAEEncoderConfigurator(ModelConfigurator[Encoder]):
-    @classmethod
-    def from_config(cls: type[Encoder], config: dict) -> Encoder:
-        audio_vae_cfg = config.get("audio_vae", {})
-        model_cfg = audio_vae_cfg.get("model", {})
-        model_params = model_cfg.get("params", {})
-        ddconfig = model_params.get("ddconfig", {})
-        preprocessing_cfg = audio_vae_cfg.get("preprocessing", {})
-        stft_cfg = preprocessing_cfg.get("stft", {})
-        mel_cfg = preprocessing_cfg.get("mel", {})
-        variables_cfg = audio_vae_cfg.get("variables", {})
-
-        sample_rate = model_params.get("sampling_rate", 16000)
-        mel_hop_length = stft_cfg.get("hop_length", 160)
-        n_fft = stft_cfg.get("filter_length", 1024)
-        is_causal = stft_cfg.get("causal", True)
-        mel_bins = ddconfig.get("mel_bins") or mel_cfg.get("n_mel_channels") or variables_cfg.get("mel_bins")
-
-        return Encoder(
-            ch=ddconfig.get("ch", 128),
-            ch_mult=tuple(ddconfig.get("ch_mult", (1, 2, 4))),
-            num_res_blocks=ddconfig.get("num_res_blocks", 2),
-            attn_resolutions=ddconfig.get("attn_resolutions", {8, 16, 32}),
-            resolution=ddconfig.get("resolution", 256),
-            z_channels=ddconfig.get("z_channels", 8),
-            double_z=ddconfig.get("double_z", True),
-            dropout=ddconfig.get("dropout", 0.0),
-            resamp_with_conv=ddconfig.get("resamp_with_conv", True),
-            in_channels=ddconfig.get("in_channels", 2),
-            attn_type=AttentionType(ddconfig.get("attn_type", "vanilla")),
-            mid_block_add_attention=ddconfig.get("mid_block_add_attention", True),
-            norm_type=NormType(ddconfig.get("norm_type", "pixel")),
-            causality_axis=CausalityAxis(ddconfig.get("causality_axis", "height")),
-            sample_rate=sample_rate,
-            mel_hop_length=mel_hop_length,
-            n_fft=n_fft,
-            is_causal=is_causal,
-            mel_bins=mel_bins,
-        )
-
-
-AUDIO_VAE_DECODER_COMFY_KEYS_FILTER = (
-    SDOps("AUDIO_VAE_DECODER_COMFY_KEYS_FILTER")
-    .with_matching(prefix="audio_vae.decoder.")
-    .with_matching(prefix="audio_vae.per_channel_statistics.")
-    .with_replacement("audio_vae.decoder.", "")
-    .with_replacement("audio_vae.per_channel_statistics.", "per_channel_statistics.")
-)
-
-
-AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER = (
-    SDOps("AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER")
-    .with_matching(prefix="audio_vae.encoder.")
-    .with_matching(prefix="audio_vae.per_channel_statistics.")
-    .with_replacement("audio_vae.encoder.", "")
-    .with_replacement("audio_vae.per_channel_statistics.", "per_channel_statistics.")
-)
+from ltx_core.loader.sd_ops import SDOps
+from ltx_core.model.audio_vae.attention import AttentionType
+from ltx_core.model.audio_vae.audio_vae import AudioDecoder, AudioEncoder
+from ltx_core.model.audio_vae.causality_axis import CausalityAxis
+from ltx_core.model.audio_vae.vocoder import Vocoder
+from ltx_core.model.common.normalization import NormType
+from ltx_core.model.model_protocol import ModelConfigurator
+
+
+class VocoderConfigurator(ModelConfigurator[Vocoder]):
+    @classmethod
+    def from_config(cls: type[Vocoder], config: dict) -> Vocoder:
+        config = config.get("vocoder", {})
+        return Vocoder(
+            resblock_kernel_sizes=config.get("resblock_kernel_sizes", [3, 7, 11]),
+            upsample_rates=config.get("upsample_rates", [6, 5, 2, 2, 2]),
+            upsample_kernel_sizes=config.get("upsample_kernel_sizes", [16, 15, 8, 4, 4]),
+            resblock_dilation_sizes=config.get("resblock_dilation_sizes", [[1, 3, 5], [1, 3, 5], [1, 3, 5]]),
+            upsample_initial_channel=config.get("upsample_initial_channel", 1024),
+            stereo=config.get("stereo", True),
+            resblock=config.get("resblock", "1"),
+            output_sample_rate=config.get("output_sample_rate", 24000),
+        )
+
+
+VOCODER_COMFY_KEYS_FILTER = (
+    SDOps("VOCODER_COMFY_KEYS_FILTER").with_matching(prefix="vocoder.").with_replacement("vocoder.", "")
+)
+
+
+class AudioDecoderConfigurator(ModelConfigurator[AudioDecoder]):
+    @classmethod
+    def from_config(cls: type[AudioDecoder], config: dict) -> AudioDecoder:
+        audio_vae_cfg = config.get("audio_vae", {})
+        model_cfg = audio_vae_cfg.get("model", {})
+        model_params = model_cfg.get("params", {})
+        ddconfig = model_params.get("ddconfig", {})
+        preprocessing_cfg = audio_vae_cfg.get("preprocessing", {})
+        stft_cfg = preprocessing_cfg.get("stft", {})
+        mel_cfg = preprocessing_cfg.get("mel", {})
+        variables_cfg = audio_vae_cfg.get("variables", {})
+
+        sample_rate = model_params.get("sampling_rate", 16000)
+        mel_hop_length = stft_cfg.get("hop_length", 160)
+        is_causal = stft_cfg.get("causal", True)
+        mel_bins = ddconfig.get("mel_bins") or mel_cfg.get("n_mel_channels") or variables_cfg.get("mel_bins")
+
+        return AudioDecoder(
+            ch=ddconfig.get("ch", 128),
+            out_ch=ddconfig.get("out_ch", 2),
+            ch_mult=tuple(ddconfig.get("ch_mult", (1, 2, 4))),
+            num_res_blocks=ddconfig.get("num_res_blocks", 2),
+            attn_resolutions=ddconfig.get("attn_resolutions", {8, 16, 32}),
+            resolution=ddconfig.get("resolution", 256),
+            z_channels=ddconfig.get("z_channels", 8),
+            norm_type=NormType(ddconfig.get("norm_type", "pixel")),
+            causality_axis=CausalityAxis(ddconfig.get("causality_axis", "height")),
+            dropout=ddconfig.get("dropout", 0.0),
+            mid_block_add_attention=ddconfig.get("mid_block_add_attention", True),
+            sample_rate=sample_rate,
+            mel_hop_length=mel_hop_length,
+            is_causal=is_causal,
+            mel_bins=mel_bins,
+        )
+
+
+class AudioEncoderConfigurator(ModelConfigurator[AudioEncoder]):
+    @classmethod
+    def from_config(cls: type[AudioEncoder], config: dict) -> AudioEncoder:
+        audio_vae_cfg = config.get("audio_vae", {})
+        model_cfg = audio_vae_cfg.get("model", {})
+        model_params = model_cfg.get("params", {})
+        ddconfig = model_params.get("ddconfig", {})
+        preprocessing_cfg = audio_vae_cfg.get("preprocessing", {})
+        stft_cfg = preprocessing_cfg.get("stft", {})
+        mel_cfg = preprocessing_cfg.get("mel", {})
+        variables_cfg = audio_vae_cfg.get("variables", {})
+
+        sample_rate = model_params.get("sampling_rate", 16000)
+        mel_hop_length = stft_cfg.get("hop_length", 160)
+        n_fft = stft_cfg.get("filter_length", 1024)
+        is_causal = stft_cfg.get("causal", True)
+        mel_bins = ddconfig.get("mel_bins") or mel_cfg.get("n_mel_channels") or variables_cfg.get("mel_bins")
+
+        return AudioEncoder(
+            ch=ddconfig.get("ch", 128),
+            ch_mult=tuple(ddconfig.get("ch_mult", (1, 2, 4))),
+            num_res_blocks=ddconfig.get("num_res_blocks", 2),
+            attn_resolutions=ddconfig.get("attn_resolutions", {8, 16, 32}),
+            resolution=ddconfig.get("resolution", 256),
+            z_channels=ddconfig.get("z_channels", 8),
+            double_z=ddconfig.get("double_z", True),
+            dropout=ddconfig.get("dropout", 0.0),
+            resamp_with_conv=ddconfig.get("resamp_with_conv", True),
+            in_channels=ddconfig.get("in_channels", 2),
+            attn_type=AttentionType(ddconfig.get("attn_type", "vanilla")),
+            mid_block_add_attention=ddconfig.get("mid_block_add_attention", True),
+            norm_type=NormType(ddconfig.get("norm_type", "pixel")),
+            causality_axis=CausalityAxis(ddconfig.get("causality_axis", "height")),
+            sample_rate=sample_rate,
+            mel_hop_length=mel_hop_length,
+            n_fft=n_fft,
+            is_causal=is_causal,
+            mel_bins=mel_bins,
+        )
+
+
+AUDIO_VAE_DECODER_COMFY_KEYS_FILTER = (
+    SDOps("AUDIO_VAE_DECODER_COMFY_KEYS_FILTER")
+    .with_matching(prefix="audio_vae.decoder.")
+    .with_matching(prefix="audio_vae.per_channel_statistics.")
+    .with_replacement("audio_vae.decoder.", "")
+    .with_replacement("audio_vae.per_channel_statistics.", "per_channel_statistics.")
+)
+
+
+AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER = (
+    SDOps("AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER")
+    .with_matching(prefix="audio_vae.encoder.")
+    .with_matching(prefix="audio_vae.per_channel_statistics.")
+    .with_replacement("audio_vae.encoder.", "")
+    .with_replacement("audio_vae.per_channel_statistics.", "per_channel_statistics.")
+)

packages/ltx-core/src/ltx_core/model/audio_vae/ops.py

@@ -1,77 +1,76 @@
-# Copyright (c) 2025 Lightricks. All rights reserved.
-# Created by Ivan Zorin
-
-import torch
-import torchaudio
-from torch import nn
-
-
-class AudioProcessor(nn.Module):
-    def __init__(
-        self,
-        sample_rate: int,
-        mel_bins: int,
-        mel_hop_length: int,
-        n_fft: int,
-    ) -> None:
-        super().__init__()
-        self.sample_rate = sample_rate
-        self.mel_transform = torchaudio.transforms.MelSpectrogram(
-            sample_rate=sample_rate,
-            n_fft=n_fft,
-            win_length=n_fft,
-            hop_length=mel_hop_length,
-            f_min=0.0,
-            f_max=sample_rate / 2.0,
-            n_mels=mel_bins,
-            window_fn=torch.hann_window,
-            center=True,
-            pad_mode="reflect",
-            power=1.0,
-            mel_scale="slaney",
-            norm="slaney",
-        )
-
-    def resample_waveform(
-        self,
-        waveform: torch.Tensor,
-        source_rate: int,
-        target_rate: int,
-    ) -> torch.Tensor:
-        """Resample waveform to target sample rate if needed."""
-        if source_rate == target_rate:
-            return waveform
-        resampled = torchaudio.functional.resample(waveform, source_rate, target_rate)
-        return resampled.to(device=waveform.device, dtype=waveform.dtype)
-
-    def waveform_to_mel(
-        self,
-        waveform: torch.Tensor,
-        waveform_sample_rate: int,
-    ) -> torch.Tensor:
-        """Convert waveform to log-mel spectrogram [batch, channels, time, n_mels]."""
-        waveform = self.resample_waveform(waveform, waveform_sample_rate, self.sample_rate)
-
-        mel = self.mel_transform(waveform)
-        mel = torch.log(torch.clamp(mel, min=1e-5))
-
-        mel = mel.to(device=waveform.device, dtype=waveform.dtype)
-        return mel.permute(0, 1, 3, 2).contiguous()
-
-
-class PerChannelStatistics(nn.Module):
-    """
-    Per-channel statistics for normalizing and denormalizing the latent representation.
-    This statics is computed over the entire dataset and stored in model's checkpoint under AudioVAE state_dict.
-    """
-
-    def __init__(self, latent_channels: int = 128) -> None:
-        super().__init__()
-        self.register_buffer("std-of-means", torch.empty(latent_channels))
-        self.register_buffer("mean-of-means", torch.empty(latent_channels))
-
-    def un_normalize(self, x: torch.Tensor) -> torch.Tensor:
-        return (x * self.get_buffer("std-of-means").to(x)) + self.get_buffer("mean-of-means").to(x)
-
-    def normalize(self, x: torch.Tensor) -> torch.Tensor:
-        return (x - self.get_buffer("mean-of-means").to(x)) / self.get_buffer("std-of-means").to(x)
+import torch
+import torchaudio
+from torch import nn
+
+
+class AudioProcessor(nn.Module):
+    """Converts audio waveforms to log-mel spectrograms with optional resampling."""
+
+    def __init__(
+        self,
+        sample_rate: int,
+        mel_bins: int,
+        mel_hop_length: int,
+        n_fft: int,
+    ) -> None:
+        super().__init__()
+        self.sample_rate = sample_rate
+        self.mel_transform = torchaudio.transforms.MelSpectrogram(
+            sample_rate=sample_rate,
+            n_fft=n_fft,
+            win_length=n_fft,
+            hop_length=mel_hop_length,
+            f_min=0.0,
+            f_max=sample_rate / 2.0,
+            n_mels=mel_bins,
+            window_fn=torch.hann_window,
+            center=True,
+            pad_mode="reflect",
+            power=1.0,
+            mel_scale="slaney",
+            norm="slaney",
+        )
+
+    def resample_waveform(
+        self,
+        waveform: torch.Tensor,
+        source_rate: int,
+        target_rate: int,
+    ) -> torch.Tensor:
+        """Resample waveform to target sample rate if needed."""
+        if source_rate == target_rate:
+            return waveform
+        resampled = torchaudio.functional.resample(waveform, source_rate, target_rate)
+        return resampled.to(device=waveform.device, dtype=waveform.dtype)
+
+    def waveform_to_mel(
+        self,
+        waveform: torch.Tensor,
+        waveform_sample_rate: int,
+    ) -> torch.Tensor:
+        """Convert waveform to log-mel spectrogram [batch, channels, time, n_mels]."""
+        waveform = self.resample_waveform(waveform, waveform_sample_rate, self.sample_rate)
+
+        mel = self.mel_transform(waveform)
+        mel = torch.log(torch.clamp(mel, min=1e-5))
+
+        mel = mel.to(device=waveform.device, dtype=waveform.dtype)
+        return mel.permute(0, 1, 3, 2).contiguous()
+
+
+class PerChannelStatistics(nn.Module):
+    """
+    Per-channel statistics for normalizing and denormalizing the latent representation.
+    This statics is computed over the entire dataset and stored in model's checkpoint under AudioVAE state_dict.
+    """
+
+    def __init__(self, latent_channels: int = 128) -> None:
+        super().__init__()
+        self.register_buffer("std-of-means", torch.empty(latent_channels))
+        self.register_buffer("mean-of-means", torch.empty(latent_channels))
+
+    def un_normalize(self, x: torch.Tensor) -> torch.Tensor:
+        return (x * self.get_buffer("std-of-means").to(x)) + self.get_buffer("mean-of-means").to(x)
+
+    def normalize(self, x: torch.Tensor) -> torch.Tensor:
+        return (x - self.get_buffer("mean-of-means").to(x)) / self.get_buffer("std-of-means").to(x)

packages/ltx-core/src/ltx_core/model/audio_vae/resnet.py

@@ -1,176 +1,176 @@
-from typing import Tuple
-
-import torch
-
-from ltx_core.model.audio_vae.causal_conv_2d import make_conv2d
-from ltx_core.model.audio_vae.causality_axis import CausalityAxis
-from ltx_core.model.common.normalization import NormType, build_normalization_layer
-
-LRELU_SLOPE = 0.1
-
-
-class ResBlock1(torch.nn.Module):
-    def __init__(self, channels: int, kernel_size: int = 3, dilation: Tuple[int, int, int] = (1, 3, 5)):
-        super(ResBlock1, self).__init__()
-        self.convs1 = torch.nn.ModuleList(
-            [
-                torch.nn.Conv1d(
-                    channels,
-                    channels,
-                    kernel_size,
-                    1,
-                    dilation=dilation[0],
-                    padding="same",
-                ),
-                torch.nn.Conv1d(
-                    channels,
-                    channels,
-                    kernel_size,
-                    1,
-                    dilation=dilation[1],
-                    padding="same",
-                ),
-                torch.nn.Conv1d(
-                    channels,
-                    channels,
-                    kernel_size,
-                    1,
-                    dilation=dilation[2],
-                    padding="same",
-                ),
-            ]
-        )
-
-        self.convs2 = torch.nn.ModuleList(
-            [
-                torch.nn.Conv1d(
-                    channels,
-                    channels,
-                    kernel_size,
-                    1,
-                    dilation=1,
-                    padding="same",
-                ),
-                torch.nn.Conv1d(
-                    channels,
-                    channels,
-                    kernel_size,
-                    1,
-                    dilation=1,
-                    padding="same",
-                ),
-                torch.nn.Conv1d(
-                    channels,
-                    channels,
-                    kernel_size,
-                    1,
-                    dilation=1,
-                    padding="same",
-                ),
-            ]
-        )
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        for conv1, conv2 in zip(self.convs1, self.convs2, strict=True):
-            xt = torch.nn.functional.leaky_relu(x, LRELU_SLOPE)
-            xt = conv1(xt)
-            xt = torch.nn.functional.leaky_relu(xt, LRELU_SLOPE)
-            xt = conv2(xt)
-            x = xt + x
-        return x
-
-
-class ResBlock2(torch.nn.Module):
-    def __init__(self, channels: int, kernel_size: int = 3, dilation: Tuple[int, int] = (1, 3)):
-        super(ResBlock2, self).__init__()
-        self.convs = torch.nn.ModuleList(
-            [
-                torch.nn.Conv1d(
-                    channels,
-                    channels,
-                    kernel_size,
-                    1,
-                    dilation=dilation[0],
-                    padding="same",
-                ),
-                torch.nn.Conv1d(
-                    channels,
-                    channels,
-                    kernel_size,
-                    1,
-                    dilation=dilation[1],
-                    padding="same",
-                ),
-            ]
-        )
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        for conv in self.convs:
-            xt = torch.nn.functional.leaky_relu(x, LRELU_SLOPE)
-            xt = conv(xt)
-            x = xt + x
-        return x
-
-
-class ResnetBlock(torch.nn.Module):
-    def __init__(
-        self,
-        *,
-        in_channels: int,
-        out_channels: int | None = None,
-        conv_shortcut: bool = False,
-        dropout: float = 0.0,
-        temb_channels: int = 512,
-        norm_type: NormType = NormType.GROUP,
-        causality_axis: CausalityAxis = CausalityAxis.HEIGHT,
-    ) -> None:
-        super().__init__()
-        self.causality_axis = causality_axis
-
-        if self.causality_axis != CausalityAxis.NONE and norm_type == NormType.GROUP:
-            raise ValueError("Causal ResnetBlock with GroupNorm is not supported.")
-        self.in_channels = in_channels
-        out_channels = in_channels if out_channels is None else out_channels
-        self.out_channels = out_channels
-        self.use_conv_shortcut = conv_shortcut
-
-        self.norm1 = build_normalization_layer(in_channels, normtype=norm_type)
-        self.non_linearity = torch.nn.SiLU()
-        self.conv1 = make_conv2d(in_channels, out_channels, kernel_size=3, stride=1, causality_axis=causality_axis)
-        if temb_channels > 0:
-            self.temb_proj = torch.nn.Linear(temb_channels, out_channels)
-        self.norm2 = build_normalization_layer(out_channels, normtype=norm_type)
-        self.dropout = torch.nn.Dropout(dropout)
-        self.conv2 = make_conv2d(out_channels, out_channels, kernel_size=3, stride=1, causality_axis=causality_axis)
-        if self.in_channels != self.out_channels:
-            if self.use_conv_shortcut:
-                self.conv_shortcut = make_conv2d(
-                    in_channels, out_channels, kernel_size=3, stride=1, causality_axis=causality_axis
-                )
-            else:
-                self.nin_shortcut = make_conv2d(
-                    in_channels, out_channels, kernel_size=1, stride=1, causality_axis=causality_axis
-                )
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        temb: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        h = x
-        h = self.norm1(h)
-        h = self.non_linearity(h)
-        h = self.conv1(h)
-
-        if temb is not None:
-            h = h + self.temb_proj(self.non_linearity(temb))[:, :, None, None]
-
-        h = self.norm2(h)
-        h = self.non_linearity(h)
-        h = self.dropout(h)
-        h = self.conv2(h)
-
-        if self.in_channels != self.out_channels:
-            x = self.conv_shortcut(x) if self.use_conv_shortcut else self.nin_shortcut(x)
-
-        return x + h
+from typing import Tuple
+
+import torch
+
+from ltx_core.model.audio_vae.causal_conv_2d import make_conv2d
+from ltx_core.model.audio_vae.causality_axis import CausalityAxis
+from ltx_core.model.common.normalization import NormType, build_normalization_layer
+
+LRELU_SLOPE = 0.1
+
+
+class ResBlock1(torch.nn.Module):
+    def __init__(self, channels: int, kernel_size: int = 3, dilation: Tuple[int, int, int] = (1, 3, 5)):
+        super(ResBlock1, self).__init__()
+        self.convs1 = torch.nn.ModuleList(
+            [
+                torch.nn.Conv1d(
+                    channels,
+                    channels,
+                    kernel_size,
+                    1,
+                    dilation=dilation[0],
+                    padding="same",
+                ),
+                torch.nn.Conv1d(
+                    channels,
+                    channels,
+                    kernel_size,
+                    1,
+                    dilation=dilation[1],
+                    padding="same",
+                ),
+                torch.nn.Conv1d(
+                    channels,
+                    channels,
+                    kernel_size,
+                    1,
+                    dilation=dilation[2],
+                    padding="same",
+                ),
+            ]
+        )
+
+        self.convs2 = torch.nn.ModuleList(
+            [
+                torch.nn.Conv1d(
+                    channels,
+                    channels,
+                    kernel_size,
+                    1,
+                    dilation=1,
+                    padding="same",
+                ),
+                torch.nn.Conv1d(
+                    channels,
+                    channels,
+                    kernel_size,
+                    1,
+                    dilation=1,
+                    padding="same",
+                ),
+                torch.nn.Conv1d(
+                    channels,
+                    channels,
+                    kernel_size,
+                    1,
+                    dilation=1,
+                    padding="same",
+                ),
+            ]
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        for conv1, conv2 in zip(self.convs1, self.convs2, strict=True):
+            xt = torch.nn.functional.leaky_relu(x, LRELU_SLOPE)
+            xt = conv1(xt)
+            xt = torch.nn.functional.leaky_relu(xt, LRELU_SLOPE)
+            xt = conv2(xt)
+            x = xt + x
+        return x
+
+
+class ResBlock2(torch.nn.Module):
+    def __init__(self, channels: int, kernel_size: int = 3, dilation: Tuple[int, int] = (1, 3)):
+        super(ResBlock2, self).__init__()
+        self.convs = torch.nn.ModuleList(
+            [
+                torch.nn.Conv1d(
+                    channels,
+                    channels,
+                    kernel_size,
+                    1,
+                    dilation=dilation[0],
+                    padding="same",
+                ),
+                torch.nn.Conv1d(
+                    channels,
+                    channels,
+                    kernel_size,
+                    1,
+                    dilation=dilation[1],
+                    padding="same",
+                ),
+            ]
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        for conv in self.convs:
+            xt = torch.nn.functional.leaky_relu(x, LRELU_SLOPE)
+            xt = conv(xt)
+            x = xt + x
+        return x
+
+
+class ResnetBlock(torch.nn.Module):
+    def __init__(
+        self,
+        *,
+        in_channels: int,
+        out_channels: int | None = None,
+        conv_shortcut: bool = False,
+        dropout: float = 0.0,
+        temb_channels: int = 512,
+        norm_type: NormType = NormType.GROUP,
+        causality_axis: CausalityAxis = CausalityAxis.HEIGHT,
+    ) -> None:
+        super().__init__()
+        self.causality_axis = causality_axis
+
+        if self.causality_axis != CausalityAxis.NONE and norm_type == NormType.GROUP:
+            raise ValueError("Causal ResnetBlock with GroupNorm is not supported.")
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+        self.use_conv_shortcut = conv_shortcut
+
+        self.norm1 = build_normalization_layer(in_channels, normtype=norm_type)
+        self.non_linearity = torch.nn.SiLU()
+        self.conv1 = make_conv2d(in_channels, out_channels, kernel_size=3, stride=1, causality_axis=causality_axis)
+        if temb_channels > 0:
+            self.temb_proj = torch.nn.Linear(temb_channels, out_channels)
+        self.norm2 = build_normalization_layer(out_channels, normtype=norm_type)
+        self.dropout = torch.nn.Dropout(dropout)
+        self.conv2 = make_conv2d(out_channels, out_channels, kernel_size=3, stride=1, causality_axis=causality_axis)
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                self.conv_shortcut = make_conv2d(
+                    in_channels, out_channels, kernel_size=3, stride=1, causality_axis=causality_axis
+                )
+            else:
+                self.nin_shortcut = make_conv2d(
+                    in_channels, out_channels, kernel_size=1, stride=1, causality_axis=causality_axis
+                )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        temb: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        h = x
+        h = self.norm1(h)
+        h = self.non_linearity(h)
+        h = self.conv1(h)
+
+        if temb is not None:
+            h = h + self.temb_proj(self.non_linearity(temb))[:, :, None, None]
+
+        h = self.norm2(h)
+        h = self.non_linearity(h)
+        h = self.dropout(h)
+        h = self.conv2(h)
+
+        if self.in_channels != self.out_channels:
+            x = self.conv_shortcut(x) if self.use_conv_shortcut else self.nin_shortcut(x)
+
+        return x + h

packages/ltx-core/src/ltx_core/model/audio_vae/upsample.py

@@ -1,106 +1,106 @@
-from typing import Set, Tuple
-
-import torch
-
-from ltx_core.model.audio_vae.attention import AttentionType, make_attn
-from ltx_core.model.audio_vae.causal_conv_2d import make_conv2d
-from ltx_core.model.audio_vae.causality_axis import CausalityAxis
-from ltx_core.model.audio_vae.resnet import ResnetBlock
-from ltx_core.model.common.normalization import NormType
-
-
-class Upsample(torch.nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        with_conv: bool,
-        causality_axis: CausalityAxis = CausalityAxis.HEIGHT,
-    ) -> None:
-        super().__init__()
-        self.with_conv = with_conv
-        self.causality_axis = causality_axis
-        if self.with_conv:
-            self.conv = make_conv2d(in_channels, in_channels, kernel_size=3, stride=1, causality_axis=causality_axis)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
-        if self.with_conv:
-            x = self.conv(x)
-            # Drop FIRST element in the causal axis to undo encoder's padding, while keeping the length 1 + 2 * n.
-            # For example, if the input is [0, 1, 2], after interpolation, the output is [0, 0, 1, 1, 2, 2].
-            # The causal convolution will pad the first element as [-, -, 0, 0, 1, 1, 2, 2],
-            # So the output elements rely on the following windows:
-            # 0: [-,-,0]
-            # 1: [-,0,0]
-            # 2: [0,0,1]
-            # 3: [0,1,1]
-            # 4: [1,1,2]
-            # 5: [1,2,2]
-            # Notice that the first and second elements in the output rely only on the first element in the input,
-            # while all other elements rely on two elements in the input.
-            # So we can drop the first element to undo the padding (rather than the last element).
-            # This is a no-op for non-causal convolutions.
-            match self.causality_axis:
-                case CausalityAxis.NONE:
-                    pass  # x remains unchanged
-                case CausalityAxis.HEIGHT:
-                    x = x[:, :, 1:, :]
-                case CausalityAxis.WIDTH:
-                    x = x[:, :, :, 1:]
-                case CausalityAxis.WIDTH_COMPATIBILITY:
-                    pass  # x remains unchanged
-                case _:
-                    raise ValueError(f"Invalid causality_axis: {self.causality_axis}")
-
-        return x
-
-
-def build_upsampling_path(  # noqa: PLR0913
-    *,
-    ch: int,
-    ch_mult: Tuple[int, ...],
-    num_resolutions: int,
-    num_res_blocks: int,
-    resolution: int,
-    temb_channels: int,
-    dropout: float,
-    norm_type: NormType,
-    causality_axis: CausalityAxis,
-    attn_type: AttentionType,
-    attn_resolutions: Set[int],
-    resamp_with_conv: bool,
-    initial_block_channels: int,
-) -> tuple[torch.nn.ModuleList, int]:
-    """Build the upsampling path with residual blocks, attention, and upsampling layers."""
-    up_modules = torch.nn.ModuleList()
-    block_in = initial_block_channels
-    curr_res = resolution // (2 ** (num_resolutions - 1))
-
-    for level in reversed(range(num_resolutions)):
-        stage = torch.nn.Module()
-        stage.block = torch.nn.ModuleList()
-        stage.attn = torch.nn.ModuleList()
-        block_out = ch * ch_mult[level]
-
-        for _ in range(num_res_blocks + 1):
-            stage.block.append(
-                ResnetBlock(
-                    in_channels=block_in,
-                    out_channels=block_out,
-                    temb_channels=temb_channels,
-                    dropout=dropout,
-                    norm_type=norm_type,
-                    causality_axis=causality_axis,
-                )
-            )
-            block_in = block_out
-            if curr_res in attn_resolutions:
-                stage.attn.append(make_attn(block_in, attn_type=attn_type, norm_type=norm_type))
-
-        if level != 0:
-            stage.upsample = Upsample(block_in, resamp_with_conv, causality_axis=causality_axis)
-            curr_res *= 2
-
-        up_modules.insert(0, stage)
-
-    return up_modules, block_in
+from typing import Set, Tuple
+
+import torch
+
+from ltx_core.model.audio_vae.attention import AttentionType, make_attn
+from ltx_core.model.audio_vae.causal_conv_2d import make_conv2d
+from ltx_core.model.audio_vae.causality_axis import CausalityAxis
+from ltx_core.model.audio_vae.resnet import ResnetBlock
+from ltx_core.model.common.normalization import NormType
+
+
+class Upsample(torch.nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        with_conv: bool,
+        causality_axis: CausalityAxis = CausalityAxis.HEIGHT,
+    ) -> None:
+        super().__init__()
+        self.with_conv = with_conv
+        self.causality_axis = causality_axis
+        if self.with_conv:
+            self.conv = make_conv2d(in_channels, in_channels, kernel_size=3, stride=1, causality_axis=causality_axis)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
+        if self.with_conv:
+            x = self.conv(x)
+            # Drop FIRST element in the causal axis to undo encoder's padding, while keeping the length 1 + 2 * n.
+            # For example, if the input is [0, 1, 2], after interpolation, the output is [0, 0, 1, 1, 2, 2].
+            # The causal convolution will pad the first element as [-, -, 0, 0, 1, 1, 2, 2],
+            # So the output elements rely on the following windows:
+            # 0: [-,-,0]
+            # 1: [-,0,0]
+            # 2: [0,0,1]
+            # 3: [0,1,1]
+            # 4: [1,1,2]
+            # 5: [1,2,2]
+            # Notice that the first and second elements in the output rely only on the first element in the input,
+            # while all other elements rely on two elements in the input.
+            # So we can drop the first element to undo the padding (rather than the last element).
+            # This is a no-op for non-causal convolutions.
+            match self.causality_axis:
+                case CausalityAxis.NONE:
+                    pass  # x remains unchanged
+                case CausalityAxis.HEIGHT:
+                    x = x[:, :, 1:, :]
+                case CausalityAxis.WIDTH:
+                    x = x[:, :, :, 1:]
+                case CausalityAxis.WIDTH_COMPATIBILITY:
+                    pass  # x remains unchanged
+                case _:
+                    raise ValueError(f"Invalid causality_axis: {self.causality_axis}")
+
+        return x
+
+
+def build_upsampling_path(  # noqa: PLR0913
+    *,
+    ch: int,
+    ch_mult: Tuple[int, ...],
+    num_resolutions: int,
+    num_res_blocks: int,
+    resolution: int,
+    temb_channels: int,
+    dropout: float,
+    norm_type: NormType,
+    causality_axis: CausalityAxis,
+    attn_type: AttentionType,
+    attn_resolutions: Set[int],
+    resamp_with_conv: bool,
+    initial_block_channels: int,
+) -> tuple[torch.nn.ModuleList, int]:
+    """Build the upsampling path with residual blocks, attention, and upsampling layers."""
+    up_modules = torch.nn.ModuleList()
+    block_in = initial_block_channels
+    curr_res = resolution // (2 ** (num_resolutions - 1))
+
+    for level in reversed(range(num_resolutions)):
+        stage = torch.nn.Module()
+        stage.block = torch.nn.ModuleList()
+        stage.attn = torch.nn.ModuleList()
+        block_out = ch * ch_mult[level]
+
+        for _ in range(num_res_blocks + 1):
+            stage.block.append(
+                ResnetBlock(
+                    in_channels=block_in,
+                    out_channels=block_out,
+                    temb_channels=temb_channels,
+                    dropout=dropout,
+                    norm_type=norm_type,
+                    causality_axis=causality_axis,
+                )
+            )
+            block_in = block_out
+            if curr_res in attn_resolutions:
+                stage.attn.append(make_attn(block_in, attn_type=attn_type, norm_type=norm_type))
+
+        if level != 0:
+            stage.upsample = Upsample(block_in, resamp_with_conv, causality_axis=causality_axis)
+            curr_res *= 2
+
+        up_modules.insert(0, stage)
+
+    return up_modules, block_in

packages/ltx-core/src/ltx_core/model/audio_vae/vocoder.py

@@ -1,129 +1,123 @@
-# Copyright (c) 2025 Lightricks. All rights reserved.
-# Created by Ivan Zorin
-
-import math
-from typing import List
-
-import einops
-import torch
-import torch.nn.functional as F
-from torch import nn
-
-from ltx_core.model.audio_vae.resnet import LRELU_SLOPE, ResBlock1, ResBlock2
-
-
-class Vocoder(torch.nn.Module):
-    """
-    Vocoder model for synthesizing audio from Mel spectrograms.
-
-    Args:
-        resblock_kernel_sizes: List of kernel sizes for the residual blocks.
-                               This value is read from the checkpoint at `config.vocoder.resblock_kernel_sizes`.
-        upsample_rates: List of upsampling rates.
-                               This value is read from the checkpoint at `config.vocoder.upsample_rates`.
-        upsample_kernel_sizes: List of kernel sizes for the upsampling layers.
-                               This value is read from the checkpoint at `config.vocoder.upsample_kernel_sizes`.
-        resblock_dilation_sizes: List of dilation sizes for the residual blocks.
-                               This value is read from the checkpoint at `config.vocoder.resblock_dilation_sizes`.
-        upsample_initial_channel: Initial number of channels for the upsampling layers.
-                               This value is read from the checkpoint at `config.vocoder.upsample_initial_channel`.
-        stereo: Whether to use stereo output.
-                               This value is read from the checkpoint at `config.vocoder.stereo`.
-        resblock: Type of residual block to use.
-                                This value is read from the checkpoint at `config.vocoder.resblock`.
-        output_sample_rate: Waveform sample rate.
-                               This value is read from the checkpoint at `config.vocoder.output_sample_rate`.
-    """
-
-    def __init__(
-        self,
-        resblock_kernel_sizes: List[int] | None = None,
-        upsample_rates: List[int] | None = None,
-        upsample_kernel_sizes: List[int] | None = None,
-        resblock_dilation_sizes: List[List[int]] | None = None,
-        upsample_initial_channel: int = 1024,
-        stereo: bool = True,
-        resblock: str = "1",
-        output_sample_rate: int = 24000,
-    ):
-        super().__init__()
-
-        # Initialize default values if not provided. Note that mutable default values are not supported.
-        if resblock_kernel_sizes is None:
-            resblock_kernel_sizes = [3, 7, 11]
-        if upsample_rates is None:
-            upsample_rates = [6, 5, 2, 2, 2]
-        if upsample_kernel_sizes is None:
-            upsample_kernel_sizes = [16, 15, 8, 4, 4]
-        if resblock_dilation_sizes is None:
-            resblock_dilation_sizes = [[1, 3, 5], [1, 3, 5], [1, 3, 5]]
-
-        self.output_sample_rate = output_sample_rate
-        self.num_kernels = len(resblock_kernel_sizes)
-        self.num_upsamples = len(upsample_rates)
-        in_channels = 128 if stereo else 64
-        self.conv_pre = nn.Conv1d(in_channels, upsample_initial_channel, 7, 1, padding=3)
-        resblock_class = ResBlock1 if resblock == "1" else ResBlock2
-
-        self.ups = nn.ModuleList()
-        for i, (stride, kernel_size) in enumerate(zip(upsample_rates, upsample_kernel_sizes, strict=True)):
-            self.ups.append(
-                nn.ConvTranspose1d(
-                    upsample_initial_channel // (2**i),
-                    upsample_initial_channel // (2 ** (i + 1)),
-                    kernel_size,
-                    stride,
-                    padding=(kernel_size - stride) // 2,
-                )
-            )
-
-        self.resblocks = nn.ModuleList()
-        for i, _ in enumerate(self.ups):
-            ch = upsample_initial_channel // (2 ** (i + 1))
-            for kernel_size, dilations in zip(resblock_kernel_sizes, resblock_dilation_sizes, strict=True):
-                self.resblocks.append(resblock_class(ch, kernel_size, dilations))
-
-        out_channels = 2 if stereo else 1
-        final_channels = upsample_initial_channel // (2**self.num_upsamples)
-        self.conv_post = nn.Conv1d(final_channels, out_channels, 7, 1, padding=3)
-
-        self.upsample_factor = math.prod(layer.stride[0] for layer in self.ups)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        Forward pass of the vocoder.
-
-        Args:
-            x: Input Mel spectrogram tensor. Can be either:
-               - 3D: (batch_size, time, mel_bins) for mono
-               - 4D: (batch_size, 2, time, mel_bins) for stereo
-
-        Returns:
-            Audio waveform tensor of shape (batch_size, out_channels, audio_length)
-        """
-        x = x.transpose(2, 3)  # (batch, channels, time, mel_bins) -> (batch, channels, mel_bins, time)
-
-        if x.dim() == 4:  # stereo
-            assert x.shape[1] == 2, "Input must have 2 channels for stereo"
-            x = einops.rearrange(x, "b s c t -> b (s c) t")
-
-        x = self.conv_pre(x)
-
-        for i in range(self.num_upsamples):
-            x = F.leaky_relu(x, LRELU_SLOPE)
-            x = self.ups[i](x)
-            start = i * self.num_kernels
-            end = start + self.num_kernels
-
-            # Evaluate all resblocks with the same input tensor so they can run
-            # independently (and thus in parallel on accelerator hardware) before
-            # aggregating their outputs via mean.
-            block_outputs = torch.stack(
-                [self.resblocks[idx](x) for idx in range(start, end)],
-                dim=0,
-            )
-
-            x = block_outputs.mean(dim=0)
-
-        x = self.conv_post(F.leaky_relu(x))
-        return torch.tanh(x)
+import math
+from typing import List
+
+import einops
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from ltx_core.model.audio_vae.resnet import LRELU_SLOPE, ResBlock1, ResBlock2
+
+
+class Vocoder(torch.nn.Module):
+    """
+    Vocoder model for synthesizing audio from Mel spectrograms.
+    Args:
+        resblock_kernel_sizes: List of kernel sizes for the residual blocks.
+                               This value is read from the checkpoint at `config.vocoder.resblock_kernel_sizes`.
+        upsample_rates: List of upsampling rates.
+                               This value is read from the checkpoint at `config.vocoder.upsample_rates`.
+        upsample_kernel_sizes: List of kernel sizes for the upsampling layers.
+                               This value is read from the checkpoint at `config.vocoder.upsample_kernel_sizes`.
+        resblock_dilation_sizes: List of dilation sizes for the residual blocks.
+                               This value is read from the checkpoint at `config.vocoder.resblock_dilation_sizes`.
+        upsample_initial_channel: Initial number of channels for the upsampling layers.
+                               This value is read from the checkpoint at `config.vocoder.upsample_initial_channel`.
+        stereo: Whether to use stereo output.
+                               This value is read from the checkpoint at `config.vocoder.stereo`.
+        resblock: Type of residual block to use.
+                                This value is read from the checkpoint at `config.vocoder.resblock`.
+        output_sample_rate: Waveform sample rate.
+                               This value is read from the checkpoint at `config.vocoder.output_sample_rate`.
+    """
+
+    def __init__(
+        self,
+        resblock_kernel_sizes: List[int] | None = None,
+        upsample_rates: List[int] | None = None,
+        upsample_kernel_sizes: List[int] | None = None,
+        resblock_dilation_sizes: List[List[int]] | None = None,
+        upsample_initial_channel: int = 1024,
+        stereo: bool = True,
+        resblock: str = "1",
+        output_sample_rate: int = 24000,
+    ):
+        super().__init__()
+
+        # Initialize default values if not provided. Note that mutable default values are not supported.
+        if resblock_kernel_sizes is None:
+            resblock_kernel_sizes = [3, 7, 11]
+        if upsample_rates is None:
+            upsample_rates = [6, 5, 2, 2, 2]
+        if upsample_kernel_sizes is None:
+            upsample_kernel_sizes = [16, 15, 8, 4, 4]
+        if resblock_dilation_sizes is None:
+            resblock_dilation_sizes = [[1, 3, 5], [1, 3, 5], [1, 3, 5]]
+
+        self.output_sample_rate = output_sample_rate
+        self.num_kernels = len(resblock_kernel_sizes)
+        self.num_upsamples = len(upsample_rates)
+        in_channels = 128 if stereo else 64
+        self.conv_pre = nn.Conv1d(in_channels, upsample_initial_channel, 7, 1, padding=3)
+        resblock_class = ResBlock1 if resblock == "1" else ResBlock2
+
+        self.ups = nn.ModuleList()
+        for i, (stride, kernel_size) in enumerate(zip(upsample_rates, upsample_kernel_sizes, strict=True)):
+            self.ups.append(
+                nn.ConvTranspose1d(
+                    upsample_initial_channel // (2**i),
+                    upsample_initial_channel // (2 ** (i + 1)),
+                    kernel_size,
+                    stride,
+                    padding=(kernel_size - stride) // 2,
+                )
+            )
+
+        self.resblocks = nn.ModuleList()
+        for i, _ in enumerate(self.ups):
+            ch = upsample_initial_channel // (2 ** (i + 1))
+            for kernel_size, dilations in zip(resblock_kernel_sizes, resblock_dilation_sizes, strict=True):
+                self.resblocks.append(resblock_class(ch, kernel_size, dilations))
+
+        out_channels = 2 if stereo else 1
+        final_channels = upsample_initial_channel // (2**self.num_upsamples)
+        self.conv_post = nn.Conv1d(final_channels, out_channels, 7, 1, padding=3)
+
+        self.upsample_factor = math.prod(layer.stride[0] for layer in self.ups)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass of the vocoder.
+        Args:
+            x: Input Mel spectrogram tensor. Can be either:
+               - 3D: (batch_size, time, mel_bins) for mono
+               - 4D: (batch_size, 2, time, mel_bins) for stereo
+        Returns:
+            Audio waveform tensor of shape (batch_size, out_channels, audio_length)
+        """
+        x = x.transpose(2, 3)  # (batch, channels, time, mel_bins) -> (batch, channels, mel_bins, time)
+
+        if x.dim() == 4:  # stereo
+            assert x.shape[1] == 2, "Input must have 2 channels for stereo"
+            x = einops.rearrange(x, "b s c t -> b (s c) t")
+
+        x = self.conv_pre(x)
+
+        for i in range(self.num_upsamples):
+            x = F.leaky_relu(x, LRELU_SLOPE)
+            x = self.ups[i](x)
+            start = i * self.num_kernels
+            end = start + self.num_kernels
+
+            # Evaluate all resblocks with the same input tensor so they can run
+            # independently (and thus in parallel on accelerator hardware) before
+            # aggregating their outputs via mean.
+            block_outputs = torch.stack(
+                [self.resblocks[idx](x) for idx in range(start, end)],
+                dim=0,
+            )
+
+            x = block_outputs.mean(dim=0)
+
+        x = self.conv_post(F.leaky_relu(x))
+        return torch.tanh(x)

packages/ltx-core/src/ltx_core/model/common/__init__.py

@@ -0,0 +1,9 @@
+"""Common model utilities."""
+
+from ltx_core.model.common.normalization import NormType, PixelNorm, build_normalization_layer
+
+__all__ = [
+    "NormType",
+    "PixelNorm",
+    "build_normalization_layer",
+]

packages/ltx-core/src/ltx_core/model/common/normalization.py

@@ -1,63 +1,59 @@
-# Copyright (c) 2025 Lightricks. All rights reserved.
-# Created by Ivan Zorin
-
-from enum import Enum
-
-import torch
-from torch import nn
-
-
-class NormType(Enum):
-    GROUP = "group"
-    PIXEL = "pixel"
-
-
-class PixelNorm(nn.Module):
-    """
-    Per-pixel (per-location) RMS normalization layer.
-
-    For each element along the chosen dimension, this layer normalizes the tensor
-    by the root-mean-square of its values across that dimension:
-
-        y = x / sqrt(mean(x^2, dim=dim, keepdim=True) + eps)
-
-    """
-
-    def __init__(self, dim: int = 1, eps: float = 1e-8) -> None:
-        """
-        Args:
-            dim: Dimension along which to compute the RMS (typically channels).
-            eps: Small constant added for numerical stability.
-        """
-        super().__init__()
-        self.dim = dim
-        self.eps = eps
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        Apply RMS normalization along the configured dimension.
-        """
-        # Compute mean of squared values along `dim`, keep dimensions for broadcasting.
-        mean_sq = torch.mean(x**2, dim=self.dim, keepdim=True)
-        # Normalize by the root-mean-square (RMS).
-        rms = torch.sqrt(mean_sq + self.eps)
-        return x / rms
-
-
-def build_normalization_layer(
-    in_channels: int, *, num_groups: int = 32, normtype: NormType = NormType.GROUP
-) -> nn.Module:
-    """
-    Create a normalization layer based on the normalization type.
-    Args:
-        in_channels: Number of input channels
-        num_groups: Number of groups for group normalization
-        normtype: Type of normalization: "group" or "pixel"
-    Returns:
-        A normalization layer
-    """
-    if normtype == NormType.GROUP:
-        return torch.nn.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True)
-    if normtype == NormType.PIXEL:
-        return PixelNorm(dim=1, eps=1e-6)
-    raise ValueError(f"Invalid normalization type: {normtype}")
+from enum import Enum
+
+import torch
+from torch import nn
+
+
+class NormType(Enum):
+    """Normalization layer types: GROUP (GroupNorm) or PIXEL (per-location RMS norm)."""
+
+    GROUP = "group"
+    PIXEL = "pixel"
+
+
+class PixelNorm(nn.Module):
+    """
+    Per-pixel (per-location) RMS normalization layer.
+    For each element along the chosen dimension, this layer normalizes the tensor
+    by the root-mean-square of its values across that dimension:
+        y = x / sqrt(mean(x^2, dim=dim, keepdim=True) + eps)
+    """
+
+    def __init__(self, dim: int = 1, eps: float = 1e-8) -> None:
+        """
+        Args:
+            dim: Dimension along which to compute the RMS (typically channels).
+            eps: Small constant added for numerical stability.
+        """
+        super().__init__()
+        self.dim = dim
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Apply RMS normalization along the configured dimension.
+        """
+        # Compute mean of squared values along `dim`, keep dimensions for broadcasting.
+        mean_sq = torch.mean(x**2, dim=self.dim, keepdim=True)
+        # Normalize by the root-mean-square (RMS).
+        rms = torch.sqrt(mean_sq + self.eps)
+        return x / rms
+
+
+def build_normalization_layer(
+    in_channels: int, *, num_groups: int = 32, normtype: NormType = NormType.GROUP
+) -> nn.Module:
+    """
+    Create a normalization layer based on the normalization type.
+    Args:
+        in_channels: Number of input channels
+        num_groups: Number of groups for group normalization
+        normtype: Type of normalization: "group" or "pixel"
+    Returns:
+        A normalization layer
+    """
+    if normtype == NormType.GROUP:
+        return torch.nn.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True)
+    if normtype == NormType.PIXEL:
+        return PixelNorm(dim=1, eps=1e-6)
+    raise ValueError(f"Invalid normalization type: {normtype}")

packages/ltx-core/src/ltx_core/model/model_protocol.py

@@ -1,13 +1,10 @@
-# Copyright (c) 2025 Lightricks. All rights reserved.
-# Created by Andrew Kvochko
-
-from typing import Protocol, TypeVar
-
-ModelType = TypeVar("ModelType")
-
-
-class ModelConfigurator(Protocol[ModelType]):
-    """Protocol for model loader classes that instantiates models from a configuration dictionary."""
-
-    @classmethod
-    def from_config(cls, config: dict) -> ModelType: ...
+from typing import Protocol, TypeVar
+
+ModelType = TypeVar("ModelType")
+
+
+class ModelConfigurator(Protocol[ModelType]):
+    """Protocol for model loader classes that instantiates models from a configuration dictionary."""
+
+    @classmethod
+    def from_config(cls, config: dict) -> ModelType: ...

packages/ltx-core/src/ltx_core/model/transformer/__init__.py

@@ -0,0 +1,24 @@
+"""Transformer model components."""
+
+from ltx_core.model.transformer.modality import Modality
+from ltx_core.model.transformer.model import LTXModel, X0Model
+from ltx_core.model.transformer.model_configurator import (
+    LTXV_MODEL_COMFY_RENAMING_MAP,
+    LTXV_MODEL_COMFY_RENAMING_WITH_TRANSFORMER_LINEAR_DOWNCAST_MAP,
+    UPCAST_DURING_INFERENCE,
+    LTXModelConfigurator,
+    LTXVideoOnlyModelConfigurator,
+    UpcastWithStochasticRounding,
+)
+
+__all__ = [
+    "LTXV_MODEL_COMFY_RENAMING_MAP",
+    "LTXV_MODEL_COMFY_RENAMING_WITH_TRANSFORMER_LINEAR_DOWNCAST_MAP",
+    "UPCAST_DURING_INFERENCE",
+    "LTXModel",
+    "LTXModelConfigurator",
+    "LTXVideoOnlyModelConfigurator",
+    "Modality",
+    "UpcastWithStochasticRounding",
+    "X0Model",
+]

packages/ltx-core/src/ltx_core/model/transformer/adaln.py

@@ -1,39 +1,34 @@
-# Copyright (c) 2025 Lightricks. All rights reserved.
-# Created by Andrew Kvochko
-
-from typing import Optional, Tuple
-
-import torch
-
-from ltx_core.model.transformer.timestep_embedding import PixArtAlphaCombinedTimestepSizeEmbeddings
-
-
-class AdaLayerNormSingle(torch.nn.Module):
-    r"""
-    Norm layer adaptive layer norm single (adaLN-single).
-
-    As proposed in PixArt-Alpha (see: https://arxiv.org/abs/2310.00426; Section 2.3).
-
-    Parameters:
-        embedding_dim (`int`): The size of each embedding vector.
-        use_additional_conditions (`bool`): To use additional conditions for normalization or not.
-    """
-
-    def __init__(self, embedding_dim: int, embedding_coefficient: int = 6):
-        super().__init__()
-
-        self.emb = PixArtAlphaCombinedTimestepSizeEmbeddings(
-            embedding_dim,
-            size_emb_dim=embedding_dim // 3,
-        )
-
-        self.silu = torch.nn.SiLU()
-        self.linear = torch.nn.Linear(embedding_dim, embedding_coefficient * embedding_dim, bias=True)
-
-    def forward(
-        self,
-        timestep: torch.Tensor,
-        hidden_dtype: Optional[torch.dtype] = None,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        embedded_timestep = self.emb(timestep, hidden_dtype=hidden_dtype)
-        return self.linear(self.silu(embedded_timestep)), embedded_timestep
+from typing import Optional, Tuple
+
+import torch
+
+from ltx_core.model.transformer.timestep_embedding import PixArtAlphaCombinedTimestepSizeEmbeddings
+
+
+class AdaLayerNormSingle(torch.nn.Module):
+    r"""
+    Norm layer adaptive layer norm single (adaLN-single).
+    As proposed in PixArt-Alpha (see: https://arxiv.org/abs/2310.00426; Section 2.3).
+    Parameters:
+        embedding_dim (`int`): The size of each embedding vector.
+        use_additional_conditions (`bool`): To use additional conditions for normalization or not.
+    """
+
+    def __init__(self, embedding_dim: int, embedding_coefficient: int = 6):
+        super().__init__()
+
+        self.emb = PixArtAlphaCombinedTimestepSizeEmbeddings(
+            embedding_dim,
+            size_emb_dim=embedding_dim // 3,
+        )
+
+        self.silu = torch.nn.SiLU()
+        self.linear = torch.nn.Linear(embedding_dim, embedding_coefficient * embedding_dim, bias=True)
+
+    def forward(
+        self,
+        timestep: torch.Tensor,
+        hidden_dtype: Optional[torch.dtype] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        embedded_timestep = self.emb(timestep, hidden_dtype=hidden_dtype)
+        return self.linear(self.silu(embedded_timestep)), embedded_timestep

packages/ltx-core/src/ltx_core/model/transformer/attention.py

@@ -1,199 +1,188 @@
-# Copyright (c) 2025 Lightricks. All rights reserved.
-# Created by Andrew Kvochko
-
-from enum import Enum
-from typing import Protocol
-
-import torch
-
-from ltx_core.model.transformer.rope import LTXRopeType, apply_rotary_emb
-
-memory_efficient_attention = None
-flash_attn_interface = None
-try:
-    from xformers.ops import memory_efficient_attention
-except ImportError:
-    memory_efficient_attention = None
-try:
-    # FlashAttention3 and XFormersAttention cannot be used together
-    if memory_efficient_attention is None:
-        import flash_attn_interface
-except ImportError:
-    flash_attn_interface = None
-
-
-class AttentionCallable(Protocol):
-    def __call__(
-        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int, mask: torch.Tensor | None = None
-    ) -> torch.Tensor: ...
-
-
-class PytorchAttention(AttentionCallable):
-    def __call__(
-        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int, mask: torch.Tensor | None = None
-    ) -> torch.Tensor:
-        b, _, dim_head = q.shape
-        dim_head //= heads
-        q, k, v = (t.view(b, -1, heads, dim_head).transpose(1, 2) for t in (q, k, v))
-
-        if mask is not None:
-            # add a batch dimension if there isn't already one
-            if mask.ndim == 2:
-                mask = mask.unsqueeze(0)
-            # add a heads dimension if there isn't already one
-            if mask.ndim == 3:
-                mask = mask.unsqueeze(1)
-
-        out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False)
-        out = out.transpose(1, 2).reshape(b, -1, heads * dim_head)
-        return out
-
-
-class XFormersAttention(AttentionCallable):
-    def __call__(
-        self,
-        q: torch.Tensor,
-        k: torch.Tensor,
-        v: torch.Tensor,
-        heads: int,
-        mask: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        if memory_efficient_attention is None:
-            raise RuntimeError("XFormersAttention was selected but `xformers` is not installed.")
-
-        b, _, dim_head = q.shape
-        dim_head //= heads
-
-        # xformers expects [B, M, H, K]
-        q, k, v = (t.view(b, -1, heads, dim_head) for t in (q, k, v))
-
-        # LT_INTERNAL: https://github.com/LightricksResearch/ComfyUI/blob/ee2a50cd8fb3544c66f8a3096390c741fff12ae3/comfy/ldm/modules/attention.py#L441-L459
-        if mask is not None:
-            # add a singleton batch dimension
-            if mask.ndim == 2:
-                mask = mask.unsqueeze(0)
-            # add a singleton heads dimension
-            if mask.ndim == 3:
-                mask = mask.unsqueeze(1)
-            # pad to a multiple of 8
-            pad = 8 - mask.shape[-1] % 8
-            # the xformers docs says that it's allowed to have a mask of shape (1, Nq, Nk)
-            # but when using separated heads, the shape has to be (B, H, Nq, Nk)
-            # in flux, this matrix ends up being over 1GB
-            # here, we create a mask with the same batch/head size as the input mask (potentially singleton or full)
-            mask_out = torch.empty(
-                [mask.shape[0], mask.shape[1], q.shape[1], mask.shape[-1] + pad], dtype=q.dtype, device=q.device
-            )
-
-            mask_out[..., : mask.shape[-1]] = mask
-            # doesn't this remove the padding again??
-            mask = mask_out[..., : mask.shape[-1]]
-            mask = mask.expand(b, heads, -1, -1)
-
-        out = memory_efficient_attention(q.to(v.dtype), k.to(v.dtype), v, attn_bias=mask, p=0.0)
-        out = out.reshape(b, -1, heads * dim_head)
-        return out
-
-
-class FlashAttention3(AttentionCallable):
-    def __call__(
-        self,
-        q: torch.Tensor,
-        k: torch.Tensor,
-        v: torch.Tensor,
-        heads: int,
-        mask: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        if flash_attn_interface is None:
-            raise RuntimeError("FlashAttention3 was selected but `FlashAttention3` is not installed.")
-
-        b, _, dim_head = q.shape
-        dim_head //= heads
-
-        q, k, v = (t.view(b, -1, heads, dim_head) for t in (q, k, v))
-
-        if mask is not None:
-            raise NotImplementedError("Mask is not supported for FlashAttention3")
-
-        out = flash_attn_interface.flash_attn_func(q.to(v.dtype), k.to(v.dtype), v)
-        out = out.reshape(b, -1, heads * dim_head)
-        return out
-
-
-class AttentionFunction(Enum):
-    PYTORCH = "pytorch"
-    XFORMERS = "xformers"
-    FLASH_ATTENTION_3 = "flash_attention_3"
-    DEFAULT = "default"
-
-    def __call__(
-        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int, mask: torch.Tensor | None = None
-    ) -> torch.Tensor:
-        if self is AttentionFunction.PYTORCH:
-            return PytorchAttention()(q, k, v, heads, mask)
-        elif self is AttentionFunction.XFORMERS:
-            return XFormersAttention()(q, k, v, heads, mask)
-        elif self is AttentionFunction.FLASH_ATTENTION_3:
-            return FlashAttention3()(q, k, v, heads, mask)
-        else:
-            # Default behavior: XFormers if installed else - PyTorch
-            return (
-                XFormersAttention()(q, k, v, heads, mask)
-                if memory_efficient_attention is not None
-                else PytorchAttention()(q, k, v, heads, mask)
-            )
-
-
-class Attention(torch.nn.Module):
-    def __init__(
-        self,
-        query_dim: int,
-        context_dim: int | None = None,
-        heads: int = 8,
-        dim_head: int = 64,
-        norm_eps: float = 1e-6,
-        rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
-        attention_function: AttentionCallable | AttentionFunction = AttentionFunction.DEFAULT,
-    ) -> None:
-        super().__init__()
-        self.rope_type = rope_type
-        self.attention_function = attention_function
-
-        inner_dim = dim_head * heads
-        context_dim = query_dim if context_dim is None else context_dim
-
-        self.heads = heads
-        self.dim_head = dim_head
-
-        self.q_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps)
-        self.k_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps)
-
-        self.to_q = torch.nn.Linear(query_dim, inner_dim, bias=True)
-        self.to_k = torch.nn.Linear(context_dim, inner_dim, bias=True)
-        self.to_v = torch.nn.Linear(context_dim, inner_dim, bias=True)
-
-        self.to_out = torch.nn.Sequential(torch.nn.Linear(inner_dim, query_dim, bias=True), torch.nn.Identity())
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        context: torch.Tensor | None = None,
-        mask: torch.Tensor | None = None,
-        pe: torch.Tensor | None = None,
-        k_pe: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        q = self.to_q(x)
-        context = x if context is None else context
-        k = self.to_k(context)
-        v = self.to_v(context)
-
-        q = self.q_norm(q)
-        k = self.k_norm(k)
-
-        if pe is not None:
-            q = apply_rotary_emb(q, pe, self.rope_type)
-            k = apply_rotary_emb(k, pe if k_pe is None else k_pe, self.rope_type)
-
-        # attention_function can be an enum *or* a custom callable
-        out = self.attention_function(q, k, v, self.heads, mask)
-        return self.to_out(out)
+from enum import Enum
+from typing import Protocol
+
+import torch
+
+from ltx_core.model.transformer.rope import LTXRopeType, apply_rotary_emb
+
+memory_efficient_attention = None
+flash_attn_interface = None
+try:
+    from xformers.ops import memory_efficient_attention
+except ImportError:
+    memory_efficient_attention = None
+
+try:  
+    import flash_attn_interface
+except ImportError:
+    flash_attn_interface = None
+
+class AttentionCallable(Protocol):
+    def __call__(
+        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int, mask: torch.Tensor | None = None
+    ) -> torch.Tensor: ...
+
+
+class PytorchAttention(AttentionCallable):
+    def __call__(
+        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int, mask: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        b, _, dim_head = q.shape
+        dim_head //= heads
+        q, k, v = (t.view(b, -1, heads, dim_head).transpose(1, 2) for t in (q, k, v))
+
+        if mask is not None:
+            # add a batch dimension if there isn't already one
+            if mask.ndim == 2:
+                mask = mask.unsqueeze(0)
+            # add a heads dimension if there isn't already one
+            if mask.ndim == 3:
+                mask = mask.unsqueeze(1)
+
+        out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False)
+        out = out.transpose(1, 2).reshape(b, -1, heads * dim_head)
+        return out
+
+
+class XFormersAttention(AttentionCallable):
+    def __call__(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        heads: int,
+        mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if memory_efficient_attention is None:
+            raise RuntimeError("XFormersAttention was selected but `xformers` is not installed.")
+
+        b, _, dim_head = q.shape
+        dim_head //= heads
+
+        # xformers expects [B, M, H, K]
+        q, k, v = (t.view(b, -1, heads, dim_head) for t in (q, k, v))
+
+        if mask is not None:
+            # add a singleton batch dimension
+            if mask.ndim == 2:
+                mask = mask.unsqueeze(0)
+            # add a singleton heads dimension
+            if mask.ndim == 3:
+                mask = mask.unsqueeze(1)
+            # pad to a multiple of 8
+            pad = 8 - mask.shape[-1] % 8
+            # the xformers docs says that it's allowed to have a mask of shape (1, Nq, Nk)
+            # but when using separated heads, the shape has to be (B, H, Nq, Nk)
+            # in flux, this matrix ends up being over 1GB
+            # here, we create a mask with the same batch/head size as the input mask (potentially singleton or full)
+            mask_out = torch.empty(
+                [mask.shape[0], mask.shape[1], q.shape[1], mask.shape[-1] + pad], dtype=q.dtype, device=q.device
+            )
+
+            mask_out[..., : mask.shape[-1]] = mask
+            # doesn't this remove the padding again??
+            mask = mask_out[..., : mask.shape[-1]]
+            mask = mask.expand(b, heads, -1, -1)
+
+        out = memory_efficient_attention(q.to(v.dtype), k.to(v.dtype), v, attn_bias=mask, p=0.0)
+        out = out.reshape(b, -1, heads * dim_head)
+        return out
+
+
+class FlashAttention3(AttentionCallable):
+    def __call__(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        heads: int,
+        mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if flash_attn_interface is None:
+            raise RuntimeError("FlashAttention3 was selected but `FlashAttention3` is not installed.")
+
+        b, _, dim_head = q.shape
+        dim_head //= heads
+
+        q, k, v = (t.view(b, -1, heads, dim_head) for t in (q, k, v))
+
+        if mask is not None:
+            raise NotImplementedError("Mask is not supported for FlashAttention3")
+
+        out = flash_attn_interface.flash_attn_func(q.to(v.dtype), k.to(v.dtype), v)
+        out = out.reshape(b, -1, heads * dim_head)
+        return out
+
+
+class AttentionFunction(Enum):
+    PYTORCH = "pytorch"
+    XFORMERS = "xformers"
+    FLASH_ATTENTION_3 = "flash_attention_3"
+    DEFAULT = "default"
+
+    def __call__(
+        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int, mask: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        if mask is None and flash_attn_interface is not None:
+            return FlashAttention3()(q, k, v, heads, mask)
+        else:
+            return (
+                XFormersAttention()(q, k, v, heads, mask)
+                if memory_efficient_attention is not None
+                else PytorchAttention()(q, k, v, heads, mask)
+            )
+
+
+class Attention(torch.nn.Module):
+    def __init__(
+        self,
+        query_dim: int,
+        context_dim: int | None = None,
+        heads: int = 8,
+        dim_head: int = 64,
+        norm_eps: float = 1e-6,
+        rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
+        attention_function: AttentionCallable | AttentionFunction = AttentionFunction.DEFAULT,
+    ) -> None:
+        super().__init__()
+        self.rope_type = rope_type
+        self.attention_function = attention_function
+
+        inner_dim = dim_head * heads
+        context_dim = query_dim if context_dim is None else context_dim
+
+        self.heads = heads
+        self.dim_head = dim_head
+
+        self.q_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps)
+        self.k_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps)
+
+        self.to_q = torch.nn.Linear(query_dim, inner_dim, bias=True)
+        self.to_k = torch.nn.Linear(context_dim, inner_dim, bias=True)
+        self.to_v = torch.nn.Linear(context_dim, inner_dim, bias=True)
+
+        self.to_out = torch.nn.Sequential(torch.nn.Linear(inner_dim, query_dim, bias=True), torch.nn.Identity())
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        context: torch.Tensor | None = None,
+        mask: torch.Tensor | None = None,
+        pe: torch.Tensor | None = None,
+        k_pe: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        q = self.to_q(x)
+        context = x if context is None else context
+        k = self.to_k(context)
+        v = self.to_v(context)
+
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+
+        if pe is not None:
+            q = apply_rotary_emb(q, pe, self.rope_type)
+            k = apply_rotary_emb(k, pe if k_pe is None else k_pe, self.rope_type)
+
+        # attention_function can be an enum *or* a custom callable
+        out = self.attention_function(q, k, v, self.heads, mask)
+        return self.to_out(out)

packages/ltx-core/src/ltx_core/model/transformer/feed_forward.py

@@ -1,18 +1,15 @@
-# Copyright (c) 2025 Lightricks. All rights reserved.
-# Created by Andrew Kvochko
-
-import torch
-
-from ltx_core.model.transformer.gelu_approx import GELUApprox
-
-
-class FeedForward(torch.nn.Module):
-    def __init__(self, dim: int, dim_out: int, mult: int = 4) -> None:
-        super().__init__()
-        inner_dim = int(dim * mult)
-        project_in = GELUApprox(dim, inner_dim)
-
-        self.net = torch.nn.Sequential(project_in, torch.nn.Identity(), torch.nn.Linear(inner_dim, dim_out))
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.net(x)
+import torch
+
+from ltx_core.model.transformer.gelu_approx import GELUApprox
+
+
+class FeedForward(torch.nn.Module):
+    def __init__(self, dim: int, dim_out: int, mult: int = 4) -> None:
+        super().__init__()
+        inner_dim = int(dim * mult)
+        project_in = GELUApprox(dim, inner_dim)
+
+        self.net = torch.nn.Sequential(project_in, torch.nn.Identity(), torch.nn.Linear(inner_dim, dim_out))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.net(x)

packages/ltx-core/src/ltx_core/model/transformer/gelu_approx.py

@@ -1,13 +1,10 @@
-# Copyright (c) 2025 Lightricks. All rights reserved.
-# Created by Andrew Kvochko
-
-import torch
-
-
-class GELUApprox(torch.nn.Module):
-    def __init__(self, dim_in: int, dim_out: int) -> None:
-        super().__init__()
-        self.proj = torch.nn.Linear(dim_in, dim_out)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return torch.nn.functional.gelu(self.proj(x), approximate="tanh")
+import torch
+
+
+class GELUApprox(torch.nn.Module):
+    def __init__(self, dim_in: int, dim_out: int) -> None:
+        super().__init__()
+        self.proj = torch.nn.Linear(dim_in, dim_out)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.nn.functional.gelu(self.proj(x), approximate="tanh")

packages/ltx-core/src/ltx_core/model/transformer/modality.py

@@ -1,20 +1,23 @@
-# Copyright (c) 2025 Lightricks. All rights reserved.
-# Created by Andrew Kvochko
-
-from dataclasses import dataclass
-
-import torch
-
-
-@dataclass(frozen=True)
-class Modality:
-    latent: (
-        torch.Tensor
-    )  # Shape: (B, T, D) where B is the batch size, T is the number of tokens, and D is input dimension
-    timesteps: torch.Tensor  # Shape: (B, T) where T is the number of timesteps
-    positions: (
-        torch.Tensor
-    )  # Shape: (B, 3, T) for video, where 3 is the number of dimensions and T is the number of tokens
-    context: torch.Tensor
-    enabled: bool = True
-    context_mask: torch.Tensor | None = None
+from dataclasses import dataclass
+
+import torch
+
+
+@dataclass(frozen=True)
+class Modality:
+    """
+    Input data for a single modality (video or audio) in the transformer.
+    Bundles the latent tokens, timestep embeddings, positional information,
+    and text conditioning context for processing by the diffusion transformer.
+    """
+
+    latent: (
+        torch.Tensor
+    )  # Shape: (B, T, D) where B is the batch size, T is the number of tokens, and D is input dimension
+    timesteps: torch.Tensor  # Shape: (B, T) where T is the number of timesteps
+    positions: (
+        torch.Tensor
+    )  # Shape: (B, 3, T) for video, where 3 is the number of dimensions and T is the number of tokens
+    context: torch.Tensor
+    enabled: bool = True
+    context_mask: torch.Tensor | None = None