Automated PR - 2026-01-05

packages/ltx-trainer/src/ltx_trainer/__init__.py

@@ -0,0 +1,44 @@
+import logging
+import os
+import sys
+from logging import getLogger
+from pathlib import Path
+
+from rich.logging import RichHandler
+
+# Get the process rank
+IS_MULTI_GPU = os.environ.get("LOCAL_RANK") is not None
+RANK = int(os.environ.get("LOCAL_RANK", "0"))
+
+# Configure with Rich
+logging.basicConfig(
+    level="INFO",
+    format=f"\\[rank {RANK}] %(message)s" if IS_MULTI_GPU else "%(message)s",
+    handlers=[
+        RichHandler(
+            rich_tracebacks=True,
+            show_time=False,
+            markup=True,
+        )
+    ],
+)
+
+# Get the logger and configure it
+logger = getLogger("ltxv_trainer")
+logger.setLevel(logging.DEBUG)
+logger.propagate = True
+
+# Set level based on process
+if RANK != 0:
+    logger.setLevel(logging.WARNING)
+
+# Expose common logging functions directly
+debug = logger.debug
+info = logger.info
+warning = logger.warning
+error = logger.error
+critical = logger.critical
+
+
+# Add the root directory to the Python path so we can import from scripts.
+sys.path.insert(0, str(Path(__file__).parent.parent.parent))

packages/ltx-trainer/src/ltx_trainer/captioning.py

@@ -0,0 +1,401 @@
+"""
+Audio-visual media captioning using multimodal models.
+This module provides captioning capabilities for videos with audio using:
+- Qwen2.5-Omni: Local model supporting text, audio, image, and video inputs (default)
+- Gemini Flash: Cloud-based API for audio-visual captioning
+Requirements:
+- Qwen2.5-Omni: transformers>=4.50, torch
+- Gemini Flash: google-generativeai (pip install google-generativeai)
+  Set GEMINI_API_KEY or GOOGLE_API_KEY environment variable
+"""
+
+import itertools
+import re
+from abc import ABC, abstractmethod
+from enum import Enum
+from pathlib import Path
+
+import torch
+
+# Instruction for audio-visual captioning (default) - includes speech transcription and sounds
+DEFAULT_CAPTION_INSTRUCTION = """\
+Analyze this media and provide a detailed caption in the following EXACT format. Fill in ALL sections:
+
+[VISUAL]: <Detailed description of people, objects, actions, settings, colors, and movements>
+[SPEECH]: <Word-for-word transcription of everything spoken.
+           Listen carefully and transcribe the exact words. If no speech, write "None">
+[SOUNDS]: <Description of music, ambient sounds, sound effects. If none, write "None">
+[TEXT]: <Any on-screen text visible. If none, write "None">
+
+You MUST fill in all four sections. For [SPEECH], transcribe the actual words spoken, not a summary."""
+
+# Instruction for video-only captioning (no audio processing)
+VIDEO_ONLY_CAPTION_INSTRUCTION = """\
+Analyze this media and provide a detailed caption in the following EXACT format. Fill in ALL sections:
+
+[VISUAL]: <Detailed description of people, objects, actions, settings, colors, and movements>
+[TEXT]: <Any on-screen text visible. If none, write "None">
+
+You MUST fill in both sections."""
+
+
+class CaptionerType(str, Enum):
+    """Enum for different types of media captioners."""
+
+    QWEN_OMNI = "qwen_omni"  # Local Qwen2.5-Omni model (audio + video)
+    GEMINI_FLASH = "gemini_flash"  # Gemini Flash API (audio + video)
+
+
+def create_captioner(captioner_type: CaptionerType, **kwargs) -> "MediaCaptioningModel":
+    """Factory function to create a media captioner.
+    Args:
+        captioner_type: The type of captioner to create
+        **kwargs: Additional arguments to pass to the captioner constructor
+    Returns:
+        An instance of a MediaCaptioningModel
+    """
+    match captioner_type:
+        case CaptionerType.QWEN_OMNI:
+            return QwenOmniCaptioner(**kwargs)
+        case CaptionerType.GEMINI_FLASH:
+            return GeminiFlashCaptioner(**kwargs)
+        case _:
+            raise ValueError(f"Unsupported captioner type: {captioner_type}")
+
+
+class MediaCaptioningModel(ABC):
+    """Abstract base class for audio-visual media captioning models."""
+
+    @abstractmethod
+    def caption(self, path: str | Path, **kwargs) -> str:
+        """Generate a caption for the given video or image.
+        Args:
+            path: Path to the video/image file to caption
+        Returns:
+            A string containing the generated caption
+        """
+
+    @property
+    @abstractmethod
+    def supports_audio(self) -> bool:
+        """Whether this captioner supports audio input."""
+
+    @staticmethod
+    def _is_image_file(path: str | Path) -> bool:
+        """Check if the file is an image based on extension."""
+        return str(path).lower().endswith((".png", ".jpg", ".jpeg", ".heic", ".heif", ".webp"))
+
+    @staticmethod
+    def _is_video_file(path: str | Path) -> bool:
+        """Check if the file is a video based on extension."""
+        return str(path).lower().endswith((".mp4", ".avi", ".mov", ".mkv", ".webm"))
+
+    @staticmethod
+    def _clean_raw_caption(caption: str) -> str:
+        """Clean up the raw caption by removing common VLM patterns."""
+        start = ["The", "This"]
+        kind = ["video", "image", "scene", "animated sequence", "clip", "footage"]
+        act = ["displays", "shows", "features", "depicts", "presents", "showcases", "captures", "contains"]
+
+        for x, y, z in itertools.product(start, kind, act):
+            caption = caption.replace(f"{x} {y} {z} ", "", 1)
+
+        return caption
+
+
+class QwenOmniCaptioner(MediaCaptioningModel):
+    """Audio-visual captioning using Alibaba's Qwen2.5-Omni model.
+    Qwen2.5-Omni is an end-to-end multimodal model that can perceive text, images, audio, and video.
+    It uses a Thinker-Talker architecture where the Thinker generates text and the Talker can
+    generate speech. For captioning, we use only the Thinker component for text generation.
+    Key features:
+    - Block-wise processing for streaming multimodal inputs
+    - TMRoPE (Time-aligned Multimodal RoPE) for synchronizing video and audio timestamps
+    - Can extract and process audio directly from video files
+    See: https://huggingface.co/docs/transformers/en/model_doc/qwen2_5_omni
+    Model: Qwen/Qwen2.5-Omni-7B (7B parameters)
+    """
+
+    MODEL_ID = "Qwen/Qwen2.5-Omni-7B"
+
+    # Default system prompt required by Qwen2.5-Omni for proper audio processing
+    DEFAULT_SYSTEM_PROMPT = (
+        "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, "
+        "capable of perceiving auditory and visual inputs, as well as generating text and speech."
+    )
+
+    def __init__(
+        self,
+        device: str | torch.device | None = None,
+        use_8bit: bool = False,
+        instruction: str | None = None,
+    ):
+        """
+        Initialize the Qwen2.5-Omni captioner.
+        Args:
+            device: Device to use for inference (e.g., 'cuda', 'cuda:0', 'cpu')
+            use_8bit: Whether to use 8-bit quantization for reduced memory usage
+            instruction: Custom instruction prompt. If None, uses the default instruction
+        """
+        self.device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
+        self.instruction = instruction
+        self._load_model(use_8bit=use_8bit)
+
+    @property
+    def supports_audio(self) -> bool:
+        return True
+
+    def caption(
+        self,
+        path: str | Path,
+        fps: int = 1,
+        include_audio: bool = True,
+        clean_caption: bool = True,
+    ) -> str:
+        """Generate a caption for the given video or image.
+        Args:
+            path: Path to the video/image file to caption
+            fps: Frames per second to sample from videos
+            include_audio: Whether to include audio in the captioning (for videos)
+            clean_caption: Whether to clean up the raw caption by removing common VLM patterns
+        Returns:
+            A string containing the generated caption
+        """
+        path = Path(path)
+        is_image = self._is_image_file(path)
+        is_video = self._is_video_file(path)
+
+        # Determine if we should process audio
+        use_audio = include_audio and is_video
+
+        # Use custom instruction if provided, otherwise pick appropriate default
+        if self.instruction is not None:
+            instruction = self.instruction
+        else:
+            instruction = DEFAULT_CAPTION_INSTRUCTION if use_audio else VIDEO_ONLY_CAPTION_INSTRUCTION
+
+        # Build the user content based on media type
+        # Based on HuggingFace docs: https://huggingface.co/docs/transformers/en/model_doc/qwen2_5_omni
+        user_content = []
+
+        if is_image:
+            user_content.append({"type": "image", "image": str(path)})
+        elif is_video:
+            user_content.append({"type": "video", "video": str(path)})
+
+        # Add the instruction text
+        user_content.append({"type": "text", "text": instruction})
+
+        # Build conversation - use the default system prompt required by Qwen2.5-Omni
+        # Using a custom system prompt causes warnings and may affect audio processing
+        messages = [
+            {
+                "role": "system",
+                "content": [{"type": "text", "text": self.DEFAULT_SYSTEM_PROMPT}],
+            },
+            {"role": "user", "content": user_content},
+        ]
+
+        # Process inputs using the processor's apply_chat_template
+        # For videos with audio, use load_audio_from_video=True and use_audio_in_video=True
+        inputs = self.processor.apply_chat_template(
+            messages,
+            load_audio_from_video=use_audio,
+            add_generation_prompt=True,
+            tokenize=True,
+            return_dict=True,
+            return_tensors="pt",
+            fps=fps,
+            padding=True,
+            use_audio_in_video=use_audio,
+        ).to(self.model.device)
+
+        # Generate caption (text only, using Thinker-only model)
+        # Note: For Qwen2_5OmniThinkerForConditionalGeneration, use standard generate params
+        # (not thinker_ prefixed ones, those are for the full Qwen2_5OmniForConditionalGeneration)
+        input_len = inputs["input_ids"].shape[1]
+
+        output_tokens = self.model.generate(
+            **inputs,
+            use_audio_in_video=use_audio,
+            do_sample=False,
+            max_new_tokens=1024,
+        )
+
+        # Extract only the generated tokens (exclude the input/prompt tokens)
+        generated_tokens = output_tokens[:, input_len:]
+
+        # Decode only the generated response
+        caption_raw = self.processor.batch_decode(
+            generated_tokens,
+            skip_special_tokens=True,
+            clean_up_tokenization_spaces=False,
+        )[0]
+
+        # Remove hallucinated conversation turns (e.g., "Human\nHuman\n..." or "Human: ...")
+        # This is a known issue with chat models continuing to generate fake turns
+        # We look for patterns that are clearly hallucinated chat turns, not legitimate uses of "human"
+
+        # Match "\nHuman" followed by ":", "\n", or end of string (chat turn patterns)
+        # This won't match "A human walks..." or "...the human body..."
+        caption_raw = re.split(r"\nHuman(?::|(?:\s*\n)|$)", caption_raw, maxsplit=1)[0]
+        caption_raw = caption_raw.strip()
+
+        # Clean up caption if requested
+        return self._clean_raw_caption(caption_raw) if clean_caption else caption_raw
+
+    def _load_model(self, use_8bit: bool) -> None:
+        """Load the Qwen2.5-Omni model and processor.
+        Uses the Thinker-only model (Qwen2_5OmniThinkerForConditionalGeneration) for text generation
+        to save compute by not loading the audio generation components.
+        """
+        from transformers import (  # noqa: PLC0415
+            BitsAndBytesConfig,
+            Qwen2_5OmniProcessor,
+            Qwen2_5OmniThinkerForConditionalGeneration,
+        )
+
+        quantization_config = BitsAndBytesConfig(load_in_8bit=True) if use_8bit else None
+
+        # Use Thinker-only model for text generation (saves memory by not loading Talker)
+        self.model = Qwen2_5OmniThinkerForConditionalGeneration.from_pretrained(
+            self.MODEL_ID,
+            dtype=torch.bfloat16,
+            low_cpu_mem_usage=True,
+            quantization_config=quantization_config,
+            device_map="auto",
+        )
+
+        self.processor = Qwen2_5OmniProcessor.from_pretrained(self.MODEL_ID)
+
+
+class GeminiFlashCaptioner(MediaCaptioningModel):
+    """Audio-visual captioning using Google's Gemini Flash API.
+    Gemini Flash is a cloud-based multimodal model that natively supports
+    audio and video understanding. Requires a Google API key.
+    Note: This captioner requires the `google-generativeai` package and a valid API key.
+    Set the GEMINI_API_KEY or GOOGLE_API_KEY environment variable, or pass the key directly.
+    """
+
+    MODEL_ID = "gemini-flash-lite-latest"
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        instruction: str | None = None,
+    ):
+        """Initialize the Gemini Flash captioner.
+        Args:
+            api_key: Google API key. If not provided, will look for
+                     GEMINI_API_KEY or GOOGLE_API_KEY environment variable.
+            instruction: Custom instruction prompt. If None, uses the default instruction
+        """
+        self.instruction = instruction
+        self._init_client(api_key)
+
+    @property
+    def supports_audio(self) -> bool:
+        return True
+
+    def caption(
+        self,
+        path: str | Path,
+        fps: int = 3,  # noqa: ARG002 - kept for API compatibility
+        include_audio: bool = True,
+        clean_caption: bool = True,
+    ) -> str:
+        """Generate a caption for the given video or image.
+        Args:
+            path: Path to the video/image file to caption
+            fps: Frames per second (not used for Gemini, kept for API compatibility)
+            include_audio: Whether to include audio content in the caption
+            clean_caption: Whether to clean up the raw caption
+        Returns:
+            A string containing the generated caption
+        """
+        import time  # noqa: PLC0415
+
+        path = Path(path)
+        is_video = self._is_video_file(path)
+        use_audio = include_audio and is_video
+
+        # Use custom instruction if provided, otherwise pick appropriate default
+        if self.instruction is not None:
+            instruction = self.instruction
+        else:
+            instruction = DEFAULT_CAPTION_INSTRUCTION if use_audio else VIDEO_ONLY_CAPTION_INSTRUCTION
+
+        # Upload the file to Gemini
+        uploaded_file = self._genai.upload_file(path)
+
+        # Wait for processing to complete (videos need time to process)
+        while uploaded_file.state.name == "PROCESSING":
+            time.sleep(1)
+            uploaded_file = self._genai.get_file(uploaded_file.name)
+
+        if uploaded_file.state.name == "FAILED":
+            raise RuntimeError(f"File processing failed: {uploaded_file.state.name}")
+
+        # Generate caption
+        response = self._model.generate_content([uploaded_file, instruction])
+
+        caption_raw = response.text
+
+        # Clean up the uploaded file
+        self._genai.delete_file(uploaded_file.name)
+
+        # Clean up caption if requested
+        return self._clean_raw_caption(caption_raw) if clean_caption else caption_raw
+
+    def _init_client(self, api_key: str | None) -> None:
+        """Initialize the Gemini API client."""
+        import os  # noqa: PLC0415
+
+        try:
+            import google.generativeai as genai  # noqa: PLC0415
+        except ImportError as e:
+            raise ImportError(
+                "The `google-generativeai` package is required for Gemini Flash captioning. "
+                "Install it with: `uv pip install google-generativeai`"
+            ) from e
+
+        # Get API key from argument or environment
+        # GEMINI_API_KEY is the recommended variable, GOOGLE_API_KEY also works
+        resolved_api_key = api_key or os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY")
+
+        if not resolved_api_key:
+            raise ValueError(
+                "Gemini API key is required. Provide it via the `api_key` argument "
+                "or set the GEMINI_API_KEY or GOOGLE_API_KEY environment variable."
+            )
+
+        # Configure the genai library with the API key
+        genai.configure(api_key=resolved_api_key)
+
+        # Store reference to genai module for file operations
+        self._genai = genai
+
+        # Initialize the model
+        self._model = genai.GenerativeModel(self.MODEL_ID)
+
+
+def example() -> None:
+    """Example usage of the captioning module."""
+    import sys  # noqa: PLC0415
+
+    if len(sys.argv) < 2:
+        print(f"Usage: python {sys.argv[0]} <video_path> [captioner_type]")  # noqa: T201
+        print("  captioner_type: qwen_omni (default) or gemini_flash")  # noqa: T201
+        sys.exit(1)
+
+    video_path = sys.argv[1]
+    captioner_type = CaptionerType(sys.argv[2]) if len(sys.argv) > 2 else CaptionerType.QWEN_OMNI
+
+    print(f"Using {captioner_type.value} captioner:")  # noqa: T201
+    captioner = create_captioner(captioner_type)
+    caption = captioner.caption(video_path)
+    print(f"CAPTION: {caption}")  # noqa: T201
+
+
+if __name__ == "__main__":
+    example()

packages/ltx-trainer/src/ltx_trainer/config.py

@@ -0,0 +1,472 @@
+from pathlib import Path
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, ConfigDict, Discriminator, Field, Tag, ValidationInfo, field_validator, model_validator
+
+from ltx_trainer.quantization import QuantizationOptions
+from ltx_trainer.training_strategies.base_strategy import TrainingStrategyConfigBase
+from ltx_trainer.training_strategies.text_to_video import TextToVideoConfig
+from ltx_trainer.training_strategies.video_to_video import VideoToVideoConfig
+
+
+class ConfigBaseModel(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+
+class ModelConfig(ConfigBaseModel):
+    """Configuration for the base model and training mode"""
+
+    model_path: str | Path = Field(
+        ...,
+        description="Model path - local path to safetensors checkpoint file",
+    )
+
+    text_encoder_path: str | Path | None = Field(
+        default=None,
+        description="Path to text encoder (required for LTX-2/Gemma models, optional for LTXV/T5 models)",
+    )
+
+    training_mode: Literal["lora", "full"] = Field(
+        default="lora",
+        description="Training mode - either LoRA fine-tuning or full model fine-tuning",
+    )
+
+    load_checkpoint: str | Path | None = Field(
+        default=None,
+        description="Path to a checkpoint file or directory to load from. "
+        "If a directory is provided, the latest checkpoint will be used.",
+    )
+
+    @field_validator("model_path")
+    @classmethod
+    def validate_model_path(cls, v: str | Path) -> str | Path:
+        """Validate that model_path is either a valid URL or an existing local path."""
+        is_url = str(v).startswith(("http://", "https://"))
+
+        if is_url:
+            raise ValueError(f"Model path cannot be a URL: {v}")
+
+        if not Path(v).exists():
+            raise ValueError(f"Model path does not exist: {v}")
+
+        return v
+
+
+class LoraConfig(ConfigBaseModel):
+    """Configuration for LoRA fine-tuning"""
+
+    rank: int = Field(
+        default=64,
+        description="Rank of LoRA adaptation",
+        ge=2,
+    )
+
+    alpha: int = Field(
+        default=64,
+        description="Alpha scaling factor for LoRA",
+        ge=1,
+    )
+
+    dropout: float = Field(
+        default=0.0,
+        description="Dropout probability for LoRA layers",
+        ge=0.0,
+        le=1.0,
+    )
+
+    target_modules: list[str] = Field(
+        default=["to_k", "to_q", "to_v", "to_out.0"],
+        description="List of modules to target with LoRA",
+    )
+
+
+def _get_strategy_discriminator(v: dict | TrainingStrategyConfigBase) -> str:
+    """Discriminator function for strategy config union."""
+    if isinstance(v, dict):
+        return v.get("name", "text_to_video")
+    return v.name
+
+
+# Union type for all strategy configs with discriminator
+TrainingStrategyConfig = Annotated[
+    Annotated[TextToVideoConfig, Tag("text_to_video")] | Annotated[VideoToVideoConfig, Tag("video_to_video")],
+    Discriminator(_get_strategy_discriminator),
+]
+
+
+class OptimizationConfig(ConfigBaseModel):
+    """Configuration for optimization parameters"""
+
+    learning_rate: float = Field(
+        default=5e-4,
+        description="Learning rate for optimization",
+    )
+
+    steps: int = Field(
+        default=3000,
+        description="Number of training steps",
+    )
+
+    batch_size: int = Field(
+        default=2,
+        description="Batch size for training",
+    )
+
+    gradient_accumulation_steps: int = Field(
+        default=1,
+        description="Number of steps to accumulate gradients",
+    )
+
+    max_grad_norm: float = Field(
+        default=1.0,
+        description="Maximum gradient norm for clipping",
+    )
+
+    optimizer_type: Literal["adamw", "adamw8bit"] = Field(
+        default="adamw",
+        description="Type of optimizer to use for training",
+    )
+
+    scheduler_type: Literal[
+        "constant",
+        "linear",
+        "cosine",
+        "cosine_with_restarts",
+        "polynomial",
+    ] = Field(
+        default="linear",
+        description="Type of scheduler to use for training",
+    )
+
+    scheduler_params: dict = Field(
+        default_factory=dict,
+        description="Parameters for the scheduler",
+    )
+
+    enable_gradient_checkpointing: bool = Field(
+        default=False,
+        description="Enable gradient checkpointing to save memory at the cost of slower training",
+    )
+
+
+class AccelerationConfig(ConfigBaseModel):
+    """Configuration for hardware acceleration and compute optimization"""
+
+    mixed_precision_mode: Literal["no", "fp16", "bf16"] | None = Field(
+        default="bf16",
+        description="Mixed precision training mode",
+    )
+
+    quantization: QuantizationOptions | None = Field(
+        default=None,
+        description="Quantization precision to use",
+    )
+
+    load_text_encoder_in_8bit: bool = Field(
+        default=False,
+        description="Whether to load the text encoder in 8-bit precision to save memory",
+    )
+
+
+class DataConfig(ConfigBaseModel):
+    """Configuration for data loading and processing"""
+
+    preprocessed_data_root: str = Field(
+        description="Path to folder containing preprocessed training data",
+    )
+
+    num_dataloader_workers: int = Field(
+        default=2,
+        description="Number of background processes for data loading (0 means synchronous loading)",
+        ge=0,
+    )
+
+
+class ValidationConfig(ConfigBaseModel):
+    """Configuration for validation during training"""
+
+    prompts: list[str] = Field(
+        default_factory=list,
+        description="List of prompts to use for validation",
+    )
+
+    negative_prompt: str = Field(
+        default="worst quality, inconsistent motion, blurry, jittery, distorted",
+        description="Negative prompt to use for validation examples",
+    )
+
+    images: list[str] | None = Field(
+        default=None,
+        description="List of image paths to use for validation. "
+        "One image path must be provided for each validation prompt",
+    )
+
+    reference_videos: list[str] | None = Field(
+        default=None,
+        description="List of reference video paths to use for validation. "
+        "One video path must be provided for each validation prompt",
+    )
+
+    video_dims: tuple[int, int, int] = Field(
+        default=(960, 544, 97),
+        description="Dimensions of validation videos (width, height, frames). "
+        "Width and height must be divisible by 32. Frames must satisfy frames % 8 == 1 for LTX-2.",
+    )
+
+    @field_validator("video_dims")
+    @classmethod
+    def validate_video_dims(cls, v: tuple[int, int, int]) -> tuple[int, int, int]:
+        """Validate video dimensions for LTX-2 compatibility."""
+        width, height, frames = v
+
+        if width % 32 != 0:
+            raise ValueError(f"Width ({width}) must be divisible by 32")
+        if height % 32 != 0:
+            raise ValueError(f"Height ({height}) must be divisible by 32")
+        if frames % 8 != 1:
+            raise ValueError(f"Frames ({frames}) must satisfy frames % 8 == 1 for LTX-2 (e.g., 1, 9, 17, 25, ...)")
+
+        return v
+
+    frame_rate: float = Field(
+        default=25.0,
+        description="Frame rate for validation videos",
+        gt=0,
+    )
+
+    seed: int = Field(
+        default=42,
+        description="Random seed used when sampling validation videos",
+    )
+
+    inference_steps: int = Field(
+        default=50,
+        description="Number of inference steps for validation",
+        gt=0,
+    )
+
+    interval: int | None = Field(
+        default=100,
+        description="Number of steps between validation runs. If None, validation is disabled.",
+        gt=0,
+    )
+
+    videos_per_prompt: int = Field(
+        default=1,
+        description="Number of videos to generate per validation prompt",
+        gt=0,
+    )
+
+    guidance_scale: float = Field(
+        default=4.0,
+        description="CFG guidance scale to use during validation",
+        ge=1.0,
+    )
+
+    stg_scale: float = Field(
+        default=1.0,
+        description="STG (Spatio-Temporal Guidance) scale. 0.0 disables STG. "
+        "Recommended value is 1.0. STG is combined with CFG for improved video quality.",
+        ge=0.0,
+    )
+
+    stg_blocks: list[int] | None = Field(
+        default=[29],
+        description="Which transformer blocks to perturb for STG. "
+        "None means all blocks are perturbed. Recommended for LTX-2: [29].",
+    )
+
+    stg_mode: Literal["stg_av", "stg_v"] = Field(
+        default="stg_av",
+        description="STG mode: 'stg_av' skips both audio and video self-attention, "
+        "'stg_v' skips only video self-attention.",
+    )
+
+    generate_audio: bool = Field(
+        default=True,
+        description="Whether to generate audio in validation samples. "
+        "Independent of training strategy setting - you can generate audio "
+        "in validation even when not training the audio branch.",
+    )
+
+    skip_initial_validation: bool = Field(
+        default=False,
+        description="Skip validation video sampling at step 0 (beginning of training)",
+    )
+
+    include_reference_in_output: bool = Field(
+        default=False,
+        description="For video-to-video training: concatenate the original reference video side-by-side "
+        "with the generated output. The reference comes from the input video, not from the model's output.",
+    )
+
+    @field_validator("images")
+    @classmethod
+    def validate_images(cls, v: list[str] | None, info: ValidationInfo) -> list[str] | None:
+        """Validate that number of images (if provided) matches number of prompts."""
+        if v is None:
+            return None
+
+        num_prompts = len(info.data.get("prompts", []))
+        if v is not None and len(v) != num_prompts:
+            raise ValueError(f"Number of images ({len(v)}) must match number of prompts ({num_prompts})")
+
+        for image_path in v:
+            if not Path(image_path).exists():
+                raise ValueError(f"Image path '{image_path}' does not exist")
+
+        return v
+
+    @field_validator("reference_videos")
+    @classmethod
+    def validate_reference_videos(cls, v: list[str] | None, info: ValidationInfo) -> list[str] | None:
+        """Validate that number of reference videos (if provided) matches number of prompts."""
+        if v is None:
+            return None
+
+        num_prompts = len(info.data.get("prompts", []))
+        if v is not None and len(v) != num_prompts:
+            raise ValueError(f"Number of reference videos ({len(v)}) must match number of prompts ({num_prompts})")
+
+        for video_path in v:
+            if not Path(video_path).exists():
+                raise ValueError(f"Reference video path '{video_path}' does not exist")
+
+        return v
+
+
+class CheckpointsConfig(ConfigBaseModel):
+    """Configuration for model checkpointing during training"""
+
+    interval: int | None = Field(
+        default=None,
+        description="Number of steps between checkpoint saves. If None, intermediate checkpoints are disabled.",
+        gt=0,
+    )
+
+    keep_last_n: int = Field(
+        default=1,
+        description="Number of most recent checkpoints to keep. Set to -1 to keep all checkpoints.",
+        ge=-1,
+    )
+
+
+class HubConfig(ConfigBaseModel):
+    """Configuration for Hugging Face Hub integration"""
+
+    push_to_hub: bool = Field(default=False, description="Whether to push the model weights to the Hugging Face Hub")
+    hub_model_id: str | None = Field(
+        default=None, description="Hugging Face Hub repository ID (e.g., 'username/repo-name')"
+    )
+
+    @model_validator(mode="after")
+    def validate_hub_config(self) -> "HubConfig":
+        """Validate that hub_model_id is not None when push_to_hub is True."""
+        if self.push_to_hub and not self.hub_model_id:
+            raise ValueError("hub_model_id must be specified when push_to_hub is True")
+        return self
+
+
+class WandbConfig(ConfigBaseModel):
+    """Configuration for Weights & Biases logging"""
+
+    enabled: bool = Field(
+        default=False,
+        description="Whether to enable W&B logging",
+    )
+
+    project: str = Field(
+        default="ltxv-trainer",
+        description="W&B project name",
+    )
+
+    entity: str | None = Field(
+        default=None,
+        description="W&B username or team",
+    )
+
+    tags: list[str] = Field(
+        default_factory=list,
+        description="Tags to add to the W&B run",
+    )
+
+    log_validation_videos: bool = Field(
+        default=True,
+        description="Whether to log validation videos to W&B",
+    )
+
+
+class FlowMatchingConfig(ConfigBaseModel):
+    """Configuration for flow matching training"""
+
+    timestep_sampling_mode: Literal["uniform", "shifted_logit_normal"] = Field(
+        default="shifted_logit_normal",
+        description="Mode to use for timestep sampling",
+    )
+
+    timestep_sampling_params: dict = Field(
+        default_factory=dict,
+        description="Parameters for timestep sampling",
+    )
+
+
+class LtxTrainerConfig(ConfigBaseModel):
+    """Unified configuration for LTXV training"""
+
+    # Sub-configurations
+    model: ModelConfig = Field(default_factory=ModelConfig)
+    lora: LoraConfig | None = Field(default=None)
+    training_strategy: TrainingStrategyConfig = Field(
+        default_factory=TextToVideoConfig,
+        description="Training strategy configuration. Determines the training mode and its parameters.",
+    )
+    optimization: OptimizationConfig = Field(default_factory=OptimizationConfig)
+    acceleration: AccelerationConfig = Field(default_factory=AccelerationConfig)
+    data: DataConfig
+    validation: ValidationConfig = Field(default_factory=ValidationConfig)
+    checkpoints: CheckpointsConfig = Field(default_factory=CheckpointsConfig)
+    hub: HubConfig = Field(default_factory=HubConfig)
+    flow_matching: FlowMatchingConfig = Field(default_factory=FlowMatchingConfig)
+    wandb: WandbConfig = Field(default_factory=WandbConfig)
+
+    # General configuration
+    seed: int = Field(
+        default=42,
+        description="Random seed for reproducibility",
+    )
+
+    output_dir: str = Field(
+        default="outputs",
+        description="Directory to save model outputs",
+    )
+
+    # noinspection PyNestedDecorators
+    @field_validator("output_dir")
+    @classmethod
+    def expand_output_path(cls, v: str) -> str:
+        """Expand user home directory in output path."""
+        return str(Path(v).expanduser().resolve())
+
+    @model_validator(mode="after")
+    def validate_strategy_compatibility(self) -> "LtxTrainerConfig":
+        """Validate that training strategy and other configurations are compatible."""
+
+        # Check that reference videos are provided when using video_to_video strategy
+        if (
+            self.training_strategy.name == "video_to_video"
+            and self.validation.interval
+            and not self.validation.reference_videos
+        ):
+            raise ValueError(
+                "reference_videos must be provided in validation config when using video_to_video strategy"
+            )
+
+        # Check that LoRA config is provided when training mode is lora
+        if self.model.training_mode == "lora" and self.lora is None:
+            raise ValueError("LoRA configuration must be provided when training_mode is 'lora'")
+
+        # Check that LoRA config is provided when using video_to_video strategy
+        if self.training_strategy.name == "video_to_video" and self.model.training_mode != "lora":
+            raise ValueError("Training mode must be 'lora' when using video_to_video strategy")
+
+        return self

packages/ltx-trainer/src/ltx_trainer/config_display.py

@@ -0,0 +1,155 @@
+"""Display utilities for training configuration.
+This module provides formatted console output for LtxTrainerConfig.
+"""
+
+from rich import box
+from rich.console import Console
+from rich.table import Table
+
+from ltx_trainer.config import LtxTrainerConfig
+
+
+def print_config(config: LtxTrainerConfig) -> None:
+    """Print configuration as a nicely formatted table with sections."""
+
+    def fmt(v: object, max_len: int = 55) -> str:
+        """Format any value for display."""
+        if v is None:
+            return "[dim]—[/]"
+        if isinstance(v, bool):
+            return "[green]✓[/]" if v else "[dim]✗[/]"
+        if isinstance(v, (list, tuple)):
+            if not v:
+                return "[dim]—[/]"
+            return ", ".join(str(x) for x in v)
+        s = str(v)
+        return s[: max_len - 3] + "..." if len(s) > max_len else s
+
+    cfg = config
+    opt = cfg.optimization
+    val = cfg.validation
+    accel = cfg.acceleration
+
+    # Build sections: list of (section_title, [(key, value), ...])
+    sections: list[tuple[str, list[tuple[str, str]]]] = [
+        (
+            "🎬 Model",
+            [
+                ("Base", fmt(cfg.model.model_path)),
+                ("Text Encoder", fmt(cfg.model.text_encoder_path) or "[dim]Built-in[/]"),
+                ("Training Mode", f"[bold green]{cfg.model.training_mode.upper()}[/]"),
+                ("Load Checkpoint", fmt(cfg.model.load_checkpoint) if cfg.model.load_checkpoint else "[dim]—[/]"),
+            ],
+        ),
+    ]
+
+    if cfg.lora:
+        sections.append(
+            (
+                "🔗 LoRA",
+                [
+                    ("Rank / Alpha", f"{cfg.lora.rank} / {cfg.lora.alpha}"),
+                    ("Dropout", str(cfg.lora.dropout)),
+                    ("Target Modules", fmt(cfg.lora.target_modules)),
+                ],
+            )
+        )
+
+    # Strategy section - include strategy-specific fields
+    strategy_items: list[tuple[str, str]] = [("Name", cfg.training_strategy.name)]
+    if hasattr(cfg.training_strategy, "with_audio"):
+        strategy_items.append(("Audio", fmt(cfg.training_strategy.with_audio)))
+    if hasattr(cfg.training_strategy, "first_frame_conditioning_p"):
+        strategy_items.append(("First Frame Cond P", str(cfg.training_strategy.first_frame_conditioning_p)))
+
+    sections.append(("🎯 Strategy", strategy_items))
+
+    sections.extend(
+        [
+            (
+                "⚡ Optimization",
+                [
+                    ("Steps", f"[bold]{opt.steps:,}[/]"),
+                    ("Learning Rate", f"{opt.learning_rate:.2e}"),
+                    ("Batch Size", str(opt.batch_size)),
+                    ("Grad Accumulation", str(opt.gradient_accumulation_steps)),
+                    ("Optimizer", opt.optimizer_type),
+                    ("Scheduler", opt.scheduler_type),
+                    ("Max Grad Norm", str(opt.max_grad_norm)),
+                    ("Grad Checkpointing", fmt(opt.enable_gradient_checkpointing)),
+                ],
+            ),
+            (
+                "🚀 Acceleration",
+                [
+                    ("Mixed Precision", accel.mixed_precision_mode or "[dim]—[/]"),
+                    ("Quantization", str(accel.quantization) if accel.quantization else "[dim]—[/]"),
+                    ("Text Encoder 8bit", fmt(accel.load_text_encoder_in_8bit)),
+                ],
+            ),
+            (
+                "🎥 Validation",
+                [
+                    ("Prompts", f"{len(val.prompts)} prompt(s)" if val.prompts else "[dim]—[/]"),
+                    ("Interval", f"Every {val.interval} steps" if val.interval else "[dim]Disabled[/]"),
+                    ("Video Dims", f"{val.video_dims[0]}x{val.video_dims[1]}, {val.video_dims[2]} frames"),
+                    ("Frame Rate", f"{val.frame_rate} fps"),
+                    ("Inference Steps", str(val.inference_steps)),
+                    ("CFG Scale", str(val.guidance_scale)),
+                    (
+                        "STG",
+                        f"scale={val.stg_scale}; blocks={fmt(val.stg_blocks)}; mode={val.stg_mode}"
+                        if val.stg_scale > 0
+                        else "[dim]Disabled[/]",
+                    ),
+                    ("Seed", str(val.seed)),
+                ],
+            ),
+            (
+                "📂 Data & Output",
+                [
+                    ("Dataset", fmt(cfg.data.preprocessed_data_root)),
+                    ("Dataloader Workers", str(cfg.data.num_dataloader_workers)),
+                    ("Output Dir", fmt(cfg.output_dir)),
+                    ("Seed", str(cfg.seed)),
+                ],
+            ),
+            (
+                "🔌 Integrations",
+                [
+                    (
+                        "Checkpoints",
+                        f"Every {cfg.checkpoints.interval} steps (keep {cfg.checkpoints.keep_last_n})"
+                        if cfg.checkpoints.interval
+                        else "[dim]Disabled[/]",
+                    ),
+                    ("W&B", f"{cfg.wandb.project}" if cfg.wandb.enabled else "[dim]Disabled[/]"),
+                    ("HF Hub", cfg.hub.hub_model_id if cfg.hub.push_to_hub else "[dim]Disabled[/]"),
+                ],
+            ),
+        ]
+    )
+
+    # Build table with section headers
+    table = Table(
+        title="[bold]⚙️  Training Configuration[/]",
+        show_header=False,
+        box=box.ROUNDED,
+        border_style="bright_blue",
+        padding=(0, 1),
+        title_style="bold bright_blue",
+    )
+    table.add_column("Key", style="white", width=20)
+    table.add_column("Value", style="cyan")
+
+    for i, (section_title, items) in enumerate(sections):
+        if i > 0:
+            table.add_row("", "")  # Blank line between sections
+        table.add_row(f"[bold yellow]{section_title}[/]", "")
+        for key, value in items:
+            table.add_row(f"  {key}", value)
+
+    console = Console()
+    console.print()
+    console.print(table)
+    console.print()

packages/ltx-trainer/src/ltx_trainer/datasets.py

@@ -0,0 +1,270 @@
+from pathlib import Path
+
+import torch
+from einops import rearrange
+from torch import Tensor
+from torch.utils.data import Dataset
+
+from ltx_trainer import logger
+
+# Constants for precomputed data directories
+PRECOMPUTED_DIR_NAME = ".precomputed"
+
+
+class DummyDataset(Dataset):
+    """Produce random latents and prompt embeddings. For minimal demonstration and benchmarking purposes"""
+
+    def __init__(
+        self,
+        width: int = 1024,
+        height: int = 1024,
+        num_frames: int = 25,
+        fps: int = 24,
+        dataset_length: int = 200,
+        latent_dim: int = 128,
+        latent_spatial_compression_ratio: int = 32,
+        latent_temporal_compression_ratio: int = 8,
+        prompt_embed_dim: int = 4096,
+        prompt_sequence_length: int = 256,
+    ) -> None:
+        if width % 32 != 0:
+            raise ValueError(f"Width must be divisible by 32, got {width=}")
+
+        if height % 32 != 0:
+            raise ValueError(f"Height must be divisible by 32, got {height=}")
+
+        if num_frames % 8 != 1:
+            raise ValueError(f"Number of frames must have a remainder of 1 when divided by 8, got {num_frames=}")
+
+        self.width = width
+        self.height = height
+        self.num_frames = num_frames
+        self.fps = fps
+        self.dataset_length = dataset_length
+        self.latent_dim = latent_dim
+        self.num_latent_frames = (num_frames - 1) // latent_temporal_compression_ratio + 1
+        self.latent_height = height // latent_spatial_compression_ratio
+        self.latent_width = width // latent_spatial_compression_ratio
+        self.latent_sequence_length = self.num_latent_frames * self.latent_height * self.latent_width
+        self.prompt_embed_dim = prompt_embed_dim
+        self.prompt_sequence_length = prompt_sequence_length
+
+    def __len__(self) -> int:
+        return self.dataset_length
+
+    def __getitem__(self, idx: int) -> dict[str, dict[str, Tensor]]:
+        return {
+            "latent_conditions": {
+                "latents": torch.randn(
+                    self.latent_dim,
+                    self.num_latent_frames,
+                    self.latent_height,
+                    self.latent_width,
+                ),
+                "num_frames": self.num_latent_frames,
+                "height": self.latent_height,
+                "width": self.latent_width,
+                "fps": self.fps,
+            },
+            "text_conditions": {
+                "prompt_embeds": torch.randn(
+                    self.prompt_sequence_length,
+                    self.prompt_embed_dim,
+                ),  # random text embeddings
+                "prompt_attention_mask": torch.ones(
+                    self.prompt_sequence_length,
+                    dtype=torch.bool,
+                ),  # random attention mask
+            },
+        }
+
+
+class PrecomputedDataset(Dataset):
+    def __init__(self, data_root: str, data_sources: dict[str, str] | list[str] | None = None) -> None:
+        """
+        Generic dataset for loading precomputed data from multiple sources.
+        Args:
+            data_root: Root directory containing preprocessed data
+            data_sources: Either:
+              - Dict mapping directory names to output keys
+              - List of directory names (keys will equal values)
+              - None (defaults to ["latents", "conditions"])
+        Example:
+            # Standard mode (list)
+            dataset = PrecomputedDataset("data/", ["latents", "conditions"])
+            # Standard mode (dict)
+            dataset = PrecomputedDataset("data/", {"latents": "latent_conditions", "conditions": "text_conditions"})
+            # IC-LoRA mode
+            dataset = PrecomputedDataset("data/", ["latents", "conditions", "reference_latents"])
+        Note:
+            Latents are always returned in non-patchified format [C, F, H, W].
+            Legacy patchified format [seq_len, C] is automatically converted.
+        """
+        super().__init__()
+
+        self.data_root = self._setup_data_root(data_root)
+        self.data_sources = self._normalize_data_sources(data_sources)
+        self.source_paths = self._setup_source_paths()
+        self.sample_files = self._discover_samples()
+        self._validate_setup()
+
+    @staticmethod
+    def _setup_data_root(data_root: str) -> Path:
+        """Setup and validate the data root directory."""
+        data_root = Path(data_root).expanduser().resolve()
+
+        if not data_root.exists():
+            raise FileNotFoundError(f"Data root directory does not exist: {data_root}")
+
+        # If the given path is the dataset root, use the precomputed subdirectory
+        if (data_root / PRECOMPUTED_DIR_NAME).exists():
+            data_root = data_root / PRECOMPUTED_DIR_NAME
+
+        return data_root
+
+    @staticmethod
+    def _normalize_data_sources(data_sources: dict[str, str] | list[str] | None) -> dict[str, str]:
+        """Normalize data_sources input to a consistent dict format."""
+        if data_sources is None:
+            # Default sources
+            return {"latents": "latent_conditions", "conditions": "text_conditions"}
+        elif isinstance(data_sources, list):
+            # Convert list to dict where keys equal values
+            return {source: source for source in data_sources}
+        elif isinstance(data_sources, dict):
+            return data_sources.copy()
+        else:
+            raise TypeError(f"data_sources must be dict, list, or None, got {type(data_sources)}")
+
+    def _setup_source_paths(self) -> dict[str, Path]:
+        """Map data source names to their actual directory paths."""
+        source_paths = {}
+
+        for dir_name in self.data_sources:
+            source_path = self.data_root / dir_name
+            source_paths[dir_name] = source_path
+
+            # Check that all sources exist.
+            if not source_path.exists():
+                raise FileNotFoundError(f"Required {dir_name} directory does not exist: {source_path}")
+
+        return source_paths
+
+    def _discover_samples(self) -> dict[str, list[Path]]:
+        """Discover all valid sample files across all data sources."""
+        # Use first data source as the reference to discover samples
+        data_key = "latents" if "latents" in self.data_sources else next(iter(self.data_sources.keys()))
+        data_path = self.source_paths[data_key]
+        data_files = list(data_path.glob("**/*.pt"))
+
+        if not data_files:
+            raise ValueError(f"No data files found in {data_path}")
+
+        # Initialize sample files dict
+        sample_files = {output_key: [] for output_key in self.data_sources.values()}
+
+        # For each data file, find corresponding files in other sources
+        for data_file in data_files:
+            rel_path = data_file.relative_to(data_path)
+
+            # Check if corresponding files exist in ALL sources
+            if self._all_source_files_exist(data_file, rel_path):
+                self._fill_sample_data_files(data_file, rel_path, sample_files)
+
+        return sample_files
+
+    def _all_source_files_exist(self, data_file: Path, rel_path: Path) -> bool:
+        """Check if corresponding files exist in all data sources."""
+        for dir_name in self.data_sources:
+            expected_path = self._get_expected_file_path(dir_name, data_file, rel_path)
+            if not expected_path.exists():
+                logger.warning(
+                    f"No matching {dir_name} file found for: {data_file.name} (expected in: {expected_path})"
+                )
+                return False
+
+        return True
+
+    def _get_expected_file_path(self, dir_name: str, data_file: Path, rel_path: Path) -> Path:
+        """Get the expected file path for a given data source."""
+        source_path = self.source_paths[dir_name]
+
+        # For conditions, handle legacy naming where latent_X.pt maps to condition_X.pt
+        if dir_name == "conditions" and data_file.name.startswith("latent_"):
+            return source_path / f"condition_{data_file.stem[7:]}.pt"
+
+        return source_path / rel_path
+
+    def _fill_sample_data_files(self, data_file: Path, rel_path: Path, sample_files: dict[str, list[Path]]) -> None:
+        """Add a valid sample to the sample_files tracking."""
+        for dir_name, output_key in self.data_sources.items():
+            expected_path = self._get_expected_file_path(dir_name, data_file, rel_path)
+            sample_files[output_key].append(expected_path.relative_to(self.source_paths[dir_name]))
+
+    def _validate_setup(self) -> None:
+        """Validate that the dataset setup is correct."""
+        if not self.sample_files:
+            raise ValueError("No valid samples found - all data sources must have matching files")
+
+        # Verify all output keys have the same number of samples
+        sample_counts = {key: len(files) for key, files in self.sample_files.items()}
+        if len(set(sample_counts.values())) > 1:
+            raise ValueError(f"Mismatched sample counts across sources: {sample_counts}")
+
+    def __len__(self) -> int:
+        # Use the first output key as reference count
+        first_key = next(iter(self.sample_files.keys()))
+        return len(self.sample_files[first_key])
+
+    def __getitem__(self, index: int) -> dict[str, torch.Tensor]:
+        result = {}
+
+        for dir_name, output_key in self.data_sources.items():
+            source_path = self.source_paths[dir_name]
+            file_rel_path = self.sample_files[output_key][index]
+            file_path = source_path / file_rel_path
+
+            try:
+                data = torch.load(file_path, map_location="cpu", weights_only=True)
+
+                # Normalize video latent format if this is a latent source
+                if "latent" in dir_name.lower():
+                    data = self._normalize_video_latents(data)
+
+                result[output_key] = data
+            except Exception as e:
+                raise RuntimeError(f"Failed to load {output_key} from {file_path}: {e}") from e
+
+        # Add index for debugging
+        result["idx"] = index
+        return result
+
+    @staticmethod
+    def _normalize_video_latents(data: dict) -> dict:
+        """
+        Normalize video latents to non-patchified format [C, F, H, W].
+        Used for keeping backward compatibility with legacy datasets.
+        """
+        latents = data["latents"]
+
+        # Check if latents are in legacy patchified format [seq_len, C]
+        if latents.dim() == 2:
+            # Legacy format: [seq_len, C] where seq_len = F * H * W
+            num_frames = data["num_frames"]
+            height = data["height"]
+            width = data["width"]
+
+            # Unpatchify: [seq_len, C] -> [C, F, H, W]
+            latents = rearrange(
+                latents,
+                "(f h w) c -> c f h w",
+                f=num_frames,
+                h=height,
+                w=width,
+            )
+
+            # Update the data dict with unpatchified latents
+            data = data.copy()
+            data["latents"] = latents
+
+        return data

packages/ltx-trainer/src/ltx_trainer/hf_hub_utils.py

@@ -0,0 +1,208 @@
+import shutil
+import tempfile
+from pathlib import Path
+from typing import List, Union
+
+import imageio
+from huggingface_hub import HfApi, create_repo
+from huggingface_hub.utils import are_progress_bars_disabled, disable_progress_bars, enable_progress_bars
+from rich.progress import Progress, SpinnerColumn, TextColumn
+
+from ltx_trainer import logger
+from ltx_trainer.config import LtxTrainerConfig
+
+
+def push_to_hub(weights_path: Path, sampled_videos_paths: List[Path], config: LtxTrainerConfig) -> None:
+    """Push the trained LoRA weights to HuggingFace Hub."""
+    if not config.hub.hub_model_id:
+        logger.warning("⚠️ HuggingFace hub_model_id not specified, skipping push to hub")
+        return
+
+    api = HfApi()
+
+    # Save original progress bar state
+    original_progress_state = are_progress_bars_disabled()
+    disable_progress_bars()  # Disable during our custom progress tracking
+
+    try:
+        # Try to create repo if it doesn't exist
+        try:
+            repo = create_repo(
+                repo_id=config.hub.hub_model_id,
+                repo_type="model",
+                exist_ok=True,  # Don't raise error if repo exists
+            )
+            repo_id = repo.repo_id
+            logger.info(f"🤗 Successfully created HuggingFace model repository at: {repo.url}")
+        except Exception as e:
+            logger.error(f"❌ Failed to create HuggingFace model repository: {e}")
+            return
+
+        # Create a single temporary directory for all files
+        with tempfile.TemporaryDirectory() as temp_dir:
+            temp_path = Path(temp_dir)
+
+            with Progress(
+                SpinnerColumn(),
+                TextColumn("[progress.description]{task.description}"),
+                transient=True,
+            ) as progress:
+                try:
+                    # Copy weights
+                    task_copy = progress.add_task("Copying weights...", total=None)
+                    weights_dest = temp_path / weights_path.name
+                    shutil.copy2(weights_path, weights_dest)
+                    progress.update(task_copy, description="✓ Weights copied")
+
+                    # Create model card and save samples
+                    task_card = progress.add_task("Creating model card and samples...", total=None)
+                    _create_model_card(
+                        output_dir=temp_path,
+                        videos=sampled_videos_paths,
+                        config=config,
+                    )
+                    progress.update(task_card, description="✓ Model card and samples created")
+
+                    # Upload everything at once
+                    task_upload = progress.add_task("Pushing files to HuggingFace Hub...", total=None)
+                    api.upload_folder(
+                        folder_path=str(temp_path),
+                        repo_id=repo_id,
+                        repo_type="model",
+                    )
+                    progress.update(task_upload, description="✓ Files pushed to HuggingFace Hub")
+                    logger.info("✅ Successfully pushed files to HuggingFace Hub")
+
+                except Exception as e:
+                    logger.error(f"❌ Failed to process and push files to HuggingFace Hub: {e}")
+                    raise  # Re-raise to handle in outer try block
+
+    finally:
+        # Restore original progress bar state
+        if not original_progress_state:
+            enable_progress_bars()
+
+
+def convert_video_to_gif(video_path: Path, output_path: Path) -> None:
+    """Convert a video file to GIF format."""
+    try:
+        # Read the video file
+        reader = imageio.get_reader(str(video_path))
+        fps = reader.get_meta_data()["fps"]
+
+        # Write GIF file with infinite loop
+        writer = imageio.get_writer(
+            str(output_path),
+            fps=min(fps, 15),  # Cap FPS at 15 for reasonable file size
+            loop=0,  # 0 means infinite loop
+        )
+
+        for frame in reader:
+            writer.append_data(frame)
+
+        writer.close()
+        reader.close()
+    except Exception as e:
+        logger.error(f"Failed to convert video to GIF: {e}")
+
+
+def _create_model_card(
+    output_dir: Union[str, Path],
+    videos: List[Path],
+    config: LtxTrainerConfig,
+) -> Path:
+    """Generate and save a model card for the trained model."""
+
+    repo_id = config.hub.hub_model_id
+    pretrained_model_name_or_path = config.model.model_path
+    validation_prompts = config.validation.prompts
+    output_dir = Path(output_dir)
+    template_path = Path(__file__).parent.parent.parent / "templates" / "model_card.md"
+
+    # Read the template
+    template = template_path.read_text()
+
+    # Get model name from repo_id
+    model_name = repo_id.split("/")[-1]
+
+    # Get base model information
+    base_model_link = str(pretrained_model_name_or_path)
+    model_path_str = str(pretrained_model_name_or_path)
+    is_url = model_path_str.startswith(("http://", "https://"))
+
+    # For URLs, extract the filename from the URL. For local paths, use the filename stem
+    base_model_name = model_path_str.split("/")[-1] if is_url else Path(pretrained_model_name_or_path).name
+
+    # Format validation prompts and create grid layout
+    prompts_text = ""
+    sample_grid = []
+
+    if validation_prompts and videos:
+        prompts_text = "Example prompts used during validation:\n\n"
+
+        # Create samples directory
+        samples_dir = output_dir / "samples"
+        samples_dir.mkdir(exist_ok=True, parents=True)
+
+        # Process videos and create cells
+        cells = []
+        for i, (prompt, video) in enumerate(zip(validation_prompts, videos, strict=False)):
+            if video.exists():
+                # Add prompt to text section
+                prompts_text += f"- `{prompt}`\n"
+
+                # Convert video to GIF
+                gif_path = samples_dir / f"sample_{i}.gif"
+                try:
+                    convert_video_to_gif(video, gif_path)
+
+                    # Create grid cell with collapsible description
+                    cell = (
+                        f"![example{i + 1}](./samples/sample_{i}.gif)"
+                        "<br>"
+                        '<details style="max-width: 300px; margin: auto;">'
+                        f"<summary>Prompt</summary>"
+                        f"{prompt}"
+                        "</details>"
+                    )
+                    cells.append(cell)
+                except Exception as e:
+                    logger.error(f"Failed to process video {video}: {e}")
+
+        # Calculate optimal grid dimensions
+        num_cells = len(cells)
+        if num_cells > 0:
+            # Aim for a roughly square grid, with max 4 columns
+            num_cols = min(4, num_cells)
+            num_rows = (num_cells + num_cols - 1) // num_cols  # Ceiling division
+
+            # Create grid rows
+            for row in range(num_rows):
+                start_idx = row * num_cols
+                end_idx = min(start_idx + num_cols, num_cells)
+                row_cells = cells[start_idx:end_idx]
+                # Properly format the row with table markers and exact number of cells
+                formatted_row = "| " + " | ".join(row_cells) + " |"
+                sample_grid.append(formatted_row)
+
+    # Join grid rows with just the content, no headers needed
+    grid_text = "\n".join(sample_grid) if sample_grid else ""
+
+    # Fill in the template
+    model_card_content = template.format(
+        base_model=base_model_name,
+        base_model_link=base_model_link,
+        model_name=model_name,
+        training_type="LoRA fine-tuning" if config.model.training_mode == "lora" else "Full model fine-tuning",
+        training_steps=config.optimization.steps,
+        learning_rate=config.optimization.learning_rate,
+        batch_size=config.optimization.batch_size,
+        validation_prompts=prompts_text,
+        sample_grid=grid_text,
+    )
+
+    # Save the model card directly
+    model_card_path = output_dir / "README.md"
+    model_card_path.write_text(model_card_content)
+
+    return model_card_path

packages/ltx-trainer/src/ltx_trainer/model_loader.py

@@ -0,0 +1,336 @@
+# ruff: noqa: PLC0415
+
+"""
+Model loader for LTX-2 trainer using the new ltx-core package.
+This module provides a unified interface for loading LTX-2 model components
+for training, using SingleGPUModelBuilder from ltx-core.
+Example usage:
+    # Load individual components
+    vae_encoder = load_video_vae_encoder("/path/to/checkpoint.safetensors", device="cuda")
+    vae_decoder = load_video_vae_decoder("/path/to/checkpoint.safetensors", device="cuda")
+    text_encoder = load_text_encoder("/path/to/checkpoint.safetensors", "/path/to/gemma", device="cuda")
+    # Load all components at once
+    components = load_model("/path/to/checkpoint.safetensors", text_encoder_path="/path/to/gemma")
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import torch
+
+from ltx_trainer import logger
+
+# Type alias for device specification
+Device = str | torch.device
+
+# Type checking imports (not loaded at runtime)
+if TYPE_CHECKING:
+    from ltx_core.components.schedulers import LTX2Scheduler
+    from ltx_core.model.audio_vae import AudioDecoder, AudioEncoder, Vocoder
+    from ltx_core.model.transformer import LTXModel
+    from ltx_core.model.video_vae import VideoDecoder, VideoEncoder
+    from ltx_core.text_encoders.gemma import AVGemmaTextEncoderModel
+
+
+def _to_torch_device(device: Device) -> torch.device:
+    """Convert device specification to torch.device."""
+    return torch.device(device) if isinstance(device, str) else device
+
+
+# =============================================================================
+# Individual Component Loaders
+# =============================================================================
+
+
+def load_transformer(
+    checkpoint_path: str | Path,
+    device: Device = "cpu",
+    dtype: torch.dtype = torch.bfloat16,
+) -> "LTXModel":
+    """Load the LTX transformer model.
+    Args:
+        checkpoint_path: Path to the safetensors checkpoint file
+        device: Device to load model on
+        dtype: Data type for model weights
+    Returns:
+        Loaded LTXModel transformer
+    """
+    from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder
+    from ltx_core.model.transformer.model_configurator import (
+        LTXV_MODEL_COMFY_RENAMING_MAP,
+        LTXModelConfigurator,
+    )
+
+    return SingleGPUModelBuilder(
+        model_path=str(checkpoint_path),
+        model_class_configurator=LTXModelConfigurator,
+        model_sd_ops=LTXV_MODEL_COMFY_RENAMING_MAP,
+    ).build(device=_to_torch_device(device), dtype=dtype)
+
+
+def load_video_vae_encoder(
+    checkpoint_path: str | Path,
+    device: Device = "cpu",
+    dtype: torch.dtype = torch.bfloat16,
+) -> "VideoEncoder":
+    """Load the video VAE encoder (for preprocessing).
+    Args:
+        checkpoint_path: Path to the safetensors checkpoint file
+        device: Device to load model on
+        dtype: Data type for model weights
+    Returns:
+        Loaded VideoEncoder
+    """
+    from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder
+    from ltx_core.model.video_vae import VAE_ENCODER_COMFY_KEYS_FILTER, VideoEncoderConfigurator
+
+    return SingleGPUModelBuilder(
+        model_path=str(checkpoint_path),
+        model_class_configurator=VideoEncoderConfigurator,
+        model_sd_ops=VAE_ENCODER_COMFY_KEYS_FILTER,
+    ).build(device=_to_torch_device(device), dtype=dtype)
+
+
+def load_video_vae_decoder(
+    checkpoint_path: str | Path,
+    device: Device = "cpu",
+    dtype: torch.dtype = torch.bfloat16,
+) -> "VideoDecoder":
+    """Load the video VAE decoder (for inference/validation).
+    Args:
+        checkpoint_path: Path to the safetensors checkpoint file
+        device: Device to load model on
+        dtype: Data type for model weights
+    Returns:
+        Loaded VideoDecoder
+    """
+    from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder
+    from ltx_core.model.video_vae import VAE_DECODER_COMFY_KEYS_FILTER, VideoDecoderConfigurator
+
+    return SingleGPUModelBuilder(
+        model_path=str(checkpoint_path),
+        model_class_configurator=VideoDecoderConfigurator,
+        model_sd_ops=VAE_DECODER_COMFY_KEYS_FILTER,
+    ).build(device=_to_torch_device(device), dtype=dtype)
+
+
+def load_audio_vae_encoder(
+    checkpoint_path: str | Path,
+    device: Device = "cpu",
+    dtype: torch.dtype = torch.bfloat16,
+) -> "AudioEncoder":
+    """Load the audio VAE encoder (for preprocessing).
+    Args:
+        checkpoint_path: Path to the safetensors checkpoint file
+        device: Device to load model on
+        dtype: Data type for model weights (default bfloat16, but float32 recommended for quality)
+    Returns:
+        Loaded AudioEncoder
+    """
+    from ltx_core.loader import SingleGPUModelBuilder
+    from ltx_core.model.audio_vae import AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER, AudioEncoderConfigurator
+
+    return SingleGPUModelBuilder(
+        model_path=str(checkpoint_path),
+        model_class_configurator=AudioEncoderConfigurator,
+        model_sd_ops=AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER,
+    ).build(device=_to_torch_device(device), dtype=dtype)
+
+
+def load_audio_vae_decoder(
+    checkpoint_path: str | Path,
+    device: Device = "cpu",
+    dtype: torch.dtype = torch.bfloat16,
+) -> "AudioDecoder":
+    """Load the audio VAE decoder.
+    Args:
+        checkpoint_path: Path to the safetensors checkpoint file
+        device: Device to load model on
+        dtype: Data type for model weights
+    Returns:
+        Loaded AudioDecoder
+    """
+    from ltx_core.loader import SingleGPUModelBuilder
+    from ltx_core.model.audio_vae import AUDIO_VAE_DECODER_COMFY_KEYS_FILTER, AudioDecoderConfigurator
+
+    return SingleGPUModelBuilder(
+        model_path=str(checkpoint_path),
+        model_class_configurator=AudioDecoderConfigurator,
+        model_sd_ops=AUDIO_VAE_DECODER_COMFY_KEYS_FILTER,
+    ).build(device=_to_torch_device(device), dtype=dtype)
+
+
+def load_vocoder(
+    checkpoint_path: str | Path,
+    device: Device = "cpu",
+    dtype: torch.dtype = torch.bfloat16,
+) -> "Vocoder":
+    """Load the vocoder (for audio waveform generation).
+    Args:
+        checkpoint_path: Path to the safetensors checkpoint file
+        device: Device to load model on
+        dtype: Data type for model weights
+    Returns:
+        Loaded Vocoder
+    """
+    from ltx_core.loader import SingleGPUModelBuilder
+    from ltx_core.model.audio_vae import VOCODER_COMFY_KEYS_FILTER, VocoderConfigurator
+
+    return SingleGPUModelBuilder(
+        model_path=str(checkpoint_path),
+        model_class_configurator=VocoderConfigurator,
+        model_sd_ops=VOCODER_COMFY_KEYS_FILTER,
+    ).build(device=_to_torch_device(device), dtype=dtype)
+
+
+def load_text_encoder(
+    checkpoint_path: str | Path,
+    gemma_model_path: str | Path,
+    device: Device = "cpu",
+    dtype: torch.dtype = torch.bfloat16,
+) -> "AVGemmaTextEncoderModel":
+    """Load the Gemma text encoder.
+    Args:
+        checkpoint_path: Path to the LTX-2 safetensors checkpoint file
+        gemma_model_path: Path to Gemma model directory
+        device: Device to load model on
+        dtype: Data type for model weights
+    Returns:
+        Loaded AVGemmaTextEncoderModel
+    """
+    from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder
+    from ltx_core.text_encoders.gemma.encoders.av_encoder import (
+        AV_GEMMA_TEXT_ENCODER_KEY_OPS,
+        AVGemmaTextEncoderModelConfigurator,
+    )
+    from ltx_core.text_encoders.gemma.encoders.base_encoder import module_ops_from_gemma_root
+
+    if not Path(gemma_model_path).is_dir():
+        raise ValueError(f"Gemma model path is not a directory: {gemma_model_path}")
+
+    torch_device = _to_torch_device(device)
+    text_encoder = SingleGPUModelBuilder(
+        model_path=str(checkpoint_path),
+        model_class_configurator=AVGemmaTextEncoderModelConfigurator,
+        model_sd_ops=AV_GEMMA_TEXT_ENCODER_KEY_OPS,
+        module_ops=module_ops_from_gemma_root(str(gemma_model_path)),
+    ).build(device=torch_device, dtype=dtype)
+
+    return text_encoder
+
+
+# =============================================================================
+# Combined Component Loader
+# =============================================================================
+
+
+@dataclass
+class LtxModelComponents:
+    """Container for all LTX-2 model components."""
+
+    transformer: "LTXModel"
+    video_vae_encoder: "VideoEncoder | None" = None
+    video_vae_decoder: "VideoDecoder | None" = None
+    audio_vae_decoder: "AudioDecoder | None" = None
+    vocoder: "Vocoder | None" = None
+    text_encoder: "AVGemmaTextEncoderModel | None" = None
+    scheduler: "LTX2Scheduler | None" = None
+
+
+def load_model(
+    checkpoint_path: str | Path,
+    text_encoder_path: str | Path | None = None,
+    device: Device = "cpu",
+    dtype: torch.dtype = torch.bfloat16,
+    with_video_vae_encoder: bool = False,
+    with_video_vae_decoder: bool = True,
+    with_audio_vae_decoder: bool = True,
+    with_vocoder: bool = True,
+    with_text_encoder: bool = True,
+) -> LtxModelComponents:
+    """
+    Load LTX-2 model components from a safetensors checkpoint.
+    This is a convenience function that loads multiple components at once.
+    For loading individual components, use the dedicated functions:
+    - load_transformer()
+    - load_video_vae_encoder()
+    - load_video_vae_decoder()
+    - load_audio_vae_decoder()
+    - load_vocoder()
+    - load_text_encoder()
+    Args:
+        checkpoint_path: Path to the safetensors checkpoint file
+        text_encoder_path: Path to Gemma model directory (required if with_text_encoder=True)
+        device: Device to load models on ("cuda", "cpu", etc.)
+        dtype: Data type for model weights
+        with_video_vae_encoder: Whether to load the video VAE encoder (for preprocessing)
+        with_video_vae_decoder: Whether to load the video VAE decoder (for inference/validation)
+        with_audio_vae_decoder: Whether to load the audio VAE decoder
+        with_vocoder: Whether to load the vocoder
+        with_text_encoder: Whether to load the text encoder
+    Returns:
+        LtxModelComponents containing all loaded model components
+    """
+    from ltx_core.components.schedulers import LTX2Scheduler
+
+    checkpoint_path = Path(checkpoint_path)
+
+    # Validate checkpoint exists
+    if not checkpoint_path.exists():
+        raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")
+
+    logger.info(f"Loading LTX-2 model from {checkpoint_path}")
+
+    torch_device = _to_torch_device(device)
+
+    # Load transformer
+    logger.debug("Loading transformer...")
+    transformer = load_transformer(checkpoint_path, torch_device, dtype)
+
+    # Load video VAE encoder
+    video_vae_encoder = None
+    if with_video_vae_encoder:
+        logger.debug("Loading video VAE encoder...")
+        video_vae_encoder = load_video_vae_encoder(checkpoint_path, torch_device, dtype)
+
+    # Load video VAE decoder
+    video_vae_decoder = None
+    if with_video_vae_decoder:
+        logger.debug("Loading video VAE decoder...")
+        video_vae_decoder = load_video_vae_decoder(checkpoint_path, torch_device, dtype)
+
+    # Load audio VAE decoder
+    audio_vae_decoder = None
+    if with_audio_vae_decoder:
+        logger.debug("Loading audio VAE decoder...")
+        audio_vae_decoder = load_audio_vae_decoder(checkpoint_path, torch_device, dtype)
+
+    # Load vocoder
+    vocoder = None
+    if with_vocoder:
+        logger.debug("Loading vocoder...")
+        vocoder = load_vocoder(checkpoint_path, torch_device, dtype)
+
+    # Load text encoder
+    text_encoder = None
+    if with_text_encoder:
+        if text_encoder_path is None:
+            raise ValueError("text_encoder_path must be provided when with_text_encoder=True")
+        logger.debug("Loading Gemma text encoder...")
+        text_encoder = load_text_encoder(checkpoint_path, text_encoder_path, torch_device, dtype)
+
+    # Create scheduler (stateless, no loading needed)
+    scheduler = LTX2Scheduler()
+
+    return LtxModelComponents(
+        transformer=transformer,
+        video_vae_encoder=video_vae_encoder,
+        video_vae_decoder=video_vae_decoder,
+        audio_vae_decoder=audio_vae_decoder,
+        vocoder=vocoder,
+        text_encoder=text_encoder,
+        scheduler=scheduler,
+    )

packages/ltx-trainer/src/ltx_trainer/progress.py

@@ -0,0 +1,236 @@
+"""Progress tracking for LTX training.
+This module provides a unified progress display for training and validation sampling,
+encapsulating all Rich progress bar logic in one place.
+"""
+
+from rich.progress import (
+    BarColumn,
+    Progress,
+    TaskID,
+    TextColumn,
+    TimeElapsedColumn,
+    TimeRemainingColumn,
+)
+
+
+class SamplingContext:
+    """Context for validation sampling progress tracking.
+    Provides a unified progress display showing current video and denoising step.
+    Display format: "Sampling X/Y [████████████] step Z/W"
+    The progress bar shows the denoising progress for the current video.
+    """
+
+    def __init__(self, progress: Progress | None, task: TaskID | None, num_prompts: int, num_steps: int):
+        self._progress = progress
+        self._task = task
+        self._num_prompts = num_prompts
+        self._num_steps = num_steps
+
+    def start_video(self, video_idx: int) -> None:
+        """Start tracking a new video (resets step progress)."""
+        if self._progress is None or self._task is None:
+            return
+        # Reset task for new video: completed=0, total=num_steps
+        self._progress.reset(self._task, total=self._num_steps)
+        self._progress.update(
+            self._task,
+            completed=0,
+            video=f"{video_idx + 1}/{self._num_prompts}",
+            info=f"step 0/{self._num_steps}",
+        )
+
+    def advance_step(self) -> None:
+        """Advance the denoising step by one."""
+        if self._progress is None or self._task is None:
+            return
+        self._progress.advance(self._task)
+        completed = int(self._progress.tasks[self._task].completed)
+        self._progress.update(self._task, info=f"step {completed}/{self._num_steps}")
+
+    def cleanup(self) -> None:
+        """Hide sampling task when done."""
+        if self._progress is None or self._task is None:
+            return
+        self._progress.update(self._task, visible=False)
+
+
+class StandaloneSamplingProgress:
+    """Standalone progress display for inference scripts.
+    Unlike SamplingContext (which integrates with TrainingProgress), this class
+    manages its own Rich Progress instance for use in standalone inference scripts.
+    Usage:
+        with StandaloneSamplingProgress(num_steps=30) as ctx:
+            for step in range(30):
+                # ... denoising step ...
+                ctx.advance_step()
+    """
+
+    def __init__(self, num_steps: int, description: str = "Generating"):
+        """Initialize standalone sampling progress.
+        Args:
+            num_steps: Total number of denoising steps
+            description: Description to show in progress bar
+        """
+        self._num_steps = num_steps
+        self._description = description
+        self._progress: Progress | None = None
+        self._task: TaskID | None = None
+
+    def __enter__(self) -> "StandaloneSamplingProgress":
+        """Start the progress display."""
+        self._progress = Progress(
+            TextColumn("[progress.description]{task.description}"),
+            BarColumn(bar_width=40, style="blue"),
+            TextColumn("{task.fields[info]}", style="cyan"),
+            TimeElapsedColumn(),
+            TextColumn("ETA:"),
+            TimeRemainingColumn(compact=True),
+        )
+        self._progress.__enter__()
+        self._task = self._progress.add_task(
+            self._description,
+            total=self._num_steps,
+            info=f"step 0/{self._num_steps}",
+        )
+        return self
+
+    def __exit__(self, *args) -> None:
+        """Stop the progress display."""
+        if self._progress is not None:
+            self._progress.__exit__(*args)
+
+    def advance_step(self) -> None:
+        """Advance the denoising step by one."""
+        if self._progress is None or self._task is None:
+            return
+        self._progress.advance(self._task)
+        completed = int(self._progress.tasks[self._task].completed)
+        self._progress.update(self._task, info=f"step {completed}/{self._num_steps}")
+
+
+class TrainingProgress:
+    """Manages Rich progress display for training and validation.
+    This class encapsulates all progress bar logic, providing a clean interface
+    for the trainer to update progress without dealing with Rich internals.
+    Usage:
+        with TrainingProgress(enabled=True, total_steps=1000) as progress:
+            for step in range(1000):
+                # ... training step ...
+                progress.update_training(loss=0.1, lr=1e-4, step_time=0.5)
+                if should_validate:
+                    sampling_ctx = progress.start_sampling(num_prompts=3, num_steps=30)
+                    sampler = ValidationSampler(..., sampling_context=sampling_ctx)
+                    for prompt_idx, prompt in enumerate(prompts):
+                        sampling_ctx.start_video(prompt_idx)
+                        sampler.generate(...)
+                    sampling_ctx.cleanup()
+    """
+
+    def __init__(self, enabled: bool, total_steps: int):
+        """Initialize progress tracking.
+        Args:
+            enabled: Whether to display progress bars (False for non-main processes)
+            total_steps: Total number of training steps
+        """
+        self._enabled = enabled
+        self._total_steps = total_steps
+        self._train_task: TaskID | None = None
+
+        if not enabled:
+            self._progress = None
+            return
+
+        # Single Progress instance with flexible columns
+        self._progress = Progress(
+            TextColumn("[progress.description]{task.description}"),
+            TextColumn("{task.fields[video]}", style="magenta"),
+            BarColumn(bar_width=40, style="blue"),
+            TextColumn("{task.fields[info]}", style="cyan"),
+            TimeElapsedColumn(),
+            TextColumn("ETA:"),
+            TimeRemainingColumn(compact=True),
+        )
+
+    def __enter__(self) -> "TrainingProgress":
+        """Enter the progress context, starting the live display."""
+        if self._progress is not None:
+            self._progress.__enter__()
+            self._train_task = self._progress.add_task(
+                "Training",
+                total=self._total_steps,
+                video=f"0/{self._total_steps}",
+                info="Starting...",
+            )
+        return self
+
+    def __exit__(self, *args) -> None:
+        """Exit the progress context, stopping the live display."""
+        if self._progress is not None:
+            self._progress.__exit__(*args)
+
+    @property
+    def enabled(self) -> bool:
+        """Whether progress display is enabled."""
+        return self._enabled
+
+    def update_training(
+        self,
+        *,
+        loss: float,
+        lr: float,
+        step_time: float,
+        advance: bool = True,
+    ) -> None:
+        """Update the training progress display.
+        Args:
+            loss: Current training loss
+            lr: Current learning rate
+            step_time: Time taken for this step in seconds
+            advance: Whether to advance the progress by one step
+        """
+        if self._progress is None or self._train_task is None:
+            return
+
+        info = f"Loss: {loss:.4f} | LR: {lr:.2e} | {step_time:.2f}s/step"
+        self._progress.update(
+            self._train_task,
+            advance=1 if advance else 0,
+            info=info,
+        )
+        # Update step count in video column
+        completed = int(self._progress.tasks[self._train_task].completed)
+        self._progress.update(self._train_task, video=f"{completed}/{self._total_steps}")
+
+    def start_sampling(self, num_prompts: int, num_steps: int) -> SamplingContext:
+        """Start validation sampling progress tracking.
+        Creates a task that shows current video and denoising step progress.
+        Format: "Sampling X/Y [████████████] step Z/W"
+        Args:
+            num_prompts: Number of validation prompts to sample
+            num_steps: Number of denoising steps per sample
+        Returns:
+            SamplingContext for tracking progress (no-op if progress is disabled)
+        """
+        if self._progress is None:
+            # Return a no-op context when progress is disabled
+            return SamplingContext(
+                progress=None,
+                task=None,
+                num_prompts=num_prompts,
+                num_steps=num_steps,
+            )
+
+        task = self._progress.add_task(
+            "Sampling",
+            total=num_steps,
+            completed=0,
+            video=f"0/{num_prompts}",
+            info=f"step 0/{num_steps}",
+        )
+
+        return SamplingContext(
+            progress=self._progress,
+            task=task,
+            num_prompts=num_prompts,
+            num_steps=num_steps,
+        )

packages/ltx-trainer/src/ltx_trainer/quantization.py

@@ -0,0 +1,90 @@
+# Adapted from: https://github.com/bghira/SimpleTuner/blob/main/helpers/training/quantisation/__init__.py
+from typing import Literal
+
+import torch
+from optimum.quanto import qtype
+
+from ltx_trainer import logger
+
+QuantizationOptions = Literal[
+    "no_change",
+    "int8-quanto",
+    "int4-quanto",
+    "int2-quanto",
+    "fp8-quanto",
+    "fp8uz-quanto",
+]
+
+
+def quantize_model(
+    model: torch.nn.Module,
+    precision: QuantizationOptions,
+    quantize_activations: bool = False,
+) -> torch.nn.Module:
+    """
+    Quantize a model using the specified precision settings.
+    Args:
+        model: The model to quantize.
+        precision: The precision level to quantize to (e.g. "int8-quanto", "fp8-quanto").
+        quantize_activations: Whether to quantize activations in addition to weights.
+    Returns:
+        The quantized model, or the original model if no quantization is performed.
+    """
+    if precision is None or precision == "no_change":
+        return model
+
+    from optimum.quanto import freeze, quantize  # noqa: PLC0415
+
+    weight_quant = _quanto_type_map(precision)
+    extra_quanto_args = {
+        "exclude": [
+            "proj_in",
+            "time_embed.*",
+            "caption_projection.*",
+            "rope",
+            "*norm*",
+            "proj_out",
+        ]
+    }
+    if quantize_activations:
+        logger.info("Freezing model weights and activations")
+        extra_quanto_args["activations"] = weight_quant
+    else:
+        logger.info("Freezing model weights only")
+
+    quantize(model, weights=weight_quant, **extra_quanto_args)
+    freeze(model)
+    return model
+
+
+def _quanto_type_map(precision: QuantizationOptions) -> torch.dtype | qtype | None:  # noqa: PLR0911
+    if precision == "no_change":
+        return None
+
+    from optimum.quanto import (  # noqa: PLC0415
+        qfloat8,
+        qfloat8_e4m3fnuz,
+        qint2,
+        qint4,
+        qint8,
+    )
+
+    if precision == "int2-quanto":
+        return qint2
+    elif precision == "int4-quanto":
+        return qint4
+    elif precision == "int8-quanto":
+        return qint8
+    elif precision in ("fp8-quanto", "fp8uz-quanto"):
+        if torch.backends.mps.is_available():
+            logger.warning(
+                "MPS doesn't support dtype float8. "
+                "you must select another precision level such as int2, int8, or int8.",
+            )
+            return None
+        if precision == "fp8-quanto":
+            return qfloat8
+        elif precision == "fp8uz-quanto":
+            return qfloat8_e4m3fnuz
+
+    raise ValueError(f"Invalid quantisation level: {precision}")

packages/ltx-trainer/src/ltx_trainer/timestep_samplers.py

@@ -0,0 +1,128 @@
+import torch
+
+
+class TimestepSampler:
+    """Base class for timestep samplers.
+    Timestep samplers are used to sample timesteps for diffusion models.
+    They should implement both sample() and sample_for() methods.
+    """
+
+    def sample(self, batch_size: int, seq_length: int | None = None, device: torch.device = None) -> torch.Tensor:
+        """Sample timesteps for a batch.
+        Args:
+            batch_size: Number of timesteps to sample
+            seq_length: (optional) Length of the sequence being processed
+            device: Device to place the samples on
+        Returns:
+            Tensor of shape (batch_size,) containing timesteps
+        """
+        raise NotImplementedError
+
+    def sample_for(self, batch: torch.Tensor) -> torch.Tensor:
+        """Sample timesteps for a specific batch tensor.
+        Args:
+            batch: Input tensor of shape (batch_size, seq_length, ...)
+        Returns:
+            Tensor of shape (batch_size,) containing timesteps
+        """
+        raise NotImplementedError
+
+
+class UniformTimestepSampler(TimestepSampler):
+    """Samples timesteps uniformly between min_value and max_value (default 0 and 1)."""
+
+    def __init__(self, min_value: float = 0.0, max_value: float = 1.0):
+        self.min_value = min_value
+        self.max_value = max_value
+
+    def sample(self, batch_size: int, seq_length: int | None = None, device: torch.device = None) -> torch.Tensor:  # noqa: ARG002
+        return torch.rand(batch_size, device=device) * (self.max_value - self.min_value) + self.min_value
+
+    def sample_for(self, batch: torch.Tensor) -> torch.Tensor:
+        if batch.ndim != 3:
+            raise ValueError(f"Batch should have 3 dimensions, got {batch.ndim}")
+
+        return self.sample(batch.shape[0], device=batch.device)
+
+
+class ShiftedLogitNormalTimestepSampler:
+    """
+    Samples timesteps from a shifted logit-normal distribution,
+    where the shift is determined by the sequence length.
+    """
+
+    def __init__(self, std: float = 1.0):
+        self.std = std
+
+    def sample(self, batch_size: int, seq_length: int, device: torch.device = None) -> torch.Tensor:
+        """Sample timesteps for a batch from a shifted logit-normal distribution.
+        Args:
+            batch_size: Number of timesteps to sample
+            seq_length: Length of the sequence being processed, used to determine the shift
+            device: Device to place the samples on
+        Returns:
+            Tensor of shape (batch_size,) containing timesteps sampled from a shifted
+            logit-normal distribution, where the shift is determined by seq_length
+        """
+        shift = self._get_shift_for_sequence_length(seq_length)
+        normal_samples = torch.randn((batch_size,), device=device) * self.std + shift
+        timesteps = torch.sigmoid(normal_samples)
+        return timesteps
+
+    def sample_for(self, batch: torch.Tensor) -> torch.Tensor:
+        """Sample timesteps for a specific batch tensor.
+        Args:
+            batch: Input tensor of shape (batch_size, seq_length, ...)
+        Returns:
+            Tensor of shape (batch_size,) containing timesteps sampled from a shifted
+            logit-normal distribution, where the shift is determined by the sequence length
+            of the input batch
+        Raises:
+            ValueError: If the input batch does not have 3 dimensions
+        """
+        if batch.ndim != 3:
+            raise ValueError(f"Batch should have 3 dimensions, got {batch.ndim}")
+
+        batch_size, seq_length, _ = batch.shape
+        return self.sample(batch_size, seq_length, device=batch.device)
+
+    @staticmethod
+    def _get_shift_for_sequence_length(
+        seq_length: int,
+        min_tokens: int = 1024,
+        max_tokens: int = 4096,
+        min_shift: float = 0.95,
+        max_shift: float = 2.05,
+    ) -> float:
+        # Calculate the shift value for a given sequence length using linear interpolation
+        # between min_shift and max_shift based on sequence length.
+        m = (max_shift - min_shift) / (max_tokens - min_tokens)  # Calculate slope
+        b = min_shift - m * min_tokens  # Calculate y-intercept
+        shift = m * seq_length + b  # Apply linear equation y = mx + b
+        return shift
+
+
+SAMPLERS = {
+    "uniform": UniformTimestepSampler,
+    "shifted_logit_normal": ShiftedLogitNormalTimestepSampler,
+}
+
+
+def example() -> None:
+    # noinspection PyUnresolvedReferences
+    import matplotlib.pyplot as plt  # noqa: PLC0415
+
+    sampler = ShiftedLogitNormalTimestepSampler()
+    for seq_length in [1024, 2048, 4096, 8192]:
+        samples = sampler.sample(batch_size=1_000_000, seq_length=seq_length)
+
+        # plot the histogram of the samples
+        plt.hist(samples.numpy(), bins=100, density=True)
+        plt.title(f"Timestep Samples for Sequence Length {seq_length}")
+        plt.xlabel("Timestep")
+        plt.ylabel("Density")
+        plt.show()
+
+
+if __name__ == "__main__":
+    example()

packages/ltx-trainer/src/ltx_trainer/trainer.py

@@ -0,0 +1,955 @@
+import os
+import time
+import warnings
+from pathlib import Path
+from typing import Callable
+
+import torch
+import wandb
+import yaml
+from accelerate import Accelerator, DistributedType
+from accelerate.utils import set_seed
+from peft import LoraConfig, get_peft_model, get_peft_model_state_dict, set_peft_model_state_dict
+from peft.tuners.tuners_utils import BaseTunerLayer
+from peft.utils import ModulesToSaveWrapper
+from pydantic import BaseModel
+from safetensors.torch import load_file, save_file
+from torch import Tensor
+from torch.optim import AdamW
+from torch.optim.lr_scheduler import (
+    CosineAnnealingLR,
+    CosineAnnealingWarmRestarts,
+    LinearLR,
+    LRScheduler,
+    PolynomialLR,
+    StepLR,
+)
+from torch.utils.data import DataLoader
+from torchvision.transforms import functional as F  # noqa: N812
+
+from ltx_trainer import logger
+from ltx_trainer.config import LtxTrainerConfig
+from ltx_trainer.config_display import print_config
+from ltx_trainer.datasets import PrecomputedDataset
+from ltx_trainer.hf_hub_utils import push_to_hub
+from ltx_trainer.model_loader import load_model as load_ltx_model
+from ltx_trainer.model_loader import load_text_encoder
+from ltx_trainer.progress import TrainingProgress
+from ltx_trainer.quantization import quantize_model
+from ltx_trainer.timestep_samplers import SAMPLERS
+from ltx_trainer.training_strategies import get_training_strategy
+from ltx_trainer.utils import get_gpu_memory_gb, open_image_as_srgb, save_image
+from ltx_trainer.validation_sampler import CachedPromptEmbeddings, GenerationConfig, ValidationSampler
+from ltx_trainer.video_utils import read_video, save_video
+
+# Disable irrelevant warnings from transformers
+os.environ["TOKENIZERS_PARALLELISM"] = "true"
+
+# Silence bitsandbytes warnings about casting
+warnings.filterwarnings(
+    "ignore", message="MatMul8bitLt: inputs will be cast from torch.bfloat16 to float16 during quantization"
+)
+
+# Disable progress bars if not main process
+IS_MAIN_PROCESS = os.environ.get("LOCAL_RANK", "0") == "0"
+if not IS_MAIN_PROCESS:
+    from transformers.utils.logging import disable_progress_bar
+
+    disable_progress_bar()
+
+StepCallback = Callable[[int, int, list[Path]], None]  # (step, total, list[sampled_video_path]) -> None
+
+MEMORY_CHECK_INTERVAL = 200
+
+
+class TrainingStats(BaseModel):
+    """Statistics collected during training"""
+
+    total_time_seconds: float
+    steps_per_second: float
+    samples_per_second: float
+    peak_gpu_memory_gb: float
+    global_batch_size: int
+    num_processes: int
+
+
+class LtxvTrainer:
+    def __init__(self, trainer_config: LtxTrainerConfig) -> None:
+        self._config = trainer_config
+        if IS_MAIN_PROCESS:
+            print_config(trainer_config)
+        self._training_strategy = get_training_strategy(self._config.training_strategy)
+        self._cached_validation_embeddings = self._load_text_encoder_and_cache_embeddings()
+        self._load_models()
+        self._setup_accelerator()
+        self._collect_trainable_params()
+        self._load_checkpoint()
+        self._prepare_models_for_training()
+        self._dataset = None
+        self._global_step = -1
+        self._checkpoint_paths = []
+        self._init_wandb()
+
+    def train(  # noqa: PLR0912, PLR0915
+        self,
+        disable_progress_bars: bool = False,
+        step_callback: StepCallback | None = None,
+    ) -> tuple[Path, TrainingStats]:
+        """
+        Start the training process.
+        Returns:
+            Tuple of (saved_model_path, training_stats)
+        """
+        device = self._accelerator.device
+        cfg = self._config
+        start_mem = get_gpu_memory_gb(device)
+
+        train_start_time = time.time()
+
+        # Use the same seed for all processes and ensure deterministic operations
+        set_seed(cfg.seed)
+        logger.debug(f"Process {self._accelerator.process_index} using seed: {cfg.seed}")
+
+        self._init_optimizer()
+        self._init_dataloader()
+        data_iter = iter(self._dataloader)
+        self._init_timestep_sampler()
+
+        # Synchronize all processes after initialization
+        self._accelerator.wait_for_everyone()
+
+        Path(cfg.output_dir).mkdir(parents=True, exist_ok=True)
+
+        # Save the training configuration as YAML
+        self._save_config()
+
+        logger.info("🚀 Starting training...")
+
+        # Create progress tracking (disabled for non-main processes or when explicitly disabled)
+        progress_enabled = IS_MAIN_PROCESS and not disable_progress_bars
+        progress = TrainingProgress(
+            enabled=progress_enabled,
+            total_steps=cfg.optimization.steps,
+        )
+
+        if IS_MAIN_PROCESS and disable_progress_bars:
+            logger.warning("Progress bars disabled. Intermediate status messages will be logged instead.")
+
+        self._transformer.train()
+        self._global_step = 0
+
+        peak_mem_during_training = start_mem
+
+        sampled_videos_paths = None
+
+        with progress:
+            # Initial validation before training starts
+            if cfg.validation.interval and not cfg.validation.skip_initial_validation:
+                sampled_videos_paths = self._sample_videos(progress)
+                if IS_MAIN_PROCESS and sampled_videos_paths and self._config.wandb.log_validation_videos:
+                    self._log_validation_samples(sampled_videos_paths, cfg.validation.prompts)
+
+            self._accelerator.wait_for_everyone()
+
+            for step in range(cfg.optimization.steps * cfg.optimization.gradient_accumulation_steps):
+                # Get next batch, reset the dataloader if needed
+                try:
+                    batch = next(data_iter)
+                except StopIteration:
+                    data_iter = iter(self._dataloader)
+                    batch = next(data_iter)
+
+                step_start_time = time.time()
+                with self._accelerator.accumulate(self._transformer):
+                    is_optimization_step = (step + 1) % cfg.optimization.gradient_accumulation_steps == 0
+                    if is_optimization_step:
+                        self._global_step += 1
+
+                    loss = self._training_step(batch)
+                    self._accelerator.backward(loss)
+
+                    if self._accelerator.sync_gradients and cfg.optimization.max_grad_norm > 0:
+                        self._accelerator.clip_grad_norm_(
+                            self._trainable_params,
+                            cfg.optimization.max_grad_norm,
+                        )
+
+                    self._optimizer.step()
+                    self._optimizer.zero_grad()
+
+                    if self._lr_scheduler is not None:
+                        self._lr_scheduler.step()
+
+                    # Run validation if needed
+                    if (
+                        cfg.validation.interval
+                        and self._global_step > 0
+                        and self._global_step % cfg.validation.interval == 0
+                        and is_optimization_step
+                    ):
+                        if self._accelerator.distributed_type == DistributedType.FSDP:
+                            # FSDP: All processes must participate in validation
+                            sampled_videos_paths = self._sample_videos(progress)
+                            if IS_MAIN_PROCESS and sampled_videos_paths and self._config.wandb.log_validation_videos:
+                                self._log_validation_samples(sampled_videos_paths, cfg.validation.prompts)
+                        # DDP: Only main process runs validation
+                        elif IS_MAIN_PROCESS:
+                            sampled_videos_paths = self._sample_videos(progress)
+                            if sampled_videos_paths and self._config.wandb.log_validation_videos:
+                                self._log_validation_samples(sampled_videos_paths, cfg.validation.prompts)
+
+                    # Save checkpoint if needed
+                    if (
+                        cfg.checkpoints.interval
+                        and self._global_step > 0
+                        and self._global_step % cfg.checkpoints.interval == 0
+                        and is_optimization_step
+                    ):
+                        self._save_checkpoint()
+
+                    self._accelerator.wait_for_everyone()
+
+                    # Call step callback if provided
+                    if step_callback and is_optimization_step:
+                        step_callback(self._global_step, cfg.optimization.steps, sampled_videos_paths)
+
+                    self._accelerator.wait_for_everyone()
+
+                    # Update progress and log metrics
+                    current_lr = self._optimizer.param_groups[0]["lr"]
+                    step_time = (time.time() - step_start_time) * cfg.optimization.gradient_accumulation_steps
+
+                    progress.update_training(
+                        loss=loss.item(),
+                        lr=current_lr,
+                        step_time=step_time,
+                        advance=is_optimization_step,
+                    )
+
+                    # Log metrics to W&B (only on main process and optimization steps)
+                    if IS_MAIN_PROCESS and is_optimization_step:
+                        self._log_metrics(
+                            {
+                                "train/loss": loss.item(),
+                                "train/learning_rate": current_lr,
+                                "train/step_time": step_time,
+                                "train/global_step": self._global_step,
+                            }
+                        )
+
+                    # Fallback logging when progress bars are disabled
+                    if disable_progress_bars and IS_MAIN_PROCESS and self._global_step % 20 == 0:
+                        elapsed = time.time() - train_start_time
+                        progress_percentage = self._global_step / cfg.optimization.steps
+                        if progress_percentage > 0:
+                            total_estimated = elapsed / progress_percentage
+                            total_time = f"{total_estimated // 3600:.0f}h {(total_estimated % 3600) // 60:.0f}m"
+                        else:
+                            total_time = "calculating..."
+                        logger.info(
+                            f"Step {self._global_step}/{cfg.optimization.steps} - "
+                            f"Loss: {loss.item():.4f}, LR: {current_lr:.2e}, "
+                            f"Time/Step: {step_time:.2f}s, Total Time: {total_time}",
+                        )
+
+                    # Sample GPU memory periodically
+                    if step % MEMORY_CHECK_INTERVAL == 0:
+                        current_mem = get_gpu_memory_gb(device)
+                        peak_mem_during_training = max(peak_mem_during_training, current_mem)
+
+        # Collect final stats
+        train_end_time = time.time()
+        end_mem = get_gpu_memory_gb(device)
+        peak_mem = max(start_mem, end_mem, peak_mem_during_training)
+
+        # Calculate steps/second over entire training
+        total_time_seconds = train_end_time - train_start_time
+        steps_per_second = cfg.optimization.steps / total_time_seconds
+
+        samples_per_second = steps_per_second * self._accelerator.num_processes * cfg.optimization.batch_size
+
+        stats = TrainingStats(
+            total_time_seconds=total_time_seconds,
+            steps_per_second=steps_per_second,
+            samples_per_second=samples_per_second,
+            peak_gpu_memory_gb=peak_mem,
+            num_processes=self._accelerator.num_processes,
+            global_batch_size=cfg.optimization.batch_size * self._accelerator.num_processes,
+        )
+
+        saved_path = self._save_checkpoint()
+
+        if IS_MAIN_PROCESS:
+            # Log the training statistics
+            self._log_training_stats(stats)
+
+            # Upload artifacts to hub if enabled
+            if cfg.hub.push_to_hub:
+                push_to_hub(saved_path, sampled_videos_paths, self._config)
+
+            # Log final stats to W&B
+            if self._wandb_run is not None:
+                self._log_metrics(
+                    {
+                        "stats/total_time_minutes": stats.total_time_seconds / 60,
+                        "stats/steps_per_second": stats.steps_per_second,
+                        "stats/samples_per_second": stats.samples_per_second,
+                        "stats/peak_gpu_memory_gb": stats.peak_gpu_memory_gb,
+                    }
+                )
+                self._wandb_run.finish()
+
+        self._accelerator.wait_for_everyone()
+        self._accelerator.end_training()
+
+        return saved_path, stats
+
+    def _training_step(self, batch: dict[str, dict[str, Tensor]]) -> Tensor:
+        """Perform a single training step using the configured strategy."""
+        # Apply embedding connectors to transform pre-computed text embeddings
+        conditions = batch["conditions"]
+        video_embeds, audio_embeds, attention_mask = self._text_encoder._run_connectors(
+            conditions["prompt_embeds"], conditions["prompt_attention_mask"]
+        )
+        conditions["video_prompt_embeds"] = video_embeds
+        conditions["audio_prompt_embeds"] = audio_embeds
+        conditions["prompt_attention_mask"] = attention_mask
+
+        # Use strategy to prepare training inputs (returns ModelInputs with Modality objects)
+        model_inputs = self._training_strategy.prepare_training_inputs(batch, self._timestep_sampler)
+
+        # Run transformer forward pass with Modality-based interface
+        video_pred, audio_pred = self._transformer(
+            video=model_inputs.video,
+            audio=model_inputs.audio,
+            perturbations=None,
+        )
+
+        # Use strategy to compute loss
+        loss = self._training_strategy.compute_loss(video_pred, audio_pred, model_inputs)
+
+        return loss
+
+    def _load_text_encoder_and_cache_embeddings(self) -> list[CachedPromptEmbeddings] | None:
+        """Load text encoder, computes and returns validation embeddings."""
+
+        # This method:
+        #   1. Loads the text encoder on GPU
+        #   2. If validation prompts are configured, computes and caches their embeddings
+        #   3. Unloads the heavy Gemma model while keeping the lightweight embedding connectors
+        #   The text encoder is kept (as self._text_encoder) but with model/tokenizer/feature_extractor
+        #   set to None. Only the embedding connectors remain for use during training.
+
+        # Load text encoder on GPU
+        logger.debug("Loading text encoder...")
+        if self._config.acceleration.load_text_encoder_in_8bit:
+            logger.warning(
+                "⚠️  load_text_encoder_in_8bit is set to True but 8-bit text encoder loading "
+                "is not currently implemented. The text encoder will be loaded in bfloat16 precision."
+            )
+
+        self._text_encoder = load_text_encoder(
+            checkpoint_path=self._config.model.model_path,
+            gemma_model_path=self._config.model.text_encoder_path,
+            device="cuda",
+            dtype=torch.bfloat16,
+        )
+
+        # Cache validation embeddings if prompts are configured
+        cached_embeddings = None
+        if self._config.validation.prompts:
+            logger.info(f"Pre-computing embeddings for {len(self._config.validation.prompts)} validation prompts...")
+            cached_embeddings = []
+            with torch.inference_mode():
+                for prompt in self._config.validation.prompts:
+                    v_ctx_pos, a_ctx_pos, _ = self._text_encoder(prompt)
+                    v_ctx_neg, a_ctx_neg, _ = self._text_encoder(self._config.validation.negative_prompt)
+
+                    cached_embeddings.append(
+                        CachedPromptEmbeddings(
+                            video_context_positive=v_ctx_pos.cpu(),
+                            audio_context_positive=a_ctx_pos.cpu(),
+                            video_context_negative=v_ctx_neg.cpu() if v_ctx_neg is not None else None,
+                            audio_context_negative=a_ctx_neg.cpu() if a_ctx_neg is not None else None,
+                        )
+                    )
+
+        # Unload heavy components to free VRAM, keeping only the embedding connectors
+        self._text_encoder.model = None
+        self._text_encoder.tokenizer = None
+        self._text_encoder.feature_extractor_linear = None
+        torch.cuda.empty_cache()
+
+        logger.debug("Validation prompt embeddings cached. Gemma model unloaded")
+        return cached_embeddings
+
+    def _load_models(self) -> None:
+        """Load the LTX-2 model components."""
+        # Load audio components if:
+        # 1. Training strategy requires audio (training the audio branch), OR
+        # 2. Validation is configured to generate audio (even if not training audio)
+        load_audio = self._training_strategy.requires_audio or self._config.validation.generate_audio
+
+        # Check if we need VAE encoder (for image or reference video conditioning)
+        need_vae_encoder = (
+            self._config.validation.images is not None or self._config.validation.reference_videos is not None
+        )
+
+        # Load all model components (except text encoder - already handled)
+        components = load_ltx_model(
+            checkpoint_path=self._config.model.model_path,
+            device="cpu",
+            dtype=torch.bfloat16,
+            with_video_vae_encoder=need_vae_encoder,  # Needed for image conditioning
+            with_video_vae_decoder=True,  # Needed for validation sampling
+            with_audio_vae_decoder=load_audio,
+            with_vocoder=load_audio,
+            with_text_encoder=False,  # Text encoder handled separately
+        )
+
+        # Extract components
+        self._transformer = components.transformer
+        self._vae_decoder = components.video_vae_decoder.to(dtype=torch.bfloat16)
+        self._vae_encoder = components.video_vae_encoder
+        if self._vae_encoder is not None:
+            self._vae_encoder = self._vae_encoder.to(dtype=torch.bfloat16)
+        self._scheduler = components.scheduler
+        self._audio_vae = components.audio_vae_decoder
+        self._vocoder = components.vocoder
+        # Note: self._text_encoder was set in _load_text_encoder_and_cache_embeddings
+
+        # Determine initial dtype based on training mode.
+        # Note: For FSDP + LoRA, we'll cast to FP32 later in _prepare_models_for_training()
+        # after the accelerator is set up, and we can detect FSDP.
+        transformer_dtype = torch.bfloat16 if self._config.model.training_mode == "lora" else torch.float32
+        self._transformer = self._transformer.to(dtype=transformer_dtype)
+
+        if self._config.acceleration.quantization is not None:
+            if self._config.model.training_mode == "full":
+                raise ValueError("Quantization is not supported in full training mode.")
+
+            logger.warning(f"Quantizing model with precision: {self._config.acceleration.quantization}")
+            self._transformer = quantize_model(
+                self._transformer,
+                precision=self._config.acceleration.quantization,
+            )
+
+        # Freeze all models. We later unfreeze the transformer based on training mode.
+        # Note: embedding_connectors are already frozen (they come from the frozen text encoder)
+        self._vae_decoder.requires_grad_(False)
+        if self._vae_encoder is not None:
+            self._vae_encoder.requires_grad_(False)
+        self._transformer.requires_grad_(False)
+        if self._audio_vae is not None:
+            self._audio_vae.requires_grad_(False)
+        if self._vocoder is not None:
+            self._vocoder.requires_grad_(False)
+
+    def _collect_trainable_params(self) -> None:
+        """Collect trainable parameters based on training mode."""
+        if self._config.model.training_mode == "lora":
+            # For LoRA training, first set up LoRA layers
+            self._setup_lora()
+        elif self._config.model.training_mode == "full":
+            # For full training, unfreeze all transformer parameters
+            self._transformer.requires_grad_(True)
+        else:
+            raise ValueError(f"Unknown training mode: {self._config.model.training_mode}")
+
+        self._trainable_params = [p for p in self._transformer.parameters() if p.requires_grad]
+        logger.debug(f"Trainable params count: {sum(p.numel() for p in self._trainable_params):,}")
+
+    def _init_timestep_sampler(self) -> None:
+        """Initialize the timestep sampler based on the config."""
+        sampler_cls = SAMPLERS[self._config.flow_matching.timestep_sampling_mode]
+        self._timestep_sampler = sampler_cls(**self._config.flow_matching.timestep_sampling_params)
+
+    def _setup_lora(self) -> None:
+        """Configure LoRA adapters for the transformer. Only called in LoRA training mode."""
+        logger.debug(f"Adding LoRA adapter with rank {self._config.lora.rank}")
+        lora_config = LoraConfig(
+            r=self._config.lora.rank,
+            lora_alpha=self._config.lora.alpha,
+            target_modules=self._config.lora.target_modules,
+            lora_dropout=self._config.lora.dropout,
+            init_lora_weights=True,
+        )
+        # Wrap the transformer with PEFT to add LoRA layers
+        # noinspection PyTypeChecker
+        self._transformer = get_peft_model(self._transformer, lora_config)
+
+    def _load_checkpoint(self) -> None:
+        """Load checkpoint if specified in config."""
+        if not self._config.model.load_checkpoint:
+            return
+
+        checkpoint_path = self._find_checkpoint(self._config.model.load_checkpoint)
+        if not checkpoint_path:
+            logger.warning(f"⚠️ Could not find checkpoint at {self._config.model.load_checkpoint}")
+            return
+
+        logger.info(f"📥 Loading checkpoint from {checkpoint_path}")
+
+        if self._config.model.training_mode == "full":
+            self._load_full_checkpoint(checkpoint_path)
+        else:  # LoRA mode
+            self._load_lora_checkpoint(checkpoint_path)
+
+    def _load_full_checkpoint(self, checkpoint_path: Path) -> None:
+        """Load full model checkpoint."""
+        state_dict = load_file(checkpoint_path)
+        self._transformer.load_state_dict(state_dict, strict=True)
+
+        logger.info("✅ Full model checkpoint loaded successfully")
+
+    def _load_lora_checkpoint(self, checkpoint_path: Path) -> None:
+        """Load LoRA checkpoint with DDP/FSDP compatibility."""
+        state_dict = load_file(checkpoint_path)
+
+        # Adjust layer names to match internal format.
+        # (Weights are saved in ComfyUI-compatible format, with "diffusion_model." prefix)
+        state_dict = {k.replace("diffusion_model.", "", 1): v for k, v in state_dict.items()}
+
+        # Load LoRA weights and verify all weights were loaded
+        base_model = self._transformer.get_base_model()
+        set_peft_model_state_dict(base_model, state_dict)
+
+        logger.info("✅ LoRA checkpoint loaded successfully")
+
+    def _prepare_models_for_training(self) -> None:
+        """Prepare models for training with Accelerate."""
+
+        # For FSDP + LoRA: Cast entire model to FP32.
+        # FSDP requires uniform dtype across all parameters in wrapped modules.
+        # In LoRA mode, PEFT creates LoRA params in FP32 while base model is BF16.
+        # We cast the base model to FP32 to match the LoRA params.
+        if self._accelerator.distributed_type == DistributedType.FSDP and self._config.model.training_mode == "lora":
+            logger.debug("FSDP: casting transformer to FP32 for uniform dtype")
+            self._transformer = self._transformer.to(dtype=torch.float32)
+
+        # Enable gradient checkpointing if requested
+        # For PeftModel, we need to access the underlying base model
+        transformer = (
+            self._transformer.get_base_model() if hasattr(self._transformer, "get_base_model") else self._transformer
+        )
+
+        transformer.set_gradient_checkpointing(self._config.optimization.enable_gradient_checkpointing)
+
+        # Keep frozen models on CPU for memory efficiency
+        self._vae_decoder = self._vae_decoder.to("cpu")
+        if self._vae_encoder is not None:
+            self._vae_encoder = self._vae_encoder.to("cpu")
+
+        # Embedding connectors are already on GPU from _load_text_encoder_and_cache_embeddings
+
+        # noinspection PyTypeChecker
+        self._transformer = self._accelerator.prepare(self._transformer)
+
+        # Log GPU memory usage after model preparation
+        vram_usage_gb = torch.cuda.memory_allocated() / 1024**3
+        logger.debug(f"GPU memory usage after models preparation: {vram_usage_gb:.2f} GB")
+
+    @staticmethod
+    def _find_checkpoint(checkpoint_path: str | Path) -> Path | None:
+        """Find the checkpoint file to load, handling both file and directory paths."""
+        checkpoint_path = Path(checkpoint_path)
+
+        if checkpoint_path.is_file():
+            if not checkpoint_path.suffix == ".safetensors":
+                raise ValueError(f"Checkpoint file must have a .safetensors extension: {checkpoint_path}")
+            return checkpoint_path
+
+        if checkpoint_path.is_dir():
+            # Look for checkpoint files in the directory
+            checkpoints = list(checkpoint_path.rglob("*step_*.safetensors"))
+
+            if not checkpoints:
+                return None
+
+            # Sort by step number and return the latest
+            def _get_step_num(p: Path) -> int:
+                try:
+                    return int(p.stem.split("step_")[1])
+                except (IndexError, ValueError):
+                    return -1
+
+            latest = max(checkpoints, key=_get_step_num)
+            return latest
+
+        else:
+            raise ValueError(f"Invalid checkpoint path: {checkpoint_path}. Must be a file or directory.")
+
+    def _init_dataloader(self) -> None:
+        """Initialize the training data loader using the strategy's data sources."""
+        if self._dataset is None:
+            # Get data sources from the training strategy
+            data_sources = self._training_strategy.get_data_sources()
+
+            self._dataset = PrecomputedDataset(self._config.data.preprocessed_data_root, data_sources=data_sources)
+            logger.debug(f"Loaded dataset with {len(self._dataset):,} samples from sources: {list(data_sources)}")
+
+        num_workers = self._config.data.num_dataloader_workers
+        dataloader = DataLoader(
+            self._dataset,
+            batch_size=self._config.optimization.batch_size,
+            shuffle=True,
+            drop_last=True,
+            num_workers=num_workers,
+            pin_memory=num_workers > 0,
+            persistent_workers=num_workers > 0,
+        )
+
+        self._dataloader = self._accelerator.prepare(dataloader)
+
+    def _init_lora_weights(self) -> None:
+        """Initialize LoRA weights for the transformer."""
+        logger.debug("Initializing LoRA weights...")
+        for _, module in self._transformer.named_modules():
+            if isinstance(module, (BaseTunerLayer, ModulesToSaveWrapper)):
+                module.reset_lora_parameters(adapter_name="default", init_lora_weights=True)
+
+    def _init_optimizer(self) -> None:
+        """Initialize the optimizer and learning rate scheduler."""
+        opt_cfg = self._config.optimization
+
+        lr = opt_cfg.learning_rate
+        if opt_cfg.optimizer_type == "adamw":
+            optimizer = AdamW(self._trainable_params, lr=lr)
+        elif opt_cfg.optimizer_type == "adamw8bit":
+            # noinspection PyUnresolvedReferences
+            from bitsandbytes.optim import AdamW8bit  # noqa: PLC0415
+
+            optimizer = AdamW8bit(self._trainable_params, lr=lr)
+        else:
+            raise ValueError(f"Unknown optimizer type: {opt_cfg.optimizer_type}")
+
+        # Add scheduler initialization
+        lr_scheduler = self._create_scheduler(optimizer)
+
+        # noinspection PyTypeChecker
+        self._optimizer, self._lr_scheduler = self._accelerator.prepare(optimizer, lr_scheduler)
+
+    def _create_scheduler(self, optimizer: torch.optim.Optimizer) -> LRScheduler | None:
+        """Create learning rate scheduler based on config."""
+        scheduler_type = self._config.optimization.scheduler_type
+        steps = self._config.optimization.steps
+        params = self._config.optimization.scheduler_params or {}
+
+        if scheduler_type is None:
+            return None
+
+        if scheduler_type == "linear":
+            scheduler = LinearLR(
+                optimizer,
+                start_factor=params.pop("start_factor", 1.0),
+                end_factor=params.pop("end_factor", 0.1),
+                total_iters=steps,
+                **params,
+            )
+        elif scheduler_type == "cosine":
+            scheduler = CosineAnnealingLR(
+                optimizer,
+                T_max=steps,
+                eta_min=params.pop("eta_min", 0),
+                **params,
+            )
+        elif scheduler_type == "cosine_with_restarts":
+            scheduler = CosineAnnealingWarmRestarts(
+                optimizer,
+                T_0=params.pop("T_0", steps // 4),  # First restart cycle length
+                T_mult=params.pop("T_mult", 1),  # Multiplicative factor for cycle lengths
+                eta_min=params.pop("eta_min", 5e-5),
+                **params,
+            )
+        elif scheduler_type == "polynomial":
+            scheduler = PolynomialLR(
+                optimizer,
+                total_iters=steps,
+                power=params.pop("power", 1.0),
+                **params,
+            )
+        elif scheduler_type == "step":
+            scheduler = StepLR(
+                optimizer,
+                step_size=params.pop("step_size", steps // 2),
+                gamma=params.pop("gamma", 0.1),
+                **params,
+            )
+        elif scheduler_type == "constant":
+            scheduler = None
+        else:
+            raise ValueError(f"Unknown scheduler type: {scheduler_type}")
+
+        return scheduler
+
+    def _setup_accelerator(self) -> None:
+        """Initialize the Accelerator with the appropriate settings."""
+
+        # All distributed setup (DDP/FSDP, number of processes, etc.) is controlled by
+        # the user's Accelerate configuration (accelerate config / accelerate launch).
+        self._accelerator = Accelerator(
+            mixed_precision=self._config.acceleration.mixed_precision_mode,
+            gradient_accumulation_steps=self._config.optimization.gradient_accumulation_steps,
+        )
+
+        if self._accelerator.num_processes > 1:
+            logger.info(
+                f"{self._accelerator.distributed_type.value} distributed training enabled "
+                f"with {self._accelerator.num_processes} processes"
+            )
+
+            local_batch = self._config.optimization.batch_size
+            global_batch = self._config.optimization.batch_size * self._accelerator.num_processes
+            logger.info(f"Local batch size: {local_batch}, global batch size: {global_batch}")
+
+        # Log torch.compile status from Accelerate's dynamo plugin
+        is_compile_enabled = (
+            hasattr(self._accelerator.state, "dynamo_plugin") and self._accelerator.state.dynamo_plugin.backend != "NO"
+        )
+        if is_compile_enabled:
+            plugin = self._accelerator.state.dynamo_plugin
+            logger.info(f"🔥 torch.compile enabled via Accelerate: backend={plugin.backend}, mode={plugin.mode}")
+
+            if self._accelerator.distributed_type == DistributedType.FSDP:
+                logger.warning(
+                    "⚠️ FSDP + torch.compile is experimental and may hang on the first training iteration. "
+                    "If this occurs, disable torch.compile by removing dynamo_config from your Accelerate config."
+                )
+
+        if self._accelerator.distributed_type == DistributedType.FSDP and self._config.acceleration.quantization:
+            logger.warning(
+                f"FSDP with quantization ({self._config.acceleration.quantization}) may have compatibility issues."
+                "Monitor training stability and consider disabling quantization if issues arise."
+            )
+
+    # Note: Use @torch.no_grad() instead of @torch.inference_mode() to avoid FSDP inplace update errors after validation
+    @torch.no_grad()
+    def _sample_videos(self, progress: TrainingProgress) -> list[Path] | None:
+        """Run validation by generating videos from validation prompts."""
+        use_images = self._config.validation.images is not None
+        use_reference_videos = self._config.validation.reference_videos is not None
+        generate_audio = self._config.validation.generate_audio
+        inference_steps = self._config.validation.inference_steps
+
+        # Free up GPU memory before validation sampling.
+        # Zero gradients and empty the cache to reclaim memory.
+        self._optimizer.zero_grad(set_to_none=True)
+        torch.cuda.empty_cache()
+
+        # Start sampling progress tracking
+        sampling_ctx = progress.start_sampling(
+            num_prompts=len(self._config.validation.prompts),
+            num_steps=inference_steps,
+        )
+
+        # Create validation sampler with loaded models and progress tracking
+        sampler = ValidationSampler(
+            transformer=self._transformer,
+            vae_decoder=self._vae_decoder,
+            vae_encoder=self._vae_encoder,
+            text_encoder=None,
+            audio_decoder=self._audio_vae if generate_audio else None,
+            vocoder=self._vocoder if generate_audio else None,
+            sampling_context=sampling_ctx,
+        )
+
+        output_dir = Path(self._config.output_dir) / "samples"
+        output_dir.mkdir(exist_ok=True, parents=True)
+
+        video_paths = []
+        width, height, num_frames = self._config.validation.video_dims
+
+        for prompt_idx, prompt in enumerate(self._config.validation.prompts):
+            # Update progress to show current video
+            sampling_ctx.start_video(prompt_idx)
+
+            # Load conditioning image if provided
+            condition_image = None
+            if use_images:
+                image_path = self._config.validation.images[prompt_idx]
+                image = open_image_as_srgb(image_path)
+                # Convert PIL image to tensor [C, H, W] in [0, 1]
+                condition_image = F.to_tensor(image)
+
+            # Load reference video if provided (for IC-LoRA)
+            reference_video = None
+            if use_reference_videos:
+                ref_video_path = self._config.validation.reference_videos[prompt_idx]
+                # read_video returns [F, C, H, W] in [0, 1]
+                reference_video, _ = read_video(ref_video_path, max_frames=num_frames)
+
+            # Get cached embeddings for this prompt if available
+            cached_embeddings = (
+                self._cached_validation_embeddings[prompt_idx]
+                if self._cached_validation_embeddings is not None
+                else None
+            )
+
+            # Create generation config
+            gen_config = GenerationConfig(
+                prompt=prompt,
+                negative_prompt=self._config.validation.negative_prompt,
+                height=height,
+                width=width,
+                num_frames=num_frames,
+                frame_rate=self._config.validation.frame_rate,
+                num_inference_steps=inference_steps,
+                guidance_scale=self._config.validation.guidance_scale,
+                seed=self._config.validation.seed,
+                condition_image=condition_image,
+                reference_video=reference_video,
+                generate_audio=generate_audio,
+                include_reference_in_output=self._config.validation.include_reference_in_output,
+                cached_embeddings=cached_embeddings,
+                stg_scale=self._config.validation.stg_scale,
+                stg_blocks=self._config.validation.stg_blocks,
+                stg_mode=self._config.validation.stg_mode,
+            )
+
+            # Generate sample
+            video, audio = sampler.generate(
+                config=gen_config,
+                device=self._accelerator.device,
+            )
+
+            # Save output (image for single frame, video otherwise)
+            if IS_MAIN_PROCESS:
+                ext = "png" if num_frames == 1 else "mp4"
+                output_path = output_dir / f"step_{self._global_step:06d}_{prompt_idx + 1}.{ext}"
+                if num_frames == 1:
+                    save_image(video, output_path)
+                else:
+                    save_video(
+                        video_tensor=video,
+                        output_path=output_path,
+                        fps=self._config.validation.frame_rate,
+                        audio=audio,
+                        audio_sample_rate=self._vocoder.output_sample_rate if audio is not None else None,
+                    )
+                video_paths.append(output_path)
+
+        # Clean up progress tasks
+        sampling_ctx.cleanup()
+
+        # Clear GPU cache after validation
+        torch.cuda.empty_cache()
+
+        rel_outputs_path = output_dir.relative_to(self._config.output_dir)
+        logger.info(f"🎥 Validation samples for step {self._global_step} saved in {rel_outputs_path}")
+        return video_paths
+
+    @staticmethod
+    def _log_training_stats(stats: TrainingStats) -> None:
+        """Log training statistics."""
+        stats_str = (
+            "📊 Training Statistics:\n"
+            f" - Total time: {stats.total_time_seconds / 60:.1f} minutes\n"
+            f" - Training speed: {stats.steps_per_second:.2f} steps/second\n"
+            f" - Samples/second: {stats.samples_per_second:.2f}\n"
+            f" - Peak GPU memory: {stats.peak_gpu_memory_gb:.2f} GB"
+        )
+        if stats.num_processes > 1:
+            stats_str += f"\n - Number of processes: {stats.num_processes}\n"
+            stats_str += f" - Global batch size: {stats.global_batch_size}"
+        logger.info(stats_str)
+
+    def _save_checkpoint(self) -> Path | None:
+        """Save the model weights."""
+        is_lora = self._config.model.training_mode == "lora"
+        is_fsdp = self._accelerator.distributed_type == DistributedType.FSDP
+
+        # Prepare paths
+        save_dir = Path(self._config.output_dir) / "checkpoints"
+        prefix = "lora" if is_lora else "model"
+        filename = f"{prefix}_weights_step_{self._global_step:05d}.safetensors"
+        saved_weights_path = save_dir / filename
+
+        # Get state dict (collective operation - all processes must participate)
+        self._accelerator.wait_for_everyone()
+        full_state_dict = self._accelerator.get_state_dict(self._transformer)
+
+        if not IS_MAIN_PROCESS:
+            return None
+
+        save_dir.mkdir(exist_ok=True, parents=True)
+
+        # For LoRA: extract only adapter weights; for full: use as-is
+        if is_lora:
+            unwrapped = self._accelerator.unwrap_model(self._transformer, keep_torch_compile=False)
+            # For FSDP, pass full_state_dict since model params aren't directly accessible
+            state_dict = get_peft_model_state_dict(unwrapped, state_dict=full_state_dict if is_fsdp else None)
+
+            # Remove "base_model.model." prefix added by PEFT
+            state_dict = {k.replace("base_model.model.", "", 1): v for k, v in state_dict.items()}
+
+            # Convert to ComfyUI-compatible format (add "diffusion_model." prefix)
+            state_dict = {f"diffusion_model.{k}": v for k, v in state_dict.items()}
+
+            # Save to disk
+            save_file(state_dict, saved_weights_path)
+        else:
+            # Save to disk
+            self._accelerator.save(full_state_dict, saved_weights_path)
+
+        rel_path = saved_weights_path.relative_to(self._config.output_dir)
+        logger.info(f"💾 {prefix.capitalize()} weights for step {self._global_step} saved in {rel_path}")
+
+        # Keep track of checkpoint paths, and cleanup old checkpoints if needed
+        self._checkpoint_paths.append(saved_weights_path)
+        self._cleanup_checkpoints()
+        return saved_weights_path
+
+    def _cleanup_checkpoints(self) -> None:
+        """Clean up old checkpoints."""
+        if 0 < self._config.checkpoints.keep_last_n < len(self._checkpoint_paths):
+            checkpoints_to_remove = self._checkpoint_paths[: -self._config.checkpoints.keep_last_n]
+            for old_checkpoint in checkpoints_to_remove:
+                if old_checkpoint.exists():
+                    old_checkpoint.unlink()
+                    logger.info(f"Removed old checkpoints: {old_checkpoint}")
+            # Update the list to only contain kept checkpoints
+            self._checkpoint_paths = self._checkpoint_paths[-self._config.checkpoints.keep_last_n :]
+
+    def _save_config(self) -> None:
+        """Save the training configuration as a YAML file in the output directory."""
+        if not IS_MAIN_PROCESS:
+            return
+
+        config_path = Path(self._config.output_dir) / "training_config.yaml"
+        with open(config_path, "w") as f:
+            yaml.dump(self._config.model_dump(), f, default_flow_style=False, indent=2)
+
+        logger.info(f"💾 Training configuration saved to: {config_path.relative_to(self._config.output_dir)}")
+
+    def _init_wandb(self) -> None:
+        """Initialize Weights & Biases run."""
+        if not self._config.wandb.enabled or not IS_MAIN_PROCESS:
+            self._wandb_run = None
+            return
+
+        wandb_config = self._config.wandb
+        run = wandb.init(
+            project=wandb_config.project,
+            entity=wandb_config.entity,
+            name=Path(self._config.output_dir).name,
+            tags=wandb_config.tags,
+            config=self._config.model_dump(),
+        )
+        self._wandb_run = run
+
+    def _log_metrics(self, metrics: dict[str, float]) -> None:
+        """Log metrics to Weights & Biases."""
+        if self._wandb_run is not None:
+            self._wandb_run.log(metrics)
+
+    def _log_validation_samples(self, sample_paths: list[Path], prompts: list[str]) -> None:
+        """Log validation samples (videos or images) to Weights & Biases."""
+        if not self._config.wandb.log_validation_videos or self._wandb_run is None:
+            return
+
+        # Determine if outputs are images or videos based on file extension
+        is_image = sample_paths and sample_paths[0].suffix.lower() in (".png", ".jpg", ".jpeg", ".heic", ".webp")
+        media_cls = wandb.Image if is_image else wandb.Video
+
+        samples = [media_cls(str(path), caption=prompt) for path, prompt in zip(sample_paths, prompts, strict=True)]
+        self._wandb_run.log({"validation_samples": samples}, step=self._global_step)

packages/ltx-trainer/src/ltx_trainer/training_strategies/__init__.py

@@ -0,0 +1,58 @@
+"""Training strategies for different conditioning modes.
+This package implements the Strategy Pattern to handle different training modes:
+- Text-to-video training (standard generation, optionally with audio)
+- Video-to-video training (IC-LoRA mode with reference videos)
+Each strategy encapsulates the specific logic for preparing model inputs and computing loss.
+"""
+
+from ltx_trainer import logger
+from ltx_trainer.training_strategies.base_strategy import (
+    DEFAULT_FPS,
+    VIDEO_SCALE_FACTORS,
+    ModelInputs,
+    TrainingStrategy,
+    TrainingStrategyConfigBase,
+)
+from ltx_trainer.training_strategies.text_to_video import TextToVideoConfig, TextToVideoStrategy
+from ltx_trainer.training_strategies.video_to_video import VideoToVideoConfig, VideoToVideoStrategy
+
+# Type alias for all strategy config types
+TrainingStrategyConfig = TextToVideoConfig | VideoToVideoConfig
+
+__all__ = [
+    "DEFAULT_FPS",
+    "VIDEO_SCALE_FACTORS",
+    "ModelInputs",
+    "TextToVideoConfig",
+    "TextToVideoStrategy",
+    "TrainingStrategy",
+    "TrainingStrategyConfig",
+    "TrainingStrategyConfigBase",
+    "VideoToVideoConfig",
+    "VideoToVideoStrategy",
+    "get_training_strategy",
+]
+
+
+def get_training_strategy(config: TrainingStrategyConfig) -> TrainingStrategy:
+    """Factory function to create the appropriate training strategy.
+    The strategy is determined by the `name` field in the configuration.
+    Args:
+        config: Strategy-specific configuration with a `name` field
+    Returns:
+        The appropriate training strategy instance
+    Raises:
+        ValueError: If strategy name is not supported
+    """
+
+    match config:
+        case TextToVideoConfig():
+            strategy = TextToVideoStrategy(config)
+        case VideoToVideoConfig():
+            strategy = VideoToVideoStrategy(config)
+        case _:
+            raise ValueError(f"Unknown training strategy config type: {type(config).__name__}")
+
+    audio_mode = "(audio enabled 🔈)" if getattr(config, "with_audio", False) else "(audio disabled 🔇)"
+    logger.debug(f"🎯 Using {strategy.__class__.__name__} training strategy {audio_mode}")
+    return strategy

packages/ltx-trainer/src/ltx_trainer/training_strategies/base_strategy.py

@@ -0,0 +1,253 @@
+"""Base class for training strategies.
+This module defines the abstract base class that all training strategies must implement,
+along with the base configuration class.
+"""
+
+import random
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any, Literal
+
+import torch
+from pydantic import BaseModel, ConfigDict, Field
+from torch import Tensor
+
+from ltx_core.components.patchifiers import (
+    AudioPatchifier,
+    VideoLatentPatchifier,
+    get_pixel_coords,
+)
+from ltx_core.model.transformer.modality import Modality
+from ltx_core.types import AudioLatentShape, SpatioTemporalScaleFactors, VideoLatentShape
+from ltx_trainer.timestep_samplers import TimestepSampler
+
+# Default frames per second for video missing in the FPS metadata
+DEFAULT_FPS = 24
+
+# VAE scale factors for LTX-2
+VIDEO_SCALE_FACTORS = SpatioTemporalScaleFactors.default()
+
+
+class TrainingStrategyConfigBase(BaseModel):
+    """Base configuration class for training strategies.
+    All strategy-specific configuration classes should inherit from this.
+    """
+
+    model_config = ConfigDict(extra="forbid")
+
+    name: Literal["text_to_video", "video_to_video"] = Field(
+        description="Unique name identifying the training strategy type"
+    )
+
+
+@dataclass
+class ModelInputs:
+    """Container for model inputs using the Modality-based interface."""
+
+    video: Modality
+    audio: Modality | None
+
+    # Training targets (for loss computation)
+    video_targets: Tensor
+    audio_targets: Tensor | None
+
+    # Masks for loss computation
+    video_loss_mask: Tensor  # Boolean mask: True = compute loss for this token
+    audio_loss_mask: Tensor | None
+
+    # Metadata needed for loss computation in some strategies
+    ref_seq_len: int | None = None  # For IC-LoRA: length of reference sequence
+
+
+class TrainingStrategy(ABC):
+    """Abstract base class for training strategies.
+    Each strategy encapsulates the logic for a specific training mode,
+    handling input preparation and loss computation.
+    """
+
+    def __init__(self, config: TrainingStrategyConfigBase):
+        """Initialize strategy with configuration.
+        Args:
+            config: Strategy-specific configuration
+        """
+        self.config = config
+        self._video_patchifier = VideoLatentPatchifier(patch_size=1)
+        self._audio_patchifier = AudioPatchifier(patch_size=1)
+
+    @property
+    def requires_audio(self) -> bool:
+        """Whether this training strategy requires audio components.
+        Override this property in subclasses that support audio training.
+        The trainer uses this to determine whether to load audio VAE and vocoder.
+        Returns:
+            True if audio components should be loaded, False otherwise.
+        """
+        return False
+
+    @abstractmethod
+    def get_data_sources(self) -> list[str] | dict[str, str]:
+        """Get the required data sources for this training strategy.
+        Returns:
+            Either a list of data directory names (where output keys match directory names)
+            or a dictionary mapping data directory names to custom output keys for the dataset
+        """
+
+    @abstractmethod
+    def prepare_training_inputs(
+        self,
+        batch: dict[str, Any],
+        timestep_sampler: TimestepSampler,
+    ) -> ModelInputs:
+        """Prepare training inputs from a raw data batch.
+        Args:
+            batch: Raw batch data from the dataset. Contains:
+                - "latents": Video latent data
+                - "conditions": Text embeddings with keys:
+                    - "video_prompt_embeds": Already processed by embedding connectors
+                    - "audio_prompt_embeds": Already processed by embedding connectors
+                    - "prompt_attention_mask": Attention mask
+                - Additional keys depending on strategy (e.g., "ref_latents" for IC-LoRA)
+            timestep_sampler: Sampler for generating timesteps and noise
+        Returns:
+            ModelInputs containing Modality objects and training targets
+        """
+
+    @abstractmethod
+    def compute_loss(
+        self,
+        video_pred: Tensor,
+        audio_pred: Tensor | None,
+        inputs: ModelInputs,
+    ) -> Tensor:
+        """Compute the training loss.
+        Args:
+            video_pred: Video prediction from the transformer model
+            audio_pred: Audio prediction from the transformer model (None for video-only)
+            inputs: The prepared model inputs containing targets and masks
+        Returns:
+            Scalar loss tensor
+        """
+
+    def _get_video_positions(
+        self,
+        num_frames: int,
+        height: int,
+        width: int,
+        batch_size: int,
+        fps: float,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> Tensor:
+        """Generate video position embeddings using ltx_core's native implementation.
+        Args:
+            num_frames: Number of latent frames
+            height: Latent height
+            width: Latent width
+            batch_size: Batch size
+            fps: Frames per second
+            device: Target device
+            dtype: Target dtype
+        Returns:
+            Position tensor of shape [B, 3, seq_len, 2]
+        """
+        latent_coords = self._video_patchifier.get_patch_grid_bounds(
+            output_shape=VideoLatentShape(
+                frames=num_frames,
+                height=height,
+                width=width,
+                batch=batch_size,
+                channels=128,  # Video latent channels
+            ),
+            device=device,
+        )
+
+        # Convert latent coords to pixel coords with causal fix
+        pixel_coords = get_pixel_coords(
+            latent_coords=latent_coords,
+            scale_factors=VIDEO_SCALE_FACTORS,
+            causal_fix=True,
+        ).to(dtype)
+
+        # Scale temporal dimension by 1/fps to get time in seconds
+        pixel_coords[:, 0, ...] = pixel_coords[:, 0, ...] / fps
+
+        return pixel_coords
+
+    def _get_audio_positions(
+        self,
+        num_time_steps: int,
+        batch_size: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> Tensor:
+        """Generate audio position embeddings using ltx_core's native implementation.
+        Args:
+            num_time_steps: Number of audio time steps (T, not T*mel_bins)
+            batch_size: Batch size
+            device: Target device
+            dtype: Target dtype
+        Returns:
+            Position tensor of shape [B, 1, num_time_steps, 2]
+        Note:
+            Audio latents should be in patchified format [B, T, C*F] = [B, T, 128]
+            where T is the number of time steps, C=8 channels, F=16 mel bins.
+            This matches the format produced by AudioPatchifier.patchify().
+        """
+        mel_bins = 16
+
+        latent_coords = self._audio_patchifier.get_patch_grid_bounds(
+            output_shape=AudioLatentShape(
+                frames=num_time_steps,
+                mel_bins=mel_bins,
+                batch=batch_size,
+                channels=8,  # Audio latent channels
+            ),
+            device=device,
+        )
+
+        return latent_coords.to(dtype)
+
+    @staticmethod
+    def _create_per_token_timesteps(conditioning_mask: Tensor, sampled_sigma: Tensor) -> Tensor:
+        """Create per-token timesteps based on conditioning mask.
+        Args:
+            conditioning_mask: Boolean mask of shape (batch_size, sequence_length),
+                where True = conditioning token (timestep=0), False = target token (use sigma)
+            sampled_sigma: Sampled sigma values of shape (batch_size,) or (batch_size, 1, 1)
+        Returns:
+            Timesteps tensor of shape [batch_size, sequence_length]
+        """
+        # Expand to match conditioning mask shape [B, seq_len]
+        expanded_sigma = sampled_sigma.view(-1, 1).expand_as(conditioning_mask)
+
+        # Conditioning tokens get 0, target tokens get the sampled sigma
+        return torch.where(conditioning_mask, torch.zeros_like(expanded_sigma), expanded_sigma)
+
+    @staticmethod
+    def _create_first_frame_conditioning_mask(
+        batch_size: int,
+        sequence_length: int,
+        height: int,
+        width: int,
+        device: torch.device,
+        first_frame_conditioning_p: float = 0.0,
+    ) -> Tensor:
+        """Create conditioning mask for first frame conditioning.
+        Args:
+            batch_size: Batch size
+            sequence_length: Total sequence length
+            height: Latent height
+            width: Latent width
+            device: Target device
+            first_frame_conditioning_p: Probability of conditioning on the first frame
+        Returns:
+            Boolean mask where True indicates first frame tokens (if conditioning is enabled)
+        """
+        conditioning_mask = torch.zeros(batch_size, sequence_length, dtype=torch.bool, device=device)
+
+        if first_frame_conditioning_p > 0 and random.random() < first_frame_conditioning_p:
+            first_frame_end_idx = height * width
+            if first_frame_end_idx < sequence_length:
+                conditioning_mask[:, :first_frame_end_idx] = True
+
+        return conditioning_mask

packages/ltx-trainer/src/ltx_trainer/training_strategies/text_to_video.py

@@ -0,0 +1,289 @@
+"""Text-to-video training strategy.
+This strategy implements standard text-to-video generation training where:
+- Only target latents are used (no reference videos)
+- Standard noise application and loss computation
+- Supports first frame conditioning
+- Optionally supports joint audio-video training
+"""
+
+from typing import Any, Literal
+
+import torch
+from pydantic import Field
+from torch import Tensor
+
+from ltx_core.model.transformer.modality import Modality
+from ltx_trainer import logger
+from ltx_trainer.timestep_samplers import TimestepSampler
+from ltx_trainer.training_strategies.base_strategy import (
+    DEFAULT_FPS,
+    ModelInputs,
+    TrainingStrategy,
+    TrainingStrategyConfigBase,
+)
+
+
+class TextToVideoConfig(TrainingStrategyConfigBase):
+    """Configuration for text-to-video training strategy."""
+
+    name: Literal["text_to_video"] = "text_to_video"
+
+    first_frame_conditioning_p: float = Field(
+        default=0.1,
+        description="Probability of conditioning on the first frame during training",
+        ge=0.0,
+        le=1.0,
+    )
+
+    with_audio: bool = Field(
+        default=False,
+        description="Whether to include audio in training (joint audio-video generation)",
+    )
+
+    audio_latents_dir: str = Field(
+        default="audio_latents",
+        description="Directory name for audio latents when with_audio is True",
+    )
+
+
+class TextToVideoStrategy(TrainingStrategy):
+    """Text-to-video training strategy.
+    This strategy implements regular video generation training where:
+    - Only target latents are used (no reference videos)
+    - Standard noise application and loss computation
+    - Supports first frame conditioning
+    - Optionally supports joint audio-video training when with_audio=True
+    """
+
+    config: TextToVideoConfig
+
+    def __init__(self, config: TextToVideoConfig):
+        """Initialize strategy with configuration.
+        Args:
+            config: Text-to-video configuration
+        """
+        super().__init__(config)
+
+    @property
+    def requires_audio(self) -> bool:
+        """Whether this training strategy requires audio components."""
+        return self.config.with_audio
+
+    def get_data_sources(self) -> list[str] | dict[str, str]:
+        """
+        Text-to-video training requires latents and text conditions.
+        When with_audio is True, also requires audio latents.
+        """
+        sources = {
+            "latents": "latents",
+            "conditions": "conditions",
+        }
+
+        if self.config.with_audio:
+            sources[self.config.audio_latents_dir] = "audio_latents"
+
+        return sources
+
+    def prepare_training_inputs(
+        self,
+        batch: dict[str, Any],
+        timestep_sampler: TimestepSampler,
+    ) -> ModelInputs:
+        """Prepare inputs for text-to-video training."""
+        # Get pre-encoded latents - dataset provides uniform non-patchified format [B, C, F, H, W]
+        latents = batch["latents"]
+        video_latents = latents["latents"]
+
+        # Get video dimensions (assume same for all batch elements)
+        num_frames = latents["num_frames"][0].item()
+        height = latents["height"][0].item()
+        width = latents["width"][0].item()
+
+        # Patchify latents: [B, C, F, H, W] -> [B, seq_len, C]
+        video_latents = self._video_patchifier.patchify(video_latents)
+
+        # Handle FPS with backward compatibility
+        fps = latents.get("fps", None)
+        if fps is not None and not torch.all(fps == fps[0]):
+            logger.warning(
+                f"Different FPS values found in the batch. Found: {fps.tolist()}, using the first one: {fps[0].item()}"
+            )
+        fps = fps[0].item() if fps is not None else DEFAULT_FPS
+
+        # Get text embeddings (already processed by embedding connectors in trainer)
+        conditions = batch["conditions"]
+        video_prompt_embeds = conditions["video_prompt_embeds"]
+        audio_prompt_embeds = conditions["audio_prompt_embeds"]
+        prompt_attention_mask = conditions["prompt_attention_mask"]
+
+        batch_size = video_latents.shape[0]
+        video_seq_len = video_latents.shape[1]
+        device = video_latents.device
+        dtype = video_latents.dtype
+
+        # Create conditioning mask (first frame conditioning)
+        video_conditioning_mask = self._create_first_frame_conditioning_mask(
+            batch_size=batch_size,
+            sequence_length=video_seq_len,
+            height=height,
+            width=width,
+            device=device,
+            first_frame_conditioning_p=self.config.first_frame_conditioning_p,
+        )
+
+        # Sample noise and sigmas
+        sigmas = timestep_sampler.sample_for(video_latents)
+        video_noise = torch.randn_like(video_latents)
+
+        # Apply noise: noisy = (1 - sigma) * clean + sigma * noise
+        sigmas_expanded = sigmas.view(-1, 1, 1)
+        noisy_video = (1 - sigmas_expanded) * video_latents + sigmas_expanded * video_noise
+
+        # For conditioning tokens, use clean latents
+        conditioning_mask_expanded = video_conditioning_mask.unsqueeze(-1)
+        noisy_video = torch.where(conditioning_mask_expanded, video_latents, noisy_video)
+
+        # Compute video targets (velocity prediction)
+        video_targets = video_noise - video_latents
+
+        # Create per-token timesteps
+        video_timesteps = self._create_per_token_timesteps(video_conditioning_mask, sigmas.squeeze())
+
+        # Generate video positions using ltx_core's native implementation
+        video_positions = self._get_video_positions(
+            num_frames=num_frames,
+            height=height,
+            width=width,
+            batch_size=batch_size,
+            fps=fps,
+            device=device,
+            dtype=dtype,
+        )
+
+        # Create video Modality
+        video_modality = Modality(
+            enabled=True,
+            latent=noisy_video,
+            timesteps=video_timesteps,
+            positions=video_positions,
+            context=video_prompt_embeds,
+            context_mask=prompt_attention_mask,
+        )
+
+        # Video loss mask: True for tokens we want to compute loss on (non-conditioning tokens)
+        video_loss_mask = ~video_conditioning_mask
+
+        # Handle audio if enabled
+        audio_modality = None
+        audio_targets = None
+        audio_loss_mask = None
+
+        if self.config.with_audio:
+            audio_modality, audio_targets, audio_loss_mask = self._prepare_audio_inputs(
+                batch=batch,
+                sigmas=sigmas,
+                audio_prompt_embeds=audio_prompt_embeds,
+                prompt_attention_mask=prompt_attention_mask,
+                batch_size=batch_size,
+                device=device,
+                dtype=dtype,
+            )
+
+        return ModelInputs(
+            video=video_modality,
+            audio=audio_modality,
+            video_targets=video_targets,
+            audio_targets=audio_targets,
+            video_loss_mask=video_loss_mask,
+            audio_loss_mask=audio_loss_mask,
+        )
+
+    def _prepare_audio_inputs(
+        self,
+        batch: dict[str, Any],
+        sigmas: Tensor,
+        audio_prompt_embeds: Tensor,
+        prompt_attention_mask: Tensor,
+        batch_size: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> tuple[Modality, Tensor, Tensor]:
+        """Prepare audio inputs for joint audio-video training.
+        Args:
+            batch: Raw batch data containing audio_latents
+            sigmas: Sampled sigma values (same as video)
+            audio_prompt_embeds: Audio context embeddings
+            prompt_attention_mask: Attention mask for context
+            batch_size: Batch size
+            device: Target device
+            dtype: Target dtype
+        Returns:
+            Tuple of (audio_modality, audio_targets, audio_loss_mask)
+        """
+        # Get audio latents - dataset provides uniform non-patchified format [B, C, T, F]
+        audio_data = batch["audio_latents"]
+        audio_latents = audio_data["latents"]
+
+        # Patchify audio latents: [B, C, T, F] -> [B, T, C*F]
+        audio_latents = self._audio_patchifier.patchify(audio_latents)
+
+        audio_seq_len = audio_latents.shape[1]
+
+        # Sample audio noise
+        audio_noise = torch.randn_like(audio_latents)
+
+        # Apply noise to audio (same sigma as video)
+        sigmas_expanded = sigmas.view(-1, 1, 1)
+        noisy_audio = (1 - sigmas_expanded) * audio_latents + sigmas_expanded * audio_noise
+
+        # Compute audio targets
+        audio_targets = audio_noise - audio_latents
+
+        # Audio timesteps: all tokens use the sampled sigma (no conditioning mask)
+        audio_timesteps = sigmas.view(-1, 1).expand(-1, audio_seq_len)
+
+        # Generate audio positions
+        audio_positions = self._get_audio_positions(
+            num_time_steps=audio_seq_len,
+            batch_size=batch_size,
+            device=device,
+            dtype=dtype,
+        )
+
+        # Create audio Modality
+        audio_modality = Modality(
+            enabled=True,
+            latent=noisy_audio,
+            timesteps=audio_timesteps,
+            positions=audio_positions,
+            context=audio_prompt_embeds,
+            context_mask=prompt_attention_mask,
+        )
+
+        # Audio loss mask: all tokens contribute to loss (no conditioning)
+        audio_loss_mask = torch.ones(batch_size, audio_seq_len, dtype=torch.bool, device=device)
+
+        return audio_modality, audio_targets, audio_loss_mask
+
+    def compute_loss(
+        self,
+        video_pred: Tensor,
+        audio_pred: Tensor | None,
+        inputs: ModelInputs,
+    ) -> Tensor:
+        """Compute masked MSE loss for video and optionally audio."""
+        # Video loss
+        video_loss = (video_pred - inputs.video_targets).pow(2)
+        video_loss_mask = inputs.video_loss_mask.unsqueeze(-1).float()
+        video_loss = video_loss.mul(video_loss_mask).div(video_loss_mask.mean())
+        video_loss = video_loss.mean()
+
+        # If no audio, return video loss only
+        if not self.config.with_audio or audio_pred is None or inputs.audio_targets is None:
+            return video_loss
+
+        # Audio loss (no conditioning mask)
+        audio_loss = (audio_pred - inputs.audio_targets).pow(2).mean()
+
+        # Combined loss
+        return video_loss + audio_loss

packages/ltx-trainer/src/ltx_trainer/training_strategies/video_to_video.py

@@ -0,0 +1,223 @@
+"""Video-to-video training strategy for IC-LoRA.
+This strategy implements training with reference video conditioning where:
+- Reference latents (clean) are concatenated with target latents (noised)
+- Video coordinates handle both reference and target sequences
+- Loss is computed only on the target portion
+"""
+
+from typing import Any, Literal
+
+import torch
+from pydantic import Field
+from torch import Tensor
+
+from ltx_core.model.transformer.modality import Modality
+from ltx_trainer import logger
+from ltx_trainer.timestep_samplers import TimestepSampler
+from ltx_trainer.training_strategies.base_strategy import (
+    DEFAULT_FPS,
+    ModelInputs,
+    TrainingStrategy,
+    TrainingStrategyConfigBase,
+)
+
+
+class VideoToVideoConfig(TrainingStrategyConfigBase):
+    """Configuration for video-to-video (IC-LoRA) training strategy."""
+
+    name: Literal["video_to_video"] = "video_to_video"
+
+    first_frame_conditioning_p: float = Field(
+        default=0.1,
+        description="Probability of conditioning on the first frame during training",
+        ge=0.0,
+        le=1.0,
+    )
+
+    reference_latents_dir: str = Field(
+        default="reference_latents",
+        description="Directory name for latents of reference videos",
+    )
+
+
+class VideoToVideoStrategy(TrainingStrategy):
+    """Video-to-video training strategy for IC-LoRA.
+    This strategy implements training with reference video conditioning where:
+    - Reference latents (clean) are concatenated with target latents (noised)
+    - Video coordinates handle both reference and target sequences
+    - Loss is computed only on the target portion
+    """
+
+    config: VideoToVideoConfig
+
+    def __init__(self, config: VideoToVideoConfig):
+        """Initialize strategy with configuration.
+        Args:
+            config: Video-to-video configuration
+        """
+        super().__init__(config)
+
+    def get_data_sources(self) -> dict[str, str]:
+        """IC-LoRA training requires latents, conditions, and reference latents."""
+        return {
+            "latents": "latents",
+            "conditions": "conditions",
+            self.config.reference_latents_dir: "ref_latents",
+        }
+
+    def prepare_training_inputs(
+        self,
+        batch: dict[str, Any],
+        timestep_sampler: TimestepSampler,
+    ) -> ModelInputs:
+        """Prepare inputs for IC-LoRA training with reference videos."""
+        # Get pre-encoded latents - dataset provides uniform non-patchified format [B, C, F, H, W]
+        latents = batch["latents"]
+        target_latents = latents["latents"]
+        ref_latents = batch["ref_latents"]["latents"]
+
+        # Get dimensions
+        num_frames = latents["num_frames"][0].item()
+        height = latents["height"][0].item()
+        width = latents["width"][0].item()
+
+        ref_latents_info = batch["ref_latents"]
+        ref_frames = ref_latents_info["num_frames"][0].item()
+        ref_height = ref_latents_info["height"][0].item()
+        ref_width = ref_latents_info["width"][0].item()
+
+        # Patchify latents: [B, C, F, H, W] -> [B, seq_len, C]
+        target_latents = self._video_patchifier.patchify(target_latents)
+        ref_latents = self._video_patchifier.patchify(ref_latents)
+
+        # Handle FPS
+        fps = latents.get("fps", None)
+        if fps is not None and not torch.all(fps == fps[0]):
+            logger.warning(
+                f"Different FPS values found in the batch. Found: {fps.tolist()}, using the first one: {fps[0].item()}"
+            )
+        fps = fps[0].item() if fps is not None else DEFAULT_FPS
+
+        # Get text embeddings (already processed by embedding connectors in trainer)
+        # Video-to-video uses only video embeddings
+        conditions = batch["conditions"]
+        prompt_embeds = conditions["video_prompt_embeds"]
+        prompt_attention_mask = conditions["prompt_attention_mask"]
+
+        batch_size = target_latents.shape[0]
+        ref_seq_len = ref_latents.shape[1]
+        target_seq_len = target_latents.shape[1]
+        device = target_latents.device
+        dtype = target_latents.dtype
+
+        # Create conditioning mask
+        # Reference tokens are always conditioning (timestep=0)
+        ref_conditioning_mask = torch.ones(batch_size, ref_seq_len, dtype=torch.bool, device=device)
+
+        # Target tokens: check for first frame conditioning
+        target_conditioning_mask = self._create_first_frame_conditioning_mask(
+            batch_size=batch_size,
+            sequence_length=target_seq_len,
+            height=height,
+            width=width,
+            device=device,
+            first_frame_conditioning_p=self.config.first_frame_conditioning_p,
+        )
+
+        # Combined conditioning mask
+        conditioning_mask = torch.cat([ref_conditioning_mask, target_conditioning_mask], dim=1)
+
+        # Sample noise and sigmas for target
+        sigmas = timestep_sampler.sample_for(target_latents)
+        noise = torch.randn_like(target_latents)
+        sigmas_expanded = sigmas.view(-1, 1, 1)
+
+        # Apply noise to target
+        noisy_target = (1 - sigmas_expanded) * target_latents + sigmas_expanded * noise
+
+        # For first frame conditioning in target, use clean latents
+        target_conditioning_mask_expanded = target_conditioning_mask.unsqueeze(-1)
+        noisy_target = torch.where(target_conditioning_mask_expanded, target_latents, noisy_target)
+
+        # Targets for loss computation
+        targets = noise - target_latents
+
+        # Concatenate reference (clean) and target (noisy)
+        combined_latents = torch.cat([ref_latents, noisy_target], dim=1)
+
+        # Create per-token timesteps
+        timesteps = self._create_per_token_timesteps(conditioning_mask, sigmas.squeeze())
+
+        # Generate positions for reference and target separately, then concatenate
+        ref_positions = self._get_video_positions(
+            num_frames=ref_frames,
+            height=ref_height,
+            width=ref_width,
+            batch_size=batch_size,
+            fps=fps,
+            device=device,
+            dtype=dtype,
+        )
+
+        target_positions = self._get_video_positions(
+            num_frames=num_frames,
+            height=height,
+            width=width,
+            batch_size=batch_size,
+            fps=fps,
+            device=device,
+            dtype=dtype,
+        )
+
+        # Concatenate positions along sequence dimension
+        positions = torch.cat([ref_positions, target_positions], dim=2)
+
+        # Create video Modality
+        video_modality = Modality(
+            enabled=True,
+            latent=combined_latents,
+            timesteps=timesteps,
+            positions=positions,
+            context=prompt_embeds,
+            context_mask=prompt_attention_mask,
+        )
+
+        # Loss mask: only compute loss on non-conditioning target tokens
+        # Reference tokens: all False (no loss)
+        # Target tokens: True where not conditioning
+        ref_loss_mask = torch.zeros(batch_size, ref_seq_len, dtype=torch.bool, device=device)
+        target_loss_mask = ~target_conditioning_mask
+        video_loss_mask = torch.cat([ref_loss_mask, target_loss_mask], dim=1)
+
+        return ModelInputs(
+            video=video_modality,
+            audio=None,
+            video_targets=targets,
+            audio_targets=None,
+            video_loss_mask=video_loss_mask,
+            audio_loss_mask=None,
+            ref_seq_len=ref_seq_len,
+        )
+
+    def compute_loss(
+        self,
+        video_pred: Tensor,
+        _audio_pred: Tensor | None,
+        inputs: ModelInputs,
+    ) -> Tensor:
+        """Compute masked loss only on target portion."""
+        # Extract target portion of prediction
+        ref_seq_len = inputs.ref_seq_len
+        target_pred = video_pred[:, ref_seq_len:, :]
+
+        # Get target portion of loss mask
+        target_loss_mask = inputs.video_loss_mask[:, ref_seq_len:]
+
+        # Compute loss
+        loss = (target_pred - inputs.video_targets).pow(2)
+
+        # Apply loss mask
+        loss_mask = target_loss_mask.unsqueeze(-1).float()
+        loss = loss.mul(loss_mask).div(loss_mask.mean())
+
+        return loss.mean()

packages/ltx-trainer/src/ltx_trainer/utils.py

@@ -0,0 +1,118 @@
+import io
+import subprocess
+from pathlib import Path
+
+import numpy as np
+import torch
+from PIL import ExifTags, Image, ImageCms, ImageOps
+from PIL.Image import Image as PilImage
+
+from ltx_trainer import logger
+
+
+def get_gpu_memory_gb(device: torch.device) -> float:
+    """
+    Get current GPU memory usage in GB using nvidia-smi
+    Args:
+        device: torch.device to get memory usage for
+    Returns:
+        Current GPU memory usage in GB
+    """
+    try:
+        device_id = device.index if device.index is not None else 0
+        result = subprocess.check_output(
+            [
+                "nvidia-smi",
+                "--query-gpu=memory.used",
+                "--format=csv,nounits,noheader",
+                "-i",
+                str(device_id),
+            ],
+            encoding="utf-8",
+        )
+        return float(result.strip()) / 1024  # Convert MB to GB
+    except (subprocess.CalledProcessError, FileNotFoundError, ValueError) as e:
+        logger.error(f"Failed to get GPU memory from nvidia-smi: {e}")
+        # Fallback to torch
+        return torch.cuda.memory_allocated(device) / 1024**3
+
+
+def open_image_as_srgb(image_path: str | Path | io.BytesIO) -> PilImage:
+    """
+    Opens an image file, applies rotation (if it's set in metadata) and converts it
+    to the sRGB color space respecting the original image color space .
+    Args:
+        image_path: Path to the image file
+    Returns:
+        PIL Image in sRGB color space
+    """
+    exif_colorspace_srgb = 1
+
+    with Image.open(image_path) as img_raw:
+        img = ImageOps.exif_transpose(img_raw)
+
+    input_icc_profile = img.info.get("icc_profile")
+
+    # Try to convert to sRGB if the image has ICC profile metadata
+    srgb_profile = ImageCms.createProfile(colorSpace="sRGB")
+    if input_icc_profile is not None:
+        input_profile = ImageCms.ImageCmsProfile(io.BytesIO(input_icc_profile))
+        srgb_img = ImageCms.profileToProfile(img, input_profile, srgb_profile, outputMode="RGB")
+    else:
+        # Try fall back to checking EXIF
+        exif_data = img.getexif()
+        if exif_data is not None:
+            # Assume sRGB if no ICC profile and EXIF has no ColorSpace tag
+            color_space_value = exif_data.get(ExifTags.Base.ColorSpace.value)
+            if color_space_value is not None and color_space_value != exif_colorspace_srgb:
+                raise ValueError(
+                    "Image has colorspace tag in EXIF but it isn't set to sRGB,"
+                    " conversion is not supported."
+                    f" EXIF ColorSpace tag value is {color_space_value}",
+                )
+
+        srgb_img = img.convert("RGB")
+
+        # Set sRGB profile in metadata since now the image is assumed to be in sRGB.
+        srgb_profile_data = ImageCms.ImageCmsProfile(srgb_profile).tobytes()
+        srgb_img.info["icc_profile"] = srgb_profile_data
+
+    return srgb_img
+
+
+def save_image(image_tensor: torch.Tensor, output_path: Path | str) -> None:
+    """Save an image tensor to a file.
+    Args:
+        image_tensor: Image tensor of shape [C, H, W] or [C, 1, H, W] in range [0, 1] or [0, 255].
+            C must be 3 (RGB).
+        output_path: Path to save the image (any PIL-supported format, e.g., .png or .jpg)
+    """
+    output_path = Path(output_path)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    # Handle [C, 1, H, W] format (single frame from video tensor)
+    if image_tensor.ndim == 4:
+        # Squeeze frame dimension: [C, 1, H, W] -> [C, H, W]
+        if image_tensor.shape[1] == 1:
+            image_tensor = image_tensor.squeeze(1)
+        else:
+            raise ValueError(f"Expected single-frame tensor with shape [C, 1, H, W], got shape {image_tensor.shape}")
+
+    if image_tensor.ndim != 3:
+        raise ValueError(f"Expected 3D tensor [C, H, W], got {image_tensor.ndim}D tensor")
+
+    if image_tensor.shape[0] != 3:
+        raise ValueError(f"Expected 3 channels (RGB), got {image_tensor.shape[0]} channels")
+
+    # Normalize to [0, 255] uint8
+    if torch.is_floating_point(image_tensor) and image_tensor.max() <= 1.0:
+        image_tensor = image_tensor * 255
+
+    # Clamp to valid uint8 range to prevent overflow
+    image_tensor = image_tensor.clamp(0, 255)
+
+    # [C, H, W] -> [H, W, C]
+    image_np: np.ndarray = image_tensor.permute(1, 2, 0).to(torch.uint8).cpu().numpy()
+
+    # Save using PIL
+    Image.fromarray(image_np).save(output_path)

packages/ltx-trainer/src/ltx_trainer/validation_sampler.py

@@ -0,0 +1,817 @@
+"""Validation sampling for LTX-2 training using ltx-core components.
+This module provides a simplified validation pipeline for generating samples during training,
+using the new ltx-core components (VideoLatentTools, AudioLatentTools, LatentState, etc.).
+"""
+
+from dataclasses import dataclass, replace
+from typing import TYPE_CHECKING, Literal
+
+import torch
+from einops import rearrange
+from torch import Tensor
+
+from ltx_core.components.diffusion_steps import EulerDiffusionStep
+from ltx_core.components.guiders import CFGGuider, STGGuider
+from ltx_core.components.noisers import GaussianNoiser
+from ltx_core.components.patchifiers import (
+    AudioPatchifier,
+    VideoLatentPatchifier,
+    get_pixel_coords,
+)
+from ltx_core.components.schedulers import LTX2Scheduler
+from ltx_core.guidance.perturbations import (
+    BatchedPerturbationConfig,
+    Perturbation,
+    PerturbationConfig,
+    PerturbationType,
+)
+from ltx_core.model.transformer.modality import Modality
+from ltx_core.model.transformer.model import X0Model
+from ltx_core.model.video_vae import SpatialTilingConfig, TemporalTilingConfig, TilingConfig
+from ltx_core.tools import AudioLatentTools, VideoLatentTools
+from ltx_core.types import AudioLatentShape, LatentState, SpatioTemporalScaleFactors, VideoLatentShape, VideoPixelShape
+from ltx_trainer.progress import SamplingContext
+
+if TYPE_CHECKING:
+    from ltx_core.model.audio_vae import AudioDecoder, Vocoder
+    from ltx_core.model.transformer import LTXModel
+    from ltx_core.model.video_vae import VideoDecoder, VideoEncoder
+    from ltx_core.text_encoders.gemma import AVGemmaTextEncoderModel
+
+VIDEO_SCALE_FACTORS = SpatioTemporalScaleFactors.default()
+
+
+@dataclass
+class CachedPromptEmbeddings:
+    """Pre-computed text embeddings for a validation prompt.
+    These embeddings are computed once at training start and reused for all validation runs,
+    avoiding the need to load the full Gemma text encoder during validation.
+    """
+
+    video_context_positive: Tensor  # [1, seq_len, hidden_dim]
+    audio_context_positive: Tensor  # [1, seq_len, hidden_dim]
+    video_context_negative: Tensor | None = None
+    audio_context_negative: Tensor | None = None
+
+
+@dataclass
+class TiledDecodingConfig:
+    """Configuration for tiled video decoding to reduce VRAM usage.
+    Tiled decoding splits the latent tensor into overlapping tiles, decodes each
+    tile individually, and blends them together. This significantly reduces peak
+    VRAM usage at the cost of slightly slower decoding.
+    Defaults match the recommended values from ltx-core tests.
+    """
+
+    enabled: bool = True  # Whether to use tiled decoding (enabled by default)
+    tile_size_pixels: int = 192  # Spatial tile size in pixels (must be ≥64 and divisible by 32)
+    tile_overlap_pixels: int = 64  # Spatial tile overlap in pixels (must be divisible by 32)
+    tile_size_frames: int = 48  # Temporal tile size in frames (must be ≥16 and divisible by 8)
+    tile_overlap_frames: int = 24  # Temporal tile overlap in frames (must be divisible by 8)
+
+
+@dataclass
+class GenerationConfig:
+    """Configuration for video/audio generation."""
+
+    prompt: str  # Text prompt for generation
+    negative_prompt: str = ""  # Negative prompt to avoid unwanted artifacts
+    height: int = 544  # Output video height in pixels
+    width: int = 960  # Output video width in pixels
+    num_frames: int = 97  # Number of frames to generate
+    frame_rate: float = 25.0  # Frame rate for temporal position scaling
+    num_inference_steps: int = 30  # Number of denoising steps
+    guidance_scale: float = 4.0  # CFG guidance scale
+    seed: int = 42  # Random seed for reproducibility
+    condition_image: Tensor | None = None  # Optional first frame image for image-to-video
+    reference_video: Tensor | None = None  # For IC-LoRA: [F, C, H, W] in [0, 1]
+    generate_audio: bool = True  # Whether to generate audio alongside video
+    include_reference_in_output: bool = False  # For IC-LoRA: concatenate original reference with generated output
+    cached_embeddings: CachedPromptEmbeddings | None = None  # Pre-computed text embeddings (avoids loading Gemma)
+    stg_scale: float = 0.0  # STG strength (0.0 = disabled, recommended: 1.0)
+    stg_blocks: list[int] | None = None  # Transformer blocks to perturb (None = all, recommended: [29])
+    stg_mode: Literal["stg_av", "stg_v"] = "stg_av"  # STG mode: "stg_av" (audio+video) or "stg_v" (video only)
+    # Tiled decoding config: None = use defaults (enabled), False = disable, or TiledDecodingConfig for custom settings
+    tiled_decoding: TiledDecodingConfig | Literal[False] | None = None
+
+    def __post_init__(self) -> None:
+        """Apply default tiled decoding config if not provided."""
+        if self.tiled_decoding is None:
+            # Use default config with tiling enabled
+            object.__setattr__(self, "tiled_decoding", TiledDecodingConfig())
+        elif self.tiled_decoding is False:
+            # Explicitly disabled - use config with enabled=False
+            object.__setattr__(self, "tiled_decoding", TiledDecodingConfig(enabled=False))
+
+
+class ValidationSampler:
+    """Generates validation samples during training using ltx-core components.
+    This class provides a simplified interface for generating video (and optionally audio)
+    samples during training validation. It supports:
+    - Text-to-video generation
+    - Image-to-video generation (first frame conditioning)
+    - Video-to-video generation (IC-LoRA reference video conditioning)
+    - Optional audio generation
+    The implementation follows the patterns from ltx_pipelines.single_stage.
+    Text embeddings can be provided either via:
+    - A full text_encoder (encodes prompts on-the-fly)
+    - Pre-computed cached_embeddings (avoids loading Gemma during validation)
+    """
+
+    def __init__(
+        self,
+        transformer: "LTXModel",
+        vae_decoder: "VideoDecoder",
+        vae_encoder: "VideoEncoder | None",
+        text_encoder: "AVGemmaTextEncoderModel | None" = None,
+        audio_decoder: "AudioDecoder | None" = None,
+        vocoder: "Vocoder | None" = None,
+        sampling_context: SamplingContext | None = None,
+    ):
+        """Initialize the validation sampler.
+        Args:
+            transformer: LTX-2 transformer model
+            vae_decoder: Video VAE decoder
+            vae_encoder: Video VAE encoder (for image/video conditioning), can be None if not needed
+            text_encoder: Gemma text encoder with embeddings connector (optional if cached_embeddings in config)
+            audio_decoder: Optional audio VAE decoder (for audio generation)
+            vocoder: Optional vocoder (for audio generation)
+            sampling_context: Optional SamplingContext for progress display during denoising
+        """
+        self._transformer = transformer
+        self._vae_decoder = vae_decoder
+        self._vae_encoder = vae_encoder
+        self._text_encoder = text_encoder
+        self._audio_decoder = audio_decoder
+        self._vocoder = vocoder
+        self._sampling_context = sampling_context
+
+        # Patchifiers
+        self._video_patchifier = VideoLatentPatchifier(patch_size=1)
+        self._audio_patchifier = AudioPatchifier(patch_size=1)
+
+    # Note: Use @torch.no_grad() instead of @torch.inference_mode() to avoid FSDP inplace update errors after validation
+    @torch.no_grad()
+    def generate(
+        self,
+        config: GenerationConfig,
+        device: torch.device | str = "cuda",
+    ) -> tuple[Tensor, Tensor | None]:
+        """Generate a video (and optionally audio) sample.
+        Args:
+            config: Generation configuration
+            device: Device to run generation on
+        Returns:
+            Tuple of:
+                - video: Video tensor [C, F, H, W] in [0, 1] (float32)
+                - audio: Audio waveform tensor [C, samples] or None
+        """
+        device = torch.device(device) if isinstance(device, str) else device
+        self._validate_config(config)
+
+        # Route to appropriate generation method
+        if config.reference_video is not None:
+            return self._generate_with_reference(config, device)
+        return self._generate_standard(config, device)
+
+    def _generate_standard(self, config: GenerationConfig, device: torch.device) -> tuple[Tensor, Tensor | None]:
+        """Standard generation (text-to-video or image-to-video)."""
+        # Get prompt embeddings (from cache or encode on-the-fly)
+        v_ctx_pos, a_ctx_pos, v_ctx_neg, a_ctx_neg = self._get_prompt_embeddings(config, device)
+
+        # Setup generator
+        generator = torch.Generator(device=device).manual_seed(config.seed)
+
+        # Create latent tools
+        video_tools = self._create_video_latent_tools(config)
+        audio_tools = self._create_audio_latent_tools(config) if config.generate_audio else None
+
+        # Create initial states
+        video_clean_state = video_tools.create_initial_state(device=device, dtype=torch.bfloat16)
+        audio_clean_state = (
+            audio_tools.create_initial_state(device=device, dtype=torch.bfloat16) if audio_tools else None
+        )
+
+        # Apply image conditioning if provided
+        if config.condition_image is not None:
+            video_clean_state = self._apply_image_conditioning(
+                video_clean_state, config.condition_image, config, device
+            )
+
+        # Add noise
+        noiser = GaussianNoiser(generator=generator)
+        video_state = noiser(latent_state=video_clean_state, noise_scale=1.0)
+        audio_state = noiser(latent_state=audio_clean_state, noise_scale=1.0) if audio_clean_state else None
+
+        # Run denoising loop
+        video_state, audio_state = self._run_denoising(
+            config=config,
+            video_state=video_state,
+            audio_state=audio_state,
+            video_clean_state=video_clean_state,
+            audio_clean_state=audio_clean_state,
+            v_ctx_pos=v_ctx_pos,
+            a_ctx_pos=a_ctx_pos,
+            v_ctx_neg=v_ctx_neg,
+            a_ctx_neg=a_ctx_neg,
+            device=device,
+        )
+
+        # Decode outputs
+        video_state = video_tools.clear_conditioning(video_state)
+        video_state = video_tools.unpatchify(video_state)
+        video_output = self._decode_video(video_state, device, config.tiled_decoding)
+
+        audio_output = None
+        if audio_state is not None and audio_tools is not None:
+            audio_state = audio_tools.clear_conditioning(audio_state)
+            audio_state = audio_tools.unpatchify(audio_state)
+            audio_output = self._decode_audio(audio_state, device)
+
+        return video_output, audio_output
+
+    def _generate_with_reference(self, config: GenerationConfig, device: torch.device) -> tuple[Tensor, Tensor | None]:
+        """Generate with reference video conditioning (IC-LoRA style).
+        For IC-LoRA:
+        - Reference video latents are concatenated with target latents
+        - Reference latents have timestep=0 (clean, not denoised)
+        - Target latents are denoised normally
+        - If condition_image is also provided, the first frame of the target is conditioned
+        - If include_reference_in_output is True, the preprocessed reference video
+          is concatenated side-by-side with the generated video
+        """
+        # Get prompt embeddings (from cache or encode on-the-fly)
+        v_ctx_pos, a_ctx_pos, v_ctx_neg, a_ctx_neg = self._get_prompt_embeddings(config, device)
+
+        # Setup generator
+        generator = torch.Generator(device=device).manual_seed(config.seed)
+
+        # Preprocess and encode reference video
+        ref_video_preprocessed = self._preprocess_reference_video(config)
+        ref_latent, ref_positions = self._encode_video(ref_video_preprocessed, config.frame_rate, device)
+        ref_seq_len = ref_latent.shape[1]
+
+        # Create target video state
+        video_tools = self._create_video_latent_tools(config)
+        target_clean_state = video_tools.create_initial_state(device=device, dtype=torch.bfloat16)
+
+        # Apply first-frame image conditioning to target if provided
+        if config.condition_image is not None:
+            target_clean_state = self._apply_image_conditioning(
+                target_clean_state, config.condition_image, config, device
+            )
+
+        # Create combined state (reference + target)
+        # denoise_mask shape is [B, seq_len, 1] after patchification
+        ref_denoise_mask = torch.zeros(1, ref_seq_len, 1, device=device, dtype=torch.float32)
+        combined_clean_state = LatentState(
+            latent=torch.cat([ref_latent, target_clean_state.latent], dim=1),
+            denoise_mask=torch.cat([ref_denoise_mask, target_clean_state.denoise_mask], dim=1),
+            positions=torch.cat([ref_positions, target_clean_state.positions], dim=2),
+            clean_latent=torch.cat([ref_latent, target_clean_state.clean_latent], dim=1),
+        )
+
+        # Add noise (only to the target portion via denoise_mask)
+        noiser = GaussianNoiser(generator=generator)
+        combined_state = noiser(latent_state=combined_clean_state, noise_scale=1.0)
+
+        # Create audio state if needed
+        audio_tools = self._create_audio_latent_tools(config) if config.generate_audio else None
+        audio_clean_state = (
+            audio_tools.create_initial_state(device=device, dtype=torch.bfloat16) if audio_tools else None
+        )
+        audio_state = noiser(latent_state=audio_clean_state, noise_scale=1.0) if audio_clean_state else None
+
+        # Run denoising loop
+        combined_state, audio_state = self._run_denoising(
+            config=config,
+            video_state=combined_state,
+            audio_state=audio_state,
+            video_clean_state=combined_clean_state,
+            audio_clean_state=audio_clean_state,
+            v_ctx_pos=v_ctx_pos,
+            a_ctx_pos=a_ctx_pos,
+            v_ctx_neg=v_ctx_neg,
+            a_ctx_neg=a_ctx_neg,
+            device=device,
+        )
+
+        # Extract target portion and decode
+        target_latent = combined_state.latent[:, ref_seq_len:]
+        video_output = self._decode_video_latent(target_latent, config, device)
+
+        # Optionally concatenate original reference video side-by-side
+        if config.include_reference_in_output:
+            # Use preprocessed reference (already resized/cropped, in pixel space)
+            # Convert from [B, C, F, H, W] to [C, F, H, W]
+            ref_video_pixels = ref_video_preprocessed[0].cpu()
+            # Normalize from [-1, 1] to [0, 1]
+            ref_video_pixels = ((ref_video_pixels + 1.0) / 2.0).clamp(0.0, 1.0)
+            video_output = self._concatenate_videos_side_by_side(ref_video_pixels, video_output)
+
+        # Decode audio
+        audio_output = None
+        if audio_state is not None and audio_tools is not None:
+            audio_state = audio_tools.clear_conditioning(audio_state)
+            audio_state = audio_tools.unpatchify(audio_state)
+            audio_output = self._decode_audio(audio_state, device)
+
+        return video_output, audio_output
+
+    def _create_video_latent_tools(self, config: GenerationConfig) -> VideoLatentTools:
+        """Create video latent tools for the given configuration."""
+        pixel_shape = VideoPixelShape(
+            batch=1,
+            frames=config.num_frames,
+            height=config.height,
+            width=config.width,
+            fps=config.frame_rate,
+        )
+        return VideoLatentTools(
+            patchifier=self._video_patchifier,
+            target_shape=VideoLatentShape.from_pixel_shape(shape=pixel_shape),
+            fps=config.frame_rate,
+            scale_factors=VIDEO_SCALE_FACTORS,
+            causal_fix=True,
+        )
+
+    def _create_audio_latent_tools(self, config: GenerationConfig) -> AudioLatentTools:
+        """Create audio latent tools for the given configuration."""
+        return AudioLatentTools(
+            patchifier=self._audio_patchifier,
+            target_shape=AudioLatentShape.from_duration(batch=1, duration=config.num_frames / config.frame_rate),
+        )
+
+    def _apply_image_conditioning(
+        self, video_state: LatentState, image: Tensor, config: GenerationConfig, device: torch.device
+    ) -> LatentState:
+        """Apply first-frame image conditioning to the video state."""
+        # Encode the image
+        encoded_image = self._encode_conditioning_image(image, config.height, config.width, device)
+
+        # Patchify the encoded image (single frame)
+        patchified_image = self._video_patchifier.patchify(encoded_image)  # [1, 1, C] -> [1, num_patches, C]
+        num_image_tokens = patchified_image.shape[1]
+
+        # Update the first frame tokens in the latent
+        new_latent = video_state.latent.clone()
+        new_latent[:, :num_image_tokens] = patchified_image.to(new_latent.dtype)
+
+        # Update clean_latent as well (conditioning image is clean)
+        new_clean_latent = video_state.clean_latent.clone()
+        new_clean_latent[:, :num_image_tokens] = patchified_image.to(new_clean_latent.dtype)
+
+        # Set denoise_mask to 0 for conditioned tokens (don't denoise them)
+        new_denoise_mask = video_state.denoise_mask.clone()
+        new_denoise_mask[:, :num_image_tokens] = 0.0
+
+        return LatentState(
+            latent=new_latent,
+            denoise_mask=new_denoise_mask,
+            positions=video_state.positions,
+            clean_latent=new_clean_latent,
+        )
+
+    @staticmethod
+    def _preprocess_reference_video(config: GenerationConfig) -> Tensor:
+        """Preprocess reference video: resize, crop, and convert to model input format.
+        Args:
+            config: Generation configuration with reference_video
+        Returns:
+            Preprocessed video tensor [B, C, F, H, W] in [-1, 1] range
+        """
+        ref_video = config.reference_video  # [F, C, H, W] in [0, 1]
+        target_height, target_width = config.height, config.width
+        current_height, current_width = ref_video.shape[2:]
+
+        # Resize maintaining aspect ratio and center crop if needed
+        if current_height != target_height or current_width != target_width:
+            aspect_ratio = current_width / current_height
+            target_aspect_ratio = target_width / target_height
+
+            if aspect_ratio > target_aspect_ratio:
+                resize_height, resize_width = target_height, int(target_height * aspect_ratio)
+            else:
+                resize_height, resize_width = int(target_width / aspect_ratio), target_width
+
+            ref_video = torch.nn.functional.interpolate(
+                ref_video, size=(resize_height, resize_width), mode="bilinear", align_corners=False
+            )
+
+            # Center crop
+            h_start = (resize_height - target_height) // 2
+            w_start = (resize_width - target_width) // 2
+            ref_video = ref_video[:, :, h_start : h_start + target_height, w_start : w_start + target_width]
+
+        # Convert to [B, C, F, H, W] and trim to valid frame count (k*8 + 1)
+        ref_video = rearrange(ref_video, "f c h w -> 1 c f h w")
+        valid_frames = (ref_video.shape[2] - 1) // 8 * 8 + 1
+        ref_video = ref_video[:, :, :valid_frames]
+
+        # Convert to [-1, 1] range
+        return ref_video * 2.0 - 1.0
+
+    def _encode_video(self, video: Tensor, fps: float, device: torch.device) -> tuple[Tensor, Tensor]:
+        """Encode video to patchified latents and compute positions.
+        Args:
+            video: Video tensor [B, C, F, H, W] in [-1, 1] range
+            fps: Frame rate for temporal position scaling
+            device: Device to run encoding on
+        Returns:
+            Tuple of (patchified_latents, positions)
+        """
+        video = video.to(device=device, dtype=torch.float32)
+
+        # Encode with VAE
+        self._vae_encoder.to(device)
+        with torch.autocast(device_type=str(device).split(":")[0], dtype=torch.bfloat16):
+            latents = self._vae_encoder(video)
+        self._vae_encoder.to("cpu")
+
+        latents = latents.to(torch.bfloat16)
+        patchified = self._video_patchifier.patchify(latents)
+
+        # Compute positions
+        latent_shape = VideoLatentShape(
+            batch=1,
+            channels=latents.shape[1],
+            frames=latents.shape[2],
+            height=latents.shape[3],
+            width=latents.shape[4],
+        )
+        latent_coords = self._video_patchifier.get_patch_grid_bounds(output_shape=latent_shape, device=device)
+        positions = get_pixel_coords(latent_coords, scale_factors=VIDEO_SCALE_FACTORS, causal_fix=True)
+        positions = positions.to(torch.bfloat16)
+        positions[:, 0, ...] = positions[:, 0, ...] / fps
+
+        return patchified, positions
+
+    def _run_denoising(
+        self,
+        config: GenerationConfig,
+        video_state: LatentState,
+        audio_state: LatentState | None,
+        video_clean_state: LatentState,
+        audio_clean_state: LatentState | None,
+        v_ctx_pos: Tensor,
+        a_ctx_pos: Tensor,
+        v_ctx_neg: Tensor | None,
+        a_ctx_neg: Tensor | None,
+        device: torch.device,
+    ) -> tuple[LatentState, LatentState | None]:
+        """Run the denoising loop using X0 prediction with CFG and optional STG."""
+        scheduler = LTX2Scheduler()
+        sigmas = scheduler.execute(steps=config.num_inference_steps).to(device).float()
+        stepper = EulerDiffusionStep()
+        cfg_guider = CFGGuider(config.guidance_scale)
+        stg_guider = STGGuider(config.stg_scale)
+
+        # Build STG perturbation config if STG is enabled
+        stg_perturbation_config = self._build_stg_perturbation_config(config) if stg_guider.enabled() else None
+
+        # Create initial modalities (will be updated each step via replace())
+        video = Modality(
+            enabled=True,
+            latent=video_state.latent,
+            timesteps=video_state.denoise_mask,
+            positions=video_state.positions,
+            context=v_ctx_pos,
+            context_mask=None,
+        )
+
+        # Audio modality is None when not generating audio
+        audio: Modality | None = None
+        if audio_state is not None:
+            audio = Modality(
+                enabled=True,
+                latent=audio_state.latent,
+                timesteps=audio_state.denoise_mask,
+                positions=audio_state.positions,
+                context=a_ctx_pos,
+                context_mask=None,
+            )
+
+        # Wrap transformer with X0Model to convert velocity predictions to denoised outputs
+        self._transformer.to(device)
+        x0_model = X0Model(self._transformer)
+
+        with torch.autocast(device_type=str(device).split(":")[0], dtype=torch.bfloat16):
+            for step_idx, sigma in enumerate(sigmas[:-1]):
+                # Update modalities with current state and timesteps
+                video = replace(
+                    video,
+                    latent=video_state.latent,
+                    timesteps=sigma * video_state.denoise_mask,
+                    positions=video_state.positions,
+                )
+
+                if audio is not None and audio_state is not None:
+                    audio = replace(
+                        audio,
+                        latent=audio_state.latent,
+                        timesteps=sigma * audio_state.denoise_mask,
+                        positions=audio_state.positions,
+                    )
+
+                # Run model (positive pass) - X0Model returns denoised outputs
+                pos_video, pos_audio = x0_model(video=video, audio=audio, perturbations=None)
+                denoised_video, denoised_audio = pos_video, pos_audio
+
+                # Apply CFG if guidance_scale != 1.0
+                if cfg_guider.enabled() and v_ctx_neg is not None:
+                    video_neg = replace(video, context=v_ctx_neg)
+                    audio_neg = replace(audio, context=a_ctx_neg) if audio is not None else None
+                    neg_video, neg_audio = x0_model(video=video_neg, audio=audio_neg, perturbations=None)
+
+                    denoised_video = denoised_video + cfg_guider.delta(pos_video, neg_video)
+                    if audio is not None and denoised_audio is not None:
+                        denoised_audio = denoised_audio + cfg_guider.delta(pos_audio, neg_audio)
+
+                # Apply STG if stg_scale != 0.0
+                if stg_guider.enabled() and stg_perturbation_config is not None:
+                    perturbed_video, perturbed_audio = x0_model(
+                        video=video, audio=audio, perturbations=stg_perturbation_config
+                    )
+                    denoised_video = denoised_video + stg_guider.delta(pos_video, perturbed_video)
+                    if audio is not None and denoised_audio is not None and perturbed_audio is not None:
+                        denoised_audio = denoised_audio + stg_guider.delta(pos_audio, perturbed_audio)
+
+                # Apply conditioning mask (keep conditioned tokens clean)
+                denoised_video = denoised_video * video_state.denoise_mask + video_clean_state.latent.float() * (
+                    1 - video_state.denoise_mask
+                )
+                if audio is not None and audio_state is not None and audio_clean_state is not None:
+                    denoised_audio = denoised_audio * audio_state.denoise_mask + audio_clean_state.latent.float() * (
+                        1 - audio_state.denoise_mask
+                    )
+
+                # Euler step
+                video_state = replace(
+                    video_state,
+                    latent=stepper.step(
+                        sample=video.latent, denoised_sample=denoised_video, sigmas=sigmas, step_index=step_idx
+                    ),
+                )
+                if audio is not None and audio_state is not None:
+                    audio_state = replace(
+                        audio_state,
+                        latent=stepper.step(
+                            sample=audio.latent, denoised_sample=denoised_audio, sigmas=sigmas, step_index=step_idx
+                        ),
+                    )
+
+                # Update progress
+                if self._sampling_context is not None:
+                    self._sampling_context.advance_step()
+
+        return video_state, audio_state
+
+    @staticmethod
+    def _build_stg_perturbation_config(config: GenerationConfig) -> BatchedPerturbationConfig:
+        """Build the perturbation config for STG based on the stg_mode."""
+        # Always skip video self-attention for STG
+        perturbations: list[Perturbation] = [
+            Perturbation(type=PerturbationType.SKIP_VIDEO_SELF_ATTN, blocks=config.stg_blocks)
+        ]
+
+        # Optionally also skip audio self-attention (stg_av mode)
+        if config.stg_mode == "stg_av":
+            perturbations.append(Perturbation(type=PerturbationType.SKIP_AUDIO_SELF_ATTN, blocks=config.stg_blocks))
+
+        perturbation_config = PerturbationConfig(perturbations=perturbations)
+        # Batch size is 1 for validation
+        return BatchedPerturbationConfig(perturbations=[perturbation_config])
+
+    def _decode_video_latent(self, latent: Tensor, config: GenerationConfig, device: torch.device) -> Tensor:
+        """Decode patchified video latent to pixel space."""
+        # Unpatchify
+        latent_frames = config.num_frames // VIDEO_SCALE_FACTORS.time + 1
+        latent_height = config.height // VIDEO_SCALE_FACTORS.height
+        latent_width = config.width // VIDEO_SCALE_FACTORS.width
+
+        unpatchified = self._video_patchifier.unpatchify(
+            latent,
+            output_shape=VideoLatentShape(
+                height=latent_height,
+                width=latent_width,
+                frames=latent_frames,
+                batch=1,
+                channels=128,
+            ),
+        )
+
+        # Decode - ensure bfloat16 to match decoder weights
+        self._vae_decoder.to(device)
+        unpatchified = unpatchified.to(dtype=torch.bfloat16)
+        tiled_config = config.tiled_decoding
+
+        if tiled_config is not None and tiled_config.enabled:
+            # Use tiled decoding for reduced VRAM
+            tiling_config = TilingConfig(
+                spatial_config=SpatialTilingConfig(
+                    tile_size_in_pixels=tiled_config.tile_size_pixels,
+                    tile_overlap_in_pixels=tiled_config.tile_overlap_pixels,
+                ),
+                temporal_config=TemporalTilingConfig(
+                    tile_size_in_frames=tiled_config.tile_size_frames,
+                    tile_overlap_in_frames=tiled_config.tile_overlap_frames,
+                ),
+            )
+            chunks = []
+            for video_chunk in self._vae_decoder.tiled_decode(
+                unpatchified,
+                tiling_config=tiling_config,
+            ):
+                chunks.append(video_chunk)
+            decoded_video = torch.cat(chunks, dim=2)
+        else:
+            # Standard full decoding
+            decoded_video = self._vae_decoder(unpatchified)
+
+        decoded_video = ((decoded_video + 1.0) / 2.0).clamp(0.0, 1.0)
+        self._vae_decoder.to("cpu")
+
+        return decoded_video[0].float().cpu()
+
+    def _validate_config(self, config: GenerationConfig) -> None:
+        """Validate generation configuration."""
+        if config.height % 32 != 0 or config.width % 32 != 0:
+            raise ValueError(f"height and width must be divisible by 32, got {config.height}x{config.width}")
+        if config.num_frames % 8 != 1:
+            raise ValueError(f"num_frames must satisfy num_frames % 8 == 1, got {config.num_frames}")
+        if config.generate_audio and (self._audio_decoder is None or self._vocoder is None):
+            raise ValueError("Audio generation requires audio_decoder and vocoder")
+        if config.condition_image is not None and self._vae_encoder is None:
+            raise ValueError("Image conditioning requires vae_encoder")
+        if config.reference_video is not None and self._vae_encoder is None:
+            raise ValueError("Reference video conditioning requires vae_encoder")
+
+        # Validate prompt embedding source
+        if config.cached_embeddings is None and self._text_encoder is None:
+            raise ValueError("Either text_encoder or config.cached_embeddings must be provided")
+
+    def _get_prompt_embeddings(
+        self, config: GenerationConfig, device: torch.device
+    ) -> tuple[Tensor, Tensor, Tensor | None, Tensor | None]:
+        """Get prompt embeddings from config cache or encode on-the-fly."""
+        if config.cached_embeddings is not None:
+            # Use pre-computed embeddings from config
+            cached = config.cached_embeddings
+            v_ctx_pos = cached.video_context_positive.to(device)
+            a_ctx_pos = cached.audio_context_positive.to(device)
+            v_ctx_neg = cached.video_context_negative.to(device) if cached.video_context_negative is not None else None
+            a_ctx_neg = cached.audio_context_negative.to(device) if cached.audio_context_negative is not None else None
+            return v_ctx_pos, a_ctx_pos, v_ctx_neg, a_ctx_neg
+
+        # Fall back to encoding on-the-fly
+        return self._encode_prompts(config, device)
+
+    def _encode_prompts(
+        self, config: GenerationConfig, device: torch.device
+    ) -> tuple[Tensor, Tensor, Tensor | None, Tensor | None]:
+        """Encode positive and negative prompts using the text encoder."""
+        self._text_encoder.to(device)
+        v_ctx_pos, a_ctx_pos, _ = self._text_encoder(config.prompt)
+        v_ctx_neg, a_ctx_neg = None, None
+        if config.guidance_scale != 1.0:
+            v_ctx_neg, a_ctx_neg, _ = self._text_encoder(config.negative_prompt)
+
+        # Move the base Gemma model to CPU but keep embeddings connectors on GPU
+        # as this module is also used during training
+        self._text_encoder.model.to("cpu")
+        self._text_encoder.feature_extractor_linear.to("cpu")
+
+        return v_ctx_pos, a_ctx_pos, v_ctx_neg, a_ctx_neg
+
+    def _decode_video(
+        self, video_state: LatentState, device: torch.device, tiled_config: TiledDecodingConfig | None = None
+    ) -> Tensor:
+        """Decode video latents to pixel space.
+        Args:
+            video_state: Video latent state to decode
+            device: Device to run decoding on
+            tiled_config: Optional tiled decoding configuration for reduced VRAM usage
+        Returns:
+            Decoded video tensor [C, F, H, W] in [0, 1] range
+        """
+        self._vae_decoder.to(device)
+        # Ensure latent is bfloat16 to match decoder weights
+        latent = video_state.latent.to(dtype=torch.bfloat16)
+
+        if tiled_config is not None and tiled_config.enabled:
+            # Use tiled decoding for reduced VRAM
+            tiling_config = TilingConfig(
+                spatial_config=SpatialTilingConfig(
+                    tile_size_in_pixels=tiled_config.tile_size_pixels,
+                    tile_overlap_in_pixels=tiled_config.tile_overlap_pixels,
+                ),
+                temporal_config=TemporalTilingConfig(
+                    tile_size_in_frames=tiled_config.tile_size_frames,
+                    tile_overlap_in_frames=tiled_config.tile_overlap_frames,
+                ),
+            )
+            chunks = []
+            for video_chunk in self._vae_decoder.tiled_decode(
+                latent,
+                tiling_config=tiling_config,
+            ):
+                chunks.append(video_chunk)
+            decoded_video = torch.cat(chunks, dim=2)
+        else:
+            # Standard full decoding
+            decoded_video = self._vae_decoder(latent)
+
+        decoded_video = ((decoded_video + 1.0) / 2.0).clamp(0.0, 1.0)
+        self._vae_decoder.to("cpu")
+        return decoded_video[0].float().cpu()
+
+    def _decode_audio(self, audio_state: LatentState, device: torch.device) -> Tensor:
+        """Decode audio latents to waveform."""
+        self._audio_decoder.to(device)
+        # Ensure latent is bfloat16 to match decoder weights
+        latent = audio_state.latent.to(dtype=torch.bfloat16)
+        decoded_audio = self._audio_decoder(latent)
+        self._audio_decoder.to("cpu")
+
+        self._vocoder.to(device)
+        audio_waveform = self._vocoder(decoded_audio)
+        self._vocoder.to("cpu")
+
+        return audio_waveform.squeeze(0).float().cpu()
+
+    @staticmethod
+    def _concatenate_videos_side_by_side(left_video: Tensor, right_video: Tensor) -> Tensor:
+        """Concatenate two videos side-by-side (horizontally).
+        If the videos have different frame counts, the shorter one is padded with
+        its last frame repeated.
+        Args:
+            left_video: Left video tensor [C, F1, H, W] in [0, 1]
+            right_video: Right video tensor [C, F2, H, W] in [0, 1]
+        Returns:
+            Concatenated video tensor [C, max(F1,F2), H, W*2] in [0, 1]
+        """
+        left_frames = left_video.shape[1]
+        right_frames = right_video.shape[1]
+
+        # Pad shorter video by repeating last frame
+        if left_frames < right_frames:
+            padding = left_video[:, -1:, :, :].expand(-1, right_frames - left_frames, -1, -1)
+            left_video = torch.cat([left_video, padding], dim=1)
+        elif right_frames < left_frames:
+            padding = right_video[:, -1:, :, :].expand(-1, left_frames - right_frames, -1, -1)
+            right_video = torch.cat([right_video, padding], dim=1)
+
+        # Concatenate along width dimension
+        return torch.cat([left_video, right_video], dim=3)
+
+    def _encode_conditioning_image(
+        self,
+        image: Tensor,
+        target_height: int,
+        target_width: int,
+        device: torch.device,
+    ) -> Tensor:
+        """Encode a conditioning image to latent space.
+        The image is resized to cover the target dimensions while preserving aspect ratio,
+        then center-cropped to exactly match the target size.
+        """
+        # image is [C, H, W] in [0, 1]  # noqa: ERA001
+        current_height, current_width = image.shape[1:]
+
+        # Resize maintaining aspect ratio (cover target, then center crop)
+        if current_height != target_height or current_width != target_width:
+            aspect_ratio = current_width / current_height
+            target_aspect_ratio = target_width / target_height
+
+            if aspect_ratio > target_aspect_ratio:
+                # Image is wider than target - resize to match height, crop width
+                resize_height = target_height
+                resize_width = int(target_height * aspect_ratio)
+            else:
+                # Image is taller than target - resize to match width, crop height
+                resize_height = int(target_width / aspect_ratio)
+                resize_width = target_width
+
+            image = rearrange(image, "c h w -> 1 c h w")
+            image = torch.nn.functional.interpolate(
+                image, size=(resize_height, resize_width), mode="bilinear", align_corners=False
+            )
+
+            # Center crop to target dimensions
+            h_start = (resize_height - target_height) // 2
+            w_start = (resize_width - target_width) // 2
+            image = image[:, :, h_start : h_start + target_height, w_start : w_start + target_width]
+        else:
+            image = rearrange(image, "c h w -> 1 c h w")
+
+        # Add frame dimension and convert to [-1, 1]
+        image = rearrange(image, "b c h w -> b c 1 h w")
+        image = (image * 2.0 - 1.0).to(device=device, dtype=torch.float32)
+
+        # Encode
+        self._vae_encoder.to(device)
+        with torch.autocast(device_type=str(device).split(":")[0], dtype=torch.bfloat16):
+            encoded = self._vae_encoder(image)
+        self._vae_encoder.to("cpu")
+
+        return encoded

packages/ltx-trainer/src/ltx_trainer/video_utils.py

@@ -0,0 +1,159 @@
+"""Video I/O utilities using PyAV.
+This module provides functions for reading and writing video files using PyAV,
+with optional audio support.
+"""
+
+from fractions import Fraction
+from pathlib import Path
+
+import av
+import numpy as np
+import torch
+from torch import Tensor
+
+
+def get_video_frame_count(video_path: str | Path) -> int:
+    """Get the number of frames in a video file.
+    Args:
+        video_path: Path to the video file
+    Returns:
+        Number of frames in the video
+    """
+    with av.open(str(video_path)) as container:
+        video_stream = container.streams.video[0]
+        frame_count = video_stream.frames
+        if frame_count == 0:
+            # Fallback: count frames by decoding
+            frame_count = sum(1 for _ in container.decode(video=0))
+    return frame_count
+
+
+def read_video(video_path: str | Path, max_frames: int | None = None) -> tuple[Tensor, float]:
+    """Load frames from a video file using PyAV.
+    Args:
+        video_path: Path to the video file
+        max_frames: Maximum number of frames to read. If None, reads all frames.
+    Returns:
+        Video tensor with shape [F, C, H, W] in range [0, 1] and frames per second (fps).
+    """
+    with av.open(str(video_path)) as container:
+        video_stream = container.streams.video[0]
+        fps = float(video_stream.average_rate or video_stream.base_rate or 24)
+
+        frames = []
+        for frame in container.decode(video=0):
+            if max_frames is not None and len(frames) >= max_frames:
+                break
+            frames.append(frame.to_ndarray(format="rgb24"))
+
+    frames_np = np.stack(frames, axis=0)  # [F, H, W, C]
+    video = torch.from_numpy(frames_np).float().div(255.0)  # [F, H, W, C] in [0, 1]
+    return video.permute(0, 3, 1, 2), fps  # [F, C, H, W]
+
+
+def save_video(
+    video_tensor: torch.Tensor,
+    output_path: Path | str,
+    fps: float = 24.0,
+    audio: torch.Tensor | None = None,
+    audio_sample_rate: int | None = None,
+) -> None:
+    """Save a video tensor to a file using PyAV, optionally with audio.
+    Args:
+        video_tensor: Video tensor of shape [C, F, H, W] or [F, C, H, W] in range [0, 1] or [0, 255]
+        output_path: Path to save the video
+        fps: Frames per second for the output video
+        audio: Optional audio tensor of shape [C, samples] or [samples, C] in range [-1, 1]
+        audio_sample_rate: Sample rate for the audio (required if audio is provided)
+    """
+    output_path = Path(output_path)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    # Normalize to [F, H, W, C] uint8 numpy array
+    video_np = _prepare_video_array(video_tensor)
+    _, height, width, _ = video_np.shape
+
+    with av.open(str(output_path), mode="w") as container:
+        # Setup video stream
+        video_stream = container.add_stream("libx264", rate=int(fps))
+        video_stream.width = width
+        video_stream.height = height
+        video_stream.pix_fmt = "yuv420p"
+        video_stream.options = {"crf": "18"}
+
+        # Setup audio stream if needed
+        if audio is not None:
+            if audio_sample_rate is None:
+                raise ValueError("audio_sample_rate must be provided when audio is given")
+            audio_stream = container.add_stream("aac", rate=audio_sample_rate)
+            audio_stream.layout = "stereo"
+            audio_stream.time_base = Fraction(1, audio_sample_rate)
+
+        # Write video frames
+        for frame_array in video_np:
+            frame = av.VideoFrame.from_ndarray(frame_array, format="rgb24")
+            for packet in video_stream.encode(frame):
+                container.mux(packet)
+        for packet in video_stream.encode():
+            container.mux(packet)
+
+        # Write audio if provided
+        if audio is not None:
+            _write_audio(container, audio_stream, audio, audio_sample_rate)
+
+
+def _prepare_video_array(video_tensor: torch.Tensor) -> np.ndarray:
+    """Convert video tensor to [F, H, W, C] uint8 numpy array."""
+    # Handle [C, F, H, W] vs [F, C, H, W] format
+    if video_tensor.shape[0] == 3 and video_tensor.shape[1] > 3:
+        video_tensor = video_tensor.permute(1, 0, 2, 3)  # [C, F, H, W] -> [F, C, H, W]
+
+    # Normalize to [0, 255] uint8
+    if video_tensor.max() <= 1.0:
+        video_tensor = video_tensor * 255
+
+    # [F, C, H, W] -> [F, H, W, C]
+    return video_tensor.permute(0, 2, 3, 1).to(torch.uint8).cpu().numpy()
+
+
+def _write_audio(
+    container: av.container.Container,
+    audio_stream: av.audio.AudioStream,
+    audio: torch.Tensor,
+    sample_rate: int,
+) -> None:
+    """Write audio tensor to container as stereo AAC."""
+    audio = audio.cpu().float()
+
+    # Normalize to [samples, 2] stereo format
+    if audio.ndim == 1:
+        audio = audio.unsqueeze(1).repeat(1, 2)  # Mono -> stereo
+    elif audio.shape[0] == 2 and audio.shape[1] != 2:
+        audio = audio.T  # [2, samples] -> [samples, 2]
+    if audio.shape[1] == 1:
+        audio = audio.repeat(1, 2)  # Mono -> stereo
+
+    # Convert to int16 interleaved: [samples, 2] -> [1, samples*2]
+    audio_int16 = (audio.clamp(-1, 1) * 32767).to(torch.int16)
+    audio_interleaved = audio_int16.contiguous().view(1, -1).numpy()
+
+    # Create audio frame
+    frame = av.AudioFrame.from_ndarray(audio_interleaved, format="s16", layout="stereo")
+    frame.sample_rate = sample_rate
+
+    # Resample to encoder format and write
+    resampler = av.audio.resampler.AudioResampler(
+        format=audio_stream.codec_context.format,
+        layout=audio_stream.codec_context.layout,
+        rate=sample_rate,
+    )
+
+    pts = 0
+    for resampled_frame in resampler.resample(frame):
+        resampled_frame.pts = pts
+        pts += resampled_frame.samples
+        for packet in audio_stream.encode(resampled_frame):
+            container.mux(packet)
+
+    for packet in audio_stream.encode():
+        container.mux(packet)