diff --git a/app.py b/app.py
new file mode 100644
index 0000000000000000000000000000000000000000..a8451f31e8a28fad7a4c90a28fb9e6d682a35965
--- /dev/null
+++ b/app.py
@@ -0,0 +1,726 @@
+import sys
+from pathlib import Path
+import uuid
+import tempfile
+
+# 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 flash_attn_interface
+import time
+import gradio as gr
+import numpy as np
+import random
+import torch
+from typing import Optional
+from pathlib import Path
+from huggingface_hub import hf_hub_download, snapshot_download
+from ltx_pipelines.distilled import DistilledPipeline
+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,
+)
+
+
+MAX_SEED = np.iinfo(np.int32).max
+# Import from public LTX-2 package
+# Install with: pip install git+https://github.com/Lightricks/LTX-2.git
+from ltx_pipelines.utils import ModelLedger
+from ltx_pipelines.utils.helpers import generate_enhanced_prompt
+
+# HuggingFace Hub defaults
+DEFAULT_REPO_ID = "Lightricks/LTX-2"
+DEFAULT_GEMMA_REPO_ID = "unsloth/gemma-3-12b-it-qat-bnb-4bit"
+DEFAULT_CHECKPOINT_FILENAME = "ltx-2-19b-dev.safetensors"
+
+
+def get_hub_or_local_checkpoint(repo_id: str, filename: str):
+ """Download from HuggingFace Hub."""
+ print(f"Downloading {filename} from {repo_id}...")
+ ckpt_path = hf_hub_download(repo_id=repo_id, filename=filename)
+ print(f"Downloaded to {ckpt_path}")
+ return ckpt_path
+
+def download_gemma_model(repo_id: str):
+ """Download the full Gemma model directory."""
+ print(f"Downloading Gemma model from {repo_id}...")
+ local_dir = snapshot_download(repo_id=repo_id)
+ print(f"Gemma model downloaded to {local_dir}")
+ return local_dir
+
+# Initialize model ledger and text encoder at startup (load once, keep in memory)
+print("=" * 80)
+print("Loading Gemma Text Encoder...")
+print("=" * 80)
+
+checkpoint_path = get_hub_or_local_checkpoint(DEFAULT_REPO_ID, DEFAULT_CHECKPOINT_FILENAME)
+gemma_local_path = download_gemma_model(DEFAULT_GEMMA_REPO_ID)
+device = "cuda"
+
+print(f"Initializing text encoder with:")
+print(f" checkpoint_path={checkpoint_path}")
+print(f" gemma_root={gemma_local_path}")
+print(f" device={device}")
+
+
+model_ledger = ModelLedger(
+ dtype=torch.bfloat16,
+ device=device,
+ checkpoint_path=checkpoint_path,
+ gemma_root_path=DEFAULT_GEMMA_REPO_ID,
+ local_files_only=False
+)
+
+
+# Load text encoder once and keep it in memory
+text_encoder = model_ledger.text_encoder()
+
+print("=" * 80)
+print("Text encoder loaded and ready!")
+print("=" * 80)
+
+def encode_text_simple(text_encoder, prompt: str):
+ """Simple text encoding without using pipeline_utils."""
+ v_context, a_context, _ = text_encoder(prompt)
+ return v_context, a_context
+
+@spaces.GPU()
+def encode_prompt(
+ prompt: str,
+ enhance_prompt: bool = True,
+ input_image=None, # this is now filepath (string) or None
+ seed: int = 42,
+ negative_prompt: str = ""
+):
+ start_time = time.time()
+ try:
+ final_prompt = prompt
+ if enhance_prompt:
+ final_prompt = generate_enhanced_prompt(
+ text_encoder=text_encoder,
+ prompt=prompt,
+ image_path=input_image if input_image is not None else None,
+ seed=seed,
+ )
+
+ with torch.inference_mode():
+ video_context, audio_context = encode_text_simple(text_encoder, final_prompt)
+
+ video_context_negative = None
+ audio_context_negative = None
+ if negative_prompt:
+ video_context_negative, audio_context_negative = encode_text_simple(text_encoder, negative_prompt)
+
+ # IMPORTANT: return tensors directly (no torch.save)
+ embedding_data = {
+ "video_context": video_context.detach().cpu(),
+ "audio_context": audio_context.detach().cpu(),
+ "prompt": final_prompt,
+ "original_prompt": prompt,
+ }
+ if video_context_negative is not None:
+ embedding_data["video_context_negative"] = video_context_negative
+ embedding_data["audio_context_negative"] = audio_context_negative
+ embedding_data["negative_prompt"] = negative_prompt
+
+ elapsed_time = time.time() - start_time
+ if torch.cuda.is_available():
+ allocated = torch.cuda.memory_allocated() / 1024**3
+ peak = torch.cuda.max_memory_allocated() / 1024**3
+ status = f"β Encoded in {elapsed_time:.2f}s | VRAM: {allocated:.2f}GB allocated, {peak:.2f}GB peak"
+ else:
+ status = f"β Encoded in {elapsed_time:.2f}s (CPU mode)"
+
+ return embedding_data, final_prompt, status
+
+ except Exception as e:
+ import traceback
+ error_msg = f"Error: {str(e)}\n{traceback.format_exc()}"
+ print(error_msg)
+ return None, prompt, error_msg
+
+
+# Default prompt from docstring example
+DEFAULT_PROMPT = "An astronaut hatches from a fragile egg on the surface of the Moon, the shell cracking and peeling apart in gentle low-gravity motion. Fine lunar dust lifts and drifts outward with each movement, floating in slow arcs before settling back onto the ground. The astronaut pushes free in a deliberate, weightless motion, small fragments of the egg tumbling and spinning through the air. In the background, the deep darkness of space subtly shifts as stars glide with the camera's movement, emphasizing vast depth and scale. The camera performs a smooth, cinematic slow push-in, with natural parallax between the foreground dust, the astronaut, and the distant starfield. Ultra-realistic detail, physically accurate low-gravity motion, cinematic lighting, and a breath-taking, movie-like shot."
+
+# HuggingFace Hub defaults
+DEFAULT_REPO_ID = "Lightricks/LTX-2"
+DEFAULT_CHECKPOINT_FILENAME = "ltx-2-19b-dev.safetensors"
+DEFAULT_DISTILLED_LORA_FILENAME = "ltx-2-19b-distilled-lora-384.safetensors"
+DEFAULT_SPATIAL_UPSAMPLER_FILENAME = "ltx-2-spatial-upscaler-x2-1.0.safetensors"
+
+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...")
+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)
+
+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}")
+
+
+# Load distilled LoRA as a regular LoRA
+loras = [
+ LoraPathStrengthAndSDOps(
+ path=distilled_lora_path,
+ strength=DEFAULT_LORA_STRENGTH,
+ 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(
+ device=torch.device("cuda"),
+ checkpoint_path=checkpoint_path,
+ spatial_upsampler_path=spatial_upsampler_path,
+ gemma_root=None, # No text encoder in this space
+ loras=loras,
+ fp8transformer=False,
+ local_files_only=False,
+)
+
+pipeline._video_encoder = pipeline.model_ledger.video_encoder()
+pipeline._transformer = pipeline.model_ledger.transformer()
+# pipeline.device = torch.device("cuda")
+# pipeline.model_ledger.device = torch.device("cuda")
+
+
+print("=" * 80)
+print("Pipeline fully loaded and ready!")
+print("=" * 80)
+
+def get_duration(
+ input_image,
+ prompt,
+ duration,
+ enhance_prompt,
+ seed,
+ randomize_seed,
+ height,
+ width,
+ progress
+):
+ if duration <= 5:
+ return 80
+ else:
+ return 120
+
+class RadioAnimated(gr.HTML):
+ """
+ Animated segmented radio (like iOS pill selector).
+ Outputs: selected option string, e.g. "768x512"
+ """
+ def __init__(self, choices, value=None, **kwargs):
+ if not choices or len(choices) < 2:
+ raise ValueError("RadioAnimated requires at least 2 choices.")
+ if value is None:
+ value = choices[0]
+
+ uid = uuid.uuid4().hex[:8] # unique per instance
+ group_name = f"ra-{uid}"
+
+ inputs_html = "\n".join(
+ f"""
+
+
+ """
+ for i, c in enumerate(choices)
+ )
+
+ # NOTE: use classes instead of duplicate IDs
+ html_template = f"""
+
+ """
+
+ js_on_load = r"""
+ (() => {
+ const wrap = element.querySelector('.ra-wrap');
+ const inner = element.querySelector('.ra-inner');
+ const highlight = element.querySelector('.ra-highlight');
+ const inputs = Array.from(element.querySelectorAll('.ra-input'));
+
+ if (!inputs.length) return;
+
+ const choices = inputs.map(i => i.value);
+
+ function setHighlightByIndex(idx) {
+ const n = choices.length;
+ const pct = 100 / n;
+ highlight.style.width = `calc(${pct}% - 6px)`;
+ highlight.style.transform = `translateX(${idx * 100}%)`;
+ }
+
+ function setCheckedByValue(val, shouldTrigger=false) {
+ const idx = Math.max(0, choices.indexOf(val));
+ inputs.forEach((inp, i) => { inp.checked = (i === idx); });
+ setHighlightByIndex(idx);
+
+ props.value = choices[idx];
+ if (shouldTrigger) trigger('change', props.value);
+ }
+
+ // Init from props.value
+ setCheckedByValue(props.value ?? choices[0], false);
+
+ // Input handlers
+ inputs.forEach((inp) => {
+ inp.addEventListener('change', () => {
+ setCheckedByValue(inp.value, true);
+ });
+ });
+ })();
+ """
+
+ super().__init__(
+ value=value,
+ html_template=html_template,
+ js_on_load=js_on_load,
+ **kwargs
+ )
+
+def generate_video_example(input_image, prompt, duration, progress=gr.Progress(track_tqdm=True)):
+ output_video, seed = generate_video(input_image, prompt, 5, True, 42, True, DEFAULT_1_STAGE_HEIGHT, DEFAULT_1_STAGE_WIDTH, progress)
+
+ return output_video
+
+@spaces.GPU(duration=get_duration)
+def generate_video(
+ input_image,
+ prompt: str,
+ duration: float,
+ enhance_prompt: bool = True,
+ seed: int = 42,
+ randomize_seed: bool = True,
+ height: int = DEFAULT_1_STAGE_HEIGHT,
+ width: int = DEFAULT_1_STAGE_WIDTH,
+ progress=gr.Progress(track_tqdm=True),
+):
+ """
+ Generate a short cinematic video from a text prompt and optional input image using the LTX-2 distilled pipeline.
+ Args:
+ input_image: Optional input image for image-to-video. If provided, it is injected at frame 0 to guide motion.
+ prompt: Text description of the scene, motion, and cinematic style to generate.
+ duration: Desired video length in seconds. Converted to frames using a fixed 24 FPS rate.
+ enhance_prompt: Whether to enhance the prompt using the prompt enhancer before encoding.
+ seed: Base random seed for reproducibility (ignored if randomize_seed is True).
+ randomize_seed: If True, a random seed is generated for each run.
+ height: Output video height in pixels.
+ width: Output video width in pixels.
+ progress: Gradio progress tracker.
+ Returns:
+ A tuple of:
+ - output_path: Path to the generated MP4 video file.
+ - seed: The seed used for generation.
+ Notes:
+ - Uses a fixed frame rate of 24 FPS.
+ - Prompt embeddings are generated externally to avoid reloading the text encoder.
+ - GPU cache is cleared after generation to reduce VRAM pressure.
+ - If an input image is provided, it is temporarily saved to disk for processing.
+ """
+ try:
+ # 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 # +1 to ensure we meet the duration
+
+ with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
+ output_path = tmpfile.name
+
+ # Handle image input
+ images = []
+ temp_image_path = None # Initialize to None
+
+ images = []
+ if input_image is not None:
+ images = [(input_image, 0, 1.0)] # input_image is already a path
+
+ # Prepare image for upload if it exists
+ image_input = None
+
+
+ embeddings, final_prompt, status = encode_prompt(
+ prompt=prompt,
+ enhance_prompt=enhance_prompt,
+ input_image=input_image,
+ seed=current_seed,
+ negative_prompt="",
+ )
+
+ video_context = embeddings["video_context"].to("cuda", non_blocking=True)
+ audio_context = embeddings["audio_context"].to("cuda", non_blocking=True)
+ print("β Embeddings loaded successfully")
+
+ # free prompt enhancer / encoder temps ASAP
+ del embeddings, final_prompt, status
+ torch.cuda.empty_cache()
+
+ # Run inference - progress automatically tracks tqdm from pipeline
+ 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,
+ )
+ del video_context, audio_context
+ torch.cuda.empty_cache()
+ print("successful generation")
+
+ 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
+
+
+def apply_resolution(resolution: str):
+ w, h = resolution.split("x")
+ return int(w), int(h)
+
+def apply_duration(duration: str):
+ duration_s = int(duration[:-1])
+ return duration_s
+
+css = """
+ #col-container {
+ margin: 0 auto;
+ max-width: 1600px;
+ }
+ #modal-container {
+ width: 100vw; /* Take full viewport width */
+ height: 100vh; /* Take full viewport height (optional) */
+ display: flex;
+ justify-content: center; /* Center content horizontally */
+ align-items: center; /* Center content vertically if desired */
+ }
+ #modal-content {
+ width: 100%;
+ max-width: 700px; /* Limit content width */
+ margin: 0 auto;
+ border-radius: 8px;
+ padding: 1.5rem;
+ }
+ #step-column {
+ padding: 10px;
+ border-radius: 8px;
+ box-shadow: var(--card-shadow);
+ margin: 10px;
+ }
+ #col-showcase {
+ margin: 0 auto;
+ max-width: 1100px;
+ }
+ .button-gradient {
+ background: linear-gradient(45deg, rgb(255, 65, 108), rgb(255, 75, 43), rgb(255, 155, 0), rgb(255, 65, 108)) 0% 0% / 400% 400%;
+ border: none;
+ padding: 14px 28px;
+ font-size: 16px;
+ font-weight: bold;
+ color: white;
+ border-radius: 10px;
+ cursor: pointer;
+ transition: 0.3s ease-in-out;
+ animation: 2s linear 0s infinite normal none running gradientAnimation;
+ box-shadow: rgba(255, 65, 108, 0.6) 0px 4px 10px;
+ }
+ .toggle-container {
+ display: inline-flex;
+ background-color: #ffd6ff; /* light pink background */
+ border-radius: 9999px;
+ padding: 4px;
+ position: relative;
+ width: fit-content;
+ font-family: sans-serif;
+ }
+ .toggle-container input[type="radio"] {
+ display: none;
+ }
+ .toggle-container label {
+ position: relative;
+ z-index: 2;
+ flex: 1;
+ text-align: center;
+ font-weight: 700;
+ color: #4b2ab5; /* dark purple text for unselected */
+ padding: 6px 22px;
+ border-radius: 9999px;
+ cursor: pointer;
+ transition: color 0.25s ease;
+ }
+ /* Moving highlight */
+ .toggle-highlight {
+ position: absolute;
+ top: 4px;
+ left: 4px;
+ width: calc(50% - 4px);
+ height: calc(100% - 8px);
+ background-color: #4b2ab5; /* dark purple background */
+ border-radius: 9999px;
+ transition: transform 0.25s ease;
+ z-index: 1;
+ }
+ /* When "True" is checked */
+ #true:checked ~ label[for="true"] {
+ color: #ffd6ff; /* light pink text */
+ }
+ /* When "False" is checked */
+ #false:checked ~ label[for="false"] {
+ color: #ffd6ff; /* light pink text */
+ }
+ /* Move highlight to right side when False is checked */
+ #false:checked ~ .toggle-highlight {
+ transform: translateX(100%);
+ }
+ """
+
+css += """
+ /* ---- radioanimated ---- */
+ .ra-wrap{
+ width: fit-content;
+ }
+ .ra-inner{
+ position: relative;
+ display: inline-flex;
+ align-items: center;
+ gap: 0;
+ padding: 6px;
+ background: #0b0b0b;
+ border-radius: 9999px;
+ overflow: hidden;
+ user-select: none;
+ }
+ .ra-input{
+ display: none;
+ }
+ .ra-label{
+ position: relative;
+ z-index: 2;
+ padding: 10px 18px;
+ font-family: ui-sans-serif, system-ui, -apple-system, Segoe UI, Roboto, Arial;
+ font-size: 14px;
+ font-weight: 600;
+ color: rgba(255,255,255,0.7);
+ cursor: pointer;
+ transition: color 180ms ease;
+ white-space: nowrap;
+ }
+ .ra-highlight{
+ position: absolute;
+ z-index: 1;
+ top: 6px;
+ left: 6px;
+ height: calc(100% - 12px);
+ border-radius: 9999px;
+ background: #8bff97; /* green knob */
+ transition: transform 200ms ease, width 200ms ease;
+ }
+ /* selected label becomes darker like your screenshot */
+ .ra-input:checked + .ra-label{
+ color: rgba(0,0,0,0.75);
+ }
+ """
+
+
+with gr.Blocks(title="LTX-2 Video Distilled π₯π") as demo:
+ gr.HTML(
+ """
+
+
+ LTX-2 Distilled DiT-based audio-video foundation model
+
+
+ [model]
+
+
+
+
+ Using FA3 and Gemma 3 12B 4bit Quantisation for Faster Inference
+
+
+
+
HF Space by:
+
+
+
+
+ """
+ )
+ with gr.Column(elem_id="col-container"):
+ with gr.Row():
+ with gr.Column(elem_id="step-column"):
+
+ input_image = gr.Image(
+ label="Input Image (Optional)",
+ type="filepath", # <-- was "pil"
+ height=512
+ )
+
+ prompt = gr.Textbox(
+ label="Prompt",
+ value="Make this image come alive with cinematic motion, smooth animation",
+ lines=3,
+ max_lines=3,
+ placeholder="Describe the motion and animation you want..."
+ )
+
+ enhance_prompt = gr.Checkbox(
+ label="Enhance Prompt",
+ value=True,
+ visible=False
+ )
+
+ with gr.Accordion("Advanced Settings", open=False, visible=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.Column(elem_id="step-column"):
+ output_video = gr.Video(label="Generated Video", autoplay=True, height=512)
+
+ with gr.Row():
+
+ with gr.Column():
+ radioanimated_duration = RadioAnimated(
+ choices=["3s", "5s", "10s"],
+ value="3s",
+ elem_id="radioanimated_duration"
+ )
+
+ duration = gr.Slider(
+ label="Duration (seconds)",
+ minimum=1.0,
+ maximum=10.0,
+ value=3.0,
+ step=0.1,
+ visible=False
+ )
+
+ with gr.Column():
+ radioanimated_resolution = RadioAnimated(
+ choices=["768x512", "512x512", "512x768"],
+ value=f"{DEFAULT_1_STAGE_WIDTH}x{DEFAULT_1_STAGE_HEIGHT}",
+ elem_id="radioanimated_resolution"
+ )
+
+ width = gr.Number(label="Width", value=DEFAULT_1_STAGE_WIDTH, precision=0, visible=False)
+ height = gr.Number(label="Height", value=DEFAULT_1_STAGE_HEIGHT, precision=0, visible=False)
+
+
+ generate_btn = gr.Button("π€© Generate Video", variant="primary", elem_classes="button-gradient")
+
+
+ radioanimated_duration.change(
+ fn=apply_duration,
+ inputs=radioanimated_duration,
+ outputs=[duration],
+ api_visibility="private"
+ )
+ radioanimated_resolution.change(
+ fn=apply_resolution,
+ inputs=radioanimated_resolution,
+ outputs=[width, height],
+ api_visibility="private"
+ )
+
+ generate_btn.click(
+ fn=generate_video,
+ inputs=[
+ input_image,
+ prompt,
+ duration,
+ enhance_prompt,
+ seed,
+ randomize_seed,
+ height,
+ width,
+ ],
+ outputs=[output_video,seed]
+ )
+
+ # Add example
+ gr.Examples(
+ examples=[
+ [
+ "supergirl.png",
+ "A fuzzy puppet superhero character resembling a female puppet with blonde hair and a blue superhero suit stands inside an icy cave made of frozen walls and icicles, she looks panicked and frantic, rapidly turning her head left and right and scanning the cave while waving her arms and shouting angrily and desperately, mouthing the words βwhere the hell is my dog,β her movements exaggerated and puppet-like with high energy and urgency, suddenly a second puppet dog bursts into frame from the side, jumping up excitedly and tackling her affectionately while licking her face repeatedly, she freezes in surprise and then breaks into relief and laughter as the dog continues licking her, the scene feels chaotic, comedic, and emotional with expressive puppet reactions, cinematic lighting, smooth camera motion, shallow depth of field, and high-quality puppet-style animation"
+ ],
+ [
+ "highland.png",
+ "Realistic POV selfie-style video in a snowy, foggy field. Two shaggy Highland cows with long curved horns stand ahead. The camera is handheld and slightly shaky. The woman filming talks nervously and excitedly in a vlog tone: \"Oh my god guysβ¦ look how big those horns areβ¦ Iβm kinda scared.\" The cow on the left walks toward the camera in a cute, bouncy, hopping way, curious and gentle. Snow crunches under its hooves, breath visible in the cold air. The horns look massive from the POV. As the cow gets very close, its wet nose with slight dripping fills part of the frame. She laughs nervously but reaches out and pets the cow. The cow makes deep, soft, interesting mooing and snorting sounds, calm and friendly. Ultra-realistic, natural lighting, immersive audio, documentary-style realism.",
+ ],
+ [
+ "wednesday.png",
+ "A cinematic close-up of Wednesday Addams frozen mid-dance on a dark, blue-lit ballroom floor as students move indistinctly behind her, their footsteps and muffled music reduced to a distant, underwater thrum; the audio foregrounds her steady breathing and the faint rustle of fabric as she slowly raises one arm, never breaking eye contact with the camera, then after a deliberately long silence she speaks in a flat, dry, perfectly controlled voice, βI donβt danceβ¦ I vibe code,β each word crisp and unemotional, followed by an abrupt cutoff of her voice as the background sound swells slightly, reinforcing the deadpan humor, with precise lip sync, minimal facial movement, stark gothic lighting, and cinematic realism.",
+ ],
+ [
+ "astronaut.png",
+ "An astronaut hatches from a fragile egg on the surface of the Moon, the shell cracking and peeling apart in gentle low-gravity motion. Fine lunar dust lifts and drifts outward with each movement, floating in slow arcs before settling back onto the ground. The astronaut pushes free in a deliberate, weightless motion, small fragments of the egg tumbling and spinning through the air. In the background, the deep darkness of space subtly shifts as stars glide with the camera's movement, emphasizing vast depth and scale. The camera performs a smooth, cinematic slow push-in, with natural parallax between the foreground dust, the astronaut, and the distant starfield. Ultra-realistic detail, physically accurate low-gravity motion, cinematic lighting, and a breath-taking, movie-like shot.",
+ ]
+
+ ],
+ fn=generate_video_example,
+ inputs=[input_image, prompt],
+ outputs = [output_video],
+ label="Example",
+ cache_examples=True,
+ )
+
+
+
+if __name__ == "__main__":
+ demo.launch(ssr_mode=False, mcp_server=True, css=css)
\ No newline at end of file
diff --git a/ltx2_two_stage.py b/ltx2_two_stage.py
new file mode 100644
index 0000000000000000000000000000000000000000..d365cc6190bb4d48592c484262d2d4c3151237c9
--- /dev/null
+++ b/ltx2_two_stage.py
@@ -0,0 +1,84 @@
+"""
+python ltx2_two_stage.py \
+ --image "astronaut.jpg" 0 1.0 \
+ --prompt="An astronaut hatches from a fragile egg on the surface of the Moon, the shell cracking and peeling apart in gentle low-gravity motion. Fine lunar dust lifts and drifts outward with each movement, floating in slow arcs before settling back onto the ground. The astronaut pushes free in a deliberate, weightless motion, small fragments of the egg tumbling and spinning through the air. In the background, the deep darkness of space subtly shifts as stars glide with the camera's movement, emphasizing vast depth and scale. The camera performs a smooth, cinematic slow push-in, with natural parallax between the foreground dust, the astronaut, and the distant starfield. Ultra-realistic detail, physically accurate low-gravity motion, cinematic lighting, and a breath-taking, movie-like shot." \
+ --output_path="t2v_2.mp4" \
+ --gemma_root="google/gemma-3-12b-it-qat-q4_0-unquantized" \
+ --checkpoint_path="rc1/ltx-2-19b-dev-rc1.safetensors" \
+ --distilled_lora_path "rc1/ltx-2-19b-distilled-lora-384-rc1.safetensors" \
+ --spatial_upsampler_path "rc1/ltx-2-spatial-upscaler-x2-1.0-rc1.safetensors"
+
+"""
+
+from huggingface_hub import hf_hub_download
+from typing import Optional
+from ltx_pipelines import utils
+from ltx_pipelines.constants import DEFAULT_LORA_STRENGTH
+from ltx_core.loader.primitives import LoraPathStrengthAndSDOps
+from ltx_core.loader.sd_ops import LTXV_LORA_COMFY_RENAMING_MAP
+from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
+from ltx_core.tiling import TilingConfig
+
+
+def get_hub_or_local_checkpoint(repo_id: Optional[str] = None, filename: Optional[str] = None):
+ 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.")
+ ckpt_path = hf_hub_download(repo_id=repo_id, filename=filename)
+ else:
+ ckpt_path = filename
+
+ return ckpt_path
+
+
+def default_2_stage_arg_parser_mod():
+ parser = utils.default_2_stage_arg_parser()
+ parser.add_argument("--local_files_only", action="store_true")
+ parser.add_argument("--checkpoint_id", type=str, default="diffusers-internal-dev/new-ltx-model")
+ return parser
+
+
+def main() -> None:
+ parser = default_2_stage_arg_parser_mod()
+ args = parser.parse_args()
+
+ checkpoint_path = get_hub_or_local_checkpoint(args.checkpoint_id, args.checkpoint_path)
+ distilled_lora_path = get_hub_or_local_checkpoint(args.checkpoint_id, args.distilled_lora_path)
+ spatial_upsampler_path = get_hub_or_local_checkpoint(args.checkpoint_id, args.spatial_upsampler_path)
+
+ lora_strengths = (args.lora_strength + [DEFAULT_LORA_STRENGTH] * len(args.lora))[: len(args.lora)]
+ loras = [
+ LoraPathStrengthAndSDOps(lora, strength, LTXV_LORA_COMFY_RENAMING_MAP)
+ for lora, strength in zip(args.lora, lora_strengths, strict=True)
+ ]
+ pipeline = TI2VidTwoStagesPipeline(
+ checkpoint_path=checkpoint_path,
+ distilled_lora_path=distilled_lora_path,
+ distilled_lora_strength=args.distilled_lora_strength,
+ spatial_upsampler_path=spatial_upsampler_path,
+ gemma_root=args.gemma_root,
+ loras=loras,
+ fp8transformer=args.enable_fp8,
+ local_files_only=args.local_files_only
+ )
+ pipeline(
+ prompt=args.prompt,
+ negative_prompt=args.negative_prompt,
+ output_path=args.output_path,
+ seed=args.seed,
+ height=args.height,
+ width=args.width,
+ num_frames=args.num_frames,
+ frame_rate=args.frame_rate,
+ num_inference_steps=args.num_inference_steps,
+ cfg_guidance_scale=args.cfg_guidance_scale,
+ images=args.images,
+ tiling_config=TilingConfig.default(),
+ )
+
+
+if __name__ == "__main__":
+ main()
diff --git a/packages/ltx-core/README.md b/packages/ltx-core/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..5835b132568bab56178763c41a7645b51b588e5e
--- /dev/null
+++ b/packages/ltx-core/README.md
@@ -0,0 +1,280 @@
+# 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
diff --git a/packages/ltx-core/pyproject.toml b/packages/ltx-core/pyproject.toml
new file mode 100644
index 0000000000000000000000000000000000000000..7ed993560d020fb0c4b5eaf98c387ff306573f23
--- /dev/null
+++ b/packages/ltx-core/pyproject.toml
@@ -0,0 +1,37 @@
+[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",
+]
diff --git a/packages/ltx-core/src/ltx_core/__init__.py b/packages/ltx-core/src/ltx_core/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/packages/ltx-core/src/ltx_core/components/__init__.py b/packages/ltx-core/src/ltx_core/components/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..038cbfc63d328cfaf451e52da3305ed9b09cbd70
--- /dev/null
+++ b/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)
+"""
diff --git a/packages/ltx-core/src/ltx_core/components/diffusion_steps.py b/packages/ltx-core/src/ltx_core/components/diffusion_steps.py
new file mode 100644
index 0000000000000000000000000000000000000000..44ad721df71feedab977c20bbd84292889bfbef1
--- /dev/null
+++ b/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)
diff --git a/packages/ltx-core/src/ltx_core/components/guiders.py b/packages/ltx-core/src/ltx_core/components/guiders.py
new file mode 100644
index 0000000000000000000000000000000000000000..8cdb1107d7a72ce258e788a3c860fcf9c3c0325f
--- /dev/null
+++ b/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
diff --git a/packages/ltx-core/src/ltx_core/components/noisers.py b/packages/ltx-core/src/ltx_core/components/noisers.py
new file mode 100644
index 0000000000000000000000000000000000000000..f93c41439658af49468bd29f0077c0881d624d15
--- /dev/null
+++ b/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),
+ )
diff --git a/packages/ltx-core/src/ltx_core/components/patchifiers.py b/packages/ltx-core/src/ltx_core/components/patchifiers.py
new file mode 100644
index 0000000000000000000000000000000000000000..9393dc32ffbb9704580c4a6358607b28f11fc19d
--- /dev/null
+++ b/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)
diff --git a/packages/ltx-core/src/ltx_core/components/protocols.py b/packages/ltx-core/src/ltx_core/components/protocols.py
new file mode 100644
index 0000000000000000000000000000000000000000..69667e048b9db35dd3d7c47393156f58eed1707d
--- /dev/null
+++ b/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: ...
diff --git a/packages/ltx-core/src/ltx_core/components/schedulers.py b/packages/ltx-core/src/ltx_core/components/schedulers.py
new file mode 100644
index 0000000000000000000000000000000000000000..b7f7725b695a234dcc0946e6f6a51d7b9f4c95ad
--- /dev/null
+++ b/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)
diff --git a/packages/ltx-core/src/ltx_core/conditioning/__init__.py b/packages/ltx-core/src/ltx_core/conditioning/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..6c0a17afabe6c52adfa0646d5c491a1ef8b80295
--- /dev/null
+++ b/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",
+]
diff --git a/packages/ltx-core/src/ltx_core/conditioning/exceptions.py b/packages/ltx-core/src/ltx_core/conditioning/exceptions.py
new file mode 100644
index 0000000000000000000000000000000000000000..adfdd131b77381330be4664d45e8245b138235e4
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/conditioning/exceptions.py
@@ -0,0 +1,4 @@
+class ConditioningError(Exception):
+ """
+ Class for conditioning-related errors.
+ """
diff --git a/packages/ltx-core/src/ltx_core/conditioning/item.py b/packages/ltx-core/src/ltx_core/conditioning/item.py
new file mode 100644
index 0000000000000000000000000000000000000000..3d68cbe1e560dd61ed320caee6f3f59a4dad2d95
--- /dev/null
+++ b/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.
+ """
+ ...
diff --git a/packages/ltx-core/src/ltx_core/conditioning/types/__init__.py b/packages/ltx-core/src/ltx_core/conditioning/types/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..f81f7a97cfc93c24d96fee4f99779575fb5e5e5f
--- /dev/null
+++ b/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",
+]
diff --git a/packages/ltx-core/src/ltx_core/conditioning/types/keyframe_cond.py b/packages/ltx-core/src/ltx_core/conditioning/types/keyframe_cond.py
new file mode 100644
index 0000000000000000000000000000000000000000..f5718ffd95efa3d9734577beadeaa28b8798e360
--- /dev/null
+++ b/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),
+ )
diff --git a/packages/ltx-core/src/ltx_core/conditioning/types/latent_cond.py b/packages/ltx-core/src/ltx_core/conditioning/types/latent_cond.py
new file mode 100644
index 0000000000000000000000000000000000000000..0257ca01531641368a5af518b4fa9ae8db22a5f6
--- /dev/null
+++ b/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
diff --git a/packages/ltx-core/src/ltx_core/guidance/__init__.py b/packages/ltx-core/src/ltx_core/guidance/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..a3141f98281b2add78a4192be2a1a547ac8e0607
--- /dev/null
+++ b/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",
+]
diff --git a/packages/ltx-core/src/ltx_core/guidance/perturbations.py b/packages/ltx-core/src/ltx_core/guidance/perturbations.py
new file mode 100644
index 0000000000000000000000000000000000000000..d5d17ae4298953c5dcd5eb6f11b60cbe3ebd7a57
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/guidance/perturbations.py
@@ -0,0 +1,79 @@
+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)])
diff --git a/packages/ltx-core/src/ltx_core/loader/__init__.py b/packages/ltx-core/src/ltx_core/loader/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..08bbb2ca7d651a89a56a7e67b83269cc07e93eb5
--- /dev/null
+++ b/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",
+]
diff --git a/packages/ltx-core/src/ltx_core/loader/fuse_loras.py b/packages/ltx-core/src/ltx_core/loader/fuse_loras.py
new file mode 100644
index 0000000000000000000000000000000000000000..4a2aaff1540b8d7a0db2e244aeea38e01f895497
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/loader/fuse_loras.py
@@ -0,0 +1,100 @@
+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)
diff --git a/packages/ltx-core/src/ltx_core/loader/kernels.py b/packages/ltx-core/src/ltx_core/loader/kernels.py
new file mode 100644
index 0000000000000000000000000000000000000000..ebeee59e56d537ed8f5e6034d6d1af330d7110c9
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/loader/kernels.py
@@ -0,0 +1,72 @@
+# 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)
diff --git a/packages/ltx-core/src/ltx_core/loader/module_ops.py b/packages/ltx-core/src/ltx_core/loader/module_ops.py
new file mode 100644
index 0000000000000000000000000000000000000000..edf49005f69da3c0bd1b6cea5e891d0cf44e79d0
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/loader/module_ops.py
@@ -0,0 +1,14 @@
+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]
diff --git a/packages/ltx-core/src/ltx_core/loader/primitives.py b/packages/ltx-core/src/ltx_core/loader/primitives.py
new file mode 100644
index 0000000000000000000000000000000000000000..65b3eab61b9859041df7e4859599d4f29fa75619
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/loader/primitives.py
@@ -0,0 +1,109 @@
+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
diff --git a/packages/ltx-core/src/ltx_core/loader/registry.py b/packages/ltx-core/src/ltx_core/loader/registry.py
new file mode 100644
index 0000000000000000000000000000000000000000..8f532ea72d34daa36df7930d44db79497753b33a
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/loader/registry.py
@@ -0,0 +1,84 @@
+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()
diff --git a/packages/ltx-core/src/ltx_core/loader/sd_ops.py b/packages/ltx-core/src/ltx_core/loader/sd_ops.py
new file mode 100644
index 0000000000000000000000000000000000000000..74281f06baebbe9e15dff984fc4b38aef38c2b95
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/loader/sd_ops.py
@@ -0,0 +1,127 @@
+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")
+)
diff --git a/packages/ltx-core/src/ltx_core/loader/sft_loader.py b/packages/ltx-core/src/ltx_core/loader/sft_loader.py
new file mode 100644
index 0000000000000000000000000000000000000000..08e0978f25b9c7c3d6c8937215b64c6f9ea587a3
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/loader/sft_loader.py
@@ -0,0 +1,63 @@
+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)
diff --git a/packages/ltx-core/src/ltx_core/loader/single_gpu_model_builder.py b/packages/ltx-core/src/ltx_core/loader/single_gpu_model_builder.py
new file mode 100644
index 0000000000000000000000000000000000000000..afd32d95d82d2a64c0bcab997d6e638d8abe8e82
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/loader/single_gpu_model_builder.py
@@ -0,0 +1,101 @@
+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):
+ """
+ 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)
+
+ 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)
diff --git a/packages/ltx-core/src/ltx_core/model/__init__.py b/packages/ltx-core/src/ltx_core/model/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..8800f4b1afe7ee36570be8b3637332283902dbc6
--- /dev/null
+++ b/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",
+]
diff --git a/packages/ltx-core/src/ltx_core/model/audio_vae/__init__.py b/packages/ltx-core/src/ltx_core/model/audio_vae/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..3bb385dfb9d7514a6a8508f0d50d9861677cf37e
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/audio_vae/__init__.py
@@ -0,0 +1,27 @@
+"""Audio VAE model components."""
+
+from ltx_core.model.audio_vae.audio_vae import AudioDecoder, AudioEncoder, decode_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",
+]
diff --git a/packages/ltx-core/src/ltx_core/model/audio_vae/attention.py b/packages/ltx-core/src/ltx_core/model/audio_vae/attention.py
new file mode 100644
index 0000000000000000000000000000000000000000..712f61d1c8c23e62624ad1500bd10c3790ff42f5
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/audio_vae/attention.py
@@ -0,0 +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}")
diff --git a/packages/ltx-core/src/ltx_core/model/audio_vae/audio_vae.py b/packages/ltx-core/src/ltx_core/model/audio_vae/audio_vae.py
new file mode 100644
index 0000000000000000000000000000000000000000..846dd2df590842a9bf48a1e13628ab6526e51f6c
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/audio_vae/audio_vae.py
@@ -0,0 +1,480 @@
+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
+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
+
+
+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
diff --git a/packages/ltx-core/src/ltx_core/model/audio_vae/causal_conv_2d.py b/packages/ltx-core/src/ltx_core/model/audio_vae/causal_conv_2d.py
new file mode 100644
index 0000000000000000000000000000000000000000..c85e7238843df041793df0be636f2a36f278f2d2
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/audio_vae/causal_conv_2d.py
@@ -0,0 +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,
+ )
diff --git a/packages/ltx-core/src/ltx_core/model/audio_vae/causality_axis.py b/packages/ltx-core/src/ltx_core/model/audio_vae/causality_axis.py
new file mode 100644
index 0000000000000000000000000000000000000000..abd634bbb03810882b69a2f5aa9d70c45e8c2afb
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/audio_vae/causality_axis.py
@@ -0,0 +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"
diff --git a/packages/ltx-core/src/ltx_core/model/audio_vae/downsample.py b/packages/ltx-core/src/ltx_core/model/audio_vae/downsample.py
new file mode 100644
index 0000000000000000000000000000000000000000..31f8e353a4051eed7f5882bf9cd548db8f4a4318
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/audio_vae/downsample.py
@@ -0,0 +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
diff --git a/packages/ltx-core/src/ltx_core/model/audio_vae/model_configurator.py b/packages/ltx-core/src/ltx_core/model/audio_vae/model_configurator.py
new file mode 100644
index 0000000000000000000000000000000000000000..e0f85c74bacc96ea568d7e642fc444d322111206
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/audio_vae/model_configurator.py
@@ -0,0 +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 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.")
+)
diff --git a/packages/ltx-core/src/ltx_core/model/audio_vae/ops.py b/packages/ltx-core/src/ltx_core/model/audio_vae/ops.py
new file mode 100644
index 0000000000000000000000000000000000000000..da8709186dcc4251f2be667af176c48d0252b75a
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/audio_vae/ops.py
@@ -0,0 +1,76 @@
+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)
diff --git a/packages/ltx-core/src/ltx_core/model/audio_vae/resnet.py b/packages/ltx-core/src/ltx_core/model/audio_vae/resnet.py
new file mode 100644
index 0000000000000000000000000000000000000000..5d5afaaa580ddc1bbca95b55294660f728fa0715
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/audio_vae/resnet.py
@@ -0,0 +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
diff --git a/packages/ltx-core/src/ltx_core/model/audio_vae/upsample.py b/packages/ltx-core/src/ltx_core/model/audio_vae/upsample.py
new file mode 100644
index 0000000000000000000000000000000000000000..01cb61437e01eb98483a5e1286d109aedb9345ef
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/audio_vae/upsample.py
@@ -0,0 +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
diff --git a/packages/ltx-core/src/ltx_core/model/audio_vae/vocoder.py b/packages/ltx-core/src/ltx_core/model/audio_vae/vocoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..9bdf9196e654d3f57e342e11ddea6c2ce5f580c7
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/audio_vae/vocoder.py
@@ -0,0 +1,123 @@
+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)
diff --git a/packages/ltx-core/src/ltx_core/model/common/__init__.py b/packages/ltx-core/src/ltx_core/model/common/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..2fc0d42348a38d4d73422be895db19df900a8f89
--- /dev/null
+++ b/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",
+]
diff --git a/packages/ltx-core/src/ltx_core/model/common/normalization.py b/packages/ltx-core/src/ltx_core/model/common/normalization.py
new file mode 100644
index 0000000000000000000000000000000000000000..ae23f2e245d7800192368fc223d49a40f9c62780
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/common/normalization.py
@@ -0,0 +1,59 @@
+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}")
diff --git a/packages/ltx-core/src/ltx_core/model/model_protocol.py b/packages/ltx-core/src/ltx_core/model/model_protocol.py
new file mode 100644
index 0000000000000000000000000000000000000000..191d33c0ed653d14e6dd313770254a725cee64d3
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/model_protocol.py
@@ -0,0 +1,10 @@
+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: ...
diff --git a/packages/ltx-core/src/ltx_core/model/transformer/__init__.py b/packages/ltx-core/src/ltx_core/model/transformer/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..bf2d70f1e014f66726ab91b5caa963f1893b39da
--- /dev/null
+++ b/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",
+]
diff --git a/packages/ltx-core/src/ltx_core/model/transformer/adaln.py b/packages/ltx-core/src/ltx_core/model/transformer/adaln.py
new file mode 100644
index 0000000000000000000000000000000000000000..97e27f62ece5e972e6f78ac454bfa095d7fc9bf0
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/transformer/adaln.py
@@ -0,0 +1,34 @@
+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
diff --git a/packages/ltx-core/src/ltx_core/model/transformer/attention.py b/packages/ltx-core/src/ltx_core/model/transformer/attention.py
new file mode 100644
index 0000000000000000000000000000000000000000..0fa21c52514467333a985e072f70d096038ae448
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/transformer/attention.py
@@ -0,0 +1,185 @@
+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
+
+import flash_attn_interface
+
+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:
+ 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)
diff --git a/packages/ltx-core/src/ltx_core/model/transformer/feed_forward.py b/packages/ltx-core/src/ltx_core/model/transformer/feed_forward.py
new file mode 100644
index 0000000000000000000000000000000000000000..fda3f2a6650a13d0e7b69ddaa331078af6306e1b
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/transformer/feed_forward.py
@@ -0,0 +1,15 @@
+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)
diff --git a/packages/ltx-core/src/ltx_core/model/transformer/gelu_approx.py b/packages/ltx-core/src/ltx_core/model/transformer/gelu_approx.py
new file mode 100644
index 0000000000000000000000000000000000000000..d86ee0a534b02054290754ea428c39d3692272b6
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/transformer/gelu_approx.py
@@ -0,0 +1,10 @@
+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")
diff --git a/packages/ltx-core/src/ltx_core/model/transformer/modality.py b/packages/ltx-core/src/ltx_core/model/transformer/modality.py
new file mode 100644
index 0000000000000000000000000000000000000000..5eaeb38b8304439ca2126e332db5bef0028aeee5
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/transformer/modality.py
@@ -0,0 +1,23 @@
+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
diff --git a/packages/ltx-core/src/ltx_core/model/transformer/model.py b/packages/ltx-core/src/ltx_core/model/transformer/model.py
new file mode 100644
index 0000000000000000000000000000000000000000..b906b62adc82138d11783fc1279aa01121853142
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/transformer/model.py
@@ -0,0 +1,468 @@
+from enum import Enum
+
+import torch
+
+from ltx_core.guidance.perturbations import BatchedPerturbationConfig
+from ltx_core.model.transformer.adaln import AdaLayerNormSingle
+from ltx_core.model.transformer.attention import AttentionCallable, AttentionFunction
+from ltx_core.model.transformer.modality import Modality
+from ltx_core.model.transformer.rope import LTXRopeType
+from ltx_core.model.transformer.text_projection import PixArtAlphaTextProjection
+from ltx_core.model.transformer.transformer import BasicAVTransformerBlock, TransformerConfig
+from ltx_core.model.transformer.transformer_args import (
+ MultiModalTransformerArgsPreprocessor,
+ TransformerArgs,
+ TransformerArgsPreprocessor,
+)
+from ltx_core.utils import to_denoised
+
+
+class LTXModelType(Enum):
+ AudioVideo = "ltx av model"
+ VideoOnly = "ltx video only model"
+ AudioOnly = "ltx audio only model"
+
+ def is_video_enabled(self) -> bool:
+ return self in (LTXModelType.AudioVideo, LTXModelType.VideoOnly)
+
+ def is_audio_enabled(self) -> bool:
+ return self in (LTXModelType.AudioVideo, LTXModelType.AudioOnly)
+
+
+class LTXModel(torch.nn.Module):
+ """
+ LTX model transformer implementation.
+ This class implements the transformer blocks for the LTX model.
+ """
+
+ def __init__( # noqa: PLR0913
+ self,
+ *,
+ model_type: LTXModelType = LTXModelType.AudioVideo,
+ num_attention_heads: int = 32,
+ attention_head_dim: int = 128,
+ in_channels: int = 128,
+ out_channels: int = 128,
+ num_layers: int = 48,
+ cross_attention_dim: int = 4096,
+ norm_eps: float = 1e-06,
+ attention_type: AttentionFunction | AttentionCallable = AttentionFunction.DEFAULT,
+ caption_channels: int = 3840,
+ positional_embedding_theta: float = 10000.0,
+ positional_embedding_max_pos: list[int] | None = None,
+ timestep_scale_multiplier: int = 1000,
+ use_middle_indices_grid: bool = True,
+ audio_num_attention_heads: int = 32,
+ audio_attention_head_dim: int = 64,
+ audio_in_channels: int = 128,
+ audio_out_channels: int = 128,
+ audio_cross_attention_dim: int = 2048,
+ audio_positional_embedding_max_pos: list[int] | None = None,
+ av_ca_timestep_scale_multiplier: int = 1,
+ rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
+ double_precision_rope: bool = False,
+ ):
+ super().__init__()
+ self._enable_gradient_checkpointing = False
+ self.use_middle_indices_grid = use_middle_indices_grid
+ self.rope_type = rope_type
+ self.double_precision_rope = double_precision_rope
+ self.timestep_scale_multiplier = timestep_scale_multiplier
+ self.positional_embedding_theta = positional_embedding_theta
+ self.model_type = model_type
+ cross_pe_max_pos = None
+ if model_type.is_video_enabled():
+ if positional_embedding_max_pos is None:
+ positional_embedding_max_pos = [20, 2048, 2048]
+ self.positional_embedding_max_pos = positional_embedding_max_pos
+ self.num_attention_heads = num_attention_heads
+ self.inner_dim = num_attention_heads * attention_head_dim
+ self._init_video(
+ in_channels=in_channels,
+ out_channels=out_channels,
+ caption_channels=caption_channels,
+ norm_eps=norm_eps,
+ )
+
+ if model_type.is_audio_enabled():
+ if audio_positional_embedding_max_pos is None:
+ audio_positional_embedding_max_pos = [20]
+ self.audio_positional_embedding_max_pos = audio_positional_embedding_max_pos
+ self.audio_num_attention_heads = audio_num_attention_heads
+ self.audio_inner_dim = self.audio_num_attention_heads * audio_attention_head_dim
+ self._init_audio(
+ in_channels=audio_in_channels,
+ out_channels=audio_out_channels,
+ caption_channels=caption_channels,
+ norm_eps=norm_eps,
+ )
+
+ if model_type.is_video_enabled() and model_type.is_audio_enabled():
+ cross_pe_max_pos = max(self.positional_embedding_max_pos[0], self.audio_positional_embedding_max_pos[0])
+ self.av_ca_timestep_scale_multiplier = av_ca_timestep_scale_multiplier
+ self.audio_cross_attention_dim = audio_cross_attention_dim
+ self._init_audio_video(num_scale_shift_values=4)
+
+ self._init_preprocessors(cross_pe_max_pos)
+ # Initialize transformer blocks
+ self._init_transformer_blocks(
+ num_layers=num_layers,
+ attention_head_dim=attention_head_dim if model_type.is_video_enabled() else 0,
+ cross_attention_dim=cross_attention_dim,
+ audio_attention_head_dim=audio_attention_head_dim if model_type.is_audio_enabled() else 0,
+ audio_cross_attention_dim=audio_cross_attention_dim,
+ norm_eps=norm_eps,
+ attention_type=attention_type,
+ )
+
+ def _init_video(
+ self,
+ in_channels: int,
+ out_channels: int,
+ caption_channels: int,
+ norm_eps: float,
+ ) -> None:
+ """Initialize video-specific components."""
+ # Video input components
+ self.patchify_proj = torch.nn.Linear(in_channels, self.inner_dim, bias=True)
+
+ self.adaln_single = AdaLayerNormSingle(self.inner_dim)
+
+ # Video caption projection
+ self.caption_projection = PixArtAlphaTextProjection(
+ in_features=caption_channels,
+ hidden_size=self.inner_dim,
+ )
+
+ # Video output components
+ self.scale_shift_table = torch.nn.Parameter(torch.empty(2, self.inner_dim))
+ self.norm_out = torch.nn.LayerNorm(self.inner_dim, elementwise_affine=False, eps=norm_eps)
+ self.proj_out = torch.nn.Linear(self.inner_dim, out_channels)
+
+ def _init_audio(
+ self,
+ in_channels: int,
+ out_channels: int,
+ caption_channels: int,
+ norm_eps: float,
+ ) -> None:
+ """Initialize audio-specific components."""
+
+ # Audio input components
+ self.audio_patchify_proj = torch.nn.Linear(in_channels, self.audio_inner_dim, bias=True)
+
+ self.audio_adaln_single = AdaLayerNormSingle(
+ self.audio_inner_dim,
+ )
+
+ # Audio caption projection
+ self.audio_caption_projection = PixArtAlphaTextProjection(
+ in_features=caption_channels,
+ hidden_size=self.audio_inner_dim,
+ )
+
+ # Audio output components
+ self.audio_scale_shift_table = torch.nn.Parameter(torch.empty(2, self.audio_inner_dim))
+ self.audio_norm_out = torch.nn.LayerNorm(self.audio_inner_dim, elementwise_affine=False, eps=norm_eps)
+ self.audio_proj_out = torch.nn.Linear(self.audio_inner_dim, out_channels)
+
+ def _init_audio_video(
+ self,
+ num_scale_shift_values: int,
+ ) -> None:
+ """Initialize audio-video cross-attention components."""
+ self.av_ca_video_scale_shift_adaln_single = AdaLayerNormSingle(
+ self.inner_dim,
+ embedding_coefficient=num_scale_shift_values,
+ )
+
+ self.av_ca_audio_scale_shift_adaln_single = AdaLayerNormSingle(
+ self.audio_inner_dim,
+ embedding_coefficient=num_scale_shift_values,
+ )
+
+ self.av_ca_a2v_gate_adaln_single = AdaLayerNormSingle(
+ self.inner_dim,
+ embedding_coefficient=1,
+ )
+
+ self.av_ca_v2a_gate_adaln_single = AdaLayerNormSingle(
+ self.audio_inner_dim,
+ embedding_coefficient=1,
+ )
+
+ def _init_preprocessors(
+ self,
+ cross_pe_max_pos: int | None = None,
+ ) -> None:
+ """Initialize preprocessors for LTX."""
+
+ if self.model_type.is_video_enabled() and self.model_type.is_audio_enabled():
+ self.video_args_preprocessor = MultiModalTransformerArgsPreprocessor(
+ patchify_proj=self.patchify_proj,
+ adaln=self.adaln_single,
+ caption_projection=self.caption_projection,
+ cross_scale_shift_adaln=self.av_ca_video_scale_shift_adaln_single,
+ cross_gate_adaln=self.av_ca_a2v_gate_adaln_single,
+ inner_dim=self.inner_dim,
+ max_pos=self.positional_embedding_max_pos,
+ num_attention_heads=self.num_attention_heads,
+ cross_pe_max_pos=cross_pe_max_pos,
+ use_middle_indices_grid=self.use_middle_indices_grid,
+ audio_cross_attention_dim=self.audio_cross_attention_dim,
+ timestep_scale_multiplier=self.timestep_scale_multiplier,
+ double_precision_rope=self.double_precision_rope,
+ positional_embedding_theta=self.positional_embedding_theta,
+ rope_type=self.rope_type,
+ av_ca_timestep_scale_multiplier=self.av_ca_timestep_scale_multiplier,
+ )
+ self.audio_args_preprocessor = MultiModalTransformerArgsPreprocessor(
+ patchify_proj=self.audio_patchify_proj,
+ adaln=self.audio_adaln_single,
+ caption_projection=self.audio_caption_projection,
+ cross_scale_shift_adaln=self.av_ca_audio_scale_shift_adaln_single,
+ cross_gate_adaln=self.av_ca_v2a_gate_adaln_single,
+ inner_dim=self.audio_inner_dim,
+ max_pos=self.audio_positional_embedding_max_pos,
+ num_attention_heads=self.audio_num_attention_heads,
+ cross_pe_max_pos=cross_pe_max_pos,
+ use_middle_indices_grid=self.use_middle_indices_grid,
+ audio_cross_attention_dim=self.audio_cross_attention_dim,
+ timestep_scale_multiplier=self.timestep_scale_multiplier,
+ double_precision_rope=self.double_precision_rope,
+ positional_embedding_theta=self.positional_embedding_theta,
+ rope_type=self.rope_type,
+ av_ca_timestep_scale_multiplier=self.av_ca_timestep_scale_multiplier,
+ )
+ elif self.model_type.is_video_enabled():
+ self.video_args_preprocessor = TransformerArgsPreprocessor(
+ patchify_proj=self.patchify_proj,
+ adaln=self.adaln_single,
+ caption_projection=self.caption_projection,
+ inner_dim=self.inner_dim,
+ max_pos=self.positional_embedding_max_pos,
+ num_attention_heads=self.num_attention_heads,
+ use_middle_indices_grid=self.use_middle_indices_grid,
+ timestep_scale_multiplier=self.timestep_scale_multiplier,
+ double_precision_rope=self.double_precision_rope,
+ positional_embedding_theta=self.positional_embedding_theta,
+ rope_type=self.rope_type,
+ )
+ elif self.model_type.is_audio_enabled():
+ self.audio_args_preprocessor = TransformerArgsPreprocessor(
+ patchify_proj=self.audio_patchify_proj,
+ adaln=self.audio_adaln_single,
+ caption_projection=self.audio_caption_projection,
+ inner_dim=self.audio_inner_dim,
+ max_pos=self.audio_positional_embedding_max_pos,
+ num_attention_heads=self.audio_num_attention_heads,
+ use_middle_indices_grid=self.use_middle_indices_grid,
+ timestep_scale_multiplier=self.timestep_scale_multiplier,
+ double_precision_rope=self.double_precision_rope,
+ positional_embedding_theta=self.positional_embedding_theta,
+ rope_type=self.rope_type,
+ )
+
+ def _init_transformer_blocks(
+ self,
+ num_layers: int,
+ attention_head_dim: int,
+ cross_attention_dim: int,
+ audio_attention_head_dim: int,
+ audio_cross_attention_dim: int,
+ norm_eps: float,
+ attention_type: AttentionFunction | AttentionCallable,
+ ) -> None:
+ """Initialize transformer blocks for LTX."""
+ video_config = (
+ TransformerConfig(
+ dim=self.inner_dim,
+ heads=self.num_attention_heads,
+ d_head=attention_head_dim,
+ context_dim=cross_attention_dim,
+ )
+ if self.model_type.is_video_enabled()
+ else None
+ )
+ audio_config = (
+ TransformerConfig(
+ dim=self.audio_inner_dim,
+ heads=self.audio_num_attention_heads,
+ d_head=audio_attention_head_dim,
+ context_dim=audio_cross_attention_dim,
+ )
+ if self.model_type.is_audio_enabled()
+ else None
+ )
+ self.transformer_blocks = torch.nn.ModuleList(
+ [
+ BasicAVTransformerBlock(
+ idx=idx,
+ video=video_config,
+ audio=audio_config,
+ rope_type=self.rope_type,
+ norm_eps=norm_eps,
+ attention_function=attention_type,
+ )
+ for idx in range(num_layers)
+ ]
+ )
+
+ def set_gradient_checkpointing(self, enable: bool) -> None:
+ """Enable or disable gradient checkpointing for transformer blocks.
+ Gradient checkpointing trades compute for memory by recomputing activations
+ during the backward pass instead of storing them. This can significantly
+ reduce memory usage at the cost of ~20-30% slower training.
+ Args:
+ enable: Whether to enable gradient checkpointing
+ """
+ self._enable_gradient_checkpointing = enable
+
+ def _process_transformer_blocks(
+ self,
+ video: TransformerArgs | None,
+ audio: TransformerArgs | None,
+ perturbations: BatchedPerturbationConfig,
+ ) -> tuple[TransformerArgs, TransformerArgs]:
+ """Process transformer blocks for LTXAV."""
+
+ # Process transformer blocks
+ for block in self.transformer_blocks:
+ if self._enable_gradient_checkpointing and self.training:
+ # Use gradient checkpointing to save memory during training.
+ # With use_reentrant=False, we can pass dataclasses directly -
+ # PyTorch will track all tensor leaves in the computation graph.
+ video, audio = torch.utils.checkpoint.checkpoint(
+ block,
+ video,
+ audio,
+ perturbations,
+ use_reentrant=False,
+ )
+ else:
+ video, audio = block(
+ video=video,
+ audio=audio,
+ perturbations=perturbations,
+ )
+
+ return video, audio
+
+ def _process_output(
+ self,
+ scale_shift_table: torch.Tensor,
+ norm_out: torch.nn.LayerNorm,
+ proj_out: torch.nn.Linear,
+ x: torch.Tensor,
+ embedded_timestep: torch.Tensor,
+ ) -> torch.Tensor:
+ """Process output for LTXV."""
+ # Apply scale-shift modulation
+ scale_shift_values = (
+ scale_shift_table[None, None].to(device=x.device, dtype=x.dtype) + embedded_timestep[:, :, None]
+ )
+ shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1]
+
+ x = norm_out(x)
+ x = x * (1 + scale) + shift
+ x = proj_out(x)
+ return x
+
+ def forward(
+ self, video: Modality | None, audio: Modality | None, perturbations: BatchedPerturbationConfig
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ Forward pass for LTX models.
+ Returns:
+ Processed output tensors
+ """
+ if not self.model_type.is_video_enabled() and video is not None:
+ raise ValueError("Video is not enabled for this model")
+ if not self.model_type.is_audio_enabled() and audio is not None:
+ raise ValueError("Audio is not enabled for this model")
+
+ video_args = self.video_args_preprocessor.prepare(video) if video is not None else None
+ audio_args = self.audio_args_preprocessor.prepare(audio) if audio is not None else None
+ # Process transformer blocks
+ video_out, audio_out = self._process_transformer_blocks(
+ video=video_args,
+ audio=audio_args,
+ perturbations=perturbations,
+ )
+
+ # Process output
+ vx = (
+ self._process_output(
+ self.scale_shift_table, self.norm_out, self.proj_out, video_out.x, video_out.embedded_timestep
+ )
+ if video_out is not None
+ else None
+ )
+ ax = (
+ self._process_output(
+ self.audio_scale_shift_table,
+ self.audio_norm_out,
+ self.audio_proj_out,
+ audio_out.x,
+ audio_out.embedded_timestep,
+ )
+ if audio_out is not None
+ else None
+ )
+ return vx, ax
+
+
+class LegacyX0Model(torch.nn.Module):
+ """
+ Legacy X0 model implementation.
+ Returns fully denoised output based on the velocities produced by the base model.
+ """
+
+ def __init__(self, velocity_model: LTXModel):
+ super().__init__()
+ self.velocity_model = velocity_model
+
+ def forward(
+ self,
+ video: Modality | None,
+ audio: Modality | None,
+ perturbations: BatchedPerturbationConfig,
+ sigma: float,
+ ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+ """
+ Denoise the video and audio according to the sigma.
+ Returns:
+ Denoised video and audio
+ """
+ vx, ax = self.velocity_model(video, audio, perturbations)
+ denoised_video = to_denoised(video.latent, vx, sigma) if vx is not None else None
+ denoised_audio = to_denoised(audio.latent, ax, sigma) if ax is not None else None
+ return denoised_video, denoised_audio
+
+
+class X0Model(torch.nn.Module):
+ """
+ X0 model implementation.
+ Returns fully denoised outputs based on the velocities produced by the base model.
+ Applies scaled denoising to the video and audio according to the timesteps = sigma * denoising_mask.
+ """
+
+ def __init__(self, velocity_model: LTXModel):
+ super().__init__()
+ self.velocity_model = velocity_model
+
+ def forward(
+ self,
+ video: Modality | None,
+ audio: Modality | None,
+ perturbations: BatchedPerturbationConfig,
+ ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+ """
+ Denoise the video and audio according to the sigma.
+ Returns:
+ Denoised video and audio
+ """
+ vx, ax = self.velocity_model(video, audio, perturbations)
+ denoised_video = to_denoised(video.latent, vx, video.timesteps) if vx is not None else None
+ denoised_audio = to_denoised(audio.latent, ax, audio.timesteps) if ax is not None else None
+ return denoised_video, denoised_audio
diff --git a/packages/ltx-core/src/ltx_core/model/transformer/model_configurator.py b/packages/ltx-core/src/ltx_core/model/transformer/model_configurator.py
new file mode 100644
index 0000000000000000000000000000000000000000..7b4d5c6b2f2d0f14aab6e43cd86f267c8326befc
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/transformer/model_configurator.py
@@ -0,0 +1,237 @@
+import torch
+
+from ltx_core.loader.fuse_loras import fused_add_round_launch
+from ltx_core.loader.module_ops import ModuleOps
+from ltx_core.loader.sd_ops import KeyValueOperationResult, SDOps
+from ltx_core.model.model_protocol import ModelConfigurator
+from ltx_core.model.transformer.attention import AttentionFunction
+from ltx_core.model.transformer.model import LTXModel, LTXModelType
+from ltx_core.model.transformer.rope import LTXRopeType
+from ltx_core.utils import check_config_value
+
+
+class LTXModelConfigurator(ModelConfigurator[LTXModel]):
+ """
+ Configurator for LTX model.
+ Used to create an LTX model from a configuration dictionary.
+ """
+
+ @classmethod
+ def from_config(cls: type[LTXModel], config: dict) -> LTXModel:
+ config = config.get("transformer", {})
+
+ check_config_value(config, "dropout", 0.0)
+ check_config_value(config, "attention_bias", True)
+ check_config_value(config, "num_vector_embeds", None)
+ check_config_value(config, "activation_fn", "gelu-approximate")
+ check_config_value(config, "num_embeds_ada_norm", 1000)
+ check_config_value(config, "use_linear_projection", False)
+ check_config_value(config, "only_cross_attention", False)
+ check_config_value(config, "cross_attention_norm", True)
+ check_config_value(config, "double_self_attention", False)
+ check_config_value(config, "upcast_attention", False)
+ check_config_value(config, "standardization_norm", "rms_norm")
+ check_config_value(config, "norm_elementwise_affine", False)
+ check_config_value(config, "qk_norm", "rms_norm")
+ check_config_value(config, "positional_embedding_type", "rope")
+ check_config_value(config, "use_audio_video_cross_attention", True)
+ check_config_value(config, "share_ff", False)
+ check_config_value(config, "av_cross_ada_norm", True)
+ check_config_value(config, "use_middle_indices_grid", True)
+
+ return LTXModel(
+ model_type=LTXModelType.AudioVideo,
+ num_attention_heads=config.get("num_attention_heads", 32),
+ attention_head_dim=config.get("attention_head_dim", 128),
+ in_channels=config.get("in_channels", 128),
+ out_channels=config.get("out_channels", 128),
+ num_layers=config.get("num_layers", 48),
+ cross_attention_dim=config.get("cross_attention_dim", 4096),
+ norm_eps=config.get("norm_eps", 1e-06),
+ attention_type=AttentionFunction(config.get("attention_type", "default")),
+ caption_channels=config.get("caption_channels", 3840),
+ positional_embedding_theta=config.get("positional_embedding_theta", 10000.0),
+ positional_embedding_max_pos=config.get("positional_embedding_max_pos", [20, 2048, 2048]),
+ timestep_scale_multiplier=config.get("timestep_scale_multiplier", 1000),
+ use_middle_indices_grid=config.get("use_middle_indices_grid", True),
+ audio_num_attention_heads=config.get("audio_num_attention_heads", 32),
+ audio_attention_head_dim=config.get("audio_attention_head_dim", 64),
+ audio_in_channels=config.get("audio_in_channels", 128),
+ audio_out_channels=config.get("audio_out_channels", 128),
+ audio_cross_attention_dim=config.get("audio_cross_attention_dim", 2048),
+ audio_positional_embedding_max_pos=config.get("audio_positional_embedding_max_pos", [20]),
+ av_ca_timestep_scale_multiplier=config.get("av_ca_timestep_scale_multiplier", 1),
+ rope_type=LTXRopeType(config.get("rope_type", "interleaved")),
+ double_precision_rope=config.get("frequencies_precision", False) == "float64",
+ )
+
+
+class LTXVideoOnlyModelConfigurator(ModelConfigurator[LTXModel]):
+ """
+ Configurator for LTX video only model.
+ Used to create an LTX video only model from a configuration dictionary.
+ """
+
+ @classmethod
+ def from_config(cls: type[LTXModel], config: dict) -> LTXModel:
+ config = config.get("transformer", {})
+
+ check_config_value(config, "dropout", 0.0)
+ check_config_value(config, "attention_bias", True)
+ check_config_value(config, "num_vector_embeds", None)
+ check_config_value(config, "activation_fn", "gelu-approximate")
+ check_config_value(config, "num_embeds_ada_norm", 1000)
+ check_config_value(config, "use_linear_projection", False)
+ check_config_value(config, "only_cross_attention", False)
+ check_config_value(config, "cross_attention_norm", True)
+ check_config_value(config, "double_self_attention", False)
+ check_config_value(config, "upcast_attention", False)
+ check_config_value(config, "standardization_norm", "rms_norm")
+ check_config_value(config, "norm_elementwise_affine", False)
+ check_config_value(config, "qk_norm", "rms_norm")
+ check_config_value(config, "positional_embedding_type", "rope")
+ check_config_value(config, "use_middle_indices_grid", True)
+
+ return LTXModel(
+ model_type=LTXModelType.VideoOnly,
+ num_attention_heads=config.get("num_attention_heads", 32),
+ attention_head_dim=config.get("attention_head_dim", 128),
+ in_channels=config.get("in_channels", 128),
+ out_channels=config.get("out_channels", 128),
+ num_layers=config.get("num_layers", 48),
+ cross_attention_dim=config.get("cross_attention_dim", 4096),
+ norm_eps=config.get("norm_eps", 1e-06),
+ attention_type=AttentionFunction(config.get("attention_type", "default")),
+ caption_channels=config.get("caption_channels", 3840),
+ positional_embedding_theta=config.get("positional_embedding_theta", 10000.0),
+ positional_embedding_max_pos=config.get("positional_embedding_max_pos", [20, 2048, 2048]),
+ timestep_scale_multiplier=config.get("timestep_scale_multiplier", 1000),
+ use_middle_indices_grid=config.get("use_middle_indices_grid", True),
+ rope_type=LTXRopeType(config.get("rope_type", "interleaved")),
+ double_precision_rope=config.get("frequencies_precision", False) == "float64",
+ )
+
+
+def _naive_weight_or_bias_downcast(key: str, value: torch.Tensor) -> list[KeyValueOperationResult]:
+ """
+ Downcast the weight or bias to the float8_e4m3fn dtype.
+ """
+ return [KeyValueOperationResult(key, value.to(dtype=torch.float8_e4m3fn))]
+
+
+def _upcast_and_round(
+ weight: torch.Tensor, dtype: torch.dtype, with_stochastic_rounding: bool = False, seed: int = 0
+) -> torch.Tensor:
+ """
+ Upcast the weight to the given dtype and optionally apply stochastic rounding.
+ Input weight needs to have float8_e4m3fn or float8_e5m2 dtype.
+ """
+ if not with_stochastic_rounding:
+ return weight.to(dtype)
+ return fused_add_round_launch(torch.zeros_like(weight, dtype=dtype), weight, seed)
+
+
+def replace_fwd_with_upcast(layer: torch.nn.Linear, with_stochastic_rounding: bool = False, seed: int = 0) -> None:
+ """
+ Replace linear.forward and rms_norm.forward with a version that:
+ - upcasts weight and bias to input's dtype
+ - returns F.linear or F.rms_norm calculated in that dtype
+ """
+
+ layer.original_forward = layer.forward
+
+ def new_linear_forward(*args, **_kwargs) -> torch.Tensor:
+ # assume first arg is the input tensor
+ x = args[0]
+ w_up = _upcast_and_round(layer.weight, x.dtype, with_stochastic_rounding, seed)
+ b_up = None
+
+ if layer.bias is not None:
+ b_up = _upcast_and_round(layer.bias, x.dtype, with_stochastic_rounding, seed)
+
+ return torch.nn.functional.linear(x, w_up, b_up)
+
+ layer.forward = new_linear_forward
+
+
+def amend_forward_with_upcast(
+ model: torch.nn.Module, with_stochastic_rounding: bool = False, seed: int = 0
+) -> torch.nn.Module:
+ """
+ Replace the forward method of the model's Linear and RMSNorm layers to forward
+ with upcast and optional stochastic rounding.
+ """
+ for m in model.modules():
+ if isinstance(m, (torch.nn.Linear)):
+ replace_fwd_with_upcast(m, with_stochastic_rounding, seed)
+ return model
+
+
+LTXV_MODEL_COMFY_RENAMING_MAP = (
+ SDOps("LTXV_MODEL_COMFY_PREFIX_MAP")
+ .with_matching(prefix="model.diffusion_model.")
+ .with_replacement("model.diffusion_model.", "")
+)
+
+LTXV_MODEL_COMFY_RENAMING_WITH_TRANSFORMER_LINEAR_DOWNCAST_MAP = (
+ SDOps("LTXV_MODEL_COMFY_PREFIX_MAP")
+ .with_matching(prefix="model.diffusion_model.")
+ .with_replacement("model.diffusion_model.", "")
+ .with_kv_operation(
+ key_prefix="transformer_blocks.", key_suffix=".to_q.weight", operation=_naive_weight_or_bias_downcast
+ )
+ .with_kv_operation(
+ key_prefix="transformer_blocks.", key_suffix=".to_q.bias", operation=_naive_weight_or_bias_downcast
+ )
+ .with_kv_operation(
+ key_prefix="transformer_blocks.", key_suffix=".to_k.weight", operation=_naive_weight_or_bias_downcast
+ )
+ .with_kv_operation(
+ key_prefix="transformer_blocks.", key_suffix=".to_k.bias", operation=_naive_weight_or_bias_downcast
+ )
+ .with_kv_operation(
+ key_prefix="transformer_blocks.", key_suffix=".to_v.weight", operation=_naive_weight_or_bias_downcast
+ )
+ .with_kv_operation(
+ key_prefix="transformer_blocks.", key_suffix=".to_v.bias", operation=_naive_weight_or_bias_downcast
+ )
+ .with_kv_operation(
+ key_prefix="transformer_blocks.", key_suffix=".to_out.0.weight", operation=_naive_weight_or_bias_downcast
+ )
+ .with_kv_operation(
+ key_prefix="transformer_blocks.", key_suffix=".to_out.0.bias", operation=_naive_weight_or_bias_downcast
+ )
+ .with_kv_operation(
+ key_prefix="transformer_blocks.", key_suffix=".ff.net.0.proj.weight", operation=_naive_weight_or_bias_downcast
+ )
+ .with_kv_operation(
+ key_prefix="transformer_blocks.", key_suffix=".ff.net.0.proj.bias", operation=_naive_weight_or_bias_downcast
+ )
+ .with_kv_operation(
+ key_prefix="transformer_blocks.", key_suffix=".ff.net.2.weight", operation=_naive_weight_or_bias_downcast
+ )
+ .with_kv_operation(
+ key_prefix="transformer_blocks.", key_suffix=".ff.net.2.bias", operation=_naive_weight_or_bias_downcast
+ )
+)
+
+UPCAST_DURING_INFERENCE = ModuleOps(
+ name="upcast_fp8_during_linear_forward",
+ matcher=lambda model: isinstance(model, LTXModel),
+ mutator=lambda model: amend_forward_with_upcast(model, False),
+)
+
+
+class UpcastWithStochasticRounding(ModuleOps):
+ """
+ ModuleOps for upcasting the model's float8_e4m3fn weights and biases to the bfloat16 dtype
+ and applying stochastic rounding during linear forward.
+ """
+
+ def __new__(cls, seed: int = 0):
+ return super().__new__(
+ cls,
+ name="upcast_fp8_during_linear_forward_with_stochastic_rounding",
+ matcher=lambda model: isinstance(model, LTXModel),
+ mutator=lambda model: amend_forward_with_upcast(model, True, seed),
+ )
diff --git a/packages/ltx-core/src/ltx_core/model/transformer/rope.py b/packages/ltx-core/src/ltx_core/model/transformer/rope.py
new file mode 100644
index 0000000000000000000000000000000000000000..b86d97e7f737187e06fc493ceb534e4197567661
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/transformer/rope.py
@@ -0,0 +1,204 @@
+import functools
+import math
+from enum import Enum
+from typing import Callable, Tuple
+
+import numpy as np
+import torch
+from einops import rearrange
+
+
+class LTXRopeType(Enum):
+ INTERLEAVED = "interleaved"
+ SPLIT = "split"
+
+
+def apply_rotary_emb(
+ input_tensor: torch.Tensor,
+ freqs_cis: Tuple[torch.Tensor, torch.Tensor],
+ rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
+) -> torch.Tensor:
+ if rope_type == LTXRopeType.INTERLEAVED:
+ return apply_interleaved_rotary_emb(input_tensor, *freqs_cis)
+ elif rope_type == LTXRopeType.SPLIT:
+ return apply_split_rotary_emb(input_tensor, *freqs_cis)
+ else:
+ raise ValueError(f"Invalid rope type: {rope_type}")
+
+
+def apply_interleaved_rotary_emb(
+ input_tensor: torch.Tensor, cos_freqs: torch.Tensor, sin_freqs: torch.Tensor
+) -> torch.Tensor:
+ t_dup = rearrange(input_tensor, "... (d r) -> ... d r", r=2)
+ t1, t2 = t_dup.unbind(dim=-1)
+ t_dup = torch.stack((-t2, t1), dim=-1)
+ input_tensor_rot = rearrange(t_dup, "... d r -> ... (d r)")
+
+ out = input_tensor * cos_freqs + input_tensor_rot * sin_freqs
+
+ return out
+
+
+def apply_split_rotary_emb(
+ input_tensor: torch.Tensor, cos_freqs: torch.Tensor, sin_freqs: torch.Tensor
+) -> torch.Tensor:
+ needs_reshape = False
+ if input_tensor.ndim != 4 and cos_freqs.ndim == 4:
+ b, h, t, _ = cos_freqs.shape
+ input_tensor = input_tensor.reshape(b, t, h, -1).swapaxes(1, 2)
+ needs_reshape = True
+
+ split_input = rearrange(input_tensor, "... (d r) -> ... d r", d=2)
+ first_half_input = split_input[..., :1, :]
+ second_half_input = split_input[..., 1:, :]
+
+ output = split_input * cos_freqs.unsqueeze(-2)
+ first_half_output = output[..., :1, :]
+ second_half_output = output[..., 1:, :]
+
+ first_half_output.addcmul_(-sin_freqs.unsqueeze(-2), second_half_input)
+ second_half_output.addcmul_(sin_freqs.unsqueeze(-2), first_half_input)
+
+ output = rearrange(output, "... d r -> ... (d r)")
+ if needs_reshape:
+ output = output.swapaxes(1, 2).reshape(b, t, -1)
+
+ return output
+
+
+@functools.lru_cache(maxsize=5)
+def generate_freq_grid_np(
+ positional_embedding_theta: float, positional_embedding_max_pos_count: int, inner_dim: int
+) -> torch.Tensor:
+ theta = positional_embedding_theta
+ start = 1
+ end = theta
+
+ n_elem = 2 * positional_embedding_max_pos_count
+ pow_indices = np.power(
+ theta,
+ np.linspace(
+ np.log(start) / np.log(theta),
+ np.log(end) / np.log(theta),
+ inner_dim // n_elem,
+ dtype=np.float64,
+ ),
+ )
+ return torch.tensor(pow_indices * math.pi / 2, dtype=torch.float32)
+
+
+@functools.lru_cache(maxsize=5)
+def generate_freq_grid_pytorch(
+ positional_embedding_theta: float, positional_embedding_max_pos_count: int, inner_dim: int
+) -> torch.Tensor:
+ theta = positional_embedding_theta
+ start = 1
+ end = theta
+ n_elem = 2 * positional_embedding_max_pos_count
+
+ indices = theta ** (
+ torch.linspace(
+ math.log(start, theta),
+ math.log(end, theta),
+ inner_dim // n_elem,
+ dtype=torch.float32,
+ )
+ )
+ indices = indices.to(dtype=torch.float32)
+
+ indices = indices * math.pi / 2
+
+ return indices
+
+
+def get_fractional_positions(indices_grid: torch.Tensor, max_pos: list[int]) -> torch.Tensor:
+ n_pos_dims = indices_grid.shape[1]
+ assert n_pos_dims == len(max_pos), (
+ f"Number of position dimensions ({n_pos_dims}) must match max_pos length ({len(max_pos)})"
+ )
+ fractional_positions = torch.stack(
+ [indices_grid[:, i] / max_pos[i] for i in range(n_pos_dims)],
+ dim=-1,
+ )
+ return fractional_positions
+
+
+def generate_freqs(
+ indices: torch.Tensor, indices_grid: torch.Tensor, max_pos: list[int], use_middle_indices_grid: bool
+) -> torch.Tensor:
+ if use_middle_indices_grid:
+ assert len(indices_grid.shape) == 4
+ assert indices_grid.shape[-1] == 2
+ indices_grid_start, indices_grid_end = indices_grid[..., 0], indices_grid[..., 1]
+ indices_grid = (indices_grid_start + indices_grid_end) / 2.0
+ elif len(indices_grid.shape) == 4:
+ indices_grid = indices_grid[..., 0]
+
+ fractional_positions = get_fractional_positions(indices_grid, max_pos)
+ indices = indices.to(device=fractional_positions.device)
+
+ freqs = (indices * (fractional_positions.unsqueeze(-1) * 2 - 1)).transpose(-1, -2).flatten(2)
+ return freqs
+
+
+def split_freqs_cis(freqs: torch.Tensor, pad_size: int, num_attention_heads: int) -> tuple[torch.Tensor, torch.Tensor]:
+ cos_freq = freqs.cos()
+ sin_freq = freqs.sin()
+
+ if pad_size != 0:
+ cos_padding = torch.ones_like(cos_freq[:, :, :pad_size])
+ sin_padding = torch.zeros_like(sin_freq[:, :, :pad_size])
+
+ cos_freq = torch.concatenate([cos_padding, cos_freq], axis=-1)
+ sin_freq = torch.concatenate([sin_padding, sin_freq], axis=-1)
+
+ # Reshape freqs to be compatible with multi-head attention
+ b = cos_freq.shape[0]
+ t = cos_freq.shape[1]
+
+ cos_freq = cos_freq.reshape(b, t, num_attention_heads, -1)
+ sin_freq = sin_freq.reshape(b, t, num_attention_heads, -1)
+
+ cos_freq = torch.swapaxes(cos_freq, 1, 2) # (B,H,T,D//2)
+ sin_freq = torch.swapaxes(sin_freq, 1, 2) # (B,H,T,D//2)
+ return cos_freq, sin_freq
+
+
+def interleaved_freqs_cis(freqs: torch.Tensor, pad_size: int) -> tuple[torch.Tensor, torch.Tensor]:
+ cos_freq = freqs.cos().repeat_interleave(2, dim=-1)
+ sin_freq = freqs.sin().repeat_interleave(2, dim=-1)
+ if pad_size != 0:
+ cos_padding = torch.ones_like(cos_freq[:, :, :pad_size])
+ sin_padding = torch.zeros_like(cos_freq[:, :, :pad_size])
+ cos_freq = torch.cat([cos_padding, cos_freq], dim=-1)
+ sin_freq = torch.cat([sin_padding, sin_freq], dim=-1)
+ return cos_freq, sin_freq
+
+
+def precompute_freqs_cis(
+ indices_grid: torch.Tensor,
+ dim: int,
+ out_dtype: torch.dtype,
+ theta: float = 10000.0,
+ max_pos: list[int] | None = None,
+ use_middle_indices_grid: bool = False,
+ num_attention_heads: int = 32,
+ rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
+ freq_grid_generator: Callable[[float, int, int, torch.device], torch.Tensor] = generate_freq_grid_pytorch,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ if max_pos is None:
+ max_pos = [20, 2048, 2048]
+
+ indices = freq_grid_generator(theta, indices_grid.shape[1], dim)
+ freqs = generate_freqs(indices, indices_grid, max_pos, use_middle_indices_grid)
+
+ if rope_type == LTXRopeType.SPLIT:
+ expected_freqs = dim // 2
+ current_freqs = freqs.shape[-1]
+ pad_size = expected_freqs - current_freqs
+ cos_freq, sin_freq = split_freqs_cis(freqs, pad_size, num_attention_heads)
+ else:
+ # 2 because of cos and sin by 3 for (t, x, y), 1 for temporal only
+ n_elem = 2 * indices_grid.shape[1]
+ cos_freq, sin_freq = interleaved_freqs_cis(freqs, dim % n_elem)
+ return cos_freq.to(out_dtype), sin_freq.to(out_dtype)
diff --git a/packages/ltx-core/src/ltx_core/model/transformer/text_projection.py b/packages/ltx-core/src/ltx_core/model/transformer/text_projection.py
new file mode 100644
index 0000000000000000000000000000000000000000..9f8fe86151d1b078a9a0bdf5b30afc7a30bdd7b8
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/transformer/text_projection.py
@@ -0,0 +1,27 @@
+import torch
+
+
+class PixArtAlphaTextProjection(torch.nn.Module):
+ """
+ Projects caption embeddings. Also handles dropout for classifier-free guidance.
+ Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py
+ """
+
+ def __init__(self, in_features: int, hidden_size: int, out_features: int | None = None, act_fn: str = "gelu_tanh"):
+ super().__init__()
+ if out_features is None:
+ out_features = hidden_size
+ self.linear_1 = torch.nn.Linear(in_features=in_features, out_features=hidden_size, bias=True)
+ if act_fn == "gelu_tanh":
+ self.act_1 = torch.nn.GELU(approximate="tanh")
+ elif act_fn == "silu":
+ self.act_1 = torch.nn.SiLU()
+ else:
+ raise ValueError(f"Unknown activation function: {act_fn}")
+ self.linear_2 = torch.nn.Linear(in_features=hidden_size, out_features=out_features, bias=True)
+
+ def forward(self, caption: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.linear_1(caption)
+ hidden_states = self.act_1(hidden_states)
+ hidden_states = self.linear_2(hidden_states)
+ return hidden_states
diff --git a/packages/ltx-core/src/ltx_core/model/transformer/timestep_embedding.py b/packages/ltx-core/src/ltx_core/model/transformer/timestep_embedding.py
new file mode 100644
index 0000000000000000000000000000000000000000..e8a5176f7295a266b223f3349c8b7e85036c20b6
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/transformer/timestep_embedding.py
@@ -0,0 +1,143 @@
+import math
+
+import torch
+
+
+def get_timestep_embedding(
+ timesteps: torch.Tensor,
+ embedding_dim: int,
+ flip_sin_to_cos: bool = False,
+ downscale_freq_shift: float = 1,
+ scale: float = 1,
+ max_period: int = 10000,
+) -> torch.Tensor:
+ """
+ This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings.
+ Args
+ timesteps (torch.Tensor):
+ a 1-D Tensor of N indices, one per batch element. These may be fractional.
+ embedding_dim (int):
+ the dimension of the output.
+ flip_sin_to_cos (bool):
+ Whether the embedding order should be `cos, sin` (if True) or `sin, cos` (if False)
+ downscale_freq_shift (float):
+ Controls the delta between frequencies between dimensions
+ scale (float):
+ Scaling factor applied to the embeddings.
+ max_period (int):
+ Controls the maximum frequency of the embeddings
+ Returns
+ torch.Tensor: an [N x dim] Tensor of positional embeddings.
+ """
+ assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"
+
+ half_dim = embedding_dim // 2
+ exponent = -math.log(max_period) * torch.arange(start=0, end=half_dim, dtype=torch.float32, device=timesteps.device)
+ exponent = exponent / (half_dim - downscale_freq_shift)
+
+ emb = torch.exp(exponent)
+ emb = timesteps[:, None].float() * emb[None, :]
+
+ # scale embeddings
+ emb = scale * emb
+
+ # concat sine and cosine embeddings
+ emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
+
+ # flip sine and cosine embeddings
+ if flip_sin_to_cos:
+ emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
+
+ # zero pad
+ if embedding_dim % 2 == 1:
+ emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
+ return emb
+
+
+class TimestepEmbedding(torch.nn.Module):
+ def __init__(
+ self,
+ in_channels: int,
+ time_embed_dim: int,
+ out_dim: int | None = None,
+ post_act_fn: str | None = None,
+ cond_proj_dim: int | None = None,
+ sample_proj_bias: bool = True,
+ ):
+ super().__init__()
+
+ self.linear_1 = torch.nn.Linear(in_channels, time_embed_dim, sample_proj_bias)
+
+ if cond_proj_dim is not None:
+ self.cond_proj = torch.nn.Linear(cond_proj_dim, in_channels, bias=False)
+ else:
+ self.cond_proj = None
+
+ self.act = torch.nn.SiLU()
+ time_embed_dim_out = out_dim if out_dim is not None else time_embed_dim
+
+ self.linear_2 = torch.nn.Linear(time_embed_dim, time_embed_dim_out, sample_proj_bias)
+
+ if post_act_fn is None:
+ self.post_act = None
+
+ def forward(self, sample: torch.Tensor, condition: torch.Tensor | None = None) -> torch.Tensor:
+ if condition is not None:
+ sample = sample + self.cond_proj(condition)
+ sample = self.linear_1(sample)
+
+ if self.act is not None:
+ sample = self.act(sample)
+
+ sample = self.linear_2(sample)
+
+ if self.post_act is not None:
+ sample = self.post_act(sample)
+ return sample
+
+
+class Timesteps(torch.nn.Module):
+ def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float, scale: int = 1):
+ super().__init__()
+ self.num_channels = num_channels
+ self.flip_sin_to_cos = flip_sin_to_cos
+ self.downscale_freq_shift = downscale_freq_shift
+ self.scale = scale
+
+ def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
+ t_emb = get_timestep_embedding(
+ timesteps,
+ self.num_channels,
+ flip_sin_to_cos=self.flip_sin_to_cos,
+ downscale_freq_shift=self.downscale_freq_shift,
+ scale=self.scale,
+ )
+ return t_emb
+
+
+class PixArtAlphaCombinedTimestepSizeEmbeddings(torch.nn.Module):
+ """
+ For PixArt-Alpha.
+ Reference:
+ https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L164C9-L168C29
+ """
+
+ def __init__(
+ self,
+ embedding_dim: int,
+ size_emb_dim: int,
+ ):
+ super().__init__()
+
+ self.outdim = size_emb_dim
+ self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)
+ self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
+
+ def forward(
+ self,
+ timestep: torch.Tensor,
+ hidden_dtype: torch.dtype,
+ ) -> torch.Tensor:
+ timesteps_proj = self.time_proj(timestep)
+ timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype)) # (N, D)
+ return timesteps_emb
diff --git a/packages/ltx-core/src/ltx_core/model/transformer/transformer.py b/packages/ltx-core/src/ltx_core/model/transformer/transformer.py
new file mode 100644
index 0000000000000000000000000000000000000000..6f188b06e432184c972eeeeb311791324848c5b4
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/transformer/transformer.py
@@ -0,0 +1,274 @@
+from dataclasses import dataclass, replace
+
+import torch
+
+from ltx_core.guidance.perturbations import BatchedPerturbationConfig, PerturbationType
+from ltx_core.model.transformer.attention import Attention, AttentionCallable, AttentionFunction
+from ltx_core.model.transformer.feed_forward import FeedForward
+from ltx_core.model.transformer.rope import LTXRopeType
+from ltx_core.model.transformer.transformer_args import TransformerArgs
+from ltx_core.utils import rms_norm
+
+
+@dataclass
+class TransformerConfig:
+ dim: int
+ heads: int
+ d_head: int
+ context_dim: int
+
+
+class BasicAVTransformerBlock(torch.nn.Module):
+ def __init__(
+ self,
+ idx: int,
+ video: TransformerConfig | None = None,
+ audio: TransformerConfig | None = None,
+ rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
+ norm_eps: float = 1e-6,
+ attention_function: AttentionFunction | AttentionCallable = AttentionFunction.DEFAULT,
+ ):
+ super().__init__()
+
+ self.idx = idx
+ if video is not None:
+ self.attn1 = Attention(
+ query_dim=video.dim,
+ heads=video.heads,
+ dim_head=video.d_head,
+ context_dim=None,
+ rope_type=rope_type,
+ norm_eps=norm_eps,
+ attention_function=attention_function,
+ )
+ self.attn2 = Attention(
+ query_dim=video.dim,
+ context_dim=video.context_dim,
+ heads=video.heads,
+ dim_head=video.d_head,
+ rope_type=rope_type,
+ norm_eps=norm_eps,
+ attention_function=attention_function,
+ )
+ self.ff = FeedForward(video.dim, dim_out=video.dim)
+ self.scale_shift_table = torch.nn.Parameter(torch.empty(6, video.dim))
+
+ if audio is not None:
+ self.audio_attn1 = Attention(
+ query_dim=audio.dim,
+ heads=audio.heads,
+ dim_head=audio.d_head,
+ context_dim=None,
+ rope_type=rope_type,
+ norm_eps=norm_eps,
+ attention_function=attention_function,
+ )
+ self.audio_attn2 = Attention(
+ query_dim=audio.dim,
+ context_dim=audio.context_dim,
+ heads=audio.heads,
+ dim_head=audio.d_head,
+ rope_type=rope_type,
+ norm_eps=norm_eps,
+ attention_function=attention_function,
+ )
+ self.audio_ff = FeedForward(audio.dim, dim_out=audio.dim)
+ self.audio_scale_shift_table = torch.nn.Parameter(torch.empty(6, audio.dim))
+
+ if audio is not None and video is not None:
+ # Q: Video, K,V: Audio
+ self.audio_to_video_attn = Attention(
+ query_dim=video.dim,
+ context_dim=audio.dim,
+ heads=audio.heads,
+ dim_head=audio.d_head,
+ rope_type=rope_type,
+ norm_eps=norm_eps,
+ attention_function=attention_function,
+ )
+
+ # Q: Audio, K,V: Video
+ self.video_to_audio_attn = Attention(
+ query_dim=audio.dim,
+ context_dim=video.dim,
+ heads=audio.heads,
+ dim_head=audio.d_head,
+ rope_type=rope_type,
+ norm_eps=norm_eps,
+ attention_function=attention_function,
+ )
+
+ self.scale_shift_table_a2v_ca_audio = torch.nn.Parameter(torch.empty(5, audio.dim))
+ self.scale_shift_table_a2v_ca_video = torch.nn.Parameter(torch.empty(5, video.dim))
+
+ self.norm_eps = norm_eps
+
+ def get_ada_values(
+ self, scale_shift_table: torch.Tensor, batch_size: int, timestep: torch.Tensor, indices: slice
+ ) -> tuple[torch.Tensor, ...]:
+ num_ada_params = scale_shift_table.shape[0]
+
+ ada_values = (
+ scale_shift_table[indices].unsqueeze(0).unsqueeze(0).to(device=timestep.device, dtype=timestep.dtype)
+ + timestep.reshape(batch_size, timestep.shape[1], num_ada_params, -1)[:, :, indices, :]
+ ).unbind(dim=2)
+ return ada_values
+
+ def get_av_ca_ada_values(
+ self,
+ scale_shift_table: torch.Tensor,
+ batch_size: int,
+ scale_shift_timestep: torch.Tensor,
+ gate_timestep: torch.Tensor,
+ num_scale_shift_values: int = 4,
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ scale_shift_ada_values = self.get_ada_values(
+ scale_shift_table[:num_scale_shift_values, :], batch_size, scale_shift_timestep, slice(None, None)
+ )
+ gate_ada_values = self.get_ada_values(
+ scale_shift_table[num_scale_shift_values:, :], batch_size, gate_timestep, slice(None, None)
+ )
+
+ scale_shift_chunks = [t.squeeze(2) for t in scale_shift_ada_values]
+ gate_ada_values = [t.squeeze(2) for t in gate_ada_values]
+
+ return (*scale_shift_chunks, *gate_ada_values)
+
+ def forward( # noqa: PLR0915
+ self,
+ video: TransformerArgs | None,
+ audio: TransformerArgs | None,
+ perturbations: BatchedPerturbationConfig | None = None,
+ ) -> tuple[TransformerArgs | None, TransformerArgs | None]:
+ batch_size = video.x.shape[0]
+ if perturbations is None:
+ perturbations = BatchedPerturbationConfig.empty(batch_size)
+
+ vx = video.x if video is not None else None
+ ax = audio.x if audio is not None else None
+
+ run_vx = video is not None and video.enabled and vx.numel() > 0
+ run_ax = audio is not None and audio.enabled and ax.numel() > 0
+
+ run_a2v = run_vx and (audio is not None and audio.enabled and ax.numel() > 0)
+ run_v2a = run_ax and (video is not None and video.enabled and vx.numel() > 0)
+
+ if run_vx:
+ vshift_msa, vscale_msa, vgate_msa = self.get_ada_values(
+ self.scale_shift_table, vx.shape[0], video.timesteps, slice(0, 3)
+ )
+ if not perturbations.all_in_batch(PerturbationType.SKIP_VIDEO_SELF_ATTN, self.idx):
+ norm_vx = rms_norm(vx, eps=self.norm_eps) * (1 + vscale_msa) + vshift_msa
+ v_mask = perturbations.mask_like(PerturbationType.SKIP_VIDEO_SELF_ATTN, self.idx, vx)
+ vx = vx + self.attn1(norm_vx, pe=video.positional_embeddings) * vgate_msa * v_mask
+
+ vx = vx + self.attn2(rms_norm(vx, eps=self.norm_eps), context=video.context, mask=video.context_mask)
+
+ del vshift_msa, vscale_msa, vgate_msa
+
+ if run_ax:
+ ashift_msa, ascale_msa, agate_msa = self.get_ada_values(
+ self.audio_scale_shift_table, ax.shape[0], audio.timesteps, slice(0, 3)
+ )
+
+ if not perturbations.all_in_batch(PerturbationType.SKIP_AUDIO_SELF_ATTN, self.idx):
+ norm_ax = rms_norm(ax, eps=self.norm_eps) * (1 + ascale_msa) + ashift_msa
+ a_mask = perturbations.mask_like(PerturbationType.SKIP_AUDIO_SELF_ATTN, self.idx, ax)
+ ax = ax + self.audio_attn1(norm_ax, pe=audio.positional_embeddings) * agate_msa * a_mask
+
+ ax = ax + self.audio_attn2(rms_norm(ax, eps=self.norm_eps), context=audio.context, mask=audio.context_mask)
+
+ del ashift_msa, ascale_msa, agate_msa
+
+ # Audio - Video cross attention.
+ if run_a2v or run_v2a:
+ vx_norm3 = rms_norm(vx, eps=self.norm_eps)
+ ax_norm3 = rms_norm(ax, eps=self.norm_eps)
+
+ (
+ scale_ca_audio_hidden_states_a2v,
+ shift_ca_audio_hidden_states_a2v,
+ scale_ca_audio_hidden_states_v2a,
+ shift_ca_audio_hidden_states_v2a,
+ gate_out_v2a,
+ ) = self.get_av_ca_ada_values(
+ self.scale_shift_table_a2v_ca_audio,
+ ax.shape[0],
+ audio.cross_scale_shift_timestep,
+ audio.cross_gate_timestep,
+ )
+
+ (
+ scale_ca_video_hidden_states_a2v,
+ shift_ca_video_hidden_states_a2v,
+ scale_ca_video_hidden_states_v2a,
+ shift_ca_video_hidden_states_v2a,
+ gate_out_a2v,
+ ) = self.get_av_ca_ada_values(
+ self.scale_shift_table_a2v_ca_video,
+ vx.shape[0],
+ video.cross_scale_shift_timestep,
+ video.cross_gate_timestep,
+ )
+
+ if run_a2v:
+ vx_scaled = vx_norm3 * (1 + scale_ca_video_hidden_states_a2v) + shift_ca_video_hidden_states_a2v
+ ax_scaled = ax_norm3 * (1 + scale_ca_audio_hidden_states_a2v) + shift_ca_audio_hidden_states_a2v
+ a2v_mask = perturbations.mask_like(PerturbationType.SKIP_A2V_CROSS_ATTN, self.idx, vx)
+ vx = vx + (
+ self.audio_to_video_attn(
+ vx_scaled,
+ context=ax_scaled,
+ pe=video.cross_positional_embeddings,
+ k_pe=audio.cross_positional_embeddings,
+ )
+ * gate_out_a2v
+ * a2v_mask
+ )
+
+ if run_v2a:
+ ax_scaled = ax_norm3 * (1 + scale_ca_audio_hidden_states_v2a) + shift_ca_audio_hidden_states_v2a
+ vx_scaled = vx_norm3 * (1 + scale_ca_video_hidden_states_v2a) + shift_ca_video_hidden_states_v2a
+ v2a_mask = perturbations.mask_like(PerturbationType.SKIP_V2A_CROSS_ATTN, self.idx, ax)
+ ax = ax + (
+ self.video_to_audio_attn(
+ ax_scaled,
+ context=vx_scaled,
+ pe=audio.cross_positional_embeddings,
+ k_pe=video.cross_positional_embeddings,
+ )
+ * gate_out_v2a
+ * v2a_mask
+ )
+
+ del gate_out_a2v, gate_out_v2a
+ del (
+ scale_ca_video_hidden_states_a2v,
+ shift_ca_video_hidden_states_a2v,
+ scale_ca_audio_hidden_states_a2v,
+ shift_ca_audio_hidden_states_a2v,
+ scale_ca_video_hidden_states_v2a,
+ shift_ca_video_hidden_states_v2a,
+ scale_ca_audio_hidden_states_v2a,
+ shift_ca_audio_hidden_states_v2a,
+ )
+
+ if run_vx:
+ vshift_mlp, vscale_mlp, vgate_mlp = self.get_ada_values(
+ self.scale_shift_table, vx.shape[0], video.timesteps, slice(3, None)
+ )
+ vx_scaled = rms_norm(vx, eps=self.norm_eps) * (1 + vscale_mlp) + vshift_mlp
+ vx = vx + self.ff(vx_scaled) * vgate_mlp
+
+ del vshift_mlp, vscale_mlp, vgate_mlp
+
+ if run_ax:
+ ashift_mlp, ascale_mlp, agate_mlp = self.get_ada_values(
+ self.audio_scale_shift_table, ax.shape[0], audio.timesteps, slice(3, None)
+ )
+ ax_scaled = rms_norm(ax, eps=self.norm_eps) * (1 + ascale_mlp) + ashift_mlp
+ ax = ax + self.audio_ff(ax_scaled) * agate_mlp
+
+ del ashift_mlp, ascale_mlp, agate_mlp
+
+ return replace(video, x=vx) if video is not None else None, replace(audio, x=ax) if audio is not None else None
diff --git a/packages/ltx-core/src/ltx_core/model/transformer/transformer_args.py b/packages/ltx-core/src/ltx_core/model/transformer/transformer_args.py
new file mode 100644
index 0000000000000000000000000000000000000000..abfedf82772668dab95ad41f3b7ce4f76b5665ea
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/transformer/transformer_args.py
@@ -0,0 +1,239 @@
+from dataclasses import dataclass, replace
+
+import torch
+
+from ltx_core.model.transformer.adaln import AdaLayerNormSingle
+from ltx_core.model.transformer.modality import Modality
+from ltx_core.model.transformer.rope import (
+ LTXRopeType,
+ generate_freq_grid_np,
+ generate_freq_grid_pytorch,
+ precompute_freqs_cis,
+)
+from ltx_core.model.transformer.text_projection import PixArtAlphaTextProjection
+
+
+@dataclass(frozen=True)
+class TransformerArgs:
+ x: torch.Tensor
+ context: torch.Tensor
+ context_mask: torch.Tensor
+ timesteps: torch.Tensor
+ embedded_timestep: torch.Tensor
+ positional_embeddings: torch.Tensor
+ cross_positional_embeddings: torch.Tensor | None
+ cross_scale_shift_timestep: torch.Tensor | None
+ cross_gate_timestep: torch.Tensor | None
+ enabled: bool
+
+
+class TransformerArgsPreprocessor:
+ def __init__( # noqa: PLR0913
+ self,
+ patchify_proj: torch.nn.Linear,
+ adaln: AdaLayerNormSingle,
+ caption_projection: PixArtAlphaTextProjection,
+ inner_dim: int,
+ max_pos: list[int],
+ num_attention_heads: int,
+ use_middle_indices_grid: bool,
+ timestep_scale_multiplier: int,
+ double_precision_rope: bool,
+ positional_embedding_theta: float,
+ rope_type: LTXRopeType,
+ ) -> None:
+ self.patchify_proj = patchify_proj
+ self.adaln = adaln
+ self.caption_projection = caption_projection
+ self.inner_dim = inner_dim
+ self.max_pos = max_pos
+ self.num_attention_heads = num_attention_heads
+ self.use_middle_indices_grid = use_middle_indices_grid
+ self.timestep_scale_multiplier = timestep_scale_multiplier
+ self.double_precision_rope = double_precision_rope
+ self.positional_embedding_theta = positional_embedding_theta
+ self.rope_type = rope_type
+
+ def _prepare_timestep(
+ self, timestep: torch.Tensor, batch_size: int, hidden_dtype: torch.dtype
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """Prepare timestep embeddings."""
+
+ timestep = timestep * self.timestep_scale_multiplier
+ timestep, embedded_timestep = self.adaln(
+ timestep.flatten(),
+ hidden_dtype=hidden_dtype,
+ )
+
+ # Second dimension is 1 or number of tokens (if timestep_per_token)
+ timestep = timestep.view(batch_size, -1, timestep.shape[-1])
+ embedded_timestep = embedded_timestep.view(batch_size, -1, embedded_timestep.shape[-1])
+ return timestep, embedded_timestep
+
+ def _prepare_context(
+ self,
+ context: torch.Tensor,
+ x: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
+ """Prepare context for transformer blocks."""
+ batch_size = x.shape[0]
+ context = self.caption_projection(context)
+ context = context.view(batch_size, -1, x.shape[-1])
+
+ return context, attention_mask
+
+ def _prepare_attention_mask(self, attention_mask: torch.Tensor | None, x_dtype: torch.dtype) -> torch.Tensor | None:
+ """Prepare attention mask."""
+ if attention_mask is None or torch.is_floating_point(attention_mask):
+ return attention_mask
+
+ return (attention_mask - 1).to(x_dtype).reshape(
+ (attention_mask.shape[0], 1, -1, attention_mask.shape[-1])
+ ) * torch.finfo(x_dtype).max
+
+ def _prepare_positional_embeddings(
+ self,
+ positions: torch.Tensor,
+ inner_dim: int,
+ max_pos: list[int],
+ use_middle_indices_grid: bool,
+ num_attention_heads: int,
+ x_dtype: torch.dtype,
+ ) -> torch.Tensor:
+ """Prepare positional embeddings."""
+ freq_grid_generator = generate_freq_grid_np if self.double_precision_rope else generate_freq_grid_pytorch
+ pe = precompute_freqs_cis(
+ positions,
+ dim=inner_dim,
+ out_dtype=x_dtype,
+ theta=self.positional_embedding_theta,
+ max_pos=max_pos,
+ use_middle_indices_grid=use_middle_indices_grid,
+ num_attention_heads=num_attention_heads,
+ rope_type=self.rope_type,
+ freq_grid_generator=freq_grid_generator,
+ )
+ return pe
+
+ def prepare(
+ self,
+ modality: Modality,
+ ) -> TransformerArgs:
+ x = self.patchify_proj(modality.latent)
+ timestep, embedded_timestep = self._prepare_timestep(modality.timesteps, x.shape[0], modality.latent.dtype)
+ context, attention_mask = self._prepare_context(modality.context, x, modality.context_mask)
+ attention_mask = self._prepare_attention_mask(attention_mask, modality.latent.dtype)
+ pe = self._prepare_positional_embeddings(
+ positions=modality.positions,
+ inner_dim=self.inner_dim,
+ max_pos=self.max_pos,
+ use_middle_indices_grid=self.use_middle_indices_grid,
+ num_attention_heads=self.num_attention_heads,
+ x_dtype=modality.latent.dtype,
+ )
+ return TransformerArgs(
+ x=x,
+ context=context,
+ context_mask=attention_mask,
+ timesteps=timestep,
+ embedded_timestep=embedded_timestep,
+ positional_embeddings=pe,
+ cross_positional_embeddings=None,
+ cross_scale_shift_timestep=None,
+ cross_gate_timestep=None,
+ enabled=modality.enabled,
+ )
+
+
+class MultiModalTransformerArgsPreprocessor:
+ def __init__( # noqa: PLR0913
+ self,
+ patchify_proj: torch.nn.Linear,
+ adaln: AdaLayerNormSingle,
+ caption_projection: PixArtAlphaTextProjection,
+ cross_scale_shift_adaln: AdaLayerNormSingle,
+ cross_gate_adaln: AdaLayerNormSingle,
+ inner_dim: int,
+ max_pos: list[int],
+ num_attention_heads: int,
+ cross_pe_max_pos: int,
+ use_middle_indices_grid: bool,
+ audio_cross_attention_dim: int,
+ timestep_scale_multiplier: int,
+ double_precision_rope: bool,
+ positional_embedding_theta: float,
+ rope_type: LTXRopeType,
+ av_ca_timestep_scale_multiplier: int,
+ ) -> None:
+ self.simple_preprocessor = TransformerArgsPreprocessor(
+ patchify_proj=patchify_proj,
+ adaln=adaln,
+ caption_projection=caption_projection,
+ inner_dim=inner_dim,
+ max_pos=max_pos,
+ num_attention_heads=num_attention_heads,
+ use_middle_indices_grid=use_middle_indices_grid,
+ timestep_scale_multiplier=timestep_scale_multiplier,
+ double_precision_rope=double_precision_rope,
+ positional_embedding_theta=positional_embedding_theta,
+ rope_type=rope_type,
+ )
+ self.cross_scale_shift_adaln = cross_scale_shift_adaln
+ self.cross_gate_adaln = cross_gate_adaln
+ self.cross_pe_max_pos = cross_pe_max_pos
+ self.audio_cross_attention_dim = audio_cross_attention_dim
+ self.av_ca_timestep_scale_multiplier = av_ca_timestep_scale_multiplier
+
+ def prepare(
+ self,
+ modality: Modality,
+ ) -> TransformerArgs:
+ transformer_args = self.simple_preprocessor.prepare(modality)
+ cross_pe = self.simple_preprocessor._prepare_positional_embeddings(
+ positions=modality.positions[:, 0:1, :],
+ inner_dim=self.audio_cross_attention_dim,
+ max_pos=[self.cross_pe_max_pos],
+ use_middle_indices_grid=True,
+ num_attention_heads=self.simple_preprocessor.num_attention_heads,
+ x_dtype=modality.latent.dtype,
+ )
+
+ cross_scale_shift_timestep, cross_gate_timestep = self._prepare_cross_attention_timestep(
+ timestep=modality.timesteps,
+ timestep_scale_multiplier=self.simple_preprocessor.timestep_scale_multiplier,
+ batch_size=transformer_args.x.shape[0],
+ hidden_dtype=modality.latent.dtype,
+ )
+
+ return replace(
+ transformer_args,
+ cross_positional_embeddings=cross_pe,
+ cross_scale_shift_timestep=cross_scale_shift_timestep,
+ cross_gate_timestep=cross_gate_timestep,
+ )
+
+ def _prepare_cross_attention_timestep(
+ self,
+ timestep: torch.Tensor,
+ timestep_scale_multiplier: int,
+ batch_size: int,
+ hidden_dtype: torch.dtype,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """Prepare cross attention timestep embeddings."""
+ timestep = timestep * timestep_scale_multiplier
+
+ av_ca_factor = self.av_ca_timestep_scale_multiplier / timestep_scale_multiplier
+
+ scale_shift_timestep, _ = self.cross_scale_shift_adaln(
+ timestep.flatten(),
+ hidden_dtype=hidden_dtype,
+ )
+ scale_shift_timestep = scale_shift_timestep.view(batch_size, -1, scale_shift_timestep.shape[-1])
+ gate_noise_timestep, _ = self.cross_gate_adaln(
+ timestep.flatten() * av_ca_factor,
+ hidden_dtype=hidden_dtype,
+ )
+ gate_noise_timestep = gate_noise_timestep.view(batch_size, -1, gate_noise_timestep.shape[-1])
+
+ return scale_shift_timestep, gate_noise_timestep
diff --git a/packages/ltx-core/src/ltx_core/model/upsampler/__init__.py b/packages/ltx-core/src/ltx_core/model/upsampler/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..723df0dd627c5cbf15ce5d67cee6ef7a30c6afd5
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/upsampler/__init__.py
@@ -0,0 +1,10 @@
+"""Latent upsampler model components."""
+
+from ltx_core.model.upsampler.model import LatentUpsampler, upsample_video
+from ltx_core.model.upsampler.model_configurator import LatentUpsamplerConfigurator
+
+__all__ = [
+ "LatentUpsampler",
+ "LatentUpsamplerConfigurator",
+ "upsample_video",
+]
diff --git a/packages/ltx-core/src/ltx_core/model/upsampler/blur_downsample.py b/packages/ltx-core/src/ltx_core/model/upsampler/blur_downsample.py
new file mode 100644
index 0000000000000000000000000000000000000000..03ace5d8cb4f94c49f5f9dcce9ec733d6bac157a
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/upsampler/blur_downsample.py
@@ -0,0 +1,53 @@
+import math
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+
+
+class BlurDownsample(torch.nn.Module):
+ """
+ Anti-aliased spatial downsampling by integer stride using a fixed separable binomial kernel.
+ Applies only on H,W. Works for dims=2 or dims=3 (per-frame).
+ """
+
+ def __init__(self, dims: int, stride: int, kernel_size: int = 5) -> None:
+ super().__init__()
+ assert dims in (2, 3)
+ assert isinstance(stride, int)
+ assert stride >= 1
+ assert kernel_size >= 3
+ assert kernel_size % 2 == 1
+ self.dims = dims
+ self.stride = stride
+ self.kernel_size = kernel_size
+
+ # 5x5 separable binomial kernel using binomial coefficients [1, 4, 6, 4, 1] from
+ # the 4th row of Pascal's triangle. This kernel is used for anti-aliasing and
+ # provides a smooth approximation of a Gaussian filter (often called a "binomial filter").
+ # The 2D kernel is constructed as the outer product and normalized.
+ k = torch.tensor([math.comb(kernel_size - 1, k) for k in range(kernel_size)])
+ k2d = k[:, None] @ k[None, :]
+ k2d = (k2d / k2d.sum()).float() # shape (kernel_size, kernel_size)
+ self.register_buffer("kernel", k2d[None, None, :, :]) # (1, 1, kernel_size, kernel_size)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ if self.stride == 1:
+ return x
+
+ if self.dims == 2:
+ return self._apply_2d(x)
+ else:
+ # dims == 3: apply per-frame on H,W
+ b, _, f, _, _ = x.shape
+ x = rearrange(x, "b c f h w -> (b f) c h w")
+ x = self._apply_2d(x)
+ h2, w2 = x.shape[-2:]
+ x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f, h=h2, w=w2)
+ return x
+
+ def _apply_2d(self, x2d: torch.Tensor) -> torch.Tensor:
+ c = x2d.shape[1]
+ weight = self.kernel.expand(c, 1, self.kernel_size, self.kernel_size) # depthwise
+ x2d = F.conv2d(x2d, weight=weight, bias=None, stride=self.stride, padding=self.kernel_size // 2, groups=c)
+ return x2d
diff --git a/packages/ltx-core/src/ltx_core/model/upsampler/model.py b/packages/ltx-core/src/ltx_core/model/upsampler/model.py
new file mode 100644
index 0000000000000000000000000000000000000000..ed4f3edc6422c957182dbfa6970e66ce9a1f7842
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/upsampler/model.py
@@ -0,0 +1,142 @@
+import torch
+from einops import rearrange
+
+from ltx_core.model.upsampler.pixel_shuffle import PixelShuffleND
+from ltx_core.model.upsampler.res_block import ResBlock
+from ltx_core.model.upsampler.spatial_rational_resampler import SpatialRationalResampler
+from ltx_core.model.video_vae import VideoEncoder
+
+
+class LatentUpsampler(torch.nn.Module):
+ """
+ Model to upsample VAE latents spatially and/or temporally.
+ Args:
+ in_channels (`int`): Number of channels in the input latent
+ mid_channels (`int`): Number of channels in the middle layers
+ num_blocks_per_stage (`int`): Number of ResBlocks to use in each stage (pre/post upsampling)
+ dims (`int`): Number of dimensions for convolutions (2 or 3)
+ spatial_upsample (`bool`): Whether to spatially upsample the latent
+ temporal_upsample (`bool`): Whether to temporally upsample the latent
+ spatial_scale (`float`): Scale factor for spatial upsampling
+ rational_resampler (`bool`): Whether to use a rational resampler for spatial upsampling
+ """
+
+ def __init__(
+ self,
+ in_channels: int = 128,
+ mid_channels: int = 512,
+ num_blocks_per_stage: int = 4,
+ dims: int = 3,
+ spatial_upsample: bool = True,
+ temporal_upsample: bool = False,
+ spatial_scale: float = 2.0,
+ rational_resampler: bool = False,
+ ):
+ super().__init__()
+
+ self.in_channels = in_channels
+ self.mid_channels = mid_channels
+ self.num_blocks_per_stage = num_blocks_per_stage
+ self.dims = dims
+ self.spatial_upsample = spatial_upsample
+ self.temporal_upsample = temporal_upsample
+ self.spatial_scale = float(spatial_scale)
+ self.rational_resampler = rational_resampler
+
+ conv = torch.nn.Conv2d if dims == 2 else torch.nn.Conv3d
+
+ self.initial_conv = conv(in_channels, mid_channels, kernel_size=3, padding=1)
+ self.initial_norm = torch.nn.GroupNorm(32, mid_channels)
+ self.initial_activation = torch.nn.SiLU()
+
+ self.res_blocks = torch.nn.ModuleList([ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)])
+
+ if spatial_upsample and temporal_upsample:
+ self.upsampler = torch.nn.Sequential(
+ torch.nn.Conv3d(mid_channels, 8 * mid_channels, kernel_size=3, padding=1),
+ PixelShuffleND(3),
+ )
+ elif spatial_upsample:
+ if rational_resampler:
+ self.upsampler = SpatialRationalResampler(mid_channels=mid_channels, scale=self.spatial_scale)
+ else:
+ self.upsampler = torch.nn.Sequential(
+ torch.nn.Conv2d(mid_channels, 4 * mid_channels, kernel_size=3, padding=1),
+ PixelShuffleND(2),
+ )
+ elif temporal_upsample:
+ self.upsampler = torch.nn.Sequential(
+ torch.nn.Conv3d(mid_channels, 2 * mid_channels, kernel_size=3, padding=1),
+ PixelShuffleND(1),
+ )
+ else:
+ raise ValueError("Either spatial_upsample or temporal_upsample must be True")
+
+ self.post_upsample_res_blocks = torch.nn.ModuleList(
+ [ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)]
+ )
+
+ self.final_conv = conv(mid_channels, in_channels, kernel_size=3, padding=1)
+
+ def forward(self, latent: torch.Tensor) -> torch.Tensor:
+ b, _, f, _, _ = latent.shape
+
+ if self.dims == 2:
+ x = rearrange(latent, "b c f h w -> (b f) c h w")
+ x = self.initial_conv(x)
+ x = self.initial_norm(x)
+ x = self.initial_activation(x)
+
+ for block in self.res_blocks:
+ x = block(x)
+
+ x = self.upsampler(x)
+
+ for block in self.post_upsample_res_blocks:
+ x = block(x)
+
+ x = self.final_conv(x)
+ x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f)
+ else:
+ x = self.initial_conv(latent)
+ x = self.initial_norm(x)
+ x = self.initial_activation(x)
+
+ for block in self.res_blocks:
+ x = block(x)
+
+ if self.temporal_upsample:
+ x = self.upsampler(x)
+ # remove the first frame after upsampling.
+ # This is done because the first frame encodes one pixel frame.
+ x = x[:, :, 1:, :, :]
+ elif isinstance(self.upsampler, SpatialRationalResampler):
+ x = self.upsampler(x)
+ else:
+ x = rearrange(x, "b c f h w -> (b f) c h w")
+ x = self.upsampler(x)
+ x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f)
+
+ for block in self.post_upsample_res_blocks:
+ x = block(x)
+
+ x = self.final_conv(x)
+
+ return x
+
+
+def upsample_video(latent: torch.Tensor, video_encoder: VideoEncoder, upsampler: "LatentUpsampler") -> torch.Tensor:
+ """
+ Apply upsampling to the latent representation using the provided upsampler,
+ with normalization and un-normalization based on the video encoder's per-channel statistics.
+ Args:
+ latent: Input latent tensor of shape [B, C, F, H, W].
+ video_encoder: VideoEncoder with per_channel_statistics for normalization.
+ upsampler: LatentUpsampler module to perform upsampling.
+ Returns:
+ torch.Tensor: Upsampled and re-normalized latent tensor.
+ """
+ latent = video_encoder.per_channel_statistics.un_normalize(latent)
+ latent = upsampler(latent)
+ latent = video_encoder.per_channel_statistics.normalize(latent)
+ return latent
diff --git a/packages/ltx-core/src/ltx_core/model/upsampler/model_configurator.py b/packages/ltx-core/src/ltx_core/model/upsampler/model_configurator.py
new file mode 100644
index 0000000000000000000000000000000000000000..e8363708e9192ff4395b4ad14667138fb57d4fc9
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/upsampler/model_configurator.py
@@ -0,0 +1,30 @@
+from ltx_core.model.model_protocol import ModelConfigurator
+from ltx_core.model.upsampler.model import LatentUpsampler
+
+
+class LatentUpsamplerConfigurator(ModelConfigurator[LatentUpsampler]):
+ """
+ Configurator for LatentUpsampler model.
+ Used to create a LatentUpsampler model from a configuration dictionary.
+ """
+
+ @classmethod
+ def from_config(cls: type[LatentUpsampler], config: dict) -> LatentUpsampler:
+ in_channels = config.get("in_channels", 128)
+ mid_channels = config.get("mid_channels", 512)
+ num_blocks_per_stage = config.get("num_blocks_per_stage", 4)
+ dims = config.get("dims", 3)
+ spatial_upsample = config.get("spatial_upsample", True)
+ temporal_upsample = config.get("temporal_upsample", False)
+ spatial_scale = config.get("spatial_scale", 2.0)
+ rational_resampler = config.get("rational_resampler", False)
+ return LatentUpsampler(
+ in_channels=in_channels,
+ mid_channels=mid_channels,
+ num_blocks_per_stage=num_blocks_per_stage,
+ dims=dims,
+ spatial_upsample=spatial_upsample,
+ temporal_upsample=temporal_upsample,
+ spatial_scale=spatial_scale,
+ rational_resampler=rational_resampler,
+ )
diff --git a/packages/ltx-core/src/ltx_core/model/upsampler/pixel_shuffle.py b/packages/ltx-core/src/ltx_core/model/upsampler/pixel_shuffle.py
new file mode 100644
index 0000000000000000000000000000000000000000..f1c4260607256faac0312407216d3f6a6ba0f5d8
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/upsampler/pixel_shuffle.py
@@ -0,0 +1,54 @@
+import torch
+from einops import rearrange
+
+
+class PixelShuffleND(torch.nn.Module):
+ """
+ N-dimensional pixel shuffle operation for upsampling tensors.
+ Args:
+ dims (int): Number of dimensions to apply pixel shuffle to.
+ - 1: Temporal (e.g., frames)
+ - 2: Spatial (e.g., height and width)
+ - 3: Spatiotemporal (e.g., depth, height, width)
+ upscale_factors (tuple[int, int, int], optional): Upscaling factors for each dimension.
+ For dims=1, only the first value is used.
+ For dims=2, the first two values are used.
+ For dims=3, all three values are used.
+ The input tensor is rearranged so that the channel dimension is split into
+ smaller channels and upscaling factors, and the upscaling factors are moved
+ into the corresponding spatial/temporal dimensions.
+ Note:
+ This operation is equivalent to the patchifier operation in for the models. Consider
+ using this class instead.
+ """
+
+ def __init__(self, dims: int, upscale_factors: tuple[int, int, int] = (2, 2, 2)):
+ super().__init__()
+ assert dims in [1, 2, 3], "dims must be 1, 2, or 3"
+ self.dims = dims
+ self.upscale_factors = upscale_factors
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ if self.dims == 3:
+ return rearrange(
+ x,
+ "b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)",
+ p1=self.upscale_factors[0],
+ p2=self.upscale_factors[1],
+ p3=self.upscale_factors[2],
+ )
+ elif self.dims == 2:
+ return rearrange(
+ x,
+ "b (c p1 p2) h w -> b c (h p1) (w p2)",
+ p1=self.upscale_factors[0],
+ p2=self.upscale_factors[1],
+ )
+ elif self.dims == 1:
+ return rearrange(
+ x,
+ "b (c p1) f h w -> b c (f p1) h w",
+ p1=self.upscale_factors[0],
+ )
+ else:
+ raise ValueError(f"Unsupported dims: {self.dims}")
diff --git a/packages/ltx-core/src/ltx_core/model/upsampler/res_block.py b/packages/ltx-core/src/ltx_core/model/upsampler/res_block.py
new file mode 100644
index 0000000000000000000000000000000000000000..5427257653f5551ef624f27e9595eb9bf35d74a9
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/upsampler/res_block.py
@@ -0,0 +1,37 @@
+from typing import Optional
+
+import torch
+
+
+class ResBlock(torch.nn.Module):
+ """
+ Residual block with two convolutional layers, group normalization, and SiLU activation.
+ Args:
+ channels (int): Number of input and output channels.
+ mid_channels (Optional[int]): Number of channels in the intermediate convolution layer. Defaults to `channels`
+ if not specified.
+ dims (int): Dimensionality of the convolution (2 for Conv2d, 3 for Conv3d). Defaults to 3.
+ """
+
+ def __init__(self, channels: int, mid_channels: Optional[int] = None, dims: int = 3):
+ super().__init__()
+ if mid_channels is None:
+ mid_channels = channels
+
+ conv = torch.nn.Conv2d if dims == 2 else torch.nn.Conv3d
+
+ self.conv1 = conv(channels, mid_channels, kernel_size=3, padding=1)
+ self.norm1 = torch.nn.GroupNorm(32, mid_channels)
+ self.conv2 = conv(mid_channels, channels, kernel_size=3, padding=1)
+ self.norm2 = torch.nn.GroupNorm(32, channels)
+ self.activation = torch.nn.SiLU()
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ residual = x
+ x = self.conv1(x)
+ x = self.norm1(x)
+ x = self.activation(x)
+ x = self.conv2(x)
+ x = self.norm2(x)
+ x = self.activation(x + residual)
+ return x
diff --git a/packages/ltx-core/src/ltx_core/model/upsampler/spatial_rational_resampler.py b/packages/ltx-core/src/ltx_core/model/upsampler/spatial_rational_resampler.py
new file mode 100644
index 0000000000000000000000000000000000000000..cc308771c665ab318b9481253ee0d7e09059654a
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/upsampler/spatial_rational_resampler.py
@@ -0,0 +1,47 @@
+from typing import Tuple
+
+import torch
+from einops import rearrange
+
+from ltx_core.model.upsampler.blur_downsample import BlurDownsample
+from ltx_core.model.upsampler.pixel_shuffle import PixelShuffleND
+
+
+def _rational_for_scale(scale: float) -> Tuple[int, int]:
+ mapping = {0.75: (3, 4), 1.5: (3, 2), 2.0: (2, 1), 4.0: (4, 1)}
+ if float(scale) not in mapping:
+ raise ValueError(f"Unsupported scale {scale}. Choose from {list(mapping.keys())}")
+ return mapping[float(scale)]
+
+
+class SpatialRationalResampler(torch.nn.Module):
+ """
+ Fully-learned rational spatial scaling: up by 'num' via PixelShuffle, then anti-aliased
+ downsample by 'den' using fixed blur + stride. Operates on H,W only.
+ For dims==3, work per-frame for spatial scaling (temporal axis untouched).
+ Args:
+ mid_channels (`int`): Number of intermediate channels for the convolution layer
+ scale (`float`): Spatial scaling factor. Supported values are:
+ - 0.75: Downsample by 3/4 (reduce spatial size)
+ - 1.5: Upsample by 3/2 (increase spatial size)
+ - 2.0: Upsample by 2x (double spatial size)
+ - 4.0: Upsample by 4x (quadruple spatial size)
+ Any other value will raise a ValueError.
+ """
+
+ def __init__(self, mid_channels: int, scale: float):
+ super().__init__()
+ self.scale = float(scale)
+ self.num, self.den = _rational_for_scale(self.scale)
+ self.conv = torch.nn.Conv2d(mid_channels, (self.num**2) * mid_channels, kernel_size=3, padding=1)
+ self.pixel_shuffle = PixelShuffleND(2, upscale_factors=(self.num, self.num))
+ self.blur_down = BlurDownsample(dims=2, stride=self.den)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ b, _, f, _, _ = x.shape
+ x = rearrange(x, "b c f h w -> (b f) c h w")
+ x = self.conv(x)
+ x = self.pixel_shuffle(x)
+ x = self.blur_down(x)
+ x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f)
+ return x
diff --git a/packages/ltx-core/src/ltx_core/model/video_vae/__init__.py b/packages/ltx-core/src/ltx_core/model/video_vae/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..dfec2fddf1640a651746c8a6ec621c0222228b39
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/video_vae/__init__.py
@@ -0,0 +1,24 @@
+"""Video VAE package."""
+
+from ltx_core.model.video_vae.model_configurator import (
+ VAE_DECODER_COMFY_KEYS_FILTER,
+ VAE_ENCODER_COMFY_KEYS_FILTER,
+ VideoDecoderConfigurator,
+ VideoEncoderConfigurator,
+)
+from ltx_core.model.video_vae.tiling import SpatialTilingConfig, TemporalTilingConfig, TilingConfig
+from ltx_core.model.video_vae.video_vae import VideoDecoder, VideoEncoder, decode_video, get_video_chunks_number
+
+__all__ = [
+ "VAE_DECODER_COMFY_KEYS_FILTER",
+ "VAE_ENCODER_COMFY_KEYS_FILTER",
+ "SpatialTilingConfig",
+ "TemporalTilingConfig",
+ "TilingConfig",
+ "VideoDecoder",
+ "VideoDecoderConfigurator",
+ "VideoEncoder",
+ "VideoEncoderConfigurator",
+ "decode_video",
+ "get_video_chunks_number",
+]
diff --git a/packages/ltx-core/src/ltx_core/model/video_vae/convolution.py b/packages/ltx-core/src/ltx_core/model/video_vae/convolution.py
new file mode 100644
index 0000000000000000000000000000000000000000..95612d2e3c2af24dca8cef3c737fc672d63727df
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/video_vae/convolution.py
@@ -0,0 +1,317 @@
+from typing import Tuple, Union
+
+import torch
+from einops import rearrange
+from torch import nn
+from torch.nn import functional as F
+
+from ltx_core.model.video_vae.enums import PaddingModeType
+
+
+def make_conv_nd( # noqa: PLR0913
+ dims: Union[int, Tuple[int, int]],
+ in_channels: int,
+ out_channels: int,
+ kernel_size: int,
+ stride: int = 1,
+ padding: int = 0,
+ dilation: int = 1,
+ groups: int = 1,
+ bias: bool = True,
+ causal: bool = False,
+ spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
+ temporal_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
+) -> nn.Module:
+ if not (spatial_padding_mode == temporal_padding_mode or causal):
+ raise NotImplementedError("spatial and temporal padding modes must be equal")
+ if dims == 2:
+ return nn.Conv2d(
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=kernel_size,
+ stride=stride,
+ padding=padding,
+ dilation=dilation,
+ groups=groups,
+ bias=bias,
+ padding_mode=spatial_padding_mode.value,
+ )
+ elif dims == 3:
+ if causal:
+ return CausalConv3d(
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=kernel_size,
+ stride=stride,
+ dilation=dilation,
+ groups=groups,
+ bias=bias,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ return nn.Conv3d(
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=kernel_size,
+ stride=stride,
+ padding=padding,
+ dilation=dilation,
+ groups=groups,
+ bias=bias,
+ padding_mode=spatial_padding_mode.value,
+ )
+ elif dims == (2, 1):
+ return DualConv3d(
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=kernel_size,
+ stride=stride,
+ padding=padding,
+ bias=bias,
+ padding_mode=spatial_padding_mode.value,
+ )
+ else:
+ raise ValueError(f"unsupported dimensions: {dims}")
+
+
+def make_linear_nd(
+ dims: int,
+ in_channels: int,
+ out_channels: int,
+ bias: bool = True,
+) -> nn.Module:
+ if dims == 2:
+ return nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=bias)
+ elif dims in (3, (2, 1)):
+ return nn.Conv3d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=bias)
+ else:
+ raise ValueError(f"unsupported dimensions: {dims}")
+
+
+class DualConv3d(nn.Module):
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ kernel_size: int,
+ stride: Union[int, Tuple[int, int, int]] = 1,
+ padding: Union[int, Tuple[int, int, int]] = 0,
+ dilation: Union[int, Tuple[int, int, int]] = 1,
+ groups: int = 1,
+ bias: bool = True,
+ padding_mode: str = "zeros",
+ ) -> None:
+ super(DualConv3d, self).__init__()
+
+ self.in_channels = in_channels
+ self.out_channels = out_channels
+ self.padding_mode = padding_mode
+ # Ensure kernel_size, stride, padding, and dilation are tuples of length 3
+ if isinstance(kernel_size, int):
+ kernel_size = (kernel_size, kernel_size, kernel_size)
+ if kernel_size == (1, 1, 1):
+ raise ValueError("kernel_size must be greater than 1. Use make_linear_nd instead.")
+ if isinstance(stride, int):
+ stride = (stride, stride, stride)
+ if isinstance(padding, int):
+ padding = (padding, padding, padding)
+ if isinstance(dilation, int):
+ dilation = (dilation, dilation, dilation)
+
+ # Set parameters for convolutions
+ self.groups = groups
+ self.bias = bias
+
+ # Define the size of the channels after the first convolution
+ intermediate_channels = out_channels if in_channels < out_channels else in_channels
+
+ # Define parameters for the first convolution
+ self.weight1 = nn.Parameter(
+ torch.Tensor(
+ intermediate_channels,
+ in_channels // groups,
+ 1,
+ kernel_size[1],
+ kernel_size[2],
+ )
+ )
+ self.stride1 = (1, stride[1], stride[2])
+ self.padding1 = (0, padding[1], padding[2])
+ self.dilation1 = (1, dilation[1], dilation[2])
+ if bias:
+ self.bias1 = nn.Parameter(torch.Tensor(intermediate_channels))
+ else:
+ self.register_parameter("bias1", None)
+
+ # Define parameters for the second convolution
+ self.weight2 = nn.Parameter(torch.Tensor(out_channels, intermediate_channels // groups, kernel_size[0], 1, 1))
+ self.stride2 = (stride[0], 1, 1)
+ self.padding2 = (padding[0], 0, 0)
+ self.dilation2 = (dilation[0], 1, 1)
+ if bias:
+ self.bias2 = nn.Parameter(torch.Tensor(out_channels))
+ else:
+ self.register_parameter("bias2", None)
+
+ # Initialize weights and biases
+ self.reset_parameters()
+
+ def reset_parameters(self) -> None:
+ nn.init.kaiming_uniform_(self.weight1, a=torch.sqrt(5))
+ nn.init.kaiming_uniform_(self.weight2, a=torch.sqrt(5))
+ if self.bias:
+ fan_in1, _ = nn.init._calculate_fan_in_and_fan_out(self.weight1)
+ bound1 = 1 / torch.sqrt(fan_in1)
+ nn.init.uniform_(self.bias1, -bound1, bound1)
+ fan_in2, _ = nn.init._calculate_fan_in_and_fan_out(self.weight2)
+ bound2 = 1 / torch.sqrt(fan_in2)
+ nn.init.uniform_(self.bias2, -bound2, bound2)
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ use_conv3d: bool = False,
+ skip_time_conv: bool = False,
+ ) -> torch.Tensor:
+ if use_conv3d:
+ return self.forward_with_3d(x=x, skip_time_conv=skip_time_conv)
+ else:
+ return self.forward_with_2d(x=x, skip_time_conv=skip_time_conv)
+
+ def forward_with_3d(self, x: torch.Tensor, skip_time_conv: bool = False) -> torch.Tensor:
+ # First convolution
+ x = F.conv3d(
+ x,
+ self.weight1,
+ self.bias1,
+ self.stride1,
+ self.padding1,
+ self.dilation1,
+ self.groups,
+ padding_mode=self.padding_mode,
+ )
+
+ if skip_time_conv:
+ return x
+
+ # Second convolution
+ x = F.conv3d(
+ x,
+ self.weight2,
+ self.bias2,
+ self.stride2,
+ self.padding2,
+ self.dilation2,
+ self.groups,
+ padding_mode=self.padding_mode,
+ )
+
+ return x
+
+ def forward_with_2d(self, x: torch.Tensor, skip_time_conv: bool = False) -> torch.Tensor:
+ b, _, _, h, w = x.shape
+
+ # First 2D convolution
+ x = rearrange(x, "b c d h w -> (b d) c h w")
+ # Squeeze the depth dimension out of weight1 since it's 1
+ weight1 = self.weight1.squeeze(2)
+ # Select stride, padding, and dilation for the 2D convolution
+ stride1 = (self.stride1[1], self.stride1[2])
+ padding1 = (self.padding1[1], self.padding1[2])
+ dilation1 = (self.dilation1[1], self.dilation1[2])
+ x = F.conv2d(
+ x,
+ weight1,
+ self.bias1,
+ stride1,
+ padding1,
+ dilation1,
+ self.groups,
+ padding_mode=self.padding_mode,
+ )
+
+ _, _, h, w = x.shape
+
+ if skip_time_conv:
+ x = rearrange(x, "(b d) c h w -> b c d h w", b=b)
+ return x
+
+ # Second convolution which is essentially treated as a 1D convolution across the 'd' dimension
+ x = rearrange(x, "(b d) c h w -> (b h w) c d", b=b)
+
+ # Reshape weight2 to match the expected dimensions for conv1d
+ weight2 = self.weight2.squeeze(-1).squeeze(-1)
+ # Use only the relevant dimension for stride, padding, and dilation for the 1D convolution
+ stride2 = self.stride2[0]
+ padding2 = self.padding2[0]
+ dilation2 = self.dilation2[0]
+ x = F.conv1d(
+ x,
+ weight2,
+ self.bias2,
+ stride2,
+ padding2,
+ dilation2,
+ self.groups,
+ padding_mode=self.padding_mode,
+ )
+ x = rearrange(x, "(b h w) c d -> b c d h w", b=b, h=h, w=w)
+
+ return x
+
+ @property
+ def weight(self) -> torch.Tensor:
+ return self.weight2
+
+
+class CausalConv3d(nn.Module):
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ kernel_size: int = 3,
+ stride: Union[int, Tuple[int]] = 1,
+ dilation: int = 1,
+ groups: int = 1,
+ bias: bool = True,
+ spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
+ ) -> None:
+ super().__init__()
+
+ self.in_channels = in_channels
+ self.out_channels = out_channels
+
+ kernel_size = (kernel_size, kernel_size, kernel_size)
+ self.time_kernel_size = kernel_size[0]
+
+ dilation = (dilation, 1, 1)
+
+ height_pad = kernel_size[1] // 2
+ width_pad = kernel_size[2] // 2
+ padding = (0, height_pad, width_pad)
+
+ self.conv = nn.Conv3d(
+ in_channels,
+ out_channels,
+ kernel_size,
+ stride=stride,
+ dilation=dilation,
+ padding=padding,
+ padding_mode=spatial_padding_mode.value,
+ groups=groups,
+ bias=bias,
+ )
+
+ def forward(self, x: torch.Tensor, causal: bool = True) -> torch.Tensor:
+ if causal:
+ first_frame_pad = x[:, :, :1, :, :].repeat((1, 1, self.time_kernel_size - 1, 1, 1))
+ x = torch.concatenate((first_frame_pad, x), dim=2)
+ else:
+ first_frame_pad = x[:, :, :1, :, :].repeat((1, 1, (self.time_kernel_size - 1) // 2, 1, 1))
+ last_frame_pad = x[:, :, -1:, :, :].repeat((1, 1, (self.time_kernel_size - 1) // 2, 1, 1))
+ x = torch.concatenate((first_frame_pad, x, last_frame_pad), dim=2)
+ x = self.conv(x)
+ return x
+
+ @property
+ def weight(self) -> torch.Tensor:
+ return self.conv.weight
diff --git a/packages/ltx-core/src/ltx_core/model/video_vae/enums.py b/packages/ltx-core/src/ltx_core/model/video_vae/enums.py
new file mode 100644
index 0000000000000000000000000000000000000000..12468f3cb7a3bfbd3c96bc0d4a66f701a0b0419b
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/video_vae/enums.py
@@ -0,0 +1,20 @@
+from enum import Enum
+
+
+class NormLayerType(Enum):
+ GROUP_NORM = "group_norm"
+ PIXEL_NORM = "pixel_norm"
+
+
+class LogVarianceType(Enum):
+ PER_CHANNEL = "per_channel"
+ UNIFORM = "uniform"
+ CONSTANT = "constant"
+ NONE = "none"
+
+
+class PaddingModeType(Enum):
+ ZEROS = "zeros"
+ REFLECT = "reflect"
+ REPLICATE = "replicate"
+ CIRCULAR = "circular"
diff --git a/packages/ltx-core/src/ltx_core/model/video_vae/model_configurator.py b/packages/ltx-core/src/ltx_core/model/video_vae/model_configurator.py
new file mode 100644
index 0000000000000000000000000000000000000000..92e28fcdb120f8ea5d61c5ef39ba894cecf4034e
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/video_vae/model_configurator.py
@@ -0,0 +1,77 @@
+from ltx_core.loader.sd_ops import SDOps
+from ltx_core.model.model_protocol import ModelConfigurator
+from ltx_core.model.video_vae.enums import LogVarianceType, NormLayerType, PaddingModeType
+from ltx_core.model.video_vae.video_vae import VideoDecoder, VideoEncoder
+
+
+class VideoEncoderConfigurator(ModelConfigurator[VideoEncoder]):
+ """Configurator for creating a video VAE Encoder from a configuration dictionary."""
+
+ @classmethod
+ def from_config(cls: type[VideoEncoder], config: dict) -> VideoEncoder:
+ config = config.get("vae", {})
+ convolution_dimensions = config.get("dims", 3)
+ in_channels = config.get("in_channels", 3)
+ latent_channels = config.get("latent_channels", 128)
+ encoder_spatial_padding_mode = PaddingModeType(config.get("encoder_spatial_padding_mode", "zeros"))
+ encoder_blocks = config.get("encoder_blocks", [])
+ patch_size = config.get("patch_size", 4)
+ norm_layer_str = config.get("norm_layer", "pixel_norm")
+ latent_log_var_str = config.get("latent_log_var", "uniform")
+
+ return VideoEncoder(
+ convolution_dimensions=convolution_dimensions,
+ in_channels=in_channels,
+ out_channels=latent_channels,
+ encoder_blocks=encoder_blocks,
+ patch_size=patch_size,
+ norm_layer=NormLayerType(norm_layer_str),
+ latent_log_var=LogVarianceType(latent_log_var_str),
+ encoder_spatial_padding_mode=encoder_spatial_padding_mode,
+ )
+
+
+class VideoDecoderConfigurator(ModelConfigurator[VideoDecoder]):
+ """Configurator for creating a video VAE Decoder from a configuration dictionary."""
+
+ @classmethod
+ def from_config(cls: type[VideoDecoder], config: dict) -> VideoDecoder:
+ config = config.get("vae", {})
+ convolution_dimensions = config.get("dims", 3)
+ latent_channels = config.get("latent_channels", 128)
+ decoder_spatial_padding_mode = PaddingModeType(config.get("decoder_spatial_padding_mode", "reflect"))
+ out_channels = config.get("out_channels", 3)
+ decoder_blocks = config.get("decoder_blocks", [])
+ patch_size = config.get("patch_size", 4)
+ norm_layer_str = config.get("norm_layer", "pixel_norm")
+ causal = config.get("causal_decoder", False)
+ timestep_conditioning = config.get("timestep_conditioning", True)
+
+ return VideoDecoder(
+ convolution_dimensions=convolution_dimensions,
+ in_channels=latent_channels,
+ out_channels=out_channels,
+ decoder_blocks=decoder_blocks,
+ patch_size=patch_size,
+ norm_layer=NormLayerType(norm_layer_str),
+ causal=causal,
+ timestep_conditioning=timestep_conditioning,
+ decoder_spatial_padding_mode=decoder_spatial_padding_mode,
+ )
+
+
+VAE_DECODER_COMFY_KEYS_FILTER = (
+ SDOps("VAE_DECODER_COMFY_KEYS_FILTER")
+ .with_matching(prefix="vae.decoder.")
+ .with_matching(prefix="vae.per_channel_statistics.")
+ .with_replacement("vae.decoder.", "")
+ .with_replacement("vae.per_channel_statistics.", "per_channel_statistics.")
+)
+
+VAE_ENCODER_COMFY_KEYS_FILTER = (
+ SDOps("VAE_ENCODER_COMFY_KEYS_FILTER")
+ .with_matching(prefix="vae.encoder.")
+ .with_matching(prefix="vae.per_channel_statistics.")
+ .with_replacement("vae.encoder.", "")
+ .with_replacement("vae.per_channel_statistics.", "per_channel_statistics.")
+)
diff --git a/packages/ltx-core/src/ltx_core/model/video_vae/normalization.py b/packages/ltx-core/src/ltx_core/model/video_vae/normalization.py
new file mode 100644
index 0000000000000000000000000000000000000000..d469a1416da35ae88d90489d318e84eda7293266
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/video_vae/normalization.py
@@ -0,0 +1,3 @@
+from ltx_core.model.common.normalization import PixelNorm, build_normalization_layer
+
+__all__ = ["PixelNorm", "build_normalization_layer"]
diff --git a/packages/ltx-core/src/ltx_core/model/video_vae/ops.py b/packages/ltx-core/src/ltx_core/model/video_vae/ops.py
new file mode 100644
index 0000000000000000000000000000000000000000..415e40f093e3204fd6345c089ec82d87b0142870
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/video_vae/ops.py
@@ -0,0 +1,85 @@
+import torch
+from einops import rearrange
+from torch import nn
+
+
+def patchify(x: torch.Tensor, patch_size_hw: int, patch_size_t: int = 1) -> torch.Tensor:
+ """
+ Rearrange spatial dimensions into channels. Divides image into patch_size x patch_size blocks
+ and moves pixels from each block into separate channels (space-to-depth).
+ Args:
+ x: Input tensor (4D or 5D)
+ patch_size_hw: Spatial patch size for height and width. With patch_size_hw=4, divides HxW into 4x4 blocks.
+ patch_size_t: Temporal patch size for frames. Default=1 (no temporal patching).
+ For 5D: (B, C, F, H, W) -> (B, Cx(patch_size_hw^2)x(patch_size_t), F/patch_size_t, H/patch_size_hw, W/patch_size_hw)
+ Example: (B, 3, 33, 512, 512) with patch_size_hw=4, patch_size_t=1 -> (B, 48, 33, 128, 128)
+ """
+ if patch_size_hw == 1 and patch_size_t == 1:
+ return x
+ if x.dim() == 4:
+ x = rearrange(x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size_hw, r=patch_size_hw)
+ elif x.dim() == 5:
+ x = rearrange(
+ x,
+ "b c (f p) (h q) (w r) -> b (c p r q) f h w",
+ p=patch_size_t,
+ q=patch_size_hw,
+ r=patch_size_hw,
+ )
+ else:
+ raise ValueError(f"Invalid input shape: {x.shape}")
+
+ return x
+
+
+def unpatchify(x: torch.Tensor, patch_size_hw: int, patch_size_t: int = 1) -> torch.Tensor:
+ """
+ Rearrange channels back into spatial dimensions. Inverse of patchify - moves pixels from
+ channels back into patch_size x patch_size blocks (depth-to-space).
+ Args:
+ x: Input tensor (4D or 5D)
+ patch_size_hw: Spatial patch size for height and width. With patch_size_hw=4, expands HxW by 4x.
+ patch_size_t: Temporal patch size for frames. Default=1 (no temporal expansion).
+ For 5D: (B, Cx(patch_size_hw^2)x(patch_size_t), F, H, W) -> (B, C, Fxpatch_size_t, Hxpatch_size_hw, Wxpatch_size_hw)
+ Example: (B, 48, 33, 128, 128) with patch_size_hw=4, patch_size_t=1 -> (B, 3, 33, 512, 512)
+ """
+ if patch_size_hw == 1 and patch_size_t == 1:
+ return x
+
+ if x.dim() == 4:
+ x = rearrange(x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size_hw, r=patch_size_hw)
+ elif x.dim() == 5:
+ x = rearrange(
+ x,
+ "b (c p r q) f h w -> b c (f p) (h q) (w r)",
+ p=patch_size_t,
+ q=patch_size_hw,
+ r=patch_size_hw,
+ )
+
+ return x
+
+
+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 VAE state_dict.
+ """
+
+ def __init__(self, latent_channels: int = 128):
+ super().__init__()
+ self.register_buffer("std-of-means", torch.empty(latent_channels))
+ self.register_buffer("mean-of-means", torch.empty(latent_channels))
+ self.register_buffer("mean-of-stds", torch.empty(latent_channels))
+ self.register_buffer("mean-of-stds_over_std-of-means", torch.empty(latent_channels))
+ self.register_buffer("channel", torch.empty(latent_channels))
+
+ def un_normalize(self, x: torch.Tensor) -> torch.Tensor:
+ return (x * self.get_buffer("std-of-means").view(1, -1, 1, 1, 1).to(x)) + self.get_buffer("mean-of-means").view(
+ 1, -1, 1, 1, 1
+ ).to(x)
+
+ def normalize(self, x: torch.Tensor) -> torch.Tensor:
+ return (x - self.get_buffer("mean-of-means").view(1, -1, 1, 1, 1).to(x)) / self.get_buffer("std-of-means").view(
+ 1, -1, 1, 1, 1
+ ).to(x)
diff --git a/packages/ltx-core/src/ltx_core/model/video_vae/resnet.py b/packages/ltx-core/src/ltx_core/model/video_vae/resnet.py
new file mode 100644
index 0000000000000000000000000000000000000000..1e087400e88468ca07f5cc586a5c76d3d1224dfd
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/video_vae/resnet.py
@@ -0,0 +1,277 @@
+from typing import Optional, Tuple, Union
+
+import torch
+from torch import nn
+
+from ltx_core.model.common.normalization import PixelNorm
+from ltx_core.model.transformer.timestep_embedding import PixArtAlphaCombinedTimestepSizeEmbeddings
+from ltx_core.model.video_vae.convolution import make_conv_nd, make_linear_nd
+from ltx_core.model.video_vae.enums import NormLayerType, PaddingModeType
+
+
+class ResnetBlock3D(nn.Module):
+ r"""
+ A Resnet block.
+ Parameters:
+ in_channels (`int`): The number of channels in the input.
+ out_channels (`int`, *optional*, default to be `None`):
+ The number of output channels for the first conv layer. If None, same as `in_channels`.
+ dropout (`float`, *optional*, defaults to `0.0`): The dropout probability to use.
+ groups (`int`, *optional*, default to `32`): The number of groups to use for the first normalization layer.
+ eps (`float`, *optional*, defaults to `1e-6`): The epsilon to use for the normalization.
+ """
+
+ def __init__(
+ self,
+ dims: Union[int, Tuple[int, int]],
+ in_channels: int,
+ out_channels: Optional[int] = None,
+ dropout: float = 0.0,
+ groups: int = 32,
+ eps: float = 1e-6,
+ norm_layer: NormLayerType = NormLayerType.PIXEL_NORM,
+ inject_noise: bool = False,
+ timestep_conditioning: bool = False,
+ spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
+ ):
+ super().__init__()
+ self.in_channels = in_channels
+ out_channels = in_channels if out_channels is None else out_channels
+ self.out_channels = out_channels
+ self.inject_noise = inject_noise
+
+ if norm_layer == NormLayerType.GROUP_NORM:
+ self.norm1 = nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True)
+ elif norm_layer == NormLayerType.PIXEL_NORM:
+ self.norm1 = PixelNorm()
+
+ self.non_linearity = nn.SiLU()
+
+ self.conv1 = make_conv_nd(
+ dims,
+ in_channels,
+ out_channels,
+ kernel_size=3,
+ stride=1,
+ padding=1,
+ causal=True,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+
+ if inject_noise:
+ self.per_channel_scale1 = nn.Parameter(torch.zeros((in_channels, 1, 1)))
+
+ if norm_layer == NormLayerType.GROUP_NORM:
+ self.norm2 = nn.GroupNorm(num_groups=groups, num_channels=out_channels, eps=eps, affine=True)
+ elif norm_layer == NormLayerType.PIXEL_NORM:
+ self.norm2 = PixelNorm()
+
+ self.dropout = torch.nn.Dropout(dropout)
+
+ self.conv2 = make_conv_nd(
+ dims,
+ out_channels,
+ out_channels,
+ kernel_size=3,
+ stride=1,
+ padding=1,
+ causal=True,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+
+ if inject_noise:
+ self.per_channel_scale2 = nn.Parameter(torch.zeros((in_channels, 1, 1)))
+
+ self.conv_shortcut = (
+ make_linear_nd(dims=dims, in_channels=in_channels, out_channels=out_channels)
+ if in_channels != out_channels
+ else nn.Identity()
+ )
+
+ # Using GroupNorm with 1 group is equivalent to LayerNorm but works with (B, C, ...) layout
+ # avoiding the need for dimension rearrangement used in standard nn.LayerNorm
+ self.norm3 = (
+ nn.GroupNorm(num_groups=1, num_channels=in_channels, eps=eps, affine=True)
+ if in_channels != out_channels
+ else nn.Identity()
+ )
+
+ self.timestep_conditioning = timestep_conditioning
+
+ if timestep_conditioning:
+ self.scale_shift_table = nn.Parameter(torch.randn(4, in_channels) / in_channels**0.5)
+
+ def _feed_spatial_noise(
+ self,
+ hidden_states: torch.Tensor,
+ per_channel_scale: torch.Tensor,
+ generator: Optional[torch.Generator] = None,
+ ) -> torch.Tensor:
+ spatial_shape = hidden_states.shape[-2:]
+ device = hidden_states.device
+ dtype = hidden_states.dtype
+
+ # similar to the "explicit noise inputs" method in style-gan
+ spatial_noise = torch.randn(spatial_shape, device=device, dtype=dtype, generator=generator)[None]
+ scaled_noise = (spatial_noise * per_channel_scale)[None, :, None, ...]
+ hidden_states = hidden_states + scaled_noise
+
+ return hidden_states
+
+ def forward(
+ self,
+ input_tensor: torch.Tensor,
+ causal: bool = True,
+ timestep: Optional[torch.Tensor] = None,
+ generator: Optional[torch.Generator] = None,
+ ) -> torch.Tensor:
+ hidden_states = input_tensor
+ batch_size = hidden_states.shape[0]
+
+ hidden_states = self.norm1(hidden_states)
+ if self.timestep_conditioning:
+ if timestep is None:
+ raise ValueError("'timestep' parameter must be provided when 'timestep_conditioning' is True")
+ ada_values = self.scale_shift_table[None, ..., None, None, None].to(
+ device=hidden_states.device, dtype=hidden_states.dtype
+ ) + timestep.reshape(
+ batch_size,
+ 4,
+ -1,
+ timestep.shape[-3],
+ timestep.shape[-2],
+ timestep.shape[-1],
+ )
+ shift1, scale1, shift2, scale2 = ada_values.unbind(dim=1)
+
+ hidden_states = hidden_states * (1 + scale1) + shift1
+
+ hidden_states = self.non_linearity(hidden_states)
+
+ hidden_states = self.conv1(hidden_states, causal=causal)
+
+ if self.inject_noise:
+ hidden_states = self._feed_spatial_noise(
+ hidden_states,
+ self.per_channel_scale1.to(device=hidden_states.device, dtype=hidden_states.dtype),
+ generator=generator,
+ )
+
+ hidden_states = self.norm2(hidden_states)
+
+ if self.timestep_conditioning:
+ hidden_states = hidden_states * (1 + scale2) + shift2
+
+ hidden_states = self.non_linearity(hidden_states)
+
+ hidden_states = self.dropout(hidden_states)
+
+ hidden_states = self.conv2(hidden_states, causal=causal)
+
+ if self.inject_noise:
+ hidden_states = self._feed_spatial_noise(
+ hidden_states,
+ self.per_channel_scale2.to(device=hidden_states.device, dtype=hidden_states.dtype),
+ generator=generator,
+ )
+
+ input_tensor = self.norm3(input_tensor)
+
+ batch_size = input_tensor.shape[0]
+
+ input_tensor = self.conv_shortcut(input_tensor)
+
+ output_tensor = input_tensor + hidden_states
+
+ return output_tensor
+
+
+class UNetMidBlock3D(nn.Module):
+ """
+ A 3D UNet mid-block [`UNetMidBlock3D`] with multiple residual blocks.
+ Args:
+ in_channels (`int`): The number of input channels.
+ dropout (`float`, *optional*, defaults to 0.0): The dropout rate.
+ num_layers (`int`, *optional*, defaults to 1): The number of residual blocks.
+ resnet_eps (`float`, *optional*, 1e-6 ): The epsilon value for the resnet blocks.
+ resnet_groups (`int`, *optional*, defaults to 32):
+ The number of groups to use in the group normalization layers of the resnet blocks.
+ norm_layer (`str`, *optional*, defaults to `group_norm`):
+ The normalization layer to use. Can be either `group_norm` or `pixel_norm`.
+ inject_noise (`bool`, *optional*, defaults to `False`):
+ Whether to inject noise into the hidden states.
+ timestep_conditioning (`bool`, *optional*, defaults to `False`):
+ Whether to condition the hidden states on the timestep.
+ Returns:
+ `torch.Tensor`: The output of the last residual block, which is a tensor of shape `(batch_size,
+ in_channels, height, width)`.
+ """
+
+ def __init__(
+ self,
+ dims: Union[int, Tuple[int, int]],
+ in_channels: int,
+ dropout: float = 0.0,
+ num_layers: int = 1,
+ resnet_eps: float = 1e-6,
+ resnet_groups: int = 32,
+ norm_layer: NormLayerType = NormLayerType.GROUP_NORM,
+ inject_noise: bool = False,
+ timestep_conditioning: bool = False,
+ spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
+ ):
+ super().__init__()
+ resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
+
+ self.timestep_conditioning = timestep_conditioning
+
+ if timestep_conditioning:
+ self.time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings(
+ embedding_dim=in_channels * 4, size_emb_dim=0
+ )
+
+ self.res_blocks = nn.ModuleList(
+ [
+ ResnetBlock3D(
+ dims=dims,
+ in_channels=in_channels,
+ out_channels=in_channels,
+ eps=resnet_eps,
+ groups=resnet_groups,
+ dropout=dropout,
+ norm_layer=norm_layer,
+ inject_noise=inject_noise,
+ timestep_conditioning=timestep_conditioning,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ for _ in range(num_layers)
+ ]
+ )
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ causal: bool = True,
+ timestep: Optional[torch.Tensor] = None,
+ generator: Optional[torch.Generator] = None,
+ ) -> torch.Tensor:
+ timestep_embed = None
+ if self.timestep_conditioning:
+ if timestep is None:
+ raise ValueError("'timestep' parameter must be provided when 'timestep_conditioning' is True")
+ batch_size = hidden_states.shape[0]
+ timestep_embed = self.time_embedder(
+ timestep=timestep.flatten(),
+ hidden_dtype=hidden_states.dtype,
+ )
+ timestep_embed = timestep_embed.view(batch_size, timestep_embed.shape[-1], 1, 1, 1)
+
+ for resnet in self.res_blocks:
+ hidden_states = resnet(
+ hidden_states,
+ causal=causal,
+ timestep=timestep_embed,
+ generator=generator,
+ )
+
+ return hidden_states
diff --git a/packages/ltx-core/src/ltx_core/model/video_vae/sampling.py b/packages/ltx-core/src/ltx_core/model/video_vae/sampling.py
new file mode 100644
index 0000000000000000000000000000000000000000..648564601d51bff9398195a337a573fa9aa599b3
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/video_vae/sampling.py
@@ -0,0 +1,123 @@
+import math
+from typing import Tuple, Union
+
+import torch
+from einops import rearrange
+from torch import nn
+
+from .convolution import make_conv_nd
+from .enums import PaddingModeType
+
+
+class SpaceToDepthDownsample(nn.Module):
+ def __init__(
+ self,
+ dims: Union[int, Tuple[int, int]],
+ in_channels: int,
+ out_channels: int,
+ stride: Tuple[int, int, int],
+ spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
+ ):
+ super().__init__()
+ self.stride = stride
+ self.group_size = in_channels * math.prod(stride) // out_channels
+ self.conv = make_conv_nd(
+ dims=dims,
+ in_channels=in_channels,
+ out_channels=out_channels // math.prod(stride),
+ kernel_size=3,
+ stride=1,
+ causal=True,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ causal: bool = True,
+ ) -> torch.Tensor:
+ if self.stride[0] == 2:
+ x = torch.cat([x[:, :, :1, :, :], x], dim=2) # duplicate first frames for padding
+
+ # skip connection
+ x_in = rearrange(
+ x,
+ "b c (d p1) (h p2) (w p3) -> b (c p1 p2 p3) d h w",
+ p1=self.stride[0],
+ p2=self.stride[1],
+ p3=self.stride[2],
+ )
+ x_in = rearrange(x_in, "b (c g) d h w -> b c g d h w", g=self.group_size)
+ x_in = x_in.mean(dim=2)
+
+ # conv
+ x = self.conv(x, causal=causal)
+ x = rearrange(
+ x,
+ "b c (d p1) (h p2) (w p3) -> b (c p1 p2 p3) d h w",
+ p1=self.stride[0],
+ p2=self.stride[1],
+ p3=self.stride[2],
+ )
+
+ x = x + x_in
+
+ return x
+
+
+class DepthToSpaceUpsample(nn.Module):
+ def __init__(
+ self,
+ dims: int | Tuple[int, int],
+ in_channels: int,
+ stride: Tuple[int, int, int],
+ residual: bool = False,
+ out_channels_reduction_factor: int = 1,
+ spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
+ ):
+ super().__init__()
+ self.stride = stride
+ self.out_channels = math.prod(stride) * in_channels // out_channels_reduction_factor
+ self.conv = make_conv_nd(
+ dims=dims,
+ in_channels=in_channels,
+ out_channels=self.out_channels,
+ kernel_size=3,
+ stride=1,
+ causal=True,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ self.residual = residual
+ self.out_channels_reduction_factor = out_channels_reduction_factor
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ causal: bool = True,
+ ) -> torch.Tensor:
+ if self.residual:
+ # Reshape and duplicate the input to match the output shape
+ x_in = rearrange(
+ x,
+ "b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)",
+ p1=self.stride[0],
+ p2=self.stride[1],
+ p3=self.stride[2],
+ )
+ num_repeat = math.prod(self.stride) // self.out_channels_reduction_factor
+ x_in = x_in.repeat(1, num_repeat, 1, 1, 1)
+ if self.stride[0] == 2:
+ x_in = x_in[:, :, 1:, :, :]
+ x = self.conv(x, causal=causal)
+ x = rearrange(
+ x,
+ "b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)",
+ p1=self.stride[0],
+ p2=self.stride[1],
+ p3=self.stride[2],
+ )
+ if self.stride[0] == 2:
+ x = x[:, :, 1:, :, :]
+ if self.residual:
+ x = x + x_in
+ return x
diff --git a/packages/ltx-core/src/ltx_core/model/video_vae/tiling.py b/packages/ltx-core/src/ltx_core/model/video_vae/tiling.py
new file mode 100644
index 0000000000000000000000000000000000000000..c97b5552df50f495decdc848c45d76f75c8cc65e
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/video_vae/tiling.py
@@ -0,0 +1,260 @@
+import itertools
+from dataclasses import dataclass
+from typing import Callable, List, NamedTuple, Tuple
+
+import torch
+
+
+def compute_trapezoidal_mask_1d(
+ length: int,
+ ramp_left: int,
+ ramp_right: int,
+ left_starts_from_0: bool = False,
+) -> torch.Tensor:
+ """
+ Generate a 1D trapezoidal blending mask with linear ramps.
+ Args:
+ length: Output length of the mask.
+ ramp_left: Fade-in length on the left.
+ ramp_right: Fade-out length on the right.
+ left_starts_from_0: Whether the ramp starts from 0 or first non-zero value.
+ Useful for temporal tiles where the first tile is causal.
+ Returns:
+ A 1D tensor of shape `(length,)` with values in [0, 1].
+ """
+ if length <= 0:
+ raise ValueError("Mask length must be positive.")
+
+ ramp_left = max(0, min(ramp_left, length))
+ ramp_right = max(0, min(ramp_right, length))
+
+ mask = torch.ones(length)
+
+ if ramp_left > 0:
+ interval_length = ramp_left + 1 if left_starts_from_0 else ramp_left + 2
+ fade_in = torch.linspace(0.0, 1.0, interval_length)[:-1]
+ if not left_starts_from_0:
+ fade_in = fade_in[1:]
+ mask[:ramp_left] *= fade_in
+
+ if ramp_right > 0:
+ fade_out = torch.linspace(1.0, 0.0, steps=ramp_right + 2)[1:-1]
+ mask[-ramp_right:] *= fade_out
+
+ return mask.clamp_(0, 1)
+
+
+@dataclass(frozen=True)
+class SpatialTilingConfig:
+ """Configuration for dividing each frame into spatial tiles with optional overlap.
+ Args:
+ tile_size_in_pixels (int): Size of each tile in pixels. Must be at least 64 and divisible by 32.
+ tile_overlap_in_pixels (int, optional): Overlap between tiles in pixels. Must be divisible by 32. Defaults to 0.
+ """
+
+ tile_size_in_pixels: int
+ tile_overlap_in_pixels: int = 0
+
+ def __post_init__(self) -> None:
+ if self.tile_size_in_pixels < 64:
+ raise ValueError(f"tile_size_in_pixels must be at least 64, got {self.tile_size_in_pixels}")
+ if self.tile_size_in_pixels % 32 != 0:
+ raise ValueError(f"tile_size_in_pixels must be divisible by 32, got {self.tile_size_in_pixels}")
+ if self.tile_overlap_in_pixels % 32 != 0:
+ raise ValueError(f"tile_overlap_in_pixels must be divisible by 32, got {self.tile_overlap_in_pixels}")
+ if self.tile_overlap_in_pixels >= self.tile_size_in_pixels:
+ raise ValueError(
+ f"Overlap must be less than tile size, got {self.tile_overlap_in_pixels} and {self.tile_size_in_pixels}"
+ )
+
+
+@dataclass(frozen=True)
+class TemporalTilingConfig:
+ """Configuration for dividing a video into temporal tiles (chunks of frames) with optional overlap.
+ Args:
+ tile_size_in_frames (int): Number of frames in each tile. Must be at least 16 and divisible by 8.
+ tile_overlap_in_frames (int, optional): Number of overlapping frames between consecutive tiles.
+ Must be divisible by 8. Defaults to 0.
+ """
+
+ tile_size_in_frames: int
+ tile_overlap_in_frames: int = 0
+
+ def __post_init__(self) -> None:
+ if self.tile_size_in_frames < 16:
+ raise ValueError(f"tile_size_in_frames must be at least 16, got {self.tile_size_in_frames}")
+ if self.tile_size_in_frames % 8 != 0:
+ raise ValueError(f"tile_size_in_frames must be divisible by 8, got {self.tile_size_in_frames}")
+ if self.tile_overlap_in_frames % 8 != 0:
+ raise ValueError(f"tile_overlap_in_frames must be divisible by 8, got {self.tile_overlap_in_frames}")
+ if self.tile_overlap_in_frames >= self.tile_size_in_frames:
+ raise ValueError(
+ f"Overlap must be less than tile size, got {self.tile_overlap_in_frames} and {self.tile_size_in_frames}"
+ )
+
+
+@dataclass(frozen=True)
+class TilingConfig:
+ """Configuration for splitting video into tiles with optional overlap.
+ Attributes:
+ spatial_config: Configuration for splitting spatial dimensions into tiles.
+ temporal_config: Configuration for splitting temporal dimension into tiles.
+ """
+
+ spatial_config: SpatialTilingConfig | None = None
+ temporal_config: TemporalTilingConfig | None = None
+
+ @classmethod
+ def default(cls) -> "TilingConfig":
+ return cls(
+ spatial_config=SpatialTilingConfig(tile_size_in_pixels=512, tile_overlap_in_pixels=64),
+ temporal_config=TemporalTilingConfig(tile_size_in_frames=64, tile_overlap_in_frames=24),
+ )
+
+
+@dataclass(frozen=True)
+class DimensionIntervals:
+ """Intervals which a single dimension of the latent space is split into.
+ Each interval is defined by its start, end, left ramp, and right ramp.
+ The start and end are the indices of the first and last element (exclusive) in the interval.
+ Ramps are regions of the interval where the value of the mask tensor is
+ interpolated between 0 and 1 for blending with neighboring intervals.
+ The left ramp and right ramp values are the lengths of the left and right ramps.
+ """
+
+ starts: List[int]
+ ends: List[int]
+ left_ramps: List[int]
+ right_ramps: List[int]
+
+
+@dataclass(frozen=True)
+class LatentIntervals:
+ """Intervals which the latent tensor of given shape is split into.
+ Each dimension of the latent space is split into intervals based on the length along said dimension.
+ """
+
+ original_shape: torch.Size
+ dimension_intervals: Tuple[DimensionIntervals, ...]
+
+
+# Operation to split a single dimension of the tensor into intervals based on the length along the dimension.
+SplitOperation = Callable[[int], DimensionIntervals]
+# Operation to map the intervals in input dimension to slices and masks along a corresponding output dimension.
+MappingOperation = Callable[[DimensionIntervals], tuple[list[slice], list[torch.Tensor | None]]]
+
+
+def default_split_operation(length: int) -> DimensionIntervals:
+ return DimensionIntervals(starts=[0], ends=[length], left_ramps=[0], right_ramps=[0])
+
+
+DEFAULT_SPLIT_OPERATION: SplitOperation = default_split_operation
+
+
+def default_mapping_operation(
+ _intervals: DimensionIntervals,
+) -> tuple[list[slice], list[torch.Tensor | None]]:
+ return [slice(0, None)], [None]
+
+
+DEFAULT_MAPPING_OPERATION: MappingOperation = default_mapping_operation
+
+
+class Tile(NamedTuple):
+ """
+ Represents a single tile.
+ Attributes:
+ in_coords:
+ Tuple of slices specifying where to cut the tile from the INPUT tensor.
+ out_coords:
+ Tuple of slices specifying where this tile's OUTPUT should be placed in the reconstructed OUTPUT tensor.
+ masks_1d:
+ Per-dimension masks in OUTPUT units.
+ These are used to create all-dimensional blending mask.
+ Methods:
+ blend_mask:
+ Create a single N-D mask from the per-dimension masks.
+ """
+
+ in_coords: Tuple[slice, ...]
+ out_coords: Tuple[slice, ...]
+ masks_1d: Tuple[Tuple[torch.Tensor, ...]]
+
+ @property
+ def blend_mask(self) -> torch.Tensor:
+ num_dims = len(self.out_coords)
+ per_dimension_masks: List[torch.Tensor] = []
+
+ for dim_idx in range(num_dims):
+ mask_1d = self.masks_1d[dim_idx]
+ view_shape = [1] * num_dims
+ if mask_1d is None:
+ # Broadcast mask along this dimension (length 1).
+ one = torch.ones(1)
+
+ view_shape[dim_idx] = 1
+ per_dimension_masks.append(one.view(*view_shape))
+ continue
+
+ # Reshape (L,) -> (1, ..., L, ..., 1) so masks across dimensions broadcast-multiply.
+ view_shape[dim_idx] = mask_1d.shape[0]
+ per_dimension_masks.append(mask_1d.view(*view_shape))
+
+ # Multiply per-dimension masks to form the full N-D mask (separable blending window).
+ combined_mask = per_dimension_masks[0]
+ for mask in per_dimension_masks[1:]:
+ combined_mask = combined_mask * mask
+
+ return combined_mask
+
+
+def create_tiles_from_intervals_and_mappers(
+ intervals: LatentIntervals,
+ mappers: List[MappingOperation],
+) -> List[Tile]:
+ full_dim_input_slices = []
+ full_dim_output_slices = []
+ full_dim_masks_1d = []
+ for axis_index in range(len(intervals.original_shape)):
+ dimension_intervals = intervals.dimension_intervals[axis_index]
+ starts = dimension_intervals.starts
+ ends = dimension_intervals.ends
+ input_slices = [slice(s, e) for s, e in zip(starts, ends, strict=True)]
+ output_slices, masks_1d = mappers[axis_index](dimension_intervals)
+ full_dim_input_slices.append(input_slices)
+ full_dim_output_slices.append(output_slices)
+ full_dim_masks_1d.append(masks_1d)
+
+ tiles = []
+ tile_in_coords = list(itertools.product(*full_dim_input_slices))
+ tile_out_coords = list(itertools.product(*full_dim_output_slices))
+ tile_mask_1ds = list(itertools.product(*full_dim_masks_1d))
+ for in_coord, out_coord, mask_1d in zip(tile_in_coords, tile_out_coords, tile_mask_1ds, strict=True):
+ tiles.append(
+ Tile(
+ in_coords=in_coord,
+ out_coords=out_coord,
+ masks_1d=mask_1d,
+ )
+ )
+ return tiles
+
+
+def create_tiles(
+ latent_shape: torch.Size,
+ splitters: List[SplitOperation],
+ mappers: List[MappingOperation],
+) -> List[Tile]:
+ if len(splitters) != len(latent_shape):
+ raise ValueError(
+ f"Number of splitters must be equal to number of dimensions in latent shape, "
+ f"got {len(splitters)} and {len(latent_shape)}"
+ )
+ if len(mappers) != len(latent_shape):
+ raise ValueError(
+ f"Number of mappers must be equal to number of dimensions in latent shape, "
+ f"got {len(mappers)} and {len(latent_shape)}"
+ )
+ intervals = [splitter(length) for splitter, length in zip(splitters, latent_shape, strict=True)]
+ latent_intervals = LatentIntervals(original_shape=latent_shape, dimension_intervals=tuple(intervals))
+ return create_tiles_from_intervals_and_mappers(latent_intervals, mappers)
diff --git a/packages/ltx-core/src/ltx_core/model/video_vae/video_vae.py b/packages/ltx-core/src/ltx_core/model/video_vae/video_vae.py
new file mode 100644
index 0000000000000000000000000000000000000000..9d12a6737f75e2b8b0204fac04b21c6904f168c5
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/model/video_vae/video_vae.py
@@ -0,0 +1,925 @@
+from dataclasses import replace
+from typing import Any, Callable, Iterator, List, Optional, Tuple
+
+import torch
+from einops import rearrange
+from torch import nn
+
+from ltx_core.model.common.normalization import PixelNorm
+from ltx_core.model.transformer.timestep_embedding import PixArtAlphaCombinedTimestepSizeEmbeddings
+from ltx_core.model.video_vae.convolution import make_conv_nd
+from ltx_core.model.video_vae.enums import LogVarianceType, NormLayerType, PaddingModeType
+from ltx_core.model.video_vae.ops import PerChannelStatistics, patchify, unpatchify
+from ltx_core.model.video_vae.resnet import ResnetBlock3D, UNetMidBlock3D
+from ltx_core.model.video_vae.sampling import DepthToSpaceUpsample, SpaceToDepthDownsample
+from ltx_core.model.video_vae.tiling import (
+ DEFAULT_MAPPING_OPERATION,
+ DEFAULT_SPLIT_OPERATION,
+ DimensionIntervals,
+ MappingOperation,
+ SplitOperation,
+ Tile,
+ TilingConfig,
+ compute_trapezoidal_mask_1d,
+ create_tiles,
+)
+from ltx_core.types import SpatioTemporalScaleFactors, VideoLatentShape
+
+
+def _make_encoder_block(
+ block_name: str,
+ block_config: dict[str, Any],
+ in_channels: int,
+ convolution_dimensions: int,
+ norm_layer: NormLayerType,
+ norm_num_groups: int,
+ spatial_padding_mode: PaddingModeType,
+) -> Tuple[nn.Module, int]:
+ out_channels = in_channels
+
+ if block_name == "res_x":
+ block = UNetMidBlock3D(
+ dims=convolution_dimensions,
+ in_channels=in_channels,
+ num_layers=block_config["num_layers"],
+ resnet_eps=1e-6,
+ resnet_groups=norm_num_groups,
+ norm_layer=norm_layer,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ elif block_name == "res_x_y":
+ out_channels = in_channels * block_config.get("multiplier", 2)
+ block = ResnetBlock3D(
+ dims=convolution_dimensions,
+ in_channels=in_channels,
+ out_channels=out_channels,
+ eps=1e-6,
+ groups=norm_num_groups,
+ norm_layer=norm_layer,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ elif block_name == "compress_time":
+ block = make_conv_nd(
+ dims=convolution_dimensions,
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=3,
+ stride=(2, 1, 1),
+ causal=True,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ elif block_name == "compress_space":
+ block = make_conv_nd(
+ dims=convolution_dimensions,
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=3,
+ stride=(1, 2, 2),
+ causal=True,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ elif block_name == "compress_all":
+ block = make_conv_nd(
+ dims=convolution_dimensions,
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=3,
+ stride=(2, 2, 2),
+ causal=True,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ elif block_name == "compress_all_x_y":
+ out_channels = in_channels * block_config.get("multiplier", 2)
+ block = make_conv_nd(
+ dims=convolution_dimensions,
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=3,
+ stride=(2, 2, 2),
+ causal=True,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ elif block_name == "compress_all_res":
+ out_channels = in_channels * block_config.get("multiplier", 2)
+ block = SpaceToDepthDownsample(
+ dims=convolution_dimensions,
+ in_channels=in_channels,
+ out_channels=out_channels,
+ stride=(2, 2, 2),
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ elif block_name == "compress_space_res":
+ out_channels = in_channels * block_config.get("multiplier", 2)
+ block = SpaceToDepthDownsample(
+ dims=convolution_dimensions,
+ in_channels=in_channels,
+ out_channels=out_channels,
+ stride=(1, 2, 2),
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ elif block_name == "compress_time_res":
+ out_channels = in_channels * block_config.get("multiplier", 2)
+ block = SpaceToDepthDownsample(
+ dims=convolution_dimensions,
+ in_channels=in_channels,
+ out_channels=out_channels,
+ stride=(2, 1, 1),
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ else:
+ raise ValueError(f"unknown block: {block_name}")
+
+ return block, out_channels
+
+
+class VideoEncoder(nn.Module):
+ _DEFAULT_NORM_NUM_GROUPS = 32
+ """
+ Variational Autoencoder Encoder. Encodes video frames into a latent representation.
+ The encoder compresses the input video through a series of downsampling operations controlled by
+ patch_size and encoder_blocks. The output is a normalized latent tensor with shape (B, 128, F', H', W').
+ Compression Behavior:
+ The total compression is determined by:
+ 1. Initial spatial compression via patchify: H -> H/4, W -> W/4 (patch_size=4)
+ 2. Sequential compression through encoder_blocks based on their stride patterns
+ Compression blocks apply 2x compression in specified dimensions:
+ - "compress_time" / "compress_time_res": temporal only
+ - "compress_space" / "compress_space_res": spatial only (H and W)
+ - "compress_all" / "compress_all_res": all dimensions (F, H, W)
+ - "res_x" / "res_x_y": no compression
+ Standard LTX Video configuration:
+ - patch_size=4
+ - encoder_blocks: 1x compress_space_res, 1x compress_time_res, 2x compress_all_res
+ - Final dimensions: F' = 1 + (F-1)/8, H' = H/32, W' = W/32
+ - Example: (B, 3, 33, 512, 512) -> (B, 128, 5, 16, 16)
+ - Note: Input must have 1 + 8*k frames (e.g., 1, 9, 17, 25, 33...)
+ Args:
+ convolution_dimensions: The number of dimensions to use in convolutions (2D or 3D).
+ in_channels: The number of input channels. For RGB images, this is 3.
+ out_channels: The number of output channels (latent channels). For latent channels, this is 128.
+ encoder_blocks: The list of blocks to construct the encoder. Each block is a tuple of (block_name, params)
+ where params is either an int (num_layers) or a dict with configuration.
+ patch_size: The patch size for initial spatial compression. Should be a power of 2.
+ norm_layer: The normalization layer to use. Can be either `group_norm` or `pixel_norm`.
+ latent_log_var: The log variance mode. Can be either `per_channel`, `uniform`, `constant` or `none`.
+ """
+
+ def __init__(
+ self,
+ convolution_dimensions: int = 3,
+ in_channels: int = 3,
+ out_channels: int = 128,
+ encoder_blocks: List[Tuple[str, int]] | List[Tuple[str, dict[str, Any]]] = [], # noqa: B006
+ patch_size: int = 4,
+ norm_layer: NormLayerType = NormLayerType.PIXEL_NORM,
+ latent_log_var: LogVarianceType = LogVarianceType.UNIFORM,
+ encoder_spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
+ ):
+ super().__init__()
+
+ self.patch_size = patch_size
+ self.norm_layer = norm_layer
+ self.latent_channels = out_channels
+ self.latent_log_var = latent_log_var
+ self._norm_num_groups = self._DEFAULT_NORM_NUM_GROUPS
+
+ # Per-channel statistics for normalizing latents
+ self.per_channel_statistics = PerChannelStatistics(latent_channels=out_channels)
+
+ in_channels = in_channels * patch_size**2
+ feature_channels = out_channels
+
+ self.conv_in = make_conv_nd(
+ dims=convolution_dimensions,
+ in_channels=in_channels,
+ out_channels=feature_channels,
+ kernel_size=3,
+ stride=1,
+ padding=1,
+ causal=True,
+ spatial_padding_mode=encoder_spatial_padding_mode,
+ )
+
+ self.down_blocks = nn.ModuleList([])
+
+ for block_name, block_params in encoder_blocks:
+ # Convert int to dict format for uniform handling
+ block_config = {"num_layers": block_params} if isinstance(block_params, int) else block_params
+
+ block, feature_channels = _make_encoder_block(
+ block_name=block_name,
+ block_config=block_config,
+ in_channels=feature_channels,
+ convolution_dimensions=convolution_dimensions,
+ norm_layer=norm_layer,
+ norm_num_groups=self._norm_num_groups,
+ spatial_padding_mode=encoder_spatial_padding_mode,
+ )
+
+ self.down_blocks.append(block)
+
+ # out
+ if norm_layer == NormLayerType.GROUP_NORM:
+ self.conv_norm_out = nn.GroupNorm(num_channels=feature_channels, num_groups=self._norm_num_groups, eps=1e-6)
+ elif norm_layer == NormLayerType.PIXEL_NORM:
+ self.conv_norm_out = PixelNorm()
+
+ self.conv_act = nn.SiLU()
+
+ conv_out_channels = out_channels
+ if latent_log_var == LogVarianceType.PER_CHANNEL:
+ conv_out_channels *= 2
+ elif latent_log_var in {LogVarianceType.UNIFORM, LogVarianceType.CONSTANT}:
+ conv_out_channels += 1
+ elif latent_log_var != LogVarianceType.NONE:
+ raise ValueError(f"Invalid latent_log_var: {latent_log_var}")
+
+ self.conv_out = make_conv_nd(
+ dims=convolution_dimensions,
+ in_channels=feature_channels,
+ out_channels=conv_out_channels,
+ kernel_size=3,
+ padding=1,
+ causal=True,
+ spatial_padding_mode=encoder_spatial_padding_mode,
+ )
+
+ def forward(self, sample: torch.Tensor) -> torch.Tensor:
+ r"""
+ Encode video frames into normalized latent representation.
+ Args:
+ sample: Input video (B, C, F, H, W). F must be 1 + 8*k (e.g., 1, 9, 17, 25, 33...).
+ Returns:
+ Normalized latent means (B, 128, F', H', W') where F' = 1+(F-1)/8, H' = H/32, W' = W/32.
+ Example: (B, 3, 33, 512, 512) -> (B, 128, 5, 16, 16).
+ """
+ # Validate frame count
+ frames_count = sample.shape[2]
+ if ((frames_count - 1) % 8) != 0:
+ raise ValueError(
+ "Invalid number of frames: Encode input must have 1 + 8 * x frames "
+ "(e.g., 1, 9, 17, ...). Please check your input."
+ )
+
+ # Initial spatial compression: trade spatial resolution for channel depth
+ # This reduces H,W by patch_size and increases channels, making convolutions more efficient
+ # Example: (B, 3, F, 512, 512) -> (B, 48, F, 128, 128) with patch_size=4
+ sample = patchify(sample, patch_size_hw=self.patch_size, patch_size_t=1)
+ sample = self.conv_in(sample)
+
+ for down_block in self.down_blocks:
+ sample = down_block(sample)
+
+ sample = self.conv_norm_out(sample)
+ sample = self.conv_act(sample)
+ sample = self.conv_out(sample)
+
+ if self.latent_log_var == LogVarianceType.UNIFORM:
+ # Uniform Variance: model outputs N means and 1 shared log-variance channel.
+ # We need to expand the single logvar to match the number of means channels
+ # to create a format compatible with PER_CHANNEL (means + logvar, each with N channels).
+ # Sample shape: (B, N+1, ...) where N = latent_channels (e.g., 128 means + 1 logvar = 129)
+ # Target shape: (B, 2*N, ...) where first N are means, last N are logvar
+
+ if sample.shape[1] < 2:
+ raise ValueError(
+ f"Invalid channel count for UNIFORM mode: expected at least 2 channels "
+ f"(N means + 1 logvar), got {sample.shape[1]}"
+ )
+
+ # Extract means (first N channels) and logvar (last 1 channel)
+ means = sample[:, :-1, ...] # (B, N, ...)
+ logvar = sample[:, -1:, ...] # (B, 1, ...)
+
+ # Repeat logvar N times to match means channels
+ # Use expand/repeat pattern that works for both 4D and 5D tensors
+ num_channels = means.shape[1]
+ repeat_shape = [1, num_channels] + [1] * (sample.ndim - 2)
+ repeated_logvar = logvar.repeat(*repeat_shape) # (B, N, ...)
+
+ # Concatenate to create (B, 2*N, ...) format: [means, repeated_logvar]
+ sample = torch.cat([means, repeated_logvar], dim=1)
+ elif self.latent_log_var == LogVarianceType.CONSTANT:
+ sample = sample[:, :-1, ...]
+ approx_ln_0 = -30 # this is the minimal clamp value in DiagonalGaussianDistribution objects
+ sample = torch.cat(
+ [sample, torch.ones_like(sample, device=sample.device) * approx_ln_0],
+ dim=1,
+ )
+
+ # Split into means and logvar, then normalize means
+ means, _ = torch.chunk(sample, 2, dim=1)
+ return self.per_channel_statistics.normalize(means)
+
+
+def _make_decoder_block(
+ block_name: str,
+ block_config: dict[str, Any],
+ in_channels: int,
+ convolution_dimensions: int,
+ norm_layer: NormLayerType,
+ timestep_conditioning: bool,
+ norm_num_groups: int,
+ spatial_padding_mode: PaddingModeType,
+) -> Tuple[nn.Module, int]:
+ out_channels = in_channels
+ if block_name == "res_x":
+ block = UNetMidBlock3D(
+ dims=convolution_dimensions,
+ in_channels=in_channels,
+ num_layers=block_config["num_layers"],
+ resnet_eps=1e-6,
+ resnet_groups=norm_num_groups,
+ norm_layer=norm_layer,
+ inject_noise=block_config.get("inject_noise", False),
+ timestep_conditioning=timestep_conditioning,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ elif block_name == "attn_res_x":
+ block = UNetMidBlock3D(
+ dims=convolution_dimensions,
+ in_channels=in_channels,
+ num_layers=block_config["num_layers"],
+ resnet_groups=norm_num_groups,
+ norm_layer=norm_layer,
+ inject_noise=block_config.get("inject_noise", False),
+ timestep_conditioning=timestep_conditioning,
+ attention_head_dim=block_config["attention_head_dim"],
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ elif block_name == "res_x_y":
+ out_channels = in_channels // block_config.get("multiplier", 2)
+ block = ResnetBlock3D(
+ dims=convolution_dimensions,
+ in_channels=in_channels,
+ out_channels=out_channels,
+ eps=1e-6,
+ groups=norm_num_groups,
+ norm_layer=norm_layer,
+ inject_noise=block_config.get("inject_noise", False),
+ timestep_conditioning=False,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ elif block_name == "compress_time":
+ block = DepthToSpaceUpsample(
+ dims=convolution_dimensions,
+ in_channels=in_channels,
+ stride=(2, 1, 1),
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ elif block_name == "compress_space":
+ block = DepthToSpaceUpsample(
+ dims=convolution_dimensions,
+ in_channels=in_channels,
+ stride=(1, 2, 2),
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ elif block_name == "compress_all":
+ out_channels = in_channels // block_config.get("multiplier", 1)
+ block = DepthToSpaceUpsample(
+ dims=convolution_dimensions,
+ in_channels=in_channels,
+ stride=(2, 2, 2),
+ residual=block_config.get("residual", False),
+ out_channels_reduction_factor=block_config.get("multiplier", 1),
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ else:
+ raise ValueError(f"unknown layer: {block_name}")
+
+ return block, out_channels
+
+
+class VideoDecoder(nn.Module):
+ _DEFAULT_NORM_NUM_GROUPS = 32
+ """
+ Variational Autoencoder Decoder. Decodes latent representation into video frames.
+ The decoder upsamples latents through a series of upsampling operations (inverse of encoder).
+ Output dimensions: F = 8x(F'-1) + 1, H = 32xH', W = 32xW' for standard LTX Video configuration.
+ Upsampling blocks expand dimensions by 2x in specified dimensions:
+ - "compress_time": temporal only
+ - "compress_space": spatial only (H and W)
+ - "compress_all": all dimensions (F, H, W)
+ - "res_x" / "res_x_y" / "attn_res_x": no upsampling
+ Causal Mode:
+ causal=False (standard): Symmetric padding, allows future frame dependencies.
+ causal=True: Causal padding, each frame depends only on past/current frames.
+ First frame removed after temporal upsampling in both modes. Output shape unchanged.
+ Example: (B, 128, 5, 16, 16) -> (B, 3, 33, 512, 512) for both modes.
+ Args:
+ convolution_dimensions: The number of dimensions to use in convolutions (2D or 3D).
+ in_channels: The number of input channels (latent channels). Default is 128.
+ out_channels: The number of output channels. For RGB images, this is 3.
+ decoder_blocks: The list of blocks to construct the decoder. Each block is a tuple of (block_name, params)
+ where params is either an int (num_layers) or a dict with configuration.
+ patch_size: Final spatial expansion factor. For standard LTX Video, use 4 for 4x spatial expansion:
+ H -> Hx4, W -> Wx4. Should be a power of 2.
+ norm_layer: The normalization layer to use. Can be either `group_norm` or `pixel_norm`.
+ causal: Whether to use causal convolutions. For standard LTX Video, use False for symmetric padding.
+ When True, uses causal padding (past/current frames only).
+ timestep_conditioning: Whether to condition the decoder on timestep for denoising.
+ """
+
+ def __init__(
+ self,
+ convolution_dimensions: int = 3,
+ in_channels: int = 128,
+ out_channels: int = 3,
+ decoder_blocks: List[Tuple[str, int | dict]] = [], # noqa: B006
+ patch_size: int = 4,
+ norm_layer: NormLayerType = NormLayerType.PIXEL_NORM,
+ causal: bool = False,
+ timestep_conditioning: bool = False,
+ decoder_spatial_padding_mode: PaddingModeType = PaddingModeType.REFLECT,
+ ):
+ super().__init__()
+
+ # Spatiotemporal downscaling between decoded video space and VAE latents.
+ # According to the LTXV paper, the standard configuration downsamples
+ # video inputs by a factor of 8 in the temporal dimension and 32 in
+ # each spatial dimension (height and width). This parameter determines how
+ # many video frames and pixels correspond to a single latent cell.
+ self.video_downscale_factors = SpatioTemporalScaleFactors(
+ time=8,
+ width=32,
+ height=32,
+ )
+
+ self.patch_size = patch_size
+ out_channels = out_channels * patch_size**2
+ self.causal = causal
+ self.timestep_conditioning = timestep_conditioning
+ self._norm_num_groups = self._DEFAULT_NORM_NUM_GROUPS
+
+ # Per-channel statistics for denormalizing latents
+ self.per_channel_statistics = PerChannelStatistics(latent_channels=in_channels)
+
+ # Noise and timestep parameters for decoder conditioning
+ self.decode_noise_scale = 0.025
+ self.decode_timestep = 0.05
+
+ # Compute initial feature_channels by going through blocks in reverse
+ # This determines the channel width at the start of the decoder
+ feature_channels = in_channels
+ for block_name, block_params in list(reversed(decoder_blocks)):
+ block_config = block_params if isinstance(block_params, dict) else {}
+ if block_name == "res_x_y":
+ feature_channels = feature_channels * block_config.get("multiplier", 2)
+ if block_name == "compress_all":
+ feature_channels = feature_channels * block_config.get("multiplier", 1)
+
+ self.conv_in = make_conv_nd(
+ dims=convolution_dimensions,
+ in_channels=in_channels,
+ out_channels=feature_channels,
+ kernel_size=3,
+ stride=1,
+ padding=1,
+ causal=True,
+ spatial_padding_mode=decoder_spatial_padding_mode,
+ )
+
+ self.up_blocks = nn.ModuleList([])
+
+ for block_name, block_params in list(reversed(decoder_blocks)):
+ # Convert int to dict format for uniform handling
+ block_config = {"num_layers": block_params} if isinstance(block_params, int) else block_params
+
+ block, feature_channels = _make_decoder_block(
+ block_name=block_name,
+ block_config=block_config,
+ in_channels=feature_channels,
+ convolution_dimensions=convolution_dimensions,
+ norm_layer=norm_layer,
+ timestep_conditioning=timestep_conditioning,
+ norm_num_groups=self._norm_num_groups,
+ spatial_padding_mode=decoder_spatial_padding_mode,
+ )
+
+ self.up_blocks.append(block)
+
+ if norm_layer == NormLayerType.GROUP_NORM:
+ self.conv_norm_out = nn.GroupNorm(num_channels=feature_channels, num_groups=self._norm_num_groups, eps=1e-6)
+ elif norm_layer == NormLayerType.PIXEL_NORM:
+ self.conv_norm_out = PixelNorm()
+
+ self.conv_act = nn.SiLU()
+ self.conv_out = make_conv_nd(
+ dims=convolution_dimensions,
+ in_channels=feature_channels,
+ out_channels=out_channels,
+ kernel_size=3,
+ padding=1,
+ causal=True,
+ spatial_padding_mode=decoder_spatial_padding_mode,
+ )
+
+ if timestep_conditioning:
+ self.timestep_scale_multiplier = nn.Parameter(torch.tensor(1000.0))
+ self.last_time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings(
+ embedding_dim=feature_channels * 2, size_emb_dim=0
+ )
+ self.last_scale_shift_table = nn.Parameter(torch.empty(2, feature_channels))
+
+ # def forward(self, sample: torch.Tensor, target_shape) -> torch.Tensor:
+ def forward(
+ self,
+ sample: torch.Tensor,
+ timestep: Optional[torch.Tensor] = None,
+ generator: Optional[torch.Generator] = None,
+ ) -> torch.Tensor:
+ r"""
+ Decode latent representation into video frames.
+ Args:
+ sample: Latent tensor (B, 128, F', H', W').
+ timestep: Timestep for conditioning (if timestep_conditioning=True). Uses default 0.05 if None.
+ generator: Random generator for deterministic noise injection (if inject_noise=True in blocks).
+ Returns:
+ Decoded video (B, 3, F, H, W) where F = 8x(F'-1) + 1, H = 32xH', W = 32xW'.
+ Example: (B, 128, 5, 16, 16) -> (B, 3, 33, 512, 512).
+ Note: First frame is removed after temporal upsampling regardless of causal mode.
+ When causal=False, allows future frame dependencies in convolutions but maintains same output shape.
+ """
+ batch_size = sample.shape[0]
+
+ # Add noise if timestep conditioning is enabled
+ if self.timestep_conditioning:
+ noise = (
+ torch.randn(
+ sample.size(),
+ generator=generator,
+ dtype=sample.dtype,
+ device=sample.device,
+ )
+ * self.decode_noise_scale
+ )
+
+ sample = noise + (1.0 - self.decode_noise_scale) * sample
+
+ # Denormalize latents
+ sample = self.per_channel_statistics.un_normalize(sample)
+
+ # Use default decode_timestep if timestep not provided
+ if timestep is None and self.timestep_conditioning:
+ timestep = torch.full((batch_size,), self.decode_timestep, device=sample.device, dtype=sample.dtype)
+
+ sample = self.conv_in(sample, causal=self.causal)
+
+ scaled_timestep = None
+ if self.timestep_conditioning:
+ if timestep is None:
+ raise ValueError("'timestep' parameter must be provided when 'timestep_conditioning' is True")
+ scaled_timestep = timestep * self.timestep_scale_multiplier.to(sample)
+
+ for up_block in self.up_blocks:
+ if isinstance(up_block, UNetMidBlock3D):
+ block_kwargs = {
+ "causal": self.causal,
+ "timestep": scaled_timestep if self.timestep_conditioning else None,
+ "generator": generator,
+ }
+ sample = up_block(sample, **block_kwargs)
+ elif isinstance(up_block, ResnetBlock3D):
+ sample = up_block(sample, causal=self.causal, generator=generator)
+ else:
+ sample = up_block(sample, causal=self.causal)
+
+ sample = self.conv_norm_out(sample)
+
+ if self.timestep_conditioning:
+ embedded_timestep = self.last_time_embedder(
+ timestep=scaled_timestep.flatten(),
+ hidden_dtype=sample.dtype,
+ )
+ embedded_timestep = embedded_timestep.view(batch_size, embedded_timestep.shape[-1], 1, 1, 1)
+ ada_values = self.last_scale_shift_table[None, ..., None, None, None].to(
+ device=sample.device, dtype=sample.dtype
+ ) + embedded_timestep.reshape(
+ batch_size,
+ 2,
+ -1,
+ embedded_timestep.shape[-3],
+ embedded_timestep.shape[-2],
+ embedded_timestep.shape[-1],
+ )
+ shift, scale = ada_values.unbind(dim=1)
+ sample = sample * (1 + scale) + shift
+
+ sample = self.conv_act(sample)
+ sample = self.conv_out(sample, causal=self.causal)
+
+ # Final spatial expansion: reverse the initial patchify from encoder
+ # Moves pixels from channels back to spatial dimensions
+ # Example: (B, 48, F, 128, 128) -> (B, 3, F, 512, 512) with patch_size=4
+ sample = unpatchify(sample, patch_size_hw=self.patch_size, patch_size_t=1)
+
+ return sample
+
+ def _prepare_tiles(
+ self,
+ latent: torch.Tensor,
+ tiling_config: TilingConfig | None = None,
+ ) -> List[Tile]:
+ splitters = [DEFAULT_SPLIT_OPERATION] * len(latent.shape)
+ mappers = [DEFAULT_MAPPING_OPERATION] * len(latent.shape)
+ if tiling_config is not None and tiling_config.spatial_config is not None:
+ cfg = tiling_config.spatial_config
+ long_side = max(latent.shape[3], latent.shape[4])
+
+ def enable_on_axis(axis_idx: int, factor: int) -> None:
+ size = cfg.tile_size_in_pixels // factor
+ overlap = cfg.tile_overlap_in_pixels // factor
+ axis_length = latent.shape[axis_idx]
+ lower_threshold = max(2, overlap + 1)
+ tile_size = max(lower_threshold, round(size * axis_length / long_side))
+ splitters[axis_idx] = split_in_spatial(tile_size, overlap)
+ mappers[axis_idx] = to_mapping_operation(map_spatial_slice, factor)
+
+ enable_on_axis(3, self.video_downscale_factors.height)
+ enable_on_axis(4, self.video_downscale_factors.width)
+
+ if tiling_config is not None and tiling_config.temporal_config is not None:
+ cfg = tiling_config.temporal_config
+ tile_size = cfg.tile_size_in_frames // self.video_downscale_factors.time
+ overlap = cfg.tile_overlap_in_frames // self.video_downscale_factors.time
+ splitters[2] = split_in_temporal(tile_size, overlap)
+ mappers[2] = to_mapping_operation(map_temporal_slice, self.video_downscale_factors.time)
+
+ return create_tiles(latent.shape, splitters, mappers)
+
+ def tiled_decode(
+ self,
+ latent: torch.Tensor,
+ tiling_config: TilingConfig | None = None,
+ timestep: Optional[torch.Tensor] = None,
+ generator: Optional[torch.Generator] = None,
+ ) -> Iterator[torch.Tensor]:
+ """
+ Decode a latent tensor into video frames using tiled processing.
+ Splits the latent tensor into tiles, decodes each tile individually,
+ and yields video chunks as they become available.
+ Args:
+ latent: Input latent tensor (B, C, F', H', W').
+ tiling_config: Tiling configuration for the latent tensor.
+ timestep: Optional timestep for decoder conditioning.
+ generator: Optional random generator for deterministic decoding.
+ Yields:
+ Video chunks (B, C, T, H, W) by temporal slices;
+ """
+
+ # Calculate full video shape from latent shape to get spatial dimensions
+ full_video_shape = VideoLatentShape.from_torch_shape(latent.shape).upscale(self.video_downscale_factors)
+ tiles = self._prepare_tiles(latent, tiling_config)
+
+ temporal_groups = self._group_tiles_by_temporal_slice(tiles)
+
+ # State for temporal overlap handling
+ previous_chunk = None
+ previous_weights = None
+ previous_temporal_slice = None
+
+ for temporal_group_tiles in temporal_groups:
+ curr_temporal_slice = temporal_group_tiles[0].out_coords[2]
+
+ # Calculate the shape of the temporal buffer for this group of tiles.
+ # The temporal length depends on whether this is the first tile (starts at 0) or not.
+ # - First tile: (frames - 1) * scale + 1
+ # - Subsequent tiles: frames * scale
+ # This logic is handled by TemporalAxisMapping and reflected in out_coords.
+ temporal_tile_buffer_shape = full_video_shape._replace(
+ frames=curr_temporal_slice.stop - curr_temporal_slice.start,
+ )
+
+ buffer = torch.zeros(
+ temporal_tile_buffer_shape.to_torch_shape(),
+ device=latent.device,
+ dtype=latent.dtype,
+ )
+
+ curr_weights = self._accumulate_temporal_group_into_buffer(
+ group_tiles=temporal_group_tiles,
+ buffer=buffer,
+ latent=latent,
+ timestep=timestep,
+ generator=generator,
+ )
+
+ # Blend with previous temporal chunk if it exists
+ if previous_chunk is not None:
+ # Check if current temporal slice overlaps with previous temporal slice
+ if previous_temporal_slice.stop > curr_temporal_slice.start:
+ overlap_len = previous_temporal_slice.stop - curr_temporal_slice.start
+ temporal_overlap_slice = slice(curr_temporal_slice.start - previous_temporal_slice.start, None)
+
+ # The overlap is already masked before it reaches this step. Each tile is accumulated into buffer
+ # with its trapezoidal mask, and curr_weights accumulates the same mask. In the overlap blend we add
+ # the masked values (buffer[...]) and the corresponding weights (curr_weights[...]) into the
+ # previous buffers, then later normalize by weights.
+ previous_chunk[:, :, temporal_overlap_slice, :, :] += buffer[:, :, slice(0, overlap_len), :, :]
+ previous_weights[:, :, temporal_overlap_slice, :, :] += curr_weights[
+ :, :, slice(0, overlap_len), :, :
+ ]
+
+ buffer[:, :, slice(0, overlap_len), :, :] = previous_chunk[:, :, temporal_overlap_slice, :, :]
+ curr_weights[:, :, slice(0, overlap_len), :, :] = previous_weights[
+ :, :, temporal_overlap_slice, :, :
+ ]
+
+ # Yield the non-overlapping part of the previous chunk
+ previous_weights = previous_weights.clamp(min=1e-8)
+ yield_len = curr_temporal_slice.start - previous_temporal_slice.start
+ yield (previous_chunk / previous_weights)[:, :, :yield_len, :, :]
+
+ # Update state for next iteration
+ previous_chunk = buffer
+ previous_weights = curr_weights
+ previous_temporal_slice = curr_temporal_slice
+
+ # Yield any remaining chunk
+ if previous_chunk is not None:
+ previous_weights = previous_weights.clamp(min=1e-8)
+ yield previous_chunk / previous_weights
+
+ def _group_tiles_by_temporal_slice(self, tiles: List[Tile]) -> List[List[Tile]]:
+ """Group tiles by their temporal output slice."""
+ if not tiles:
+ return []
+
+ groups = []
+ current_slice = tiles[0].out_coords[2]
+ current_group = []
+
+ for tile in tiles:
+ tile_slice = tile.out_coords[2]
+ if tile_slice == current_slice:
+ current_group.append(tile)
+ else:
+ groups.append(current_group)
+ current_slice = tile_slice
+ current_group = [tile]
+
+ # Add the final group
+ if current_group:
+ groups.append(current_group)
+
+ return groups
+
+ def _accumulate_temporal_group_into_buffer(
+ self,
+ group_tiles: List[Tile],
+ buffer: torch.Tensor,
+ latent: torch.Tensor,
+ timestep: Optional[torch.Tensor],
+ generator: Optional[torch.Generator],
+ ) -> torch.Tensor:
+ """
+ Decode and accumulate all tiles of a temporal group into a local buffer.
+ The buffer is local to the group and always starts at time 0; temporal coordinates
+ are rebased by subtracting temporal_slice.start.
+ """
+ temporal_slice = group_tiles[0].out_coords[2]
+
+ weights = torch.zeros_like(buffer)
+
+ for tile in group_tiles:
+ decoded_tile = self.forward(latent[tile.in_coords], timestep, generator)
+ mask = tile.blend_mask.to(device=buffer.device, dtype=buffer.dtype)
+ temporal_offset = tile.out_coords[2].start - temporal_slice.start
+ # Use the tile's output coordinate length, not the decoded tile's length,
+ # as the decoder may produce a different number of frames than expected
+ expected_temporal_len = tile.out_coords[2].stop - tile.out_coords[2].start
+ decoded_temporal_len = decoded_tile.shape[2]
+
+ # Ensure we don't exceed the buffer or decoded tile bounds
+ actual_temporal_len = min(expected_temporal_len, decoded_temporal_len, buffer.shape[2] - temporal_offset)
+
+ chunk_coords = (
+ slice(None), # batch
+ slice(None), # channels
+ slice(temporal_offset, temporal_offset + actual_temporal_len),
+ tile.out_coords[3], # height
+ tile.out_coords[4], # width
+ )
+
+ # Slice decoded_tile and mask to match the actual length we're writing
+ decoded_slice = decoded_tile[:, :, :actual_temporal_len, :, :]
+ mask_slice = mask[:, :, :actual_temporal_len, :, :] if mask.shape[2] > 1 else mask
+
+ buffer[chunk_coords] += decoded_slice * mask_slice
+ weights[chunk_coords] += mask_slice
+
+ return weights
+
+
+def decode_video(
+ latent: torch.Tensor,
+ video_decoder: VideoDecoder,
+ tiling_config: TilingConfig | None = None,
+) -> Iterator[torch.Tensor]:
+ """
+ Decode a video latent tensor with the given decoder.
+ Args:
+ latent: Tensor [c, f, h, w]
+ video_decoder: Decoder module.
+ tiling_config: Optional tiling settings.
+ Yields:
+ Decoded chunk [f, h, w, c], uint8 in [0, 255].
+ """
+
+ def convert_to_uint8(frames: torch.Tensor) -> torch.Tensor:
+ frames = (((frames + 1.0) / 2.0).clamp(0.0, 1.0) * 255.0).to(torch.uint8)
+ frames = rearrange(frames[0], "c f h w -> f h w c")
+ return frames
+
+ if tiling_config is not None:
+ for frames in video_decoder.tiled_decode(latent, tiling_config):
+ yield convert_to_uint8(frames)
+ else:
+ decoded_video = video_decoder(latent)
+ yield convert_to_uint8(decoded_video)
+
+
+def get_video_chunks_number(num_frames: int, tiling_config: TilingConfig | None = None) -> int:
+ """
+ Get the number of video chunks for a given number of frames and tiling configuration.
+ Args:
+ num_frames: Number of frames in the video.
+ tiling_config: Tiling configuration.
+ Returns:
+ Number of video chunks.
+ """
+ if not tiling_config or not tiling_config.temporal_config:
+ return 1
+ cfg = tiling_config.temporal_config
+ frame_stride = cfg.tile_size_in_frames - cfg.tile_overlap_in_frames
+ return (num_frames - 1 + frame_stride - 1) // frame_stride
+
+
+def split_in_spatial(size: int, overlap: int) -> SplitOperation:
+ def split(dimension_size: int) -> DimensionIntervals:
+ if dimension_size <= size:
+ return DEFAULT_SPLIT_OPERATION(dimension_size)
+ amount = (dimension_size + size - 2 * overlap - 1) // (size - overlap)
+ starts = [i * (size - overlap) for i in range(amount)]
+ ends = [start + size for start in starts]
+ ends[-1] = dimension_size
+ left_ramps = [0] + [overlap] * (amount - 1)
+ right_ramps = [overlap] * (amount - 1) + [0]
+ return DimensionIntervals(starts=starts, ends=ends, left_ramps=left_ramps, right_ramps=right_ramps)
+
+ return split
+
+
+def split_in_temporal(size: int, overlap: int) -> SplitOperation:
+ non_causal_split = split_in_spatial(size, overlap)
+
+ def split(dimension_size: int) -> DimensionIntervals:
+ if dimension_size <= size:
+ return DEFAULT_SPLIT_OPERATION(dimension_size)
+ intervals = non_causal_split(dimension_size)
+ starts = intervals.starts
+ starts[1:] = [s - 1 for s in starts[1:]]
+ left_ramps = intervals.left_ramps
+ left_ramps[1:] = [r + 1 for r in left_ramps[1:]]
+ return replace(intervals, starts=starts, left_ramps=left_ramps)
+
+ return split
+
+
+def to_mapping_operation(
+ map_func: Callable[[int, int, int, int, int], Tuple[slice, torch.Tensor]],
+ scale: int,
+) -> MappingOperation:
+ def map_op(intervals: DimensionIntervals) -> tuple[list[slice], list[torch.Tensor | None]]:
+ output_slices: list[slice] = []
+ masks_1d: list[torch.Tensor | None] = []
+ number_of_slices = len(intervals.starts)
+ for i in range(number_of_slices):
+ start = intervals.starts[i]
+ end = intervals.ends[i]
+ left_ramp = intervals.left_ramps[i]
+ right_ramp = intervals.right_ramps[i]
+ output_slice, mask_1d = map_func(start, end, left_ramp, right_ramp, scale)
+ output_slices.append(output_slice)
+ masks_1d.append(mask_1d)
+ return output_slices, masks_1d
+
+ return map_op
+
+
+def map_temporal_slice(begin: int, end: int, left_ramp: int, right_ramp: int, scale: int) -> Tuple[slice, torch.Tensor]:
+ start = begin * scale
+ stop = 1 + (end - 1) * scale
+ left_ramp = 1 + (left_ramp - 1) * scale
+ right_ramp = right_ramp * scale
+
+ return slice(start, stop), compute_trapezoidal_mask_1d(stop - start, left_ramp, right_ramp, True)
+
+
+def map_spatial_slice(begin: int, end: int, left_ramp: int, right_ramp: int, scale: int) -> Tuple[slice, torch.Tensor]:
+ start = begin * scale
+ stop = end * scale
+ left_ramp = left_ramp * scale
+ right_ramp = right_ramp * scale
+
+ return slice(start, stop), compute_trapezoidal_mask_1d(stop - start, left_ramp, right_ramp, False)
diff --git a/packages/ltx-core/src/ltx_core/text_encoders/__init__.py b/packages/ltx-core/src/ltx_core/text_encoders/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..847312421bc51b19c7941acc71ebea84c0a0f0f2
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/text_encoders/__init__.py
@@ -0,0 +1 @@
+"""CLIP/text encoder model components."""
diff --git a/packages/ltx-core/src/ltx_core/text_encoders/gemma/__init__.py b/packages/ltx-core/src/ltx_core/text_encoders/gemma/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..35658c0c0fb3bcef9418bbd108568d5863ddeae2
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/text_encoders/gemma/__init__.py
@@ -0,0 +1,31 @@
+"""Gemma text encoder components."""
+
+from ltx_core.text_encoders.gemma.encoders.av_encoder import (
+ AV_GEMMA_TEXT_ENCODER_KEY_OPS,
+ AVGemmaEncoderOutput,
+ AVGemmaTextEncoderModel,
+ AVGemmaTextEncoderModelConfigurator,
+)
+from ltx_core.text_encoders.gemma.encoders.base_encoder import (
+ GemmaTextEncoderModelBase,
+ encode_text,
+ module_ops_from_gemma_root,
+)
+from ltx_core.text_encoders.gemma.encoders.video_only_encoder import (
+ VideoGemmaEncoderOutput,
+ VideoGemmaTextEncoderModel,
+ VideoGemmaTextEncoderModelConfigurator,
+)
+
+__all__ = [
+ "AV_GEMMA_TEXT_ENCODER_KEY_OPS",
+ "AVGemmaEncoderOutput",
+ "AVGemmaTextEncoderModel",
+ "AVGemmaTextEncoderModelConfigurator",
+ "GemmaTextEncoderModelBase",
+ "VideoGemmaEncoderOutput",
+ "VideoGemmaTextEncoderModel",
+ "VideoGemmaTextEncoderModelConfigurator",
+ "encode_text",
+ "module_ops_from_gemma_root",
+]
diff --git a/packages/ltx-core/src/ltx_core/text_encoders/gemma/embeddings_connector.py b/packages/ltx-core/src/ltx_core/text_encoders/gemma/embeddings_connector.py
new file mode 100644
index 0000000000000000000000000000000000000000..de3f2efd9b339d6000b623eeb0512983f4bd8b64
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/text_encoders/gemma/embeddings_connector.py
@@ -0,0 +1,210 @@
+import torch
+
+from ltx_core.model.model_protocol import ModelConfigurator
+from ltx_core.model.transformer.attention import Attention
+from ltx_core.model.transformer.feed_forward import FeedForward
+from ltx_core.model.transformer.rope import (
+ LTXRopeType,
+ generate_freq_grid_np,
+ generate_freq_grid_pytorch,
+ precompute_freqs_cis,
+)
+from ltx_core.utils import rms_norm
+
+
+class _BasicTransformerBlock1D(torch.nn.Module):
+ def __init__(
+ self,
+ dim: int,
+ heads: int,
+ dim_head: int,
+ rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
+ ):
+ super().__init__()
+
+ self.attn1 = Attention(
+ query_dim=dim,
+ heads=heads,
+ dim_head=dim_head,
+ rope_type=rope_type,
+ )
+
+ self.ff = FeedForward(
+ dim,
+ dim_out=dim,
+ )
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ pe: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ # Notice that normalization is always applied before the real computation in the following blocks.
+
+ # 1. Normalization Before Self-Attention
+ norm_hidden_states = rms_norm(hidden_states)
+
+ norm_hidden_states = norm_hidden_states.squeeze(1)
+
+ # 2. Self-Attention
+ attn_output = self.attn1(norm_hidden_states, mask=attention_mask, pe=pe)
+
+ hidden_states = attn_output + hidden_states
+ if hidden_states.ndim == 4:
+ hidden_states = hidden_states.squeeze(1)
+
+ # 3. Normalization before Feed-Forward
+ norm_hidden_states = rms_norm(hidden_states)
+
+ # 4. Feed-forward
+ ff_output = self.ff(norm_hidden_states)
+
+ hidden_states = ff_output + hidden_states
+ if hidden_states.ndim == 4:
+ hidden_states = hidden_states.squeeze(1)
+
+ return hidden_states
+
+
+class Embeddings1DConnector(torch.nn.Module):
+ """
+ Embeddings1DConnector applies a 1D transformer-based processing to sequential embeddings (e.g., for video, audio, or
+ other modalities). It supports rotary positional encoding (rope), optional causal temporal positioning, and can
+ substitute padded positions with learnable registers. The module is highly configurable for head size, number of
+ layers, and register usage.
+ Args:
+ attention_head_dim (int): Dimension of each attention head (default=128).
+ num_attention_heads (int): Number of attention heads (default=30).
+ num_layers (int): Number of transformer layers (default=2).
+ positional_embedding_theta (float): Scaling factor for position embedding (default=10000.0).
+ positional_embedding_max_pos (list[int] | None): Max positions for positional embeddings (default=[1]).
+ causal_temporal_positioning (bool): If True, uses causal attention (default=False).
+ num_learnable_registers (int | None): Number of learnable registers to replace padded tokens. If None, disables
+ register replacement. (default=128)
+ rope_type (LTXRopeType): The RoPE variant to use (default=DEFAULT_ROPE_TYPE).
+ double_precision_rope (bool): Use double precision rope calculation (default=False).
+ """
+
+ _supports_gradient_checkpointing = True
+
+ def __init__(
+ self,
+ attention_head_dim: int = 128,
+ num_attention_heads: int = 30,
+ num_layers: int = 2,
+ positional_embedding_theta: float = 10000.0,
+ positional_embedding_max_pos: list[int] | None = None,
+ causal_temporal_positioning: bool = False,
+ num_learnable_registers: int | None = 128,
+ rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
+ double_precision_rope: bool = False,
+ ):
+ super().__init__()
+ self.num_attention_heads = num_attention_heads
+ self.inner_dim = num_attention_heads * attention_head_dim
+ self.causal_temporal_positioning = causal_temporal_positioning
+ self.positional_embedding_theta = positional_embedding_theta
+ self.positional_embedding_max_pos = (
+ positional_embedding_max_pos if positional_embedding_max_pos is not None else [1]
+ )
+ self.rope_type = rope_type
+ self.double_precision_rope = double_precision_rope
+ self.transformer_1d_blocks = torch.nn.ModuleList(
+ [
+ _BasicTransformerBlock1D(
+ dim=self.inner_dim,
+ heads=num_attention_heads,
+ dim_head=attention_head_dim,
+ rope_type=rope_type,
+ )
+ for _ in range(num_layers)
+ ]
+ )
+
+ self.num_learnable_registers = num_learnable_registers
+ if self.num_learnable_registers:
+ self.learnable_registers = torch.nn.Parameter(
+ torch.rand(self.num_learnable_registers, self.inner_dim, dtype=torch.bfloat16) * 2.0 - 1.0
+ )
+
+ def _replace_padded_with_learnable_registers(
+ self, hidden_states: torch.Tensor, attention_mask: torch.Tensor
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ assert hidden_states.shape[1] % self.num_learnable_registers == 0, (
+ f"Hidden states sequence length {hidden_states.shape[1]} must be divisible by num_learnable_registers "
+ f"{self.num_learnable_registers}."
+ )
+
+ num_registers_duplications = hidden_states.shape[1] // self.num_learnable_registers
+ learnable_registers = torch.tile(self.learnable_registers, (num_registers_duplications, 1))
+ attention_mask_binary = (attention_mask.squeeze(1).squeeze(1).unsqueeze(-1) >= -9000.0).int()
+
+ non_zero_hidden_states = hidden_states[:, attention_mask_binary.squeeze().bool(), :]
+ non_zero_nums = non_zero_hidden_states.shape[1]
+ pad_length = hidden_states.shape[1] - non_zero_nums
+ adjusted_hidden_states = torch.nn.functional.pad(non_zero_hidden_states, pad=(0, 0, 0, pad_length), value=0)
+ flipped_mask = torch.flip(attention_mask_binary, dims=[1])
+ hidden_states = flipped_mask * adjusted_hidden_states + (1 - flipped_mask) * learnable_registers
+
+ attention_mask = torch.full_like(
+ attention_mask,
+ 0.0,
+ dtype=attention_mask.dtype,
+ device=attention_mask.device,
+ )
+
+ return hidden_states, attention_mask
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ Forward pass of Embeddings1DConnector.
+ Args:
+ hidden_states (torch.Tensor): Input tensor of embeddings (shape [batch, seq_len, feature_dim]).
+ attention_mask (torch.Tensor|None): Optional mask for valid tokens (shape compatible with hidden_states).
+ Returns:
+ tuple[torch.Tensor, torch.Tensor]: Processed features and the corresponding (possibly modified) mask.
+ """
+ if self.num_learnable_registers:
+ hidden_states, attention_mask = self._replace_padded_with_learnable_registers(hidden_states, attention_mask)
+
+ indices_grid = torch.arange(hidden_states.shape[1], dtype=torch.float32, device=hidden_states.device)
+ indices_grid = indices_grid[None, None, :]
+ freq_grid_generator = generate_freq_grid_np if self.double_precision_rope else generate_freq_grid_pytorch
+ freqs_cis = precompute_freqs_cis(
+ indices_grid=indices_grid,
+ dim=self.inner_dim,
+ out_dtype=hidden_states.dtype,
+ theta=self.positional_embedding_theta,
+ max_pos=self.positional_embedding_max_pos,
+ num_attention_heads=self.num_attention_heads,
+ rope_type=self.rope_type,
+ freq_grid_generator=freq_grid_generator,
+ )
+
+ for block in self.transformer_1d_blocks:
+ hidden_states = block(hidden_states, attention_mask=attention_mask, pe=freqs_cis)
+
+ hidden_states = rms_norm(hidden_states)
+
+ return hidden_states, attention_mask
+
+
+class Embeddings1DConnectorConfigurator(ModelConfigurator[Embeddings1DConnector]):
+ @classmethod
+ def from_config(cls: type[Embeddings1DConnector], config: dict) -> Embeddings1DConnector:
+ config = config.get("transformer", {})
+ rope_type = LTXRopeType(config.get("rope_type", "interleaved"))
+ double_precision_rope = config.get("frequencies_precision", False) == "float64"
+ pe_max_pos = config.get("connector_positional_embedding_max_pos", [1])
+
+ connector = Embeddings1DConnector(
+ positional_embedding_max_pos=pe_max_pos,
+ rope_type=rope_type,
+ double_precision_rope=double_precision_rope,
+ )
+ return connector
diff --git a/packages/ltx-core/src/ltx_core/text_encoders/gemma/encoders/av_encoder.py b/packages/ltx-core/src/ltx_core/text_encoders/gemma/encoders/av_encoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..63728d4e6f839bd24a0a3a85c57d9839d853cc61
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/text_encoders/gemma/encoders/av_encoder.py
@@ -0,0 +1,94 @@
+from typing import NamedTuple
+
+import torch
+from transformers.models.gemma3 import Gemma3ForConditionalGeneration
+
+from ltx_core.loader.sd_ops import SDOps
+from ltx_core.model.model_protocol import ModelConfigurator
+from ltx_core.text_encoders.gemma.embeddings_connector import (
+ Embeddings1DConnector,
+ Embeddings1DConnectorConfigurator,
+)
+from ltx_core.text_encoders.gemma.encoders.base_encoder import GemmaTextEncoderModelBase
+from ltx_core.text_encoders.gemma.feature_extractor import GemmaFeaturesExtractorProjLinear
+from ltx_core.text_encoders.gemma.tokenizer import LTXVGemmaTokenizer
+
+
+class AVGemmaEncoderOutput(NamedTuple):
+ video_encoding: torch.Tensor
+ audio_encoding: torch.Tensor
+ attention_mask: torch.Tensor
+
+
+class AVGemmaTextEncoderModel(GemmaTextEncoderModelBase):
+ """
+ AVGemma Text Encoder Model.
+ This class combines the tokenizer, Gemma model, feature extractor from base class and a
+ video and audio embeddings connectors to provide a preprocessing for audio-visual pipeline.
+ """
+
+ def __init__(
+ self,
+ feature_extractor_linear: GemmaFeaturesExtractorProjLinear,
+ embeddings_connector: Embeddings1DConnector,
+ audio_embeddings_connector: Embeddings1DConnector,
+ tokenizer: LTXVGemmaTokenizer | None = None,
+ model: Gemma3ForConditionalGeneration | None = None,
+ dtype: torch.dtype = torch.bfloat16,
+ ) -> None:
+ super().__init__(
+ feature_extractor_linear=feature_extractor_linear,
+ tokenizer=tokenizer,
+ model=model,
+ dtype=dtype,
+ )
+ self.embeddings_connector = embeddings_connector.to(dtype=dtype)
+ self.audio_embeddings_connector = audio_embeddings_connector.to(dtype=dtype)
+
+ def _run_connectors(
+ self, encoded_input: torch.Tensor, attention_mask: torch.Tensor
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ connector_attention_mask = self._convert_to_additive_mask(attention_mask, encoded_input.dtype)
+
+ encoded, encoded_connector_attention_mask = self.embeddings_connector(
+ encoded_input,
+ connector_attention_mask,
+ )
+
+ # restore the mask values to int64
+ attention_mask = (encoded_connector_attention_mask < 0.000001).to(torch.int64)
+ attention_mask = attention_mask.reshape([encoded.shape[0], encoded.shape[1], 1])
+ encoded = encoded * attention_mask
+
+ encoded_for_audio, _ = self.audio_embeddings_connector(encoded_input, connector_attention_mask)
+
+ return encoded, encoded_for_audio, attention_mask.squeeze(-1)
+
+ def forward(self, text: str, padding_side: str = "left") -> AVGemmaEncoderOutput:
+ encoded_inputs, attention_mask = self._preprocess_text(text, padding_side)
+ video_encoding, audio_encoding, attention_mask = self._run_connectors(encoded_inputs, attention_mask)
+ return AVGemmaEncoderOutput(video_encoding, audio_encoding, attention_mask)
+
+
+class AVGemmaTextEncoderModelConfigurator(ModelConfigurator[AVGemmaTextEncoderModel]):
+ @classmethod
+ def from_config(cls: type["AVGemmaTextEncoderModel"], config: dict) -> "AVGemmaTextEncoderModel":
+ feature_extractor_linear = GemmaFeaturesExtractorProjLinear.from_config(config)
+ embeddings_connector = Embeddings1DConnectorConfigurator.from_config(config)
+ audio_embeddings_connector = Embeddings1DConnectorConfigurator.from_config(config)
+ return AVGemmaTextEncoderModel(
+ feature_extractor_linear=feature_extractor_linear,
+ embeddings_connector=embeddings_connector,
+ audio_embeddings_connector=audio_embeddings_connector,
+ )
+
+
+AV_GEMMA_TEXT_ENCODER_KEY_OPS = (
+ SDOps("AV_GEMMA_TEXT_ENCODER_KEY_OPS")
+ .with_matching(prefix="text_embedding_projection.")
+ .with_matching(prefix="model.diffusion_model.audio_embeddings_connector.")
+ .with_matching(prefix="model.diffusion_model.video_embeddings_connector.")
+ .with_replacement("text_embedding_projection.", "feature_extractor_linear.")
+ .with_replacement("model.diffusion_model.video_embeddings_connector.", "embeddings_connector.")
+ .with_replacement("model.diffusion_model.audio_embeddings_connector.", "audio_embeddings_connector.")
+)
diff --git a/packages/ltx-core/src/ltx_core/text_encoders/gemma/encoders/base_encoder.py b/packages/ltx-core/src/ltx_core/text_encoders/gemma/encoders/base_encoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..0e2780e7701959231d5205ed9673332e8fb9f7ad
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/text_encoders/gemma/encoders/base_encoder.py
@@ -0,0 +1,343 @@
+import functools
+from pathlib import Path
+
+import torch
+from einops import rearrange
+from transformers import AutoImageProcessor, Gemma3ForConditionalGeneration, Gemma3Processor
+
+from ltx_core.loader.module_ops import ModuleOps
+from ltx_core.text_encoders.gemma.feature_extractor import GemmaFeaturesExtractorProjLinear
+from ltx_core.text_encoders.gemma.tokenizer import LTXVGemmaTokenizer
+
+
+class GemmaTextEncoderModelBase(torch.nn.Module):
+ """
+ Gemma Text Encoder Model.
+ This base class combines the tokenizer, Gemma model and feature extractor to provide a preprocessing
+ for implementation classes for multimodal pipelines. It processes input text through tokenization,
+ obtains hidden states from the base language model, applies a linear feature extractor.
+ Args:
+ tokenizer (LTXVGemmaTokenizer): The tokenizer used for text preprocessing.
+ model (Gemma3ForConditionalGeneration): The base Gemma LLM.
+ feature_extractor_linear (GemmaFeaturesExtractorProjLinear): Linear projection for hidden state aggregation.
+ dtype (torch.dtype, optional): The data type for model parameters (default: torch.bfloat16).
+ """
+
+ def __init__(
+ self,
+ feature_extractor_linear: GemmaFeaturesExtractorProjLinear,
+ tokenizer: LTXVGemmaTokenizer | None = None,
+ model: Gemma3ForConditionalGeneration | None = None,
+ img_processor: Gemma3Processor | None = None,
+ dtype: torch.dtype = torch.bfloat16,
+ ) -> None:
+ super().__init__()
+ self._gemma_root = None
+ self.tokenizer = tokenizer
+ self.model = model
+ self.processor = img_processor
+ self.feature_extractor_linear = feature_extractor_linear.to(dtype=dtype)
+
+ def _run_feature_extractor(
+ self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, padding_side: str = "right"
+ ) -> torch.Tensor:
+ encoded_text_features = torch.stack(hidden_states, dim=-1)
+ encoded_text_features_dtype = encoded_text_features.dtype
+
+ sequence_lengths = attention_mask.sum(dim=-1)
+ normed_concated_encoded_text_features = _norm_and_concat_padded_batch(
+ encoded_text_features, sequence_lengths, padding_side=padding_side
+ )
+
+ return self.feature_extractor_linear(normed_concated_encoded_text_features.to(encoded_text_features_dtype))
+
+ def _convert_to_additive_mask(self, attention_mask: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+ return (attention_mask - 1).to(dtype).reshape(
+ (attention_mask.shape[0], 1, -1, attention_mask.shape[-1])
+ ) * torch.finfo(dtype).max
+
+ def _preprocess_text(self, text: str, padding_side: str = "left") -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
+ """
+ Encode a given string into feature tensors suitable for downstream tasks.
+ Args:
+ text (str): Input string to encode.
+ Returns:
+ tuple[torch.Tensor, dict[str, torch.Tensor]]: Encoded features and a dictionary with attention mask.
+ """
+ token_pairs = self.tokenizer.tokenize_with_weights(text)["gemma"]
+ input_ids = torch.tensor([[t[0] for t in token_pairs]], device=self.model.device)
+ attention_mask = torch.tensor([[w[1] for w in token_pairs]], device=self.model.device)
+ outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True)
+ projected = self._run_feature_extractor(
+ hidden_states=outputs.hidden_states, attention_mask=attention_mask, padding_side=padding_side
+ )
+ return projected, attention_mask
+
+ def _init_image_processor(self) -> None:
+ img_processor = AutoImageProcessor.from_pretrained(self._gemma_root, local_files_only=True)
+ if not self.tokenizer:
+ raise ValueError("Tokenizer is not loaded, cannot load image processor")
+ self.processor = Gemma3Processor(image_processor=img_processor, tokenizer=self.tokenizer.tokenizer)
+
+ def _enhance(
+ self,
+ messages: list[dict[str, str]],
+ image: torch.Tensor | None = None,
+ max_new_tokens: int = 512,
+ seed: int = 42,
+ ) -> str:
+ if self.processor is None:
+ self._init_image_processor()
+ text = self.processor.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+
+ model_inputs = self.processor(
+ text=text,
+ images=image,
+ return_tensors="pt",
+ ).to(self.model.device)
+ pad_token_id = self.processor.tokenizer.pad_token_id if self.processor.tokenizer.pad_token_id is not None else 0
+ model_inputs = _pad_inputs_for_attention_alignment(model_inputs, pad_token_id=pad_token_id)
+
+ with torch.inference_mode(), torch.random.fork_rng(devices=[self.model.device]):
+ torch.manual_seed(seed)
+ outputs = self.model.generate(
+ **model_inputs,
+ max_new_tokens=max_new_tokens,
+ do_sample=True,
+ temperature=0.7,
+ )
+ generated_ids = outputs[0][len(model_inputs.input_ids[0]) :]
+ enhanced_prompt = self.processor.tokenizer.decode(generated_ids, skip_special_tokens=True)
+
+ return enhanced_prompt
+
+ def enhance_t2v(
+ self,
+ prompt: str,
+ max_new_tokens: int = 512,
+ system_prompt: str | None = None,
+ seed: int = 42,
+ ) -> str:
+ """Enhance a text prompt for T2V generation."""
+
+ system_prompt = system_prompt or self.default_gemma_t2v_system_prompt
+
+ messages = [
+ {"role": "system", "content": system_prompt},
+ {"role": "user", "content": f"user prompt: {prompt}"},
+ ]
+
+ return self._enhance(messages, max_new_tokens=max_new_tokens, seed=seed)
+
+ def enhance_i2v(
+ self,
+ prompt: str,
+ image: torch.Tensor,
+ max_new_tokens: int = 512,
+ system_prompt: str | None = None,
+ seed: int = 42,
+ ) -> str:
+ """Enhance a text prompt for I2V generation using a reference image."""
+ system_prompt = system_prompt or self.default_gemma_i2v_system_prompt
+ messages = [
+ {"role": "system", "content": system_prompt},
+ {
+ "role": "user",
+ "content": [
+ {"type": "image"},
+ {"type": "text", "text": f"User Raw Input Prompt: {prompt}."},
+ ],
+ },
+ ]
+ return self._enhance(messages, image=image, max_new_tokens=max_new_tokens, seed=seed)
+
+ @functools.cached_property
+ def default_gemma_i2v_system_prompt(self) -> str:
+ return _load_system_prompt("gemma_i2v_system_prompt.txt")
+
+ @functools.cached_property
+ def default_gemma_t2v_system_prompt(self) -> str:
+ return _load_system_prompt("gemma_t2v_system_prompt.txt")
+
+ def forward(self, text: str, padding_side: str = "left") -> tuple[torch.Tensor, torch.Tensor]:
+ raise NotImplementedError("This method is not implemented for the base class")
+
+
+def _norm_and_concat_padded_batch(
+ encoded_text: torch.Tensor,
+ sequence_lengths: torch.Tensor,
+ padding_side: str = "right",
+) -> torch.Tensor:
+ """Normalize and flatten multi-layer hidden states, respecting padding.
+ Performs per-batch, per-layer normalization using masked mean and range,
+ then concatenates across the layer dimension.
+ Args:
+ encoded_text: Hidden states of shape [batch, seq_len, hidden_dim, num_layers].
+ sequence_lengths: Number of valid (non-padded) tokens per batch item.
+ padding_side: Whether padding is on "left" or "right".
+ Returns:
+ Normalized tensor of shape [batch, seq_len, hidden_dim * num_layers],
+ with padded positions zeroed out.
+ """
+ b, t, d, l = encoded_text.shape # noqa: E741
+ device = encoded_text.device
+
+ # Build mask: [B, T, 1, 1]
+ token_indices = torch.arange(t, device=device)[None, :] # [1, T]
+
+ if padding_side == "right":
+ # For right padding, valid tokens are from 0 to sequence_length-1
+ mask = token_indices < sequence_lengths[:, None] # [B, T]
+ elif padding_side == "left":
+ # For left padding, valid tokens are from (T - sequence_length) to T-1
+ start_indices = t - sequence_lengths[:, None] # [B, 1]
+ mask = token_indices >= start_indices # [B, T]
+ else:
+ raise ValueError(f"padding_side must be 'left' or 'right', got {padding_side}")
+
+ mask = rearrange(mask, "b t -> b t 1 1")
+
+ eps = 1e-6
+
+ # Compute masked mean: [B, 1, 1, L]
+ masked = encoded_text.masked_fill(~mask, 0.0)
+ denom = (sequence_lengths * d).view(b, 1, 1, 1)
+ mean = masked.sum(dim=(1, 2), keepdim=True) / (denom + eps)
+
+ # Compute masked min/max: [B, 1, 1, L]
+ x_min = encoded_text.masked_fill(~mask, float("inf")).amin(dim=(1, 2), keepdim=True)
+ x_max = encoded_text.masked_fill(~mask, float("-inf")).amax(dim=(1, 2), keepdim=True)
+ range_ = x_max - x_min
+
+ # Normalize only the valid tokens
+ normed = 8 * (encoded_text - mean) / (range_ + eps)
+
+ # concat to be [Batch, T, D * L] - this preserves the original structure
+ normed = normed.reshape(b, t, -1) # [B, T, D * L]
+
+ # Apply mask to preserve original padding (set padded positions to 0)
+ mask_flattened = rearrange(mask, "b t 1 1 -> b t 1").expand(-1, -1, d * l)
+ normed = normed.masked_fill(~mask_flattened, 0.0)
+
+ return normed
+
+
+@functools.lru_cache(maxsize=2)
+def _load_system_prompt(prompt_name: str) -> str:
+ with open(Path(__file__).parent / "prompts" / f"{prompt_name}", "r") as f:
+ return f.read()
+
+
+def _find_matching_dir(root_path: str, pattern: str) -> str:
+ """
+ Recursively search for files matching a glob pattern and return the parent directory of the first match.
+
+ LT_INTERNAL_BEGIN
+ Handles both LT internal storage and HuggingFace directory structures for Gemma model files.
+ See: https://huggingface.co/google/gemma-3-12b-it-qat-q4_0-unquantized
+ LT_INTERNAL_END
+ """
+
+ matches = list(Path(root_path).rglob(pattern))
+ if not matches:
+ raise FileNotFoundError(f"No files matching pattern '{pattern}' found under {root_path}")
+ return str(matches[0].parent)
+
+
+def module_ops_from_gemma_root(gemma_root: str, local_files_only: bool = True) -> tuple[ModuleOps, ...]:
+ if len(gemma_root.split("/")) != 2:
+ gemma_path = _find_matching_dir(gemma_root, "model*.safetensors")
+ tokenizer_path = _find_matching_dir(gemma_root, "tokenizer.model")
+ else:
+ # Hub ID: google/gemma-3-12b-it-qat-q4_0-unquantized
+ gemma_path = tokenizer_path = gemma_root
+
+ # LT_INTERNAL_BEGIN
+ # Note: We pass torch_dtype to from_pretrained here to maintain backward compatibility with older versions of
+ # Transformers. This is necessary to compare results with ComfyUI, which uses an older version that raises an error
+ # when dtype is passed. Current solution only logs a warning.
+ # LT_INTERNAL_END
+ def load_gemma(module: GemmaTextEncoderModelBase) -> GemmaTextEncoderModelBase:
+ module.model = Gemma3ForConditionalGeneration.from_pretrained(
+ gemma_path, local_files_only=local_files_only, torch_dtype=torch.bfloat16
+ )
+ module._gemma_root = module._gemma_root or gemma_root
+ return module
+
+ def load_tokenizer(module: GemmaTextEncoderModelBase) -> GemmaTextEncoderModelBase:
+ module.tokenizer = LTXVGemmaTokenizer(tokenizer_path, 1024, local_files_only)
+ module._gemma_root = module._gemma_root or gemma_root
+ return module
+
+ gemma_load_ops = ModuleOps(
+ "GemmaLoad",
+ matcher=lambda module: isinstance(module, GemmaTextEncoderModelBase) and module.model is None,
+ mutator=load_gemma,
+ )
+ tokenizer_load_ops = ModuleOps(
+ "TokenizerLoad",
+ matcher=lambda module: isinstance(module, GemmaTextEncoderModelBase) and module.tokenizer is None,
+ mutator=load_tokenizer,
+ )
+ return (gemma_load_ops, tokenizer_load_ops)
+
+
+def encode_text(text_encoder: GemmaTextEncoderModelBase, prompts: list[str]) -> list[tuple[torch.Tensor, torch.Tensor]]:
+ """
+ Encode a list of prompts using the provided Gemma text encoder.
+ Args:
+ text_encoder: The Gemma text encoder instance.
+ prompts: List of prompt strings to encode.
+ Returns:
+ List of tuples, each containing (v_context, a_context) tensors for each prompt.
+ """
+ result = []
+ for prompt in prompts:
+ v_context, a_context, _ = text_encoder(prompt)
+ result.append((v_context, a_context))
+ return result
+
+
+def _cat_with_padding(
+ tensor: torch.Tensor,
+ padding_length: int,
+ value: int | float,
+) -> torch.Tensor:
+ """Concatenate a tensor with a padding tensor of the given value."""
+ return torch.cat(
+ [
+ tensor,
+ torch.full(
+ (1, padding_length),
+ value,
+ dtype=tensor.dtype,
+ device=tensor.device,
+ ),
+ ],
+ dim=1,
+ )
+
+
+def _pad_inputs_for_attention_alignment(
+ model_inputs: dict[str, torch.Tensor],
+ pad_token_id: int = 0,
+ alignment: int = 8,
+) -> dict[str, torch.Tensor]:
+ """Pad sequence length to multiple of alignment for Flash Attention compatibility.
+ Flash Attention within SDPA requires sequence lengths aligned to 8 bytes.
+ This pads input_ids, attention_mask, and token_type_ids (if present) to prevent
+ 'p.attn_bias_ptr is not correctly aligned' errors.
+ """
+ seq_len = model_inputs.input_ids.shape[1]
+ padded_len = ((seq_len + alignment - 1) // alignment) * alignment
+ padding_length = padded_len - seq_len
+
+ if padding_length > 0:
+ model_inputs["input_ids"] = _cat_with_padding(model_inputs.input_ids, padding_length, pad_token_id)
+
+ model_inputs["attention_mask"] = _cat_with_padding(model_inputs.attention_mask, padding_length, 0)
+
+ if "token_type_ids" in model_inputs and model_inputs["token_type_ids"] is not None:
+ model_inputs["token_type_ids"] = _cat_with_padding(model_inputs["token_type_ids"], padding_length, 0)
+
+ return model_inputs
diff --git a/packages/ltx-core/src/ltx_core/text_encoders/gemma/encoders/prompts/gemma_i2v_system_prompt.txt b/packages/ltx-core/src/ltx_core/text_encoders/gemma/encoders/prompts/gemma_i2v_system_prompt.txt
new file mode 100644
index 0000000000000000000000000000000000000000..bfe674606b121474f567c3a91cab61926268f7b3
--- /dev/null
+++ b/packages/ltx-core/src/ltx_core/text_encoders/gemma/encoders/prompts/gemma_i2v_system_prompt.txt
@@ -0,0 +1,30 @@
+You are a Creative Assistant writing concise, action-focused image-to-video prompts. Given an image (first frame) and user Raw Input Prompt, generate a prompt to guide video generation from that image.
+
+#### Guidelines:
+- Analyze the Image: Identify Subject, Setting, Elements, Style and Mood.
+- Follow user Raw Input Prompt: Include all requested motion, actions, camera movements, audio, and details. If in conflict with the image, prioritize user request while maintaining visual consistency (describe transition from image to user's scene).
+- Describe only changes from the image: Don't reiterate established visual details. Inaccurate descriptions may cause scene cuts.
+- Active language: Use present-progressive verbs ("is walking," "speaking"). If no action specified, describe natural movements.
+- Chronological flow: Use temporal connectors ("as," "then," "while").
+- Audio layer: Describe complete soundscape throughout the prompt alongside actionsβNOT at the end. Align audio intensity with action tempo. Include natural background audio, ambient sounds, effects, speech or music (when requested). Be specific (e.g., "soft footsteps on tile") not vague (e.g., "ambient sound").
+- Speech (only when requested): Provide exact words in quotes with character's visual/voice characteristics (e.g., "The tall man speaks in a low, gravelly voice"), language if not English and accent if relevant. If general conversation mentioned without text, generate contextual quoted dialogue. (i.e., "The man is talking" input -> the output should include exact spoken words, like: "The man is talking in an excited voice saying: 'You won't believe what I just saw!' His hands gesture expressively as he speaks, eyebrows raised with enthusiasm. The ambient sound of a quiet room underscores his animated speech.")
+- Style: Include visual style at beginning: "Style: