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