Automated PR - 2026-06-17
This commit is contained in:
github-actions[bot]
@@ -1,59 +1,117 @@
|
||||
"""
|
||||
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 (uv pip install google-generativeai)
|
||||
Set GEMINI_API_KEY or GOOGLE_API_KEY environment variable
|
||||
- Qwen3-Omni via a local vLLM server (default)
|
||||
- Gemini Flash 3.5 (cloud API)
|
||||
Both produce a single combined English caption per video as a single
|
||||
continuous paragraph of prose.
|
||||
The Qwen3-Omni backend runs in a separately-launched vLLM server rather than
|
||||
in-process, so vLLM's heavy CUDA dependencies stay out of this package. The
|
||||
captioner talks to it over the OpenAI-compatible HTTP API.
|
||||
Launch the server once (in an isolated environment) with:
|
||||
.. code-block:: bash
|
||||
uv run python scripts/serve_captioner.py
|
||||
That helper picks BF16 vs FP8 dynamic quantization based on the GPU's free
|
||||
memory and forwards everything else to ``vllm serve``. To check the recommended
|
||||
command without running it, pass ``--print-cmd``.
|
||||
To use Gemini instead, install ``google-genai`` and either set ``GEMINI_API_KEY``
|
||||
(Gemini Developer API) or have Google Cloud credentials available (gcloud / an
|
||||
attached service account), in which case it uses Vertex AI automatically.
|
||||
"""
|
||||
|
||||
import itertools
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import tempfile
|
||||
from abc import ABC, abstractmethod
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
from typing import ClassVar
|
||||
|
||||
import torch
|
||||
DEFAULT_VIDEO_CAPTION_INSTRUCTION = """\
|
||||
Analyze this video and produce a single detailed caption covering both its visual content and its audio. Be \
|
||||
detailed enough that someone reading the caption could form an accurate mental picture of what happens on screen \
|
||||
and what can be heard. Be exhaustive: include every meaningful detail you can see and hear, including small \
|
||||
objects, textures, secondary movements, and minor background sounds.
|
||||
|
||||
# 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:
|
||||
Begin the caption directly with the action or visual detail; do not preface it with phrases like \
|
||||
"The video opens with...", "The scene shows...", "We see...", or "There is...".
|
||||
|
||||
[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">
|
||||
For every shot, include:
|
||||
- The shot type and framing (extreme wide / wide / medium / medium close-up / close-up / extreme close-up) and any \
|
||||
camera motion.
|
||||
- Characters' clothing, appearance, posture, and movement (direction, speed, quality).
|
||||
- The environment's materials, textures, lighting, and colors.
|
||||
- All audio: spoken dialogue (quoted exactly in the original language), tone of voice, music (style, mood, \
|
||||
volume changes), and environmental sounds. If a category is absent -- for example no music is playing, or no one is \
|
||||
speaking -- state that explicitly. Do not invent specific instruments, music genres, moods, or ambient sounds \
|
||||
that are not actually present.
|
||||
- Any on-screen text (signs, titles, labels).
|
||||
|
||||
You MUST fill in all four sections. For [SPEECH], transcribe the actual words spoken, not a summary."""
|
||||
Describe only what is visible or audible. Do not infer emotions, intentions, or anything outside the segment. \
|
||||
Refer to people descriptively (e.g., "the man in the blue jacket"). Narrate strictly in chronological order; if \
|
||||
the video contains multiple shots, describe each one in turn.
|
||||
|
||||
# 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:
|
||||
Write everything as a single continuous paragraph of prose. Do not use section headers, bullet points, or labels \
|
||||
like "Audio:" / "Visual:" / "Shot:". Integrate visual and audio details naturally within the same sentences.
|
||||
|
||||
[VISUAL]: <Detailed description of people, objects, actions, settings, colors, and movements>
|
||||
[TEXT]: <Any on-screen text visible. If none, write "None">
|
||||
Return a JSON object with exactly one key:
|
||||
|
||||
You MUST fill in both sections."""
|
||||
{"combined_caption_english": "<your caption here>"}"""
|
||||
|
||||
|
||||
DEFAULT_IMAGE_CAPTION_INSTRUCTION = """\
|
||||
Analyze this image and produce a single detailed caption of its visual content. Be detailed enough that \
|
||||
someone reading the caption could form an accurate mental picture of the image. Be thorough: include every meaningful \
|
||||
detail that is actually present, including small objects, textures, and background elements.
|
||||
|
||||
Begin the caption directly with the main subject or a visual detail; do not preface it with phrases like \
|
||||
"The image shows...", "This is a photo of...", "We see...", or "There is...".
|
||||
|
||||
Include:
|
||||
- The framing and composition (close-up / medium / wide / overhead, etc.) and the vantage point.
|
||||
- The medium or style if distinctive (photograph, illustration, 3D render, painting).
|
||||
- People's clothing, appearance, and posture, and what they are doing.
|
||||
- The setting's materials, textures, lighting, and colors.
|
||||
- Transcribe any visible text verbatim (signs, labels, titles, captions).
|
||||
|
||||
Describe only what is visible. Do not infer emotions or intentions, and do not describe sounds, motion, or \
|
||||
events before or after the moment shown -- this is a single still image. When something is ambiguous, describe \
|
||||
the visible cue (e.g., "warm low-angle light") rather than guessing the underlying fact (e.g., "sunrise"). \
|
||||
Refer to people descriptively (e.g., "the man in the blue jacket").
|
||||
|
||||
Only describe what is present. Never state that something is absent or missing -- do not write phrases like \
|
||||
"there is no text", "no people are present", or "no other objects". If a category such as people or text does \
|
||||
not appear, simply leave it out.
|
||||
|
||||
Write everything as a single continuous paragraph of prose. Do not use section headers, bullet points, or \
|
||||
labels.
|
||||
|
||||
Return a JSON object with exactly one key:
|
||||
|
||||
{"combined_caption_english": "<your caption here>"}"""
|
||||
|
||||
|
||||
# Default model served by ``scripts/serve_captioner.py``. The captioner does not
|
||||
# download or load this model itself -- it just sends requests to the vLLM
|
||||
# server, which already has the model loaded.
|
||||
DEFAULT_QWEN_MODEL = "Qwen/Qwen3-Omni-30B-A3B-Thinking"
|
||||
DEFAULT_VLLM_BASE_URL = "http://127.0.0.1:8001/v1"
|
||||
|
||||
# Key the combined-caption prompt asks the model to return its caption under.
|
||||
_CAPTION_JSON_KEY = "combined_caption_english"
|
||||
|
||||
|
||||
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)
|
||||
QWEN_OMNI = "qwen_omni" # Qwen3-Omni via local vLLM HTTP server
|
||||
GEMINI_FLASH = "gemini_flash" # Gemini Flash 3.5 cloud API
|
||||
|
||||
|
||||
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
|
||||
"""
|
||||
"""Factory function to create a media captioner."""
|
||||
match captioner_type:
|
||||
case CaptionerType.QWEN_OMNI:
|
||||
return QwenOmniCaptioner(**kwargs)
|
||||
@@ -66,336 +124,332 @@ def create_captioner(captioner_type: CaptionerType, **kwargs) -> "MediaCaptionin
|
||||
class MediaCaptioningModel(ABC):
|
||||
"""Abstract base class for audio-visual media captioning models."""
|
||||
|
||||
instruction: str | None = None
|
||||
|
||||
@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
|
||||
"""
|
||||
"""Generate a caption for the given video or image."""
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def supports_audio(self) -> bool:
|
||||
"""Whether this captioner supports audio input."""
|
||||
def _resolve_instruction(self, path: str | Path) -> str:
|
||||
"""Return the custom instruction, or the image/video default for this input."""
|
||||
if self.instruction is not None:
|
||||
return self.instruction
|
||||
return DEFAULT_IMAGE_CAPTION_INSTRUCTION if self._is_image_file(path) else DEFAULT_VIDEO_CAPTION_INSTRUCTION
|
||||
|
||||
@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)
|
||||
"""Audio-visual captioning via a local vLLM server running Qwen3-Omni.
|
||||
The vLLM server must already be running. See ``scripts/serve_captioner.py``
|
||||
for a helper that launches one in an isolated environment (no impact on
|
||||
this package's dependency tree).
|
||||
The captioner uses the OpenAI-compatible chat completions API. It sends
|
||||
a ``file://`` URL pointing at the local video, the default combined-caption
|
||||
prompt, and parses the JSON-wrapped response.
|
||||
"""
|
||||
|
||||
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,
|
||||
base_url: str = DEFAULT_VLLM_BASE_URL,
|
||||
model: str = DEFAULT_QWEN_MODEL,
|
||||
api_key: str = "EMPTY",
|
||||
instruction: str | None = None,
|
||||
max_tokens: int = 4096,
|
||||
enable_thinking: bool = False,
|
||||
timeout_s: float = 600.0,
|
||||
):
|
||||
"""
|
||||
Initialize the Qwen2.5-Omni captioner.
|
||||
"""Initialize the Qwen3-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
|
||||
base_url: Base URL of the vLLM OpenAI-compatible server (default
|
||||
``http://127.0.0.1:8001/v1``).
|
||||
model: Model identifier the server is serving. Must match the
|
||||
server's ``--served-model-name`` (defaults to the HuggingFace
|
||||
model ID).
|
||||
api_key: Token sent in the ``Authorization`` header. vLLM accepts
|
||||
any value by default.
|
||||
instruction: Custom instruction prompt. If ``None``, uses the
|
||||
default combined-caption prompt.
|
||||
max_tokens: Maximum new tokens to generate per caption. 4096 leaves
|
||||
comfortable headroom for both ``enable_thinking`` modes.
|
||||
enable_thinking: Whether to let the Thinking model produce a
|
||||
``<think>...</think>`` chain-of-thought before the caption.
|
||||
Off by default: it makes captioning ~5x slower with little
|
||||
quality benefit and occasionally introduces hallucinations
|
||||
(e.g., inventing dialogue or background music).
|
||||
timeout_s: Per-request HTTP timeout.
|
||||
"""
|
||||
self.device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
|
||||
self.instruction = instruction
|
||||
self._load_model(use_8bit=use_8bit)
|
||||
from openai import OpenAI # noqa: PLC0415
|
||||
|
||||
@property
|
||||
def supports_audio(self) -> bool:
|
||||
return True
|
||||
self.model = model
|
||||
self.instruction = instruction
|
||||
self.max_tokens = max_tokens
|
||||
self.enable_thinking = enable_thinking
|
||||
self._client = OpenAI(base_url=base_url, api_key=api_key, timeout=timeout_s)
|
||||
|
||||
def caption(
|
||||
self,
|
||||
path: str | Path,
|
||||
fps: int = 1,
|
||||
include_audio: bool = True,
|
||||
clean_caption: bool = True,
|
||||
fps: int = 2,
|
||||
) -> 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
|
||||
path: Path to the video/image file to caption.
|
||||
fps: Frames per second to sample from the video. Passed through to
|
||||
vLLM's multimodal processor (``mm_processor_kwargs.fps``).
|
||||
Default 2 is a typical choice for video MLLMs at this resolution.
|
||||
Ignored for image inputs.
|
||||
Returns:
|
||||
A string containing the generated caption
|
||||
The extracted caption string.
|
||||
"""
|
||||
path = Path(path)
|
||||
is_image = self._is_image_file(path)
|
||||
is_video = self._is_video_file(path)
|
||||
if not (is_image or is_video):
|
||||
raise ValueError(f"Unsupported media 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 = []
|
||||
instruction = self._resolve_instruction(path)
|
||||
|
||||
if is_image:
|
||||
user_content.append({"type": "image", "image": str(path)})
|
||||
elif is_video:
|
||||
user_content.append({"type": "video", "video": str(path)})
|
||||
content = [
|
||||
{"type": "image_url", "image_url": {"url": f"file://{path.resolve()}"}},
|
||||
{"type": "text", "text": instruction},
|
||||
]
|
||||
return _parse_caption_response(self._chat(content)).strip()
|
||||
|
||||
# Add the instruction text
|
||||
user_content.append({"type": "text", "text": instruction})
|
||||
return self._caption_video(path, instruction, fps)
|
||||
|
||||
# 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,
|
||||
def _chat(self, content: list[dict], mm_kwargs: dict | None = None) -> str:
|
||||
"""Send one chat-completions request and return the raw response text."""
|
||||
extra_body: dict = {
|
||||
"repetition_penalty": 1.05,
|
||||
"chat_template_kwargs": {"enable_thinking": self.enable_thinking},
|
||||
}
|
||||
if mm_kwargs:
|
||||
extra_body["mm_processor_kwargs"] = mm_kwargs
|
||||
response = self._client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=[{"role": "user", "content": content}],
|
||||
max_tokens=self.max_tokens,
|
||||
temperature=0.0,
|
||||
extra_body=extra_body,
|
||||
)
|
||||
return response.choices[0].message.content or ""
|
||||
|
||||
# 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.
|
||||
def _caption_video(self, path: Path, instruction: str, fps: int) -> str:
|
||||
"""Caption a video, sending its audio track as a separate modality.
|
||||
vLLM does not extract a video's audio on its own (and its
|
||||
``use_audio_in_video`` path is broken server-side), so we pull the audio
|
||||
into a 16 kHz mono WAV and send it alongside the video -- otherwise the
|
||||
model only sees frames and fabricates any spoken content.
|
||||
"""
|
||||
from transformers import ( # noqa: PLC0415
|
||||
BitsAndBytesConfig,
|
||||
Qwen2_5OmniProcessor,
|
||||
Qwen2_5OmniThinkerForConditionalGeneration,
|
||||
)
|
||||
with tempfile.TemporaryDirectory(prefix="qwencap_") as tmp:
|
||||
work = Path(tmp)
|
||||
|
||||
quantization_config = BitsAndBytesConfig(load_in_8bit=True) if use_8bit else None
|
||||
# Best-effort: ffmpeg fails (and we send video only) if there's no audio.
|
||||
audio_url: str | None = None
|
||||
try:
|
||||
wav = work / "audio.wav"
|
||||
_extract_audio_wav(path, wav)
|
||||
audio_url = f"file://{wav.resolve()}"
|
||||
except subprocess.CalledProcessError:
|
||||
pass
|
||||
|
||||
# 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",
|
||||
)
|
||||
def content(video: Path) -> list[dict]:
|
||||
parts: list[dict] = [{"type": "video_url", "video_url": {"url": f"file://{video.resolve()}"}}]
|
||||
if audio_url:
|
||||
parts.append({"type": "audio_url", "audio_url": {"url": audio_url}})
|
||||
parts.append({"type": "text", "text": instruction})
|
||||
return parts
|
||||
|
||||
self.processor = Qwen2_5OmniProcessor.from_pretrained(self.MODEL_ID)
|
||||
mm_kwargs = {"fps": fps}
|
||||
try:
|
||||
raw = self._chat(content(path), mm_kwargs)
|
||||
except Exception as e:
|
||||
# Raw / variable-frame-rate videos over-report their frame count, which
|
||||
# breaks the server's frame sampler ("... frames from video"). Re-encode
|
||||
# to a constant frame rate and retry once.
|
||||
if "frames from video" not in str(e):
|
||||
raise
|
||||
cfr = work / "video_cfr.mp4"
|
||||
_transcode_cfr(path, cfr)
|
||||
raw = self._chat(content(cfr), mm_kwargs)
|
||||
|
||||
return _parse_caption_response(raw).strip()
|
||||
|
||||
|
||||
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.
|
||||
"""Audio-visual captioning using Google's Gemini via the Google Gen AI SDK.
|
||||
Uses the ``google-genai`` package (the current SDK; ``google-generativeai``
|
||||
is deprecated). Auth is resolved automatically:
|
||||
1. If an API key is given (``api_key`` argument, or ``GEMINI_API_KEY`` /
|
||||
``GOOGLE_API_KEY`` in the environment) -> the Gemini Developer API (AI Studio).
|
||||
2. Otherwise, if Google Cloud Application Default Credentials are available
|
||||
(an attached service account or ``gcloud auth application-default login``)
|
||||
-> Vertex AI. The project comes from ADC (or ``GOOGLE_CLOUD_PROJECT``) and
|
||||
the location defaults to ``global`` (override with ``GOOGLE_CLOUD_LOCATION``).
|
||||
This means it "just works" on a gcloud-authed GCP VM with no env vars.
|
||||
If neither is available, a clear error explains how to authenticate.
|
||||
Media is sent inline (``Part.from_bytes``), which works on both backends.
|
||||
"""
|
||||
|
||||
MODEL_ID = "gemini-flash-lite-latest"
|
||||
MODEL_ID = "gemini-3.5-flash"
|
||||
|
||||
_MIME_TYPES: ClassVar[dict[str, str]] = {
|
||||
".jpg": "image/jpeg",
|
||||
".jpeg": "image/jpeg",
|
||||
".png": "image/png",
|
||||
".webp": "image/webp",
|
||||
".heic": "image/heic",
|
||||
".heif": "image/heif",
|
||||
".mp4": "video/mp4",
|
||||
".mov": "video/quicktime",
|
||||
".avi": "video/x-msvideo",
|
||||
".mkv": "video/x-matroska",
|
||||
".webm": "video/webm",
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str | None = None,
|
||||
instruction: str | None = None,
|
||||
model: str | None = None,
|
||||
):
|
||||
"""Initialize the Gemini Flash captioner.
|
||||
"""Initialize the Gemini 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
|
||||
api_key: Gemini Developer API key. If ``None``, falls back to
|
||||
``GEMINI_API_KEY`` / ``GOOGLE_API_KEY``; if no key is set at all,
|
||||
uses Vertex AI via Application Default Credentials.
|
||||
instruction: Custom instruction prompt. If ``None``, uses the default
|
||||
image or video prompt depending on the input.
|
||||
model: Override the served model id (defaults to ``MODEL_ID``).
|
||||
"""
|
||||
self.instruction = instruction
|
||||
self._init_client(api_key)
|
||||
|
||||
@property
|
||||
def supports_audio(self) -> bool:
|
||||
return True
|
||||
self.model = model or self.MODEL_ID
|
||||
self._client = self._make_client(api_key)
|
||||
|
||||
def caption(
|
||||
self,
|
||||
path: str | Path,
|
||||
fps: int = 3, # noqa: ARG002 - kept for API compatibility
|
||||
include_audio: bool = True,
|
||||
clean_caption: bool = True,
|
||||
fps: int = 2, # noqa: ARG002 - kept for API compatibility
|
||||
) -> 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
|
||||
from google.genai import types # noqa: PLC0415
|
||||
|
||||
path = Path(path)
|
||||
is_video = self._is_video_file(path)
|
||||
use_audio = include_audio and is_video
|
||||
instruction = self._resolve_instruction(path)
|
||||
media = types.Part.from_bytes(data=path.read_bytes(), mime_type=self._mime_type(path))
|
||||
response = self._client.models.generate_content(
|
||||
model=self.model,
|
||||
contents=[media, instruction],
|
||||
config=types.GenerateContentConfig(temperature=0.0),
|
||||
)
|
||||
|
||||
# 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
|
||||
# Gemini may also return JSON if it followed our prompt format.
|
||||
return _parse_caption_response(response.text or "").strip()
|
||||
|
||||
# Upload the file to Gemini
|
||||
uploaded_file = self._genai.upload_file(path)
|
||||
@classmethod
|
||||
def _mime_type(cls, path: Path) -> str:
|
||||
try:
|
||||
return cls._MIME_TYPES[path.suffix.lower()]
|
||||
except KeyError:
|
||||
raise ValueError(f"Unsupported media type for Gemini: {path.suffix}") from None
|
||||
|
||||
# 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)
|
||||
def _make_client(self, api_key: str | None): # noqa: ANN202 - genai.Client type is lazy-imported
|
||||
from google import genai # noqa: PLC0415
|
||||
|
||||
if uploaded_file.state.name == "FAILED":
|
||||
raise RuntimeError(f"File processing failed: {uploaded_file.state.name}")
|
||||
# 1. API key (explicit arg or env) -> Gemini Developer API.
|
||||
key = api_key or os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY")
|
||||
if key:
|
||||
return genai.Client(api_key=key)
|
||||
|
||||
# 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
|
||||
# 2. No key -> Vertex AI via Application Default Credentials (gcloud / service account).
|
||||
import google.auth # 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`"
|
||||
_, adc_project = google.auth.default()
|
||||
except Exception as e:
|
||||
raise ValueError(
|
||||
"No Gemini credentials found. Provide an API key (--api-key, or "
|
||||
"GEMINI_API_KEY / GOOGLE_API_KEY), or set up Google Cloud credentials "
|
||||
"for Vertex AI (e.g. `gcloud auth application-default login` or an "
|
||||
"attached service account)."
|
||||
) 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)
|
||||
project = os.environ.get("GOOGLE_CLOUD_PROJECT") or adc_project
|
||||
location = os.environ.get("GOOGLE_CLOUD_LOCATION", "global")
|
||||
return genai.Client(vertexai=True, project=project, location=location)
|
||||
|
||||
|
||||
def example() -> None:
|
||||
"""Example usage of the captioning module."""
|
||||
import sys # noqa: PLC0415
|
||||
def _parse_caption_response(raw: str) -> str:
|
||||
"""Extract the caption text from a model response.
|
||||
Backend-agnostic: works for any model that follows the combined-caption
|
||||
prompt. Handles the formats a model may produce:
|
||||
- Plain caption text
|
||||
- JSON ``{"combined_caption_english": "..."}``
|
||||
- ``<think>...</think>`` chain-of-thought followed by either of the above
|
||||
- Truncated JSON (when generation hits a token limit mid-string)
|
||||
"""
|
||||
text = re.sub(r"<think>[\s\S]*?</think>", "", raw).strip()
|
||||
|
||||
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)
|
||||
# Thinking models (e.g. Qwen3-Omni-*-Thinking) emit the reasoning trace
|
||||
# without an opening ``<think>`` tag, because the chat template injects it
|
||||
# for them -- so the response starts mid-thought and is terminated by a lone
|
||||
# ``</think>`` before the real answer. Drop everything up to that closer.
|
||||
if "</think>" in text:
|
||||
text = text.rsplit("</think>", 1)[1].strip()
|
||||
|
||||
video_path = sys.argv[1]
|
||||
captioner_type = CaptionerType(sys.argv[2]) if len(sys.argv) > 2 else CaptionerType.QWEN_OMNI
|
||||
if not text:
|
||||
return raw.strip()
|
||||
|
||||
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
|
||||
try:
|
||||
parsed = json.loads(text)
|
||||
if isinstance(parsed, dict) and _CAPTION_JSON_KEY in parsed:
|
||||
return parsed[_CAPTION_JSON_KEY]
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
pass
|
||||
|
||||
match = re.search(rf"\{{[^{{}}]*\"{_CAPTION_JSON_KEY}\"[^{{}}]*\}}", text)
|
||||
if match:
|
||||
try:
|
||||
parsed = json.loads(match.group())
|
||||
if isinstance(parsed, dict) and _CAPTION_JSON_KEY in parsed:
|
||||
return parsed[_CAPTION_JSON_KEY]
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
pass
|
||||
|
||||
# Truncated JSON: extract the string value even if the closing quote/brace is missing.
|
||||
match = re.search(rf'"{_CAPTION_JSON_KEY}"\s*:\s*"((?:[^"\\]|\\.)*)', text)
|
||||
if match:
|
||||
try:
|
||||
return json.loads('"' + match.group(1) + '"')
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
return match.group(1)
|
||||
|
||||
return text
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
example()
|
||||
def _run_ffmpeg(args: list[str]) -> None:
|
||||
"""Run the ffmpeg binary bundled with ``imageio-ffmpeg`` (a dependency)."""
|
||||
import imageio_ffmpeg # noqa: PLC0415
|
||||
|
||||
cmd = [imageio_ffmpeg.get_ffmpeg_exe(), "-y", "-loglevel", "error", *args]
|
||||
subprocess.run(cmd, check=True, capture_output=True)
|
||||
|
||||
|
||||
def _extract_audio_wav(src: Path, dest: Path) -> None:
|
||||
"""Extract the audio track to a 16 kHz mono PCM WAV (matches pretraining).
|
||||
Raises ``CalledProcessError`` when the video has no audio stream.
|
||||
"""
|
||||
_run_ffmpeg(["-i", str(src), "-vn", "-ac", "1", "-ar", "16000", "-c:a", "pcm_s16le", str(dest)])
|
||||
|
||||
|
||||
def _transcode_cfr(src: Path, dest: Path) -> None:
|
||||
"""Re-encode the video to a constant frame rate so the server's frame sampler can
|
||||
read every requested index (raw / variable-frame-rate videos over-report frames)."""
|
||||
_run_ffmpeg(["-i", str(src), "-fps_mode", "cfr", "-c:v", "libx264", "-pix_fmt", "yuv420p", "-an", str(dest)])
|
||||
|
||||
@@ -1,10 +1,11 @@
|
||||
from pathlib import Path
|
||||
from typing import Annotated, Literal
|
||||
from typing import Annotated, Literal, Union
|
||||
|
||||
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.flexible import FlexibleStrategyConfig
|
||||
from ltx_trainer.training_strategies.text_to_video import TextToVideoConfig
|
||||
from ltx_trainer.training_strategies.video_to_video import VideoToVideoConfig
|
||||
|
||||
@@ -13,6 +14,226 @@ class ConfigBaseModel(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Validation Condition Types
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class FirstFrameConditionConfig(ConfigBaseModel):
|
||||
"""First-frame conditioning (intrinsic, latent_idx=0). Always targets video.
|
||||
If image_or_video points to a video file, the first frame is automatically extracted.
|
||||
"""
|
||||
|
||||
type: Literal["first_frame"] = "first_frame"
|
||||
image_or_video: str | Path
|
||||
|
||||
|
||||
class PrefixConditionConfig(ConfigBaseModel):
|
||||
"""Prefix conditioning for temporal extension (intrinsic). Exactly one of video/audio must be set."""
|
||||
|
||||
type: Literal["prefix"] = "prefix"
|
||||
video: str | None = None
|
||||
audio: str | None = None
|
||||
num_frames: int | None = Field(
|
||||
default=None,
|
||||
ge=1,
|
||||
description="Number of pixel frames for video prefix. Must satisfy num_frames %% 8 == 1.",
|
||||
)
|
||||
duration: float | None = Field(default=None, gt=0, description="Duration in seconds for audio prefix")
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_exactly_one_modality(self) -> "PrefixConditionConfig":
|
||||
if (self.video is None) == (self.audio is None):
|
||||
raise ValueError("Exactly one of 'video' or 'audio' must be set for prefix condition")
|
||||
return self
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_num_frames_constraint(self) -> "PrefixConditionConfig":
|
||||
if self.video is not None and self.num_frames is not None and self.num_frames % 8 != 1:
|
||||
raise ValueError(
|
||||
f"num_frames ({self.num_frames}) must satisfy num_frames % 8 == 1 "
|
||||
f"for video prefix (e.g., 1, 9, 17, 25, ...)"
|
||||
)
|
||||
return self
|
||||
|
||||
|
||||
class SuffixConditionConfig(ConfigBaseModel):
|
||||
"""Suffix conditioning for temporal extension (intrinsic). Exactly one of video/audio must be set."""
|
||||
|
||||
type: Literal["suffix"] = "suffix"
|
||||
video: str | None = None
|
||||
audio: str | None = None
|
||||
num_frames: int | None = Field(
|
||||
default=None,
|
||||
ge=1,
|
||||
description="Number of pixel frames for video suffix. Must satisfy num_frames %% 8 == 0.",
|
||||
)
|
||||
duration: float | None = Field(default=None, gt=0, description="Duration in seconds for audio suffix")
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_exactly_one_modality(self) -> "SuffixConditionConfig":
|
||||
if (self.video is None) == (self.audio is None):
|
||||
raise ValueError("Exactly one of 'video' or 'audio' must be set for suffix condition")
|
||||
return self
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_num_frames_constraint(self) -> "SuffixConditionConfig":
|
||||
if self.video is not None and self.num_frames is not None and self.num_frames % 8 != 0:
|
||||
raise ValueError(
|
||||
f"num_frames ({self.num_frames}) must satisfy num_frames % 8 == 0 "
|
||||
f"for video suffix (e.g., 8, 16, 24, 32, ...)"
|
||||
)
|
||||
return self
|
||||
|
||||
|
||||
class SpatialCropConditionConfig(ConfigBaseModel):
|
||||
"""Spatial crop conditioning for outpainting (intrinsic, video only)."""
|
||||
|
||||
type: Literal["spatial_crop"] = "spatial_crop"
|
||||
video: str
|
||||
spatial_region: tuple[int, int, int, int] = Field(
|
||||
..., description="Spatial crop region as (y1, x1, y2, x2) in pixel coordinates"
|
||||
)
|
||||
|
||||
|
||||
class MaskConditionConfig(ConfigBaseModel):
|
||||
"""Mask-based conditioning for inpainting (intrinsic). Exactly one of video/audio must be set."""
|
||||
|
||||
type: Literal["mask"] = "mask"
|
||||
video: str | None = None
|
||||
audio: str | None = None
|
||||
mask: str
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_exactly_one_modality(self) -> "MaskConditionConfig":
|
||||
if (self.video is None) == (self.audio is None):
|
||||
raise ValueError("Exactly one of 'video' or 'audio' must be set for mask condition")
|
||||
return self
|
||||
|
||||
|
||||
class ReferenceConditionConfig(ConfigBaseModel):
|
||||
"""Reference conditioning (IC-LoRA style concatenation). Exactly one of video/audio must be set."""
|
||||
|
||||
type: Literal["reference"] = "reference"
|
||||
video: str | None = None
|
||||
audio: str | None = None
|
||||
downscale_factor: int = Field(default=1, ge=1)
|
||||
temporal_scale_factor: int = Field(default=1, ge=1)
|
||||
include_in_output: bool = False
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_exactly_one_modality(self) -> "ReferenceConditionConfig":
|
||||
if (self.video is None) == (self.audio is None):
|
||||
raise ValueError("Exactly one of 'video' or 'audio' must be set for reference condition")
|
||||
return self
|
||||
|
||||
|
||||
class VideoToAudioConditionConfig(ConfigBaseModel):
|
||||
"""Video-to-audio — video is provided as frozen cross-modal conditioning.
|
||||
The video is kept clean (sigma=0) and influences audio generation via cross-modal attention.
|
||||
"""
|
||||
|
||||
type: Literal["video_to_audio"] = "video_to_audio"
|
||||
video: str
|
||||
|
||||
|
||||
class AudioToVideoConditionConfig(ConfigBaseModel):
|
||||
"""Audio-to-video — audio is provided as frozen cross-modal conditioning.
|
||||
The audio is kept clean (sigma=0) and influences video generation via cross-modal attention.
|
||||
"""
|
||||
|
||||
type: Literal["audio_to_video"] = "audio_to_video"
|
||||
audio: str
|
||||
|
||||
|
||||
ValidationCondition = Annotated[
|
||||
Union[
|
||||
FirstFrameConditionConfig,
|
||||
PrefixConditionConfig,
|
||||
SuffixConditionConfig,
|
||||
SpatialCropConditionConfig,
|
||||
MaskConditionConfig,
|
||||
ReferenceConditionConfig,
|
||||
VideoToAudioConditionConfig,
|
||||
AudioToVideoConditionConfig,
|
||||
],
|
||||
Field(discriminator="type"),
|
||||
]
|
||||
|
||||
|
||||
def _condition_targets_video(cond: ValidationCondition) -> bool:
|
||||
"""Check if a validation condition targets the video modality."""
|
||||
if cond.type in ("first_frame", "spatial_crop", "video_to_audio"):
|
||||
return True
|
||||
if cond.type in ("prefix", "suffix", "mask", "reference"):
|
||||
return getattr(cond, "video", None) is not None
|
||||
return False
|
||||
|
||||
|
||||
def _condition_targets_audio(cond: ValidationCondition) -> bool:
|
||||
"""Check if a validation condition targets the audio modality."""
|
||||
if cond.type == "audio_to_video":
|
||||
return True
|
||||
if cond.type in ("prefix", "suffix", "mask", "reference"):
|
||||
return getattr(cond, "audio", None) is not None
|
||||
return False
|
||||
|
||||
|
||||
class ValidationSample(ConfigBaseModel):
|
||||
"""Configuration for a single validation sample — fully self-describing."""
|
||||
|
||||
prompt: str
|
||||
conditions: list[ValidationCondition] = Field(default_factory=list)
|
||||
|
||||
video_dims: tuple[int, int, int] | None = Field(
|
||||
default=None,
|
||||
description="Per-sample override for (width, height, frames). None = inherit from ValidationConfig.",
|
||||
)
|
||||
seed: int | None = Field(
|
||||
default=None,
|
||||
description="Per-sample override for random seed. None = inherit from ValidationConfig.",
|
||||
)
|
||||
|
||||
@field_validator("video_dims")
|
||||
@classmethod
|
||||
def validate_video_dims(cls, v: tuple[int, int, int] | None) -> tuple[int, int, int] | None:
|
||||
if v is None:
|
||||
return v
|
||||
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
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_frozen_modality_conflicts(self) -> "ValidationSample":
|
||||
frozen_types = {c.type for c in self.conditions if c.type in ("video_to_audio", "audio_to_video")}
|
||||
|
||||
if "video_to_audio" in frozen_types and "audio_to_video" in frozen_types:
|
||||
raise ValueError(
|
||||
"Cannot have both video_to_audio and audio_to_video conditions — nothing would be generated"
|
||||
)
|
||||
|
||||
if "video_to_audio" in frozen_types:
|
||||
for c in self.conditions:
|
||||
if c.type != "video_to_audio" and _condition_targets_video(c):
|
||||
raise ValueError(
|
||||
f"Cannot use video-targeting '{c.type}' condition when video is frozen (video_to_audio)"
|
||||
)
|
||||
|
||||
if "audio_to_video" in frozen_types:
|
||||
for c in self.conditions:
|
||||
if c.type != "audio_to_video" and _condition_targets_audio(c):
|
||||
raise ValueError(
|
||||
f"Cannot use audio-targeting '{c.type}' condition when audio is frozen (audio_to_video)"
|
||||
)
|
||||
|
||||
return self
|
||||
|
||||
|
||||
class ModelConfig(ConfigBaseModel):
|
||||
"""Configuration for the base model and training mode"""
|
||||
|
||||
@@ -89,7 +310,9 @@ def _get_strategy_discriminator(v: dict | TrainingStrategyConfigBase) -> str:
|
||||
|
||||
# Union type for all strategy configs with discriminator
|
||||
TrainingStrategyConfig = Annotated[
|
||||
Annotated[TextToVideoConfig, Tag("text_to_video")] | Annotated[VideoToVideoConfig, Tag("video_to_video")],
|
||||
Annotated[TextToVideoConfig, Tag("text_to_video")]
|
||||
| Annotated[VideoToVideoConfig, Tag("video_to_video")]
|
||||
| Annotated[FlexibleStrategyConfig, Tag("flexible")],
|
||||
Discriminator(_get_strategy_discriminator),
|
||||
]
|
||||
|
||||
@@ -191,13 +414,32 @@ class DataConfig(ConfigBaseModel):
|
||||
ge=0,
|
||||
)
|
||||
|
||||
@field_validator("preprocessed_data_root")
|
||||
@classmethod
|
||||
def validate_preprocessed_data_root(cls, v: str) -> str:
|
||||
"""Validate that preprocessed_data_root exists."""
|
||||
path = Path(v).expanduser().resolve()
|
||||
if not path.exists():
|
||||
raise ValueError(f"Dataset path does not exist: {v}")
|
||||
if not path.is_dir():
|
||||
raise ValueError(f"Dataset path is not a directory: {v}")
|
||||
return str(path)
|
||||
|
||||
|
||||
class ValidationConfig(ConfigBaseModel):
|
||||
"""Configuration for validation during training"""
|
||||
|
||||
# Per-sample configuration (new format — preferred)
|
||||
samples: list[ValidationSample] = Field(
|
||||
default_factory=list,
|
||||
description="List of validation samples. Each sample is fully self-describing with its own "
|
||||
"prompt, conditions, and optional overrides. Replaces prompts/images/reference_videos.",
|
||||
)
|
||||
|
||||
# Legacy fields (deprecated — converted to samples internally via convert_legacy_format)
|
||||
prompts: list[str] = Field(
|
||||
default_factory=list,
|
||||
description="List of prompts to use for validation",
|
||||
description="[DEPRECATED: use 'samples' instead] List of prompts to use for validation",
|
||||
)
|
||||
|
||||
negative_prompt: str = Field(
|
||||
@@ -207,19 +449,22 @@ class ValidationConfig(ConfigBaseModel):
|
||||
|
||||
images: list[str] | None = Field(
|
||||
default=None,
|
||||
description="List of image paths to use for validation. "
|
||||
description="[DEPRECATED: use 'samples' with first_frame conditions] "
|
||||
"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. "
|
||||
description="[DEPRECATED: use 'samples' with reference conditions] "
|
||||
"List of reference video paths to use for validation. "
|
||||
"One video path must be provided for each validation prompt",
|
||||
)
|
||||
|
||||
reference_downscale_factor: int = Field(
|
||||
default=1,
|
||||
description="Downscale factor for reference videos in IC-LoRA validation. "
|
||||
description="[DEPRECATED: use downscale_factor on ReferenceCondition] "
|
||||
"Downscale factor for reference videos in IC-LoRA validation. "
|
||||
"When > 1, reference videos are processed at 1/n resolution (e.g., 2 means half resolution). "
|
||||
"Must match the factor used during dataset preprocessing.",
|
||||
ge=1,
|
||||
@@ -301,6 +546,13 @@ class ValidationConfig(ConfigBaseModel):
|
||||
"in validation even when not training the audio branch.",
|
||||
)
|
||||
|
||||
generate_video: bool = Field(
|
||||
default=True,
|
||||
description="Whether to generate video in validation samples. "
|
||||
"Set to False for audio-only or v2a validation to save VRAM by skipping video VAE decoder loading. "
|
||||
"When False, validation will only generate audio (requires generate_audio=True).",
|
||||
)
|
||||
|
||||
skip_initial_validation: bool = Field(
|
||||
default=False,
|
||||
description="Skip validation video sampling at step 0 (beginning of training)",
|
||||
@@ -308,7 +560,8 @@ class ValidationConfig(ConfigBaseModel):
|
||||
|
||||
include_reference_in_output: bool = Field(
|
||||
default=False,
|
||||
description="For video-to-video training: concatenate the original reference video side-by-side "
|
||||
description="[DEPRECATED: use include_in_output on ReferenceCondition] "
|
||||
"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.",
|
||||
)
|
||||
|
||||
@@ -346,13 +599,33 @@ class ValidationConfig(ConfigBaseModel):
|
||||
|
||||
return v
|
||||
|
||||
@model_validator(mode="after")
|
||||
def convert_legacy_format(self) -> "ValidationConfig":
|
||||
"""Convert deprecated prompts/images/reference_videos to the new samples format."""
|
||||
if self.prompts and not self.samples:
|
||||
samples = []
|
||||
for i, prompt in enumerate(self.prompts):
|
||||
conditions: list[ValidationCondition] = []
|
||||
if self.images and i < len(self.images):
|
||||
conditions.append(FirstFrameConditionConfig(image_or_video=self.images[i]))
|
||||
if self.reference_videos and i < len(self.reference_videos):
|
||||
conditions.append(
|
||||
ReferenceConditionConfig(
|
||||
video=self.reference_videos[i],
|
||||
downscale_factor=self.reference_downscale_factor,
|
||||
include_in_output=self.include_reference_in_output,
|
||||
)
|
||||
)
|
||||
samples.append(ValidationSample(prompt=prompt, conditions=conditions))
|
||||
self.samples = samples
|
||||
return self
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_scaled_reference_dimensions(self) -> "ValidationConfig":
|
||||
"""Validate that scaled reference dimensions are valid when reference_downscale_factor > 1."""
|
||||
if self.reference_downscale_factor > 1:
|
||||
width, height, _frames = self.video_dims
|
||||
|
||||
# Validate that downscale factor evenly divides the target dimensions
|
||||
if width % self.reference_downscale_factor != 0:
|
||||
raise ValueError(
|
||||
f"Width {width} is not evenly divisible by reference_downscale_factor "
|
||||
@@ -367,7 +640,6 @@ class ValidationConfig(ConfigBaseModel):
|
||||
scaled_width = width // self.reference_downscale_factor
|
||||
scaled_height = height // self.reference_downscale_factor
|
||||
|
||||
# Validate scaled dimensions are divisible by 32
|
||||
if scaled_width % 32 != 0:
|
||||
raise ValueError(
|
||||
f"Scaled reference width {scaled_width} (from {width} / {self.reference_downscale_factor}) "
|
||||
@@ -381,6 +653,16 @@ class ValidationConfig(ConfigBaseModel):
|
||||
|
||||
return self
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_output_modality_requirements(self) -> "ValidationConfig":
|
||||
"""Validate output modality settings when validation is configured."""
|
||||
has_validation = bool(self.prompts) or bool(self.samples)
|
||||
if has_validation and not self.generate_video and not self.generate_audio:
|
||||
raise ValueError(
|
||||
"At least one of generate_video or generate_audio must be True when validation is configured."
|
||||
)
|
||||
return self
|
||||
|
||||
|
||||
class CheckpointsConfig(ConfigBaseModel):
|
||||
"""Configuration for model checkpointing during training"""
|
||||
@@ -514,19 +796,31 @@ class LtxTrainerConfig(ConfigBaseModel):
|
||||
"""Expand user home directory in output path."""
|
||||
return str(Path(v).expanduser().resolve())
|
||||
|
||||
def _validate_data_dirs_exist(self) -> None:
|
||||
"""Verify that every directory declared by the training strategy exists under the data root."""
|
||||
data_root = Path(self.data.preprocessed_data_root)
|
||||
for dir_name in self.training_strategy.get_data_sources():
|
||||
dir_path = data_root / dir_name
|
||||
if not dir_path.is_dir():
|
||||
raise ValueError(
|
||||
f"Required data directory '{dir_name}' does not exist under preprocessed_data_root: {dir_path}"
|
||||
)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_strategy_compatibility(self) -> "LtxTrainerConfig":
|
||||
"""Validate that training strategy and other configurations are compatible."""
|
||||
self._validate_data_dirs_exist()
|
||||
|
||||
# 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"
|
||||
if self.training_strategy.name == "video_to_video" and self.validation.interval:
|
||||
has_reference = bool(self.validation.reference_videos) or any(
|
||||
cond.type == "reference" for sample in self.validation.samples for cond in sample.conditions
|
||||
)
|
||||
if not has_reference:
|
||||
raise ValueError(
|
||||
"reference_videos or samples with reference conditions 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:
|
||||
|
||||
@@ -54,60 +54,6 @@ class SamplingContext:
|
||||
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
|
||||
|
||||
@@ -12,8 +12,8 @@ from typing import Any, Callable
|
||||
import torch
|
||||
import wandb
|
||||
import yaml
|
||||
from accelerate import Accelerator, DistributedDataParallelKwargs, DistributedType
|
||||
from accelerate.utils import gather_object, set_seed
|
||||
from accelerate import Accelerator, DistributedType
|
||||
from accelerate.utils import DistributedDataParallelKwargs, gather_object, 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
|
||||
@@ -30,26 +30,22 @@ from torch.optim.lr_scheduler import (
|
||||
StepLR,
|
||||
)
|
||||
from torch.utils.data import DataLoader
|
||||
from torchvision.transforms import functional as F # noqa: N812
|
||||
|
||||
from ltx_core.text_encoders.gemma import convert_to_additive_mask
|
||||
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.gpu_utils import free_gpu_memory, free_gpu_memory_context, get_gpu_memory_gb
|
||||
from ltx_trainer.gpu_utils import free_gpu_memory, get_gpu_memory_gb
|
||||
from ltx_trainer.hf_hub_utils import push_to_hub
|
||||
from ltx_trainer.model_loader import load_embeddings_processor, load_text_encoder
|
||||
from ltx_trainer.model_loader import load_model as load_ltx_model
|
||||
from ltx_trainer.model_loader import load_embeddings_processor, load_transformer
|
||||
from ltx_trainer.progress import TrainingProgress
|
||||
from ltx_trainer.quantization import quantize_model
|
||||
from ltx_trainer.sigma_tracker import SigmaBucketTracker
|
||||
from ltx_trainer.timestep_samplers import SAMPLERS
|
||||
from ltx_trainer.training_state import ConfigFingerprint, RngStates, TrainingState
|
||||
from ltx_trainer.training_strategies import get_training_strategy
|
||||
from ltx_trainer.utils import 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
|
||||
from ltx_trainer.validation_runner import ValidationRunner
|
||||
|
||||
# Disable irrelevant warnings from transformers
|
||||
os.environ["TOKENIZERS_PARALLELISM"] = "true"
|
||||
@@ -66,7 +62,7 @@ if not IS_MAIN_PROCESS:
|
||||
|
||||
disable_progress_bar()
|
||||
|
||||
StepCallback = Callable[[int, int, list[Path] | None], None] # (step, total, sampled paths or None) -> None
|
||||
StepCallback = Callable[[int, int, list[Path]], None] # (step, total, list[sampled_video_path]) -> None
|
||||
|
||||
MEMORY_CHECK_INTERVAL = 200
|
||||
|
||||
@@ -96,7 +92,16 @@ class LtxvTrainer:
|
||||
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()
|
||||
|
||||
# ValidationRunner loads its own models (text encoder, VAE encoder/decoder, etc.),
|
||||
# caches prompt embeddings and conditioning media, then unloads encoders.
|
||||
self._validation_runner = ValidationRunner(
|
||||
config=self._config.validation,
|
||||
model_path=self._config.model.model_path,
|
||||
text_encoder_path=self._config.model.text_encoder_path,
|
||||
load_text_encoder_in_8bit=self._config.acceleration.load_text_encoder_in_8bit,
|
||||
)
|
||||
|
||||
self._load_models()
|
||||
self._setup_accelerator()
|
||||
self._collect_trainable_params()
|
||||
@@ -108,8 +113,8 @@ class LtxvTrainer:
|
||||
self._checkpoint_paths: list[Path] = []
|
||||
self._training_state_paths: list[Path] = []
|
||||
self._training_state_size_warned = False
|
||||
self._wandb_run = None
|
||||
self._sigma_tracker = SigmaBucketTracker()
|
||||
self._wandb_run = None
|
||||
|
||||
def train( # noqa: PLR0912, PLR0915
|
||||
self,
|
||||
@@ -190,7 +195,7 @@ class LtxvTrainer:
|
||||
with progress:
|
||||
if cfg.validation.interval and not cfg.validation.skip_initial_validation:
|
||||
with self._offloaded_optimizer_state():
|
||||
sampled_videos_paths = self._run_distributed_validation(progress)
|
||||
sampled_videos_paths = self._run_validation(progress)
|
||||
|
||||
self._accelerator.wait_for_everyone()
|
||||
|
||||
@@ -223,7 +228,7 @@ class LtxvTrainer:
|
||||
if self._lr_scheduler is not None:
|
||||
self._lr_scheduler.step()
|
||||
|
||||
# Run validation if needed
|
||||
# Run validation if needed (handles DDP/FSDP work distribution internally)
|
||||
if (
|
||||
cfg.validation.interval
|
||||
and self._global_step > 0
|
||||
@@ -231,7 +236,7 @@ class LtxvTrainer:
|
||||
and is_optimization_step
|
||||
):
|
||||
with self._offloaded_optimizer_state():
|
||||
sampled_videos_paths = self._run_distributed_validation(progress)
|
||||
sampled_videos_paths = self._run_validation(progress)
|
||||
|
||||
# Save checkpoint if needed
|
||||
if (
|
||||
@@ -376,111 +381,39 @@ class LtxvTrainer:
|
||||
perturbations=None,
|
||||
)
|
||||
|
||||
# Use strategy to compute loss
|
||||
# Use strategy to compute loss (returns per-element [B,] for sigma-bucket tracking)
|
||||
loss = self._training_strategy.compute_loss(video_pred, audio_pred, model_inputs)
|
||||
sigma = model_inputs.video.sigma.detach() if model_inputs.video.enabled else model_inputs.audio.sigma.detach()
|
||||
|
||||
# Sigma comes from whichever modality is generated (video preferred, else audio).
|
||||
if model_inputs.video is not None and model_inputs.video.enabled:
|
||||
sigma = model_inputs.video.sigma.detach()
|
||||
else:
|
||||
sigma = model_inputs.audio.sigma.detach()
|
||||
|
||||
return TrainingStepOutput(loss=loss, sigma=sigma)
|
||||
|
||||
@free_gpu_memory_context(after=True)
|
||||
def _load_text_encoder_and_cache_embeddings(self) -> list[CachedPromptEmbeddings] | None:
|
||||
"""Load text encoder + embeddings processor, compute and cache validation embeddings."""
|
||||
def _load_models(self) -> None:
|
||||
"""Load the transformer and embeddings processor for training."""
|
||||
logger.debug("Loading transformer...")
|
||||
self._transformer = load_transformer(
|
||||
checkpoint_path=self._config.model.model_path,
|
||||
device="cpu",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
|
||||
# This method:
|
||||
# 1. Loads the pure Gemma text encoder on GPU
|
||||
# 2. Loads the embeddings processor (feature extractor + connectors)
|
||||
# 3. If validation prompts are configured, computes and caches their embeddings
|
||||
# 4. Unloads the Gemma model entirely, keeps the embeddings processor for training
|
||||
|
||||
# Load text encoder (pure Gemma LLM) on GPU — LOCAL_RANK before Accelerator exists
|
||||
# DDP-safe: LOCAL_RANK is set by accelerate before trainer init. Loading on bare
|
||||
# "cuda" would resolve to cuda:0 on every rank and crash with a device mismatch.
|
||||
local_rank = int(os.environ.get("LOCAL_RANK", "0"))
|
||||
init_device = torch.device(f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
logger.debug("Loading text encoder...")
|
||||
text_encoder = load_text_encoder(
|
||||
gemma_model_path=self._config.model.text_encoder_path,
|
||||
device=init_device,
|
||||
dtype=torch.bfloat16,
|
||||
load_in_8bit=self._config.acceleration.load_text_encoder_in_8bit,
|
||||
)
|
||||
|
||||
# Load embeddings processor (feature extractor + connectors)
|
||||
logger.debug("Loading embeddings processor...")
|
||||
self._embeddings_processor = load_embeddings_processor(
|
||||
checkpoint_path=self._config.model.model_path,
|
||||
device=init_device,
|
||||
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:
|
||||
pos_hs, pos_mask = text_encoder.encode(prompt)
|
||||
pos_out = self._embeddings_processor.process_hidden_states(pos_hs, pos_mask)
|
||||
|
||||
neg_hs, neg_mask = text_encoder.encode(self._config.validation.negative_prompt)
|
||||
neg_out = self._embeddings_processor.process_hidden_states(neg_hs, neg_mask)
|
||||
|
||||
cached_embeddings.append(
|
||||
CachedPromptEmbeddings(
|
||||
video_context_positive=pos_out.video_encoding.cpu(),
|
||||
audio_context_positive=pos_out.audio_encoding.cpu(),
|
||||
video_context_negative=neg_out.video_encoding.cpu(),
|
||||
audio_context_negative=(
|
||||
neg_out.audio_encoding.cpu() if neg_out.audio_encoding is not None else None
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
# Unload Gemma model and feature extractor, keep only connectors for training
|
||||
del text_encoder
|
||||
self._embeddings_processor.feature_extractor = None
|
||||
|
||||
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._embeddings_processor 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)
|
||||
|
||||
@@ -494,16 +427,7 @@ class LtxvTrainer:
|
||||
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."""
|
||||
@@ -692,13 +616,6 @@ class LtxvTrainer:
|
||||
|
||||
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)
|
||||
|
||||
@@ -740,7 +657,7 @@ class LtxvTrainer:
|
||||
"""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()
|
||||
data_sources = self._config.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)}")
|
||||
@@ -785,41 +702,6 @@ class LtxvTrainer:
|
||||
# noinspection PyTypeChecker
|
||||
self._optimizer, self._lr_scheduler = self._accelerator.prepare(optimizer, lr_scheduler)
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _offloaded_optimizer_state(self) -> Iterator[None]:
|
||||
"""Context manager that offloads optimizer state to CPU during validation.
|
||||
Opt-in via `acceleration.offload_optimizer_during_validation`. Frees VRAM for
|
||||
validation video generation when optimizer state is large (e.g. full fine-tune
|
||||
AdamW, high-rank LoRA). No-op for FSDP (sharded state -- manual `.cpu()` breaks
|
||||
metadata).
|
||||
"""
|
||||
enabled = (
|
||||
self._config.acceleration.offload_optimizer_during_validation
|
||||
and self._accelerator.distributed_type != DistributedType.FSDP
|
||||
)
|
||||
|
||||
# Track exactly which tensors we move so we don't promote ones that were
|
||||
# intentionally on CPU (e.g. AdamW's `step` scalar on recent PyTorch).
|
||||
offloaded: list[tuple[dict, str]] = []
|
||||
if enabled:
|
||||
offloaded_bytes = 0
|
||||
for state in self._optimizer.state.values():
|
||||
for k, v in state.items():
|
||||
if isinstance(v, torch.Tensor) and v.is_cuda:
|
||||
offloaded.append((state, k))
|
||||
offloaded_bytes += v.nbytes
|
||||
if offloaded:
|
||||
logger.info(f"Offloading optimizer state to CPU ({offloaded_bytes / 1e9:.1f} GB)")
|
||||
for state, k in offloaded:
|
||||
state[k] = state[k].cpu()
|
||||
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
device = self._accelerator.device
|
||||
for state, k in offloaded:
|
||||
state[k] = state[k].to(device)
|
||||
|
||||
def _create_scheduler(self, optimizer: torch.optim.Optimizer) -> LRScheduler | None:
|
||||
"""Create learning rate scheduler based on config."""
|
||||
scheduler_type = self._config.optimization.scheduler_type
|
||||
@@ -920,164 +802,104 @@ class LtxvTrainer:
|
||||
"Monitor training stability and consider disabling quantization if issues arise."
|
||||
)
|
||||
|
||||
def _run_distributed_validation(self, progress: TrainingProgress) -> list[Path]:
|
||||
"""Run validation across all ranks and log gathered results on rank 0.
|
||||
Each rank generates only its assigned subset of prompts (see `_sample_videos`),
|
||||
so all GPUs stay busy and no rank idles long enough to trigger NCCL timeouts.
|
||||
Paths are gathered across ranks so rank 0 has the full list for W&B logging.
|
||||
@contextlib.contextmanager
|
||||
def _offloaded_optimizer_state(self) -> Iterator[None]:
|
||||
"""Context manager that offloads optimizer state to CPU during validation.
|
||||
Opt-in via `acceleration.offload_optimizer_during_validation`. Frees VRAM for
|
||||
validation video generation when optimizer state is large (e.g. full fine-tune
|
||||
AdamW, high-rank LoRA). No-op for FSDP (sharded state -- manual `.cpu()` breaks
|
||||
metadata).
|
||||
"""
|
||||
enabled = (
|
||||
self._config.acceleration.offload_optimizer_during_validation
|
||||
and self._accelerator.distributed_type != DistributedType.FSDP
|
||||
)
|
||||
|
||||
# Track exactly which tensors we move so we don't promote ones that were
|
||||
# intentionally on CPU (e.g. AdamW's `step` scalar on recent PyTorch).
|
||||
offloaded: list[tuple[dict, str]] = []
|
||||
if enabled:
|
||||
offloaded_bytes = 0
|
||||
for state in self._optimizer.state.values():
|
||||
for k, v in state.items():
|
||||
if isinstance(v, torch.Tensor) and v.is_cuda:
|
||||
offloaded.append((state, k))
|
||||
offloaded_bytes += v.nbytes
|
||||
if offloaded:
|
||||
logger.info(f"Offloading optimizer state to CPU ({offloaded_bytes / 1e9:.1f} GB)")
|
||||
for state, k in offloaded:
|
||||
state[k] = state[k].cpu()
|
||||
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
device = self._accelerator.device
|
||||
for state, k in offloaded:
|
||||
state[k] = state[k].to(device)
|
||||
|
||||
def _run_validation(self, progress: TrainingProgress) -> list[Path]:
|
||||
"""Run distributed validation by delegating to the ValidationRunner.
|
||||
Each rank generates its assigned subset of validation samples (round-robin by
|
||||
`process_index`/`num_processes`), so all GPUs stay busy and no rank idles long
|
||||
enough to trigger NCCL timeouts. Paths are gathered across ranks so rank 0 has
|
||||
the full list for W&B logging.
|
||||
Under FSDP with multiple processes, ranks pad with extra generate passes
|
||||
(same sample, no disk write) so every rank runs the same number of forwards --
|
||||
avoids collective mismatch.
|
||||
Note: Multi-node training requires a shared filesystem so rank 0 can read
|
||||
videos written by other ranks.
|
||||
"""
|
||||
sampled = self._sample_videos(progress)
|
||||
self._optimizer.zero_grad(set_to_none=True)
|
||||
free_gpu_memory()
|
||||
|
||||
if self._accelerator.num_processes > 1:
|
||||
# gather_object returns a flat list from all ranks
|
||||
num_samples = len(self._config.validation.samples)
|
||||
if num_samples == 0:
|
||||
return []
|
||||
|
||||
rank = self._accelerator.process_index
|
||||
world_size = self._accelerator.num_processes
|
||||
|
||||
rank_indices = list(range(rank, num_samples, world_size))
|
||||
work_items: list[tuple[int, bool]] = [(i, True) for i in rank_indices]
|
||||
if self._accelerator.distributed_type == DistributedType.FSDP and world_size > 1:
|
||||
# FSDP forwards run collective ops; pad short ranks with no-save duplicates so
|
||||
# every rank executes the same number of forwards. A rank with empty
|
||||
# rank_indices (world_size > num_samples) still pads with sample 0 to stay in
|
||||
# sync with the others.
|
||||
max_per_rank = math.ceil(num_samples / world_size)
|
||||
pad_seed = rank_indices[-1] if rank_indices else 0
|
||||
work_items += [(pad_seed, False)] * (max_per_rank - len(work_items))
|
||||
|
||||
# W&B logging is handled by the trainer (after gathering across ranks),
|
||||
# so we always pass wandb_run=None to the runner.
|
||||
sampled = self._validation_runner.run(
|
||||
transformer=self._transformer,
|
||||
step=self._global_step,
|
||||
output_dir=Path(self._config.output_dir),
|
||||
device=self._accelerator.device,
|
||||
progress=progress,
|
||||
wandb_run=None,
|
||||
work_items=work_items,
|
||||
)
|
||||
|
||||
if world_size > 1:
|
||||
sampled = sorted(gather_object(sampled), key=lambda x: x[0])
|
||||
|
||||
paths = [p for _, p in sampled]
|
||||
|
||||
if self._accelerator.is_main_process and paths:
|
||||
self._log_validation_samples(paths, self._config.validation.prompts)
|
||||
if (
|
||||
self._accelerator.is_main_process
|
||||
and paths
|
||||
and self._config.wandb.log_validation_videos
|
||||
and self._wandb_run is not None
|
||||
):
|
||||
self._validation_runner.log_to_wandb(self._wandb_run, paths, self._global_step)
|
||||
|
||||
# Non-main ranks must not reach checkpoint collectives while main is still logging to W&B.
|
||||
self._accelerator.wait_for_everyone()
|
||||
|
||||
return paths
|
||||
|
||||
# Note: Use @torch.no_grad() instead of @torch.inference_mode() to avoid FSDP inplace update errors after validation
|
||||
@torch.no_grad()
|
||||
@free_gpu_memory_context(after=True)
|
||||
def _sample_videos(self, progress: TrainingProgress) -> list[tuple[int, Path]]:
|
||||
"""Run validation by generating videos from this rank's share of the validation prompts.
|
||||
Prompts are split round-robin across ranks via `process_index` / `num_processes`,
|
||||
which collapses to "all prompts" when running on a single GPU. Returns
|
||||
(prompt_idx, path) tuples so the caller can reconstruct global order without
|
||||
relying on filename conventions.
|
||||
Under FSDP with multiple processes, ranks pad with extra generate passes (same prompt,
|
||||
no disk write) so every rank runs the same number of forwards — avoids collective mismatch.
|
||||
"""
|
||||
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
|
||||
|
||||
# Zero gradients and free GPU memory to reclaim memory before validation sampling
|
||||
self._optimizer.zero_grad(set_to_none=True)
|
||||
free_gpu_memory()
|
||||
|
||||
prompts = self._config.validation.prompts
|
||||
rank = self._accelerator.process_index
|
||||
world_size = self._accelerator.num_processes
|
||||
rank_indices = list(range(rank, len(prompts), world_size))
|
||||
|
||||
# FSDP: every rank must run the same number of forwards; pad with duplicate generates (no save).
|
||||
work: list[tuple[int, bool]] = [(i, True) for i in rank_indices]
|
||||
if self._accelerator.distributed_type == DistributedType.FSDP and world_size > 1:
|
||||
max_per_rank = math.ceil(len(prompts) / world_size)
|
||||
pad_seed = rank_indices[-1] if rank_indices else 0
|
||||
work += [(pad_seed, False)] * (max_per_rank - len(work))
|
||||
|
||||
sampling_ctx = progress.start_sampling(
|
||||
num_prompts=len(work),
|
||||
num_steps=inference_steps,
|
||||
)
|
||||
|
||||
# Create a 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)
|
||||
|
||||
results: list[tuple[int, Path]] = []
|
||||
width, height, num_frames = self._config.validation.video_dims
|
||||
|
||||
for local_i, (prompt_idx, save_output) in enumerate(work):
|
||||
prompt = prompts[prompt_idx]
|
||||
sampling_ctx.start_video(local_i)
|
||||
|
||||
# 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,
|
||||
reference_downscale_factor=self._config.validation.reference_downscale_factor,
|
||||
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,
|
||||
)
|
||||
|
||||
if not save_output:
|
||||
continue
|
||||
|
||||
# Save output (image for single frame, video otherwise)
|
||||
ext = "png" if num_frames == 1 else "mp4"
|
||||
output_path = output_dir / f"step_{self._global_step:06d}_{prompt_idx + 1:02d}.{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_sampling_rate if audio is not None else None,
|
||||
)
|
||||
results.append((prompt_idx, output_path))
|
||||
|
||||
# Clean up progress tasks
|
||||
sampling_ctx.cleanup()
|
||||
|
||||
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 results
|
||||
|
||||
@staticmethod
|
||||
def _log_training_stats(stats: TrainingStats) -> None:
|
||||
"""Log training statistics."""
|
||||
@@ -1166,8 +988,8 @@ class LtxvTrainer:
|
||||
def _save_training_state(self, save_dir: Path) -> None:
|
||||
"""Save training state alongside checkpoint for resume.
|
||||
Respects checkpoints.save_training_state config:
|
||||
- "full": optimizer + scheduler + RNG + step + wandb_run_id
|
||||
- "minimal": scheduler + RNG + step + wandb_run_id
|
||||
- "full": optimizer + scheduler + RNG + step
|
||||
- "minimal": scheduler + RNG + step only
|
||||
- "off": skip entirely
|
||||
"""
|
||||
if not IS_MAIN_PROCESS:
|
||||
@@ -1270,7 +1092,7 @@ class LtxvTrainer:
|
||||
logger.info(f"💾 Training configuration saved to: {config_path.relative_to(self._config.output_dir)}")
|
||||
|
||||
def _init_wandb(self, resume_run_id: str | None = None) -> None:
|
||||
"""Initialize Weights & Biases run."""
|
||||
"""Initialize Weights & Biases run, resuming an existing run if its id is provided."""
|
||||
if not self._config.wandb.enabled or not IS_MAIN_PROCESS:
|
||||
self._wandb_run = None
|
||||
return
|
||||
@@ -1285,7 +1107,7 @@ class LtxvTrainer:
|
||||
}
|
||||
if resume_run_id is not None:
|
||||
init_kwargs["id"] = resume_run_id
|
||||
init_kwargs["resume"] = "allow"
|
||||
init_kwargs["resume"] = "must"
|
||||
run = wandb.init(**init_kwargs)
|
||||
self._wandb_run = run
|
||||
|
||||
@@ -1293,22 +1115,3 @@ class LtxvTrainer:
|
||||
"""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")
|
||||
|
||||
if is_image:
|
||||
samples = [
|
||||
wandb.Image(str(path), caption=prompt) for path, prompt in zip(sample_paths, prompts, strict=True)
|
||||
]
|
||||
else:
|
||||
samples = [
|
||||
wandb.Video(str(path), caption=prompt, format=path.suffix.lower().lstrip("."))
|
||||
for path, prompt in zip(sample_paths, prompts, strict=True)
|
||||
]
|
||||
self._wandb_run.log({"validation_samples": samples}, step=self._global_step)
|
||||
|
||||
@@ -1,10 +1,13 @@
|
||||
"""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)
|
||||
- Text-to-video training (standard generation, optionally with audio) [DEPRECATED]
|
||||
- Video-to-video training (IC-LoRA mode with reference videos) [DEPRECATED]
|
||||
- Flexible training (unified conditioning framework supporting all scenarios) [RECOMMENDED]
|
||||
Each strategy encapsulates the specific logic for preparing model inputs and computing loss.
|
||||
"""
|
||||
|
||||
import warnings
|
||||
|
||||
from ltx_trainer import logger
|
||||
from ltx_trainer.training_strategies.base_strategy import (
|
||||
DEFAULT_FPS,
|
||||
@@ -13,15 +16,18 @@ from ltx_trainer.training_strategies.base_strategy import (
|
||||
TrainingStrategy,
|
||||
TrainingStrategyConfigBase,
|
||||
)
|
||||
from ltx_trainer.training_strategies.flexible import FlexibleStrategy, FlexibleStrategyConfig
|
||||
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
|
||||
TrainingStrategyConfig = TextToVideoConfig | VideoToVideoConfig | FlexibleStrategyConfig
|
||||
|
||||
__all__ = [
|
||||
"DEFAULT_FPS",
|
||||
"VIDEO_SCALE_FACTORS",
|
||||
"FlexibleStrategy",
|
||||
"FlexibleStrategyConfig",
|
||||
"ModelInputs",
|
||||
"TextToVideoConfig",
|
||||
"TextToVideoStrategy",
|
||||
@@ -43,16 +49,42 @@ def get_training_strategy(config: TrainingStrategyConfig) -> TrainingStrategy:
|
||||
The appropriate training strategy instance
|
||||
Raises:
|
||||
ValueError: If strategy name is not supported
|
||||
Note:
|
||||
The `text_to_video` and `video_to_video` strategies are deprecated.
|
||||
Please use the `flexible` strategy instead.
|
||||
"""
|
||||
|
||||
match config:
|
||||
case TextToVideoConfig():
|
||||
warnings.warn(
|
||||
"The 'text_to_video' training strategy is deprecated and will be removed "
|
||||
"in a future version. Please migrate to the 'flexible' strategy. "
|
||||
"See the migration guide in the documentation.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
strategy = TextToVideoStrategy(config)
|
||||
case VideoToVideoConfig():
|
||||
warnings.warn(
|
||||
"The 'video_to_video' training strategy is deprecated and will be removed "
|
||||
"in a future version. Please migrate to the 'flexible' strategy. "
|
||||
"See the migration guide in the documentation.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
strategy = VideoToVideoStrategy(config)
|
||||
case FlexibleStrategyConfig():
|
||||
strategy = FlexibleStrategy(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 🔇)"
|
||||
# Determine audio mode for logging
|
||||
if hasattr(config, "with_audio"):
|
||||
audio_mode = "(audio enabled 🔈)" if config.with_audio else "(audio disabled 🔇)"
|
||||
elif hasattr(config, "audio") and config.audio is not None:
|
||||
audio_mode = "(audio enabled 🔈)"
|
||||
else:
|
||||
audio_mode = "(audio disabled 🔇)"
|
||||
|
||||
logger.debug(f"🎯 Using {strategy.__class__.__name__} training strategy {audio_mode}")
|
||||
return strategy
|
||||
|
||||
@@ -34,29 +34,36 @@ class TrainingStrategyConfigBase(BaseModel):
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
name: Literal["text_to_video", "video_to_video"] = Field(
|
||||
name: Literal["text_to_video", "video_to_video", "flexible"] = Field(
|
||||
description="Unique name identifying the training strategy type"
|
||||
)
|
||||
|
||||
@abstractmethod
|
||||
def get_data_sources(self) -> dict[str, str]:
|
||||
"""Get the required data sources for this strategy.
|
||||
Returns a mapping of directory name (relative to ``preprocessed_data_root``)
|
||||
to the dataset output key under which that directory's contents are exposed.
|
||||
This is the single source of truth for which directories the strategy needs:
|
||||
it drives both dataset wiring (in the trainer) and existence validation
|
||||
(in ``LtxTrainerConfig``).
|
||||
"""
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelInputs:
|
||||
"""Container for model inputs using the Modality-based interface."""
|
||||
|
||||
video: Modality
|
||||
video: Modality | None
|
||||
audio: Modality | None
|
||||
|
||||
# Training targets (for loss computation)
|
||||
video_targets: Tensor
|
||||
video_targets: Tensor | None
|
||||
audio_targets: Tensor | None
|
||||
|
||||
# Masks for loss computation
|
||||
video_loss_mask: Tensor # Boolean mask: True = compute loss for this token
|
||||
# Masks for loss computation (True = compute loss for this token)
|
||||
video_loss_mask: Tensor | None
|
||||
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.
|
||||
@@ -73,24 +80,6 @@ class TrainingStrategy(ABC):
|
||||
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,
|
||||
@@ -145,7 +134,6 @@ class TrainingStrategy(ABC):
|
||||
batch_size: int,
|
||||
fps: float,
|
||||
device: torch.device,
|
||||
dtype: torch.dtype,
|
||||
) -> Tensor:
|
||||
"""Generate video position embeddings using ltx_core's native implementation.
|
||||
Args:
|
||||
@@ -155,9 +143,8 @@ class TrainingStrategy(ABC):
|
||||
batch_size: Batch size
|
||||
fps: Frames per second
|
||||
device: Target device
|
||||
dtype: Target dtype
|
||||
Returns:
|
||||
Position tensor of shape [B, 3, seq_len, 2]
|
||||
Position tensor of shape [B, 3, seq_len, 2] (float32)
|
||||
"""
|
||||
latent_coords = self._video_patchifier.get_patch_grid_bounds(
|
||||
output_shape=VideoLatentShape(
|
||||
@@ -175,7 +162,7 @@ class TrainingStrategy(ABC):
|
||||
latent_coords=latent_coords,
|
||||
scale_factors=VIDEO_SCALE_FACTORS,
|
||||
causal_fix=True,
|
||||
).to(dtype)
|
||||
).float()
|
||||
|
||||
# Scale temporal dimension by 1/fps to get time in seconds
|
||||
pixel_coords[:, 0, ...] = pixel_coords[:, 0, ...] / fps
|
||||
@@ -187,14 +174,12 @@ class TrainingStrategy(ABC):
|
||||
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:
|
||||
@@ -204,7 +189,7 @@ class TrainingStrategy(ABC):
|
||||
"""
|
||||
mel_bins = 16
|
||||
|
||||
latent_coords = self._audio_patchifier.get_patch_grid_bounds(
|
||||
return self._audio_patchifier.get_patch_grid_bounds(
|
||||
output_shape=AudioLatentShape(
|
||||
frames=num_time_steps,
|
||||
mel_bins=mel_bins,
|
||||
@@ -214,8 +199,6 @@ class TrainingStrategy(ABC):
|
||||
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.
|
||||
|
||||
@@ -0,0 +1,718 @@
|
||||
"""Flexible training strategy for a unified conditioning framework.
|
||||
This strategy implements the Unified Conditioning Framework that supports:
|
||||
- Simple fine-tuning with text conditioning (text-to-video/audio)
|
||||
- Intrinsic conditioning (first_frame, prefix, suffix, spatial_crop, mask)
|
||||
- Extrinsic conditioning (concatenation-based, IC-LoRA style)
|
||||
The flexible strategy replaces TextToVideoStrategy and VideoToVideoStrategy by expressing
|
||||
all conditioning scenarios through configuration rather than code.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Annotated, Any, Literal, Union
|
||||
|
||||
import torch
|
||||
from pydantic import BaseModel, ConfigDict, Field, model_validator
|
||||
from torch import Tensor
|
||||
|
||||
from ltx_core.model.transformer.modality import Modality
|
||||
from ltx_trainer.timestep_samplers import TimestepSampler
|
||||
from ltx_trainer.training_strategies.base_strategy import (
|
||||
DEFAULT_FPS,
|
||||
VIDEO_SCALE_FACTORS,
|
||||
ModelInputs,
|
||||
TrainingStrategy,
|
||||
TrainingStrategyConfigBase,
|
||||
)
|
||||
|
||||
# =============================================================================
|
||||
# Configuration Classes
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class IntrinsicConditionBase(BaseModel):
|
||||
"""Base for intrinsic conditioning — tokens get clean latents, timestep=0, no loss."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
probability: float = Field(
|
||||
default=1.0,
|
||||
ge=0.0,
|
||||
le=1.0,
|
||||
description="Probability of applying this condition",
|
||||
)
|
||||
|
||||
|
||||
class FirstFrameConditionConfig(IntrinsicConditionBase):
|
||||
"""First frame conditioning — frame 0 is clean, excluded from loss."""
|
||||
|
||||
type: Literal["first_frame"] = "first_frame"
|
||||
|
||||
|
||||
class PrefixConditionConfig(IntrinsicConditionBase):
|
||||
"""Prefix conditioning — first N temporal units are clean, excluded from loss."""
|
||||
|
||||
type: Literal["prefix"] = "prefix"
|
||||
temporal_boundary: int = Field(
|
||||
...,
|
||||
ge=1,
|
||||
description="Number of temporal units for prefix region. "
|
||||
"For video: number of latent frames. For audio: number of audio latent timesteps.",
|
||||
)
|
||||
|
||||
|
||||
class SuffixConditionConfig(IntrinsicConditionBase):
|
||||
"""Suffix conditioning — last N temporal units are clean, excluded from loss."""
|
||||
|
||||
type: Literal["suffix"] = "suffix"
|
||||
temporal_boundary: int = Field(
|
||||
...,
|
||||
ge=1,
|
||||
description="Number of temporal units for suffix region. "
|
||||
"For video: number of latent frames. For audio: number of audio latent timesteps.",
|
||||
)
|
||||
|
||||
|
||||
class SpatialCropConditionConfig(IntrinsicConditionBase):
|
||||
"""Spatial crop conditioning — rectangular pixel region is clean, excluded from loss."""
|
||||
|
||||
type: Literal["spatial_crop"] = "spatial_crop"
|
||||
spatial_region: tuple[int, int, int, int] = Field(
|
||||
...,
|
||||
description="Spatial crop region as (y1, x1, y2, x2) in pixel coordinates",
|
||||
)
|
||||
|
||||
|
||||
class MaskConditionConfig(IntrinsicConditionBase):
|
||||
"""Mask conditioning — per-sample binary mask determines conditioning tokens."""
|
||||
|
||||
type: Literal["mask"] = "mask"
|
||||
mask_dir: str = Field(
|
||||
...,
|
||||
description="Directory containing per-sample masks",
|
||||
)
|
||||
|
||||
|
||||
class ReferenceConditionConfig(BaseModel):
|
||||
"""Reference conditioning (IC-LoRA style concatenation).
|
||||
External reference latents are concatenated to the target sequence.
|
||||
Reference tokens are clean (timestep=0), excluded from loss, and
|
||||
participate in bidirectional self-attention.
|
||||
"""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
type: Literal["reference"] = "reference"
|
||||
latents_dir: str = Field(..., description="Directory for reference latents")
|
||||
probability: float = Field(default=1.0, ge=0.0, le=1.0, description="Probability of applying this condition")
|
||||
|
||||
|
||||
# Discriminated union for condition configs
|
||||
ConditionConfig = Annotated[
|
||||
Union[
|
||||
FirstFrameConditionConfig,
|
||||
PrefixConditionConfig,
|
||||
SuffixConditionConfig,
|
||||
SpatialCropConditionConfig,
|
||||
MaskConditionConfig,
|
||||
ReferenceConditionConfig,
|
||||
],
|
||||
Field(discriminator="type"),
|
||||
]
|
||||
|
||||
|
||||
class ModalityConfig(BaseModel):
|
||||
"""Configuration for a single modality (video or audio)."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
is_generated: bool = Field(
|
||||
...,
|
||||
description="True = generated modality (denoised, contributes to loss), False = conditioning-only modality",
|
||||
)
|
||||
|
||||
latents_dir: str = Field(
|
||||
...,
|
||||
description="Directory for latents",
|
||||
)
|
||||
|
||||
conditions: list[ConditionConfig] = Field(
|
||||
default_factory=list,
|
||||
description="List of conditions (e.g. first_frame, prefix, reference). Text conditioning is always applied.",
|
||||
)
|
||||
|
||||
|
||||
class FlexibleStrategyConfig(TrainingStrategyConfigBase):
|
||||
"""Configuration for the flexible training strategy.
|
||||
This strategy supports all conditioning scenarios through configuration:
|
||||
- Text-to-video/audio with simple fine-tuning
|
||||
- Intrinsic conditioning like first-frame, extension, outpainting
|
||||
- Reference conditioning like IC-LoRA (concatenation-based reference)
|
||||
"""
|
||||
|
||||
name: Literal["flexible"] = "flexible"
|
||||
|
||||
video: ModalityConfig | None = Field(
|
||||
default=None,
|
||||
description="Video modality configuration",
|
||||
)
|
||||
|
||||
audio: ModalityConfig | None = Field(
|
||||
default=None,
|
||||
description="Audio modality configuration",
|
||||
)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_at_least_one_generated(self) -> "FlexibleStrategyConfig":
|
||||
"""Ensure at least one modality has is_generated=true."""
|
||||
has_video_target = self.video is not None and self.video.is_generated
|
||||
has_audio_target = self.audio is not None and self.audio.is_generated
|
||||
if not has_video_target and not has_audio_target:
|
||||
raise ValueError("At least one modality must have is_generated=true")
|
||||
return self
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_audio_intrinsic_regions(self) -> "FlexibleStrategyConfig":
|
||||
"""Reject video-only intrinsic regions on the audio modality."""
|
||||
if self.audio is None:
|
||||
return self
|
||||
for cond in self.audio.conditions:
|
||||
if isinstance(cond, (FirstFrameConditionConfig, SpatialCropConditionConfig)):
|
||||
raise ValueError(
|
||||
f"Intrinsic condition '{cond.type}' is not supported for audio. "
|
||||
f"Audio supports: prefix, suffix, mask."
|
||||
)
|
||||
return self
|
||||
|
||||
def get_data_sources(self) -> dict[str, str]:
|
||||
"""Dynamically determine required data sources from config.
|
||||
Returns a mapping of directory name (under ``preprocessed_data_root``) to
|
||||
the dataset output key.
|
||||
"""
|
||||
sources: dict[str, str] = {"conditions": "conditions"}
|
||||
|
||||
if self.video is not None:
|
||||
sources[self.video.latents_dir] = "video_latents"
|
||||
if self.audio is not None:
|
||||
sources[self.audio.latents_dir] = "audio_latents"
|
||||
|
||||
for modality_config in (self.video, self.audio):
|
||||
if modality_config is None:
|
||||
continue
|
||||
for cond in modality_config.conditions:
|
||||
if isinstance(cond, ReferenceConditionConfig):
|
||||
sources[cond.latents_dir] = cond.latents_dir
|
||||
elif isinstance(cond, MaskConditionConfig):
|
||||
sources[cond.mask_dir] = cond.mask_dir
|
||||
|
||||
return sources
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Helper Data Structures
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModalityProcessingResult:
|
||||
"""Result of processing a single modality."""
|
||||
|
||||
modality: Modality
|
||||
targets: Tensor | None
|
||||
loss_mask: Tensor | None
|
||||
|
||||
|
||||
@dataclass
|
||||
class LatentData:
|
||||
"""Loaded and patchified latents with metadata."""
|
||||
|
||||
latents: Tensor # [B, seq_len, C]
|
||||
num_frames: int
|
||||
height: int
|
||||
width: int
|
||||
fps: float
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# FlexibleStrategy Implementation
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class FlexibleStrategy(TrainingStrategy):
|
||||
"""Unified training strategy supporting all conditioning scenarios.
|
||||
This strategy implements the Unified Conditioning Framework, allowing
|
||||
any training scenario to be expressed through configuration.
|
||||
"""
|
||||
|
||||
config: FlexibleStrategyConfig
|
||||
|
||||
def __init__(self, config: FlexibleStrategyConfig):
|
||||
"""Initialize strategy with configuration.
|
||||
Args:
|
||||
config: Flexible strategy configuration
|
||||
"""
|
||||
super().__init__(config)
|
||||
self.config = config
|
||||
self.reference_spatial_scale_factor, self.reference_temporal_scale_factor = (
|
||||
self._infer_reference_scale_factors_from_config()
|
||||
)
|
||||
|
||||
def prepare_training_inputs(
|
||||
self,
|
||||
batch: dict[str, Any],
|
||||
timestep_sampler: TimestepSampler,
|
||||
) -> ModelInputs:
|
||||
"""Prepare training inputs by processing video and audio modalities."""
|
||||
video_result = self._process_modality(self.config.video, batch, "video", timestep_sampler)
|
||||
audio_result = self._process_modality(self.config.audio, batch, "audio", timestep_sampler)
|
||||
|
||||
return ModelInputs(
|
||||
video=video_result.modality if video_result else None,
|
||||
audio=audio_result.modality if audio_result else None,
|
||||
video_targets=video_result.targets if video_result else None,
|
||||
audio_targets=audio_result.targets if audio_result else None,
|
||||
video_loss_mask=video_result.loss_mask if video_result else None,
|
||||
audio_loss_mask=audio_result.loss_mask if audio_result else None,
|
||||
)
|
||||
|
||||
def compute_loss(
|
||||
self,
|
||||
video_pred: Tensor | None,
|
||||
audio_pred: Tensor | None,
|
||||
inputs: ModelInputs,
|
||||
) -> Tensor:
|
||||
"""Compute masked MSE loss for video and audio predictions. Returns [B,]."""
|
||||
total_loss = None
|
||||
|
||||
if video_pred is not None and inputs.video_targets is not None:
|
||||
video_loss = self._compute_modality_loss(
|
||||
pred=video_pred,
|
||||
targets=inputs.video_targets,
|
||||
loss_mask=inputs.video_loss_mask,
|
||||
)
|
||||
total_loss = video_loss
|
||||
|
||||
if audio_pred is not None and inputs.audio_targets is not None:
|
||||
audio_loss = self._compute_modality_loss(
|
||||
pred=audio_pred,
|
||||
targets=inputs.audio_targets,
|
||||
loss_mask=inputs.audio_loss_mask,
|
||||
)
|
||||
total_loss = audio_loss if total_loss is None else total_loss + audio_loss
|
||||
|
||||
if total_loss is None:
|
||||
raise ValueError("No valid predictions and targets provided for loss computation")
|
||||
|
||||
return total_loss
|
||||
|
||||
def get_checkpoint_metadata(self) -> dict[str, Any]:
|
||||
"""Include reference scale factors in checkpoint metadata for inference pipelines."""
|
||||
metadata: dict[str, Any] = {}
|
||||
spatial = self.reference_spatial_scale_factor
|
||||
temporal = self.reference_temporal_scale_factor
|
||||
if spatial is not None and spatial != 1:
|
||||
metadata["reference_spatial_scale_factor"] = spatial
|
||||
metadata["reference_downscale_factor"] = spatial # backward compat
|
||||
if temporal is not None and temporal != 1:
|
||||
metadata["reference_temporal_scale_factor"] = temporal
|
||||
return metadata
|
||||
|
||||
def _infer_reference_scale_factors_from_config(self) -> tuple[int | None, int | None]:
|
||||
"""Infer spatial and temporal scale factors by peeking at one sample pair."""
|
||||
if self.config.video is None:
|
||||
return None, None
|
||||
for cond in self.config.video.conditions:
|
||||
if not isinstance(cond, ReferenceConditionConfig):
|
||||
continue
|
||||
target_dir = Path(self.config.video.latents_dir)
|
||||
ref_dir = Path(cond.latents_dir)
|
||||
for sample_file in target_dir.rglob("*.pt"):
|
||||
ref_file = ref_dir / sample_file.relative_to(target_dir)
|
||||
if not ref_file.exists():
|
||||
continue
|
||||
target_data = torch.load(sample_file, map_location="cpu", weights_only=True)
|
||||
ref_data = torch.load(ref_file, map_location="cpu", weights_only=True)
|
||||
if "height" not in ref_data or "height" not in target_data:
|
||||
continue
|
||||
spatial = self._infer_scale_factor(
|
||||
ref_data["height"],
|
||||
ref_data["width"],
|
||||
target_data["height"],
|
||||
target_data["width"],
|
||||
)
|
||||
temporal = self._infer_temporal_scale_factor(
|
||||
ref_data["num_frames"],
|
||||
target_data["num_frames"],
|
||||
)
|
||||
return spatial, temporal
|
||||
return None, None
|
||||
|
||||
def _process_modality(
|
||||
self,
|
||||
modality_config: ModalityConfig | None,
|
||||
batch: dict[str, Any],
|
||||
modality_key: str,
|
||||
timestep_sampler: TimestepSampler,
|
||||
) -> ModalityProcessingResult | None:
|
||||
"""Process a single modality: load latents, add noise, apply conditions, build Modality."""
|
||||
if modality_config is None:
|
||||
return None
|
||||
|
||||
# Step 1: Load and patchify latents
|
||||
data = self._patchify_latent_data(batch[f"{modality_key}_latents"], modality_key)
|
||||
latents = data.latents
|
||||
|
||||
batch_size, seq_len, _ = latents.shape
|
||||
device = latents.device
|
||||
dtype = latents.dtype
|
||||
|
||||
# Step 2: Get text embeddings
|
||||
conditions = batch["conditions"]
|
||||
prompt_embeds = conditions[f"{modality_key}_prompt_embeds"]
|
||||
prompt_attention_mask = conditions["prompt_attention_mask"]
|
||||
|
||||
# Step 3: Initialize noise, timesteps, and loss mask based on is_generated flag
|
||||
if modality_config.is_generated:
|
||||
noisy_latents, targets, timesteps, loss_mask, sigmas = self._initialize_noisy_target(
|
||||
latents, timestep_sampler
|
||||
)
|
||||
else:
|
||||
# Conditioning modality: keep clean (sigma=0), no loss
|
||||
noisy_latents = latents
|
||||
targets = None
|
||||
timesteps = torch.zeros(batch_size, seq_len, device=device, dtype=dtype)
|
||||
loss_mask = None
|
||||
sigmas = torch.zeros(batch_size, device=device, dtype=dtype)
|
||||
|
||||
# Step 4: Generate positions
|
||||
if modality_key == "video":
|
||||
positions = self._get_video_positions(
|
||||
num_frames=data.num_frames,
|
||||
height=data.height,
|
||||
width=data.width,
|
||||
batch_size=batch_size,
|
||||
fps=data.fps,
|
||||
device=device,
|
||||
)
|
||||
else:
|
||||
positions = self._get_audio_positions(
|
||||
num_time_steps=seq_len,
|
||||
batch_size=batch_size,
|
||||
device=device,
|
||||
)
|
||||
|
||||
# Step 5: Apply conditions (intrinsic first, then extrinsic)
|
||||
for cond in modality_config.conditions:
|
||||
if isinstance(cond, IntrinsicConditionBase) and modality_config.is_generated:
|
||||
noisy_latents, timesteps, loss_mask = self._apply_intrinsic_condition(
|
||||
noisy_latents=noisy_latents,
|
||||
clean_latents=latents,
|
||||
timesteps=timesteps,
|
||||
loss_mask=loss_mask,
|
||||
config=cond,
|
||||
height=data.height,
|
||||
width=data.width,
|
||||
batch=batch,
|
||||
)
|
||||
|
||||
for cond in modality_config.conditions:
|
||||
if isinstance(cond, ReferenceConditionConfig):
|
||||
noisy_latents, positions, timesteps, loss_mask, targets = self._apply_reference_condition(
|
||||
noisy_latents=noisy_latents,
|
||||
positions=positions,
|
||||
timesteps=timesteps,
|
||||
loss_mask=loss_mask,
|
||||
targets=targets,
|
||||
batch=batch,
|
||||
config=cond,
|
||||
modality_key=modality_key,
|
||||
)
|
||||
|
||||
# Step 6: Build Modality
|
||||
modality = Modality(
|
||||
enabled=True,
|
||||
latent=noisy_latents,
|
||||
sigma=sigmas,
|
||||
timesteps=timesteps,
|
||||
positions=positions,
|
||||
context=prompt_embeds,
|
||||
context_mask=prompt_attention_mask,
|
||||
)
|
||||
|
||||
return ModalityProcessingResult(
|
||||
modality=modality,
|
||||
targets=targets,
|
||||
loss_mask=loss_mask,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _initialize_noisy_target(
|
||||
latents: Tensor,
|
||||
timestep_sampler: TimestepSampler,
|
||||
) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
|
||||
"""Add noise to latents and create training targets. Returns (noisy, targets, timesteps, mask, sigmas)."""
|
||||
batch_size, seq_len, _ = latents.shape
|
||||
sigmas = timestep_sampler.sample_for(latents)
|
||||
noise = torch.randn_like(latents)
|
||||
sigmas_expanded = sigmas.view(-1, 1, 1)
|
||||
noisy_latents = (1 - sigmas_expanded) * latents + sigmas_expanded * noise
|
||||
targets = noise - latents # velocity prediction
|
||||
timesteps = sigmas.view(-1, 1).expand(batch_size, seq_len).clone()
|
||||
loss_mask = torch.ones(batch_size, seq_len, dtype=torch.bool, device=latents.device)
|
||||
return noisy_latents, targets, timesteps, loss_mask, sigmas
|
||||
|
||||
def _apply_intrinsic_condition(
|
||||
self,
|
||||
noisy_latents: Tensor,
|
||||
clean_latents: Tensor,
|
||||
timesteps: Tensor,
|
||||
loss_mask: Tensor,
|
||||
config: IntrinsicConditionBase,
|
||||
height: int,
|
||||
width: int,
|
||||
batch: dict[str, Any],
|
||||
) -> tuple[Tensor, Tensor, Tensor]:
|
||||
"""Apply intrinsic conditioning using a binary mask.
|
||||
For each token, the mask value determines conditioning strength:
|
||||
- mask=1: conditioned (clean latent, timestep=0, excluded from loss)
|
||||
- mask=0: generated (noisy latent, original timestep, contributes to loss)
|
||||
The conditioning decision is drawn independently per batch element so the training
|
||||
signal across samples in a batch is i.i.d. -- a single batch-wide draw would
|
||||
correlate gradient updates across the batch.
|
||||
"""
|
||||
batch_size, seq_len, _ = noisy_latents.shape
|
||||
device = noisy_latents.device
|
||||
|
||||
# Per-sample Bernoulli draw -- each element is independently conditioned.
|
||||
apply_per_sample = torch.rand(batch_size, device=device) < config.probability
|
||||
if not apply_per_sample.any():
|
||||
return noisy_latents, timesteps, loss_mask
|
||||
|
||||
if isinstance(config, FirstFrameConditionConfig):
|
||||
mask = self._compute_temporal_mask(batch_size, seq_len, height, width, 1, False, device)
|
||||
elif isinstance(config, PrefixConditionConfig):
|
||||
mask = self._compute_temporal_mask(
|
||||
batch_size, seq_len, height, width, config.temporal_boundary, False, device
|
||||
)
|
||||
elif isinstance(config, SuffixConditionConfig):
|
||||
mask = self._compute_temporal_mask(
|
||||
batch_size, seq_len, height, width, config.temporal_boundary, True, device
|
||||
)
|
||||
elif isinstance(config, SpatialCropConditionConfig):
|
||||
mask = self._compute_spatial_crop_mask(batch_size, seq_len, height, width, config.spatial_region, device)
|
||||
elif isinstance(config, MaskConditionConfig):
|
||||
# Binarize to match inference, which thresholds masks at load time
|
||||
# (validation_runner._load_and_downsample_mask / _load_audio_mask).
|
||||
mask = (batch[config.mask_dir]["mask"].reshape(batch_size, seq_len) > 0.5).float()
|
||||
else:
|
||||
raise ValueError(f"Unknown intrinsic condition type: {type(config).__name__}")
|
||||
|
||||
# Zero the mask for samples the per-sample draw did not select.
|
||||
mask = mask * apply_per_sample.view(-1, 1).to(mask.dtype)
|
||||
|
||||
# Apply binary mask: clean conditioned tokens, noisy generated tokens.
|
||||
m = mask.unsqueeze(-1)
|
||||
noisy_latents = m * clean_latents + (1 - m) * noisy_latents
|
||||
timesteps = (1 - mask) * timesteps
|
||||
loss_mask = loss_mask & (mask == 0)
|
||||
|
||||
return noisy_latents, timesteps, loss_mask
|
||||
|
||||
@staticmethod
|
||||
def _compute_temporal_mask(
|
||||
batch_size: int,
|
||||
seq_len: int,
|
||||
height: int,
|
||||
width: int,
|
||||
num_frames: int,
|
||||
from_end: bool,
|
||||
device: torch.device,
|
||||
) -> Tensor:
|
||||
"""Compute float mask for temporal region (prefix or suffix). Returns [B, seq_len] in {0, 1}."""
|
||||
tokens_per_frame = height * width
|
||||
num_tokens = num_frames * tokens_per_frame
|
||||
mask = torch.zeros(batch_size, seq_len, device=device)
|
||||
if from_end:
|
||||
mask[:, -num_tokens:] = 1.0
|
||||
else:
|
||||
mask[:, :num_tokens] = 1.0
|
||||
return mask
|
||||
|
||||
@staticmethod
|
||||
def _compute_spatial_crop_mask(
|
||||
batch_size: int,
|
||||
seq_len: int,
|
||||
height: int,
|
||||
width: int,
|
||||
region: tuple[int, int, int, int],
|
||||
device: torch.device,
|
||||
) -> Tensor:
|
||||
"""Compute float mask for spatial crop region (y1, x1, y2, x2) in pixel coords.
|
||||
Returns [B, seq_len] in {0, 1}.
|
||||
"""
|
||||
y1, x1, y2, x2 = region
|
||||
num_frames = seq_len // (height * width)
|
||||
|
||||
# Convert pixel to latent coordinates and clamp (per-axis VAE scale factor).
|
||||
def to_latent(v: int, scale: int, max_v: int) -> int:
|
||||
return max(0, min(v // scale, max_v))
|
||||
|
||||
ly1 = to_latent(y1, VIDEO_SCALE_FACTORS.height, height)
|
||||
ly2 = to_latent(y2, VIDEO_SCALE_FACTORS.height, height)
|
||||
lx1 = to_latent(x1, VIDEO_SCALE_FACTORS.width, width)
|
||||
lx2 = to_latent(x2, VIDEO_SCALE_FACTORS.width, width)
|
||||
|
||||
# Create spatial mask and tile across frames
|
||||
spatial_mask = torch.zeros(height, width, device=device)
|
||||
spatial_mask[ly1:ly2, lx1:lx2] = 1.0
|
||||
full_mask = spatial_mask.flatten().repeat(num_frames)
|
||||
|
||||
return full_mask.unsqueeze(0).expand(batch_size, -1)
|
||||
|
||||
def _patchify_latent_data(self, latent_data: dict[str, Any], modality_key: str) -> LatentData:
|
||||
"""Patchify latent data and extract metadata."""
|
||||
latents = latent_data["latents"]
|
||||
|
||||
if modality_key == "video":
|
||||
num_frames = latent_data["num_frames"][0].item()
|
||||
height = latent_data["height"][0].item()
|
||||
width = latent_data["width"][0].item()
|
||||
fps = latent_data.get("fps")
|
||||
fps = fps[0].item() if fps is not None else DEFAULT_FPS
|
||||
latents = self._video_patchifier.patchify(latents)
|
||||
else:
|
||||
num_frames = latent_data.get("num_frames", [latents.shape[2]])[0]
|
||||
if isinstance(num_frames, Tensor):
|
||||
num_frames = num_frames.item()
|
||||
height = 1
|
||||
width = 1
|
||||
fps = 1.0
|
||||
latents = self._audio_patchifier.patchify(latents)
|
||||
|
||||
return LatentData(latents=latents, num_frames=num_frames, height=height, width=width, fps=fps)
|
||||
|
||||
def _apply_reference_condition(
|
||||
self,
|
||||
noisy_latents: Tensor,
|
||||
positions: Tensor,
|
||||
timesteps: Tensor,
|
||||
loss_mask: Tensor | None,
|
||||
targets: Tensor | None,
|
||||
batch: dict[str, Any],
|
||||
config: ReferenceConditionConfig,
|
||||
modality_key: str,
|
||||
) -> tuple[Tensor, Tensor, Tensor, Tensor | None, Tensor | None]:
|
||||
"""Concatenate reference latents to target sequence for reference conditioning (IC-LoRA style).
|
||||
The apply/skip decision is batch-wide (reference conditioning changes the sequence
|
||||
length, so it cannot be applied to only part of a batch) but is drawn from the torch
|
||||
RNG so runs are reproducible under ``torch.manual_seed`` — mirroring the intrinsic
|
||||
per-sample draw rather than Python's unseeded ``random``.
|
||||
"""
|
||||
if torch.rand((), device=noisy_latents.device).item() >= config.probability:
|
||||
return noisy_latents, positions, timesteps, loss_mask, targets
|
||||
|
||||
# Load and patchify condition latents
|
||||
cond = self._patchify_latent_data(batch[config.latents_dir], modality_key)
|
||||
cond_latents = cond.latents
|
||||
|
||||
batch_size, cond_seq_len, _ = cond_latents.shape
|
||||
device = cond_latents.device
|
||||
dtype = cond_latents.dtype
|
||||
|
||||
# Generate condition positions
|
||||
if modality_key == "video":
|
||||
cond_positions = self._get_video_positions(
|
||||
num_frames=cond.num_frames,
|
||||
height=cond.height,
|
||||
width=cond.width,
|
||||
batch_size=batch_size,
|
||||
fps=cond.fps,
|
||||
device=device,
|
||||
)
|
||||
else:
|
||||
cond_positions = self._get_audio_positions(
|
||||
num_time_steps=cond_seq_len,
|
||||
batch_size=batch_size,
|
||||
device=device,
|
||||
)
|
||||
|
||||
# Translate / rescale ref positions into the target's frame (video only).
|
||||
if modality_key == "video":
|
||||
spatial_sf = self.reference_spatial_scale_factor or 1
|
||||
temporal_sf = self.reference_temporal_scale_factor or 1
|
||||
if spatial_sf != 1 or temporal_sf != 1:
|
||||
cond_positions = cond_positions.clone()
|
||||
if temporal_sf != 1:
|
||||
# Ref positions are already at the ref's effective fps (source_fps / S,
|
||||
# stored by process_videos.py). Shift by (S - 1) / target_fps so ref's
|
||||
# last patch aligns with target's last; clamp the causal patch at 0.
|
||||
t_target = positions[:, 0, 0:1, 1:2] # = 1 / target_fps
|
||||
cond_positions[:, 0, ...] = torch.clamp(
|
||||
cond_positions[:, 0, ...] - (temporal_sf - 1) * t_target, min=0
|
||||
)
|
||||
if spatial_sf != 1:
|
||||
cond_positions[:, 1, ...] *= spatial_sf
|
||||
cond_positions[:, 2, ...] *= spatial_sf
|
||||
|
||||
# Condition tokens: clean, timestep=0, no loss
|
||||
cond_timesteps = torch.zeros(batch_size, cond_seq_len, device=device, dtype=dtype)
|
||||
cond_loss_mask = torch.zeros(batch_size, cond_seq_len, dtype=torch.bool, device=device)
|
||||
|
||||
# Concatenate condition and target sequences (condition first, then target)
|
||||
combined_latents = torch.cat([cond_latents, noisy_latents], dim=1)
|
||||
combined_positions = torch.cat([cond_positions, positions], dim=2)
|
||||
combined_timesteps = torch.cat([cond_timesteps, timesteps], dim=1)
|
||||
|
||||
combined_loss_mask = torch.cat([cond_loss_mask, loss_mask], dim=1) if loss_mask is not None else None
|
||||
|
||||
# Targets remain unchanged (only for target portion, not condition portion)
|
||||
|
||||
return combined_latents, combined_positions, combined_timesteps, combined_loss_mask, targets
|
||||
|
||||
@staticmethod
|
||||
def _compute_modality_loss(pred: Tensor, targets: Tensor, loss_mask: Tensor) -> Tensor:
|
||||
"""Compute per-element MSE loss for a single modality. Returns [B,]."""
|
||||
# Slice prediction to match targets length (removes any prepended condition tokens)
|
||||
target_len = targets.shape[1]
|
||||
pred = pred[:, -target_len:, :]
|
||||
mask = loss_mask[:, -target_len:]
|
||||
|
||||
# Compute masked MSE loss, reduce per-element [B,] over (seq, channels)
|
||||
mask_expanded = mask.unsqueeze(-1).float()
|
||||
squared_error = (pred - targets).pow(2)
|
||||
masked_loss = squared_error * mask_expanded
|
||||
return masked_loss.mean(dim=[-2, -1]) / mask_expanded.mean(dim=[-2, -1]).clamp(min=1e-8)
|
||||
|
||||
@staticmethod
|
||||
def _infer_scale_factor(cond_height: int, cond_width: int, target_height: int, target_width: int) -> int:
|
||||
"""Infer spatial scale factor between condition and target resolutions."""
|
||||
if target_height == cond_height and target_width == cond_width:
|
||||
return 1
|
||||
scale_h = target_height // cond_height if cond_height > 0 else 1
|
||||
scale_w = target_width // cond_width if cond_width > 0 else 1
|
||||
if scale_h != scale_w:
|
||||
raise ValueError(
|
||||
f"Non-uniform scale factors between condition and target: height={scale_h}, width={scale_w}. "
|
||||
"Condition and target resolutions must scale uniformly."
|
||||
)
|
||||
return scale_h
|
||||
|
||||
@staticmethod
|
||||
def _infer_temporal_scale_factor(cond_num_frames: int, target_num_frames: int) -> int:
|
||||
"""Infer temporal scale factor between condition and target latent frame counts.
|
||||
The first latent frame encodes a single pixel frame (the VAE's causal structure),
|
||||
so the temporal groups count is (num_frames - 1). The scale factor is the ratio
|
||||
of target groups to condition groups.
|
||||
"""
|
||||
if target_num_frames == cond_num_frames:
|
||||
return 1
|
||||
target_groups = target_num_frames - 1
|
||||
cond_groups = cond_num_frames - 1
|
||||
if cond_groups <= 0 or target_groups <= 0:
|
||||
return 1
|
||||
if target_groups % cond_groups != 0:
|
||||
raise ValueError(
|
||||
f"Target temporal groups ({target_groups}) is not evenly divisible by "
|
||||
f"condition temporal groups ({cond_groups})."
|
||||
)
|
||||
return target_groups // cond_groups
|
||||
@@ -45,6 +45,18 @@ class TextToVideoConfig(TrainingStrategyConfigBase):
|
||||
description="Directory name for audio latents when with_audio is True",
|
||||
)
|
||||
|
||||
def get_data_sources(self) -> 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.with_audio:
|
||||
sources[self.audio_latents_dir] = "audio_latents"
|
||||
return sources
|
||||
|
||||
|
||||
class TextToVideoStrategy(TrainingStrategy):
|
||||
"""Text-to-video training strategy.
|
||||
@@ -64,26 +76,6 @@ class TextToVideoStrategy(TrainingStrategy):
|
||||
"""
|
||||
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],
|
||||
@@ -119,7 +111,6 @@ class TextToVideoStrategy(TrainingStrategy):
|
||||
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(
|
||||
@@ -157,7 +148,6 @@ class TextToVideoStrategy(TrainingStrategy):
|
||||
batch_size=batch_size,
|
||||
fps=fps,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
# Create video Modality
|
||||
@@ -187,7 +177,6 @@ class TextToVideoStrategy(TrainingStrategy):
|
||||
prompt_attention_mask=prompt_attention_mask,
|
||||
batch_size=batch_size,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
return ModelInputs(
|
||||
@@ -207,7 +196,6 @@ class TextToVideoStrategy(TrainingStrategy):
|
||||
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:
|
||||
@@ -217,7 +205,6 @@ class TextToVideoStrategy(TrainingStrategy):
|
||||
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)
|
||||
"""
|
||||
@@ -248,7 +235,6 @@ class TextToVideoStrategy(TrainingStrategy):
|
||||
num_time_steps=audio_seq_len,
|
||||
batch_size=batch_size,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
# Create audio Modality
|
||||
|
||||
@@ -39,6 +39,14 @@ class VideoToVideoConfig(TrainingStrategyConfigBase):
|
||||
description="Directory name for latents of reference videos",
|
||||
)
|
||||
|
||||
def get_data_sources(self) -> dict[str, str]:
|
||||
"""IC-LoRA training requires latents, conditions, and reference latents."""
|
||||
return {
|
||||
"latents": "latents",
|
||||
"conditions": "conditions",
|
||||
self.reference_latents_dir: "ref_latents",
|
||||
}
|
||||
|
||||
|
||||
class VideoToVideoStrategy(TrainingStrategy):
|
||||
"""Video-to-video training strategy for IC-LoRA.
|
||||
@@ -62,14 +70,6 @@ class VideoToVideoStrategy(TrainingStrategy):
|
||||
super().__init__(config)
|
||||
self.reference_downscale_factor = None # Will be inferred from first batch
|
||||
|
||||
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( # noqa: PLR0915
|
||||
self,
|
||||
batch: dict[str, Any],
|
||||
@@ -133,7 +133,6 @@ class VideoToVideoStrategy(TrainingStrategy):
|
||||
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)
|
||||
@@ -164,7 +163,7 @@ class VideoToVideoStrategy(TrainingStrategy):
|
||||
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 for loss computation (velocity prediction) - only for target portion
|
||||
targets = noise - target_latents
|
||||
|
||||
# Concatenate reference (clean) and target (noisy)
|
||||
@@ -181,7 +180,6 @@ class VideoToVideoStrategy(TrainingStrategy):
|
||||
batch_size=batch_size,
|
||||
fps=fps,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
# Scale reference positions to match target coordinate space
|
||||
@@ -200,7 +198,6 @@ class VideoToVideoStrategy(TrainingStrategy):
|
||||
batch_size=batch_size,
|
||||
fps=fps,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
# Concatenate positions along sequence dimension
|
||||
@@ -231,7 +228,6 @@ class VideoToVideoStrategy(TrainingStrategy):
|
||||
audio_targets=None,
|
||||
video_loss_mask=video_loss_mask,
|
||||
audio_loss_mask=None,
|
||||
ref_seq_len=ref_seq_len,
|
||||
)
|
||||
|
||||
def compute_loss(
|
||||
@@ -240,13 +236,11 @@ class VideoToVideoStrategy(TrainingStrategy):
|
||||
_audio_pred: Tensor | None,
|
||||
inputs: ModelInputs,
|
||||
) -> Tensor:
|
||||
"""Compute masked loss only on target portion. Returns [B,]."""
|
||||
# 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 masked loss on target portion only. Returns [B,]."""
|
||||
# Slice prediction to match targets length (removes prepended reference tokens)
|
||||
target_len = inputs.video_targets.shape[1]
|
||||
target_pred = video_pred[:, -target_len:, :]
|
||||
target_loss_mask = inputs.video_loss_mask[:, -target_len:]
|
||||
|
||||
# Compute per-element loss [B,]
|
||||
loss = (target_pred - inputs.video_targets).pow(2)
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,874 +0,0 @@
|
||||
"""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 GemmaTextEncoder
|
||||
from ltx_core.text_encoders.gemma.embeddings_processor import EmbeddingsProcessor
|
||||
|
||||
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]
|
||||
reference_downscale_factor: int = 1 # For IC-LoRA: downscale factor (1 = same resolution, 2 = half resolution)
|
||||
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: "GemmaTextEncoder | None" = None,
|
||||
audio_decoder: "AudioDecoder | None" = None,
|
||||
vocoder: "Vocoder | None" = None,
|
||||
sampling_context: SamplingContext | None = None,
|
||||
embeddings_processor: "EmbeddingsProcessor | 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 (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
|
||||
embeddings_processor: Optional embeddings processor (required if text_encoder provided)
|
||||
"""
|
||||
self._transformer = transformer
|
||||
self._vae_decoder = vae_decoder
|
||||
self._vae_encoder = vae_encoder
|
||||
self._text_encoder = text_encoder
|
||||
self._embeddings_processor = embeddings_processor
|
||||
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]
|
||||
|
||||
# Scale reference positions to match target coordinate space
|
||||
# Position tensor shape: [B, 3, seq_len, 2] where dim 1 is (time, height, width)
|
||||
if config.reference_downscale_factor != 1:
|
||||
ref_positions = ref_positions.clone()
|
||||
ref_positions[:, 1, ...] *= config.reference_downscale_factor # height axis
|
||||
ref_positions[:, 2, ...] *= config.reference_downscale_factor # width axis
|
||||
# Time axis (index 0) remains unchanged
|
||||
|
||||
# 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.
|
||||
When reference_downscale_factor > 1, the reference video is downscaled to a smaller
|
||||
resolution for more efficient inference. The positions will be scaled up later
|
||||
to match the target coordinate space.
|
||||
Args:
|
||||
config: Generation configuration
|
||||
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]
|
||||
scale_factor = config.reference_downscale_factor
|
||||
|
||||
# Target dimensions for reference (scaled down if scale_factor > 1)
|
||||
target_height = config.height // scale_factor
|
||||
target_width = config.width // scale_factor
|
||||
|
||||
# Validate scaled dimensions
|
||||
if target_height % 32 != 0 or target_width % 32 != 0:
|
||||
raise ValueError(
|
||||
f"Scaled reference dimensions ({target_height}x{target_width}) must be divisible by 32. "
|
||||
f"Original: {config.height}x{config.width}, scale_factor: {scale_factor}"
|
||||
)
|
||||
|
||||
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,
|
||||
sigma=sigmas[0].repeat(video_state.latent.shape[0]),
|
||||
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,
|
||||
sigma=sigmas[0].repeat(audio_state.latent.shape[0]),
|
||||
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,
|
||||
sigma=sigma.repeat(video_state.latent.shape[0]),
|
||||
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,
|
||||
sigma=sigma.repeat(audio_state.latent.shape[0]),
|
||||
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")
|
||||
if config.cached_embeddings is None and self._embeddings_processor is None:
|
||||
raise ValueError("embeddings_processor is required when encoding prompts on-the-fly")
|
||||
|
||||
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 + embeddings processor."""
|
||||
self._text_encoder.to(device)
|
||||
self._embeddings_processor.to(device)
|
||||
|
||||
pos_hs, pos_mask = self._text_encoder.encode(config.prompt)
|
||||
pos_out = self._embeddings_processor.process_hidden_states(pos_hs, pos_mask)
|
||||
v_ctx_pos, a_ctx_pos = pos_out.video_encoding, pos_out.audio_encoding
|
||||
|
||||
v_ctx_neg, a_ctx_neg = None, None
|
||||
if config.guidance_scale != 1.0:
|
||||
neg_hs, neg_mask = self._text_encoder.encode(config.negative_prompt)
|
||||
neg_out = self._embeddings_processor.process_hidden_states(neg_hs, neg_mask)
|
||||
v_ctx_neg, a_ctx_neg = neg_out.video_encoding, neg_out.audio_encoding
|
||||
|
||||
# Move the base Gemma model to CPU
|
||||
self._text_encoder.model.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)
|
||||
first_param = next(self._audio_decoder.parameters(), None)
|
||||
decoder_dtype = first_param.dtype if first_param is not None else audio_state.latent.dtype
|
||||
latent = audio_state.latent.to(dtype=decoder_dtype, device=device)
|
||||
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, H1, W1] in [0, 1]
|
||||
right_video: Right video tensor [C, F2, H2, W2] in [0, 1]
|
||||
Returns:
|
||||
Concatenated video tensor [C, max(F1,F2), H2, W1_scaled+W2] in [0, 1]
|
||||
"""
|
||||
left_height, left_width = left_video.shape[2], left_video.shape[3]
|
||||
right_height = right_video.shape[2]
|
||||
|
||||
# Resize left video to match right video's height if needed
|
||||
if left_height != right_height:
|
||||
# Scale width proportionally to maintain aspect ratio
|
||||
scale = right_height / left_height
|
||||
new_width = int(left_width * scale)
|
||||
# Interpolate expects [N, C, H, W], we have [C, F, H, W]
|
||||
# Reshape to [C*F, 1, H, W] -> interpolate -> reshape back
|
||||
c, f, h, w = left_video.shape
|
||||
left_video = left_video.reshape(c * f, 1, h, w)
|
||||
left_video = torch.nn.functional.interpolate(
|
||||
left_video, size=(right_height, new_width), mode="bilinear", align_corners=False
|
||||
)
|
||||
left_video = left_video.reshape(c, f, right_height, new_width)
|
||||
|
||||
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
|
||||
Reference in New Issue
Block a user