Automated PR - 2026-01-05
This commit is contained in:
sync-bot
@@ -0,0 +1,58 @@
|
||||
"""Training strategies for different conditioning modes.
|
||||
This package implements the Strategy Pattern to handle different training modes:
|
||||
- Text-to-video training (standard generation, optionally with audio)
|
||||
- Video-to-video training (IC-LoRA mode with reference videos)
|
||||
Each strategy encapsulates the specific logic for preparing model inputs and computing loss.
|
||||
"""
|
||||
|
||||
from ltx_trainer import logger
|
||||
from ltx_trainer.training_strategies.base_strategy import (
|
||||
DEFAULT_FPS,
|
||||
VIDEO_SCALE_FACTORS,
|
||||
ModelInputs,
|
||||
TrainingStrategy,
|
||||
TrainingStrategyConfigBase,
|
||||
)
|
||||
from ltx_trainer.training_strategies.text_to_video import TextToVideoConfig, TextToVideoStrategy
|
||||
from ltx_trainer.training_strategies.video_to_video import VideoToVideoConfig, VideoToVideoStrategy
|
||||
|
||||
# Type alias for all strategy config types
|
||||
TrainingStrategyConfig = TextToVideoConfig | VideoToVideoConfig
|
||||
|
||||
__all__ = [
|
||||
"DEFAULT_FPS",
|
||||
"VIDEO_SCALE_FACTORS",
|
||||
"ModelInputs",
|
||||
"TextToVideoConfig",
|
||||
"TextToVideoStrategy",
|
||||
"TrainingStrategy",
|
||||
"TrainingStrategyConfig",
|
||||
"TrainingStrategyConfigBase",
|
||||
"VideoToVideoConfig",
|
||||
"VideoToVideoStrategy",
|
||||
"get_training_strategy",
|
||||
]
|
||||
|
||||
|
||||
def get_training_strategy(config: TrainingStrategyConfig) -> TrainingStrategy:
|
||||
"""Factory function to create the appropriate training strategy.
|
||||
The strategy is determined by the `name` field in the configuration.
|
||||
Args:
|
||||
config: Strategy-specific configuration with a `name` field
|
||||
Returns:
|
||||
The appropriate training strategy instance
|
||||
Raises:
|
||||
ValueError: If strategy name is not supported
|
||||
"""
|
||||
|
||||
match config:
|
||||
case TextToVideoConfig():
|
||||
strategy = TextToVideoStrategy(config)
|
||||
case VideoToVideoConfig():
|
||||
strategy = VideoToVideoStrategy(config)
|
||||
case _:
|
||||
raise ValueError(f"Unknown training strategy config type: {type(config).__name__}")
|
||||
|
||||
audio_mode = "(audio enabled 🔈)" if getattr(config, "with_audio", False) else "(audio disabled 🔇)"
|
||||
logger.debug(f"🎯 Using {strategy.__class__.__name__} training strategy {audio_mode}")
|
||||
return strategy
|
||||
@@ -0,0 +1,253 @@
|
||||
"""Base class for training strategies.
|
||||
This module defines the abstract base class that all training strategies must implement,
|
||||
along with the base configuration class.
|
||||
"""
|
||||
|
||||
import random
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Literal
|
||||
|
||||
import torch
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
from torch import Tensor
|
||||
|
||||
from ltx_core.components.patchifiers import (
|
||||
AudioPatchifier,
|
||||
VideoLatentPatchifier,
|
||||
get_pixel_coords,
|
||||
)
|
||||
from ltx_core.model.transformer.modality import Modality
|
||||
from ltx_core.types import AudioLatentShape, SpatioTemporalScaleFactors, VideoLatentShape
|
||||
from ltx_trainer.timestep_samplers import TimestepSampler
|
||||
|
||||
# Default frames per second for video missing in the FPS metadata
|
||||
DEFAULT_FPS = 24
|
||||
|
||||
# VAE scale factors for LTX-2
|
||||
VIDEO_SCALE_FACTORS = SpatioTemporalScaleFactors.default()
|
||||
|
||||
|
||||
class TrainingStrategyConfigBase(BaseModel):
|
||||
"""Base configuration class for training strategies.
|
||||
All strategy-specific configuration classes should inherit from this.
|
||||
"""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
name: Literal["text_to_video", "video_to_video"] = Field(
|
||||
description="Unique name identifying the training strategy type"
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelInputs:
|
||||
"""Container for model inputs using the Modality-based interface."""
|
||||
|
||||
video: Modality
|
||||
audio: Modality | None
|
||||
|
||||
# Training targets (for loss computation)
|
||||
video_targets: Tensor
|
||||
audio_targets: Tensor | None
|
||||
|
||||
# Masks for loss computation
|
||||
video_loss_mask: Tensor # Boolean mask: True = compute loss for this token
|
||||
audio_loss_mask: Tensor | None
|
||||
|
||||
# Metadata needed for loss computation in some strategies
|
||||
ref_seq_len: int | None = None # For IC-LoRA: length of reference sequence
|
||||
|
||||
|
||||
class TrainingStrategy(ABC):
|
||||
"""Abstract base class for training strategies.
|
||||
Each strategy encapsulates the logic for a specific training mode,
|
||||
handling input preparation and loss computation.
|
||||
"""
|
||||
|
||||
def __init__(self, config: TrainingStrategyConfigBase):
|
||||
"""Initialize strategy with configuration.
|
||||
Args:
|
||||
config: Strategy-specific configuration
|
||||
"""
|
||||
self.config = config
|
||||
self._video_patchifier = VideoLatentPatchifier(patch_size=1)
|
||||
self._audio_patchifier = AudioPatchifier(patch_size=1)
|
||||
|
||||
@property
|
||||
def requires_audio(self) -> bool:
|
||||
"""Whether this training strategy requires audio components.
|
||||
Override this property in subclasses that support audio training.
|
||||
The trainer uses this to determine whether to load audio VAE and vocoder.
|
||||
Returns:
|
||||
True if audio components should be loaded, False otherwise.
|
||||
"""
|
||||
return False
|
||||
|
||||
@abstractmethod
|
||||
def get_data_sources(self) -> list[str] | dict[str, str]:
|
||||
"""Get the required data sources for this training strategy.
|
||||
Returns:
|
||||
Either a list of data directory names (where output keys match directory names)
|
||||
or a dictionary mapping data directory names to custom output keys for the dataset
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def prepare_training_inputs(
|
||||
self,
|
||||
batch: dict[str, Any],
|
||||
timestep_sampler: TimestepSampler,
|
||||
) -> ModelInputs:
|
||||
"""Prepare training inputs from a raw data batch.
|
||||
Args:
|
||||
batch: Raw batch data from the dataset. Contains:
|
||||
- "latents": Video latent data
|
||||
- "conditions": Text embeddings with keys:
|
||||
- "video_prompt_embeds": Already processed by embedding connectors
|
||||
- "audio_prompt_embeds": Already processed by embedding connectors
|
||||
- "prompt_attention_mask": Attention mask
|
||||
- Additional keys depending on strategy (e.g., "ref_latents" for IC-LoRA)
|
||||
timestep_sampler: Sampler for generating timesteps and noise
|
||||
Returns:
|
||||
ModelInputs containing Modality objects and training targets
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def compute_loss(
|
||||
self,
|
||||
video_pred: Tensor,
|
||||
audio_pred: Tensor | None,
|
||||
inputs: ModelInputs,
|
||||
) -> Tensor:
|
||||
"""Compute the training loss.
|
||||
Args:
|
||||
video_pred: Video prediction from the transformer model
|
||||
audio_pred: Audio prediction from the transformer model (None for video-only)
|
||||
inputs: The prepared model inputs containing targets and masks
|
||||
Returns:
|
||||
Scalar loss tensor
|
||||
"""
|
||||
|
||||
def _get_video_positions(
|
||||
self,
|
||||
num_frames: int,
|
||||
height: int,
|
||||
width: int,
|
||||
batch_size: int,
|
||||
fps: float,
|
||||
device: torch.device,
|
||||
dtype: torch.dtype,
|
||||
) -> Tensor:
|
||||
"""Generate video position embeddings using ltx_core's native implementation.
|
||||
Args:
|
||||
num_frames: Number of latent frames
|
||||
height: Latent height
|
||||
width: Latent width
|
||||
batch_size: Batch size
|
||||
fps: Frames per second
|
||||
device: Target device
|
||||
dtype: Target dtype
|
||||
Returns:
|
||||
Position tensor of shape [B, 3, seq_len, 2]
|
||||
"""
|
||||
latent_coords = self._video_patchifier.get_patch_grid_bounds(
|
||||
output_shape=VideoLatentShape(
|
||||
frames=num_frames,
|
||||
height=height,
|
||||
width=width,
|
||||
batch=batch_size,
|
||||
channels=128, # Video latent channels
|
||||
),
|
||||
device=device,
|
||||
)
|
||||
|
||||
# Convert latent coords to pixel coords with causal fix
|
||||
pixel_coords = get_pixel_coords(
|
||||
latent_coords=latent_coords,
|
||||
scale_factors=VIDEO_SCALE_FACTORS,
|
||||
causal_fix=True,
|
||||
).to(dtype)
|
||||
|
||||
# Scale temporal dimension by 1/fps to get time in seconds
|
||||
pixel_coords[:, 0, ...] = pixel_coords[:, 0, ...] / fps
|
||||
|
||||
return pixel_coords
|
||||
|
||||
def _get_audio_positions(
|
||||
self,
|
||||
num_time_steps: int,
|
||||
batch_size: int,
|
||||
device: torch.device,
|
||||
dtype: torch.dtype,
|
||||
) -> Tensor:
|
||||
"""Generate audio position embeddings using ltx_core's native implementation.
|
||||
Args:
|
||||
num_time_steps: Number of audio time steps (T, not T*mel_bins)
|
||||
batch_size: Batch size
|
||||
device: Target device
|
||||
dtype: Target dtype
|
||||
Returns:
|
||||
Position tensor of shape [B, 1, num_time_steps, 2]
|
||||
Note:
|
||||
Audio latents should be in patchified format [B, T, C*F] = [B, T, 128]
|
||||
where T is the number of time steps, C=8 channels, F=16 mel bins.
|
||||
This matches the format produced by AudioPatchifier.patchify().
|
||||
"""
|
||||
mel_bins = 16
|
||||
|
||||
latent_coords = self._audio_patchifier.get_patch_grid_bounds(
|
||||
output_shape=AudioLatentShape(
|
||||
frames=num_time_steps,
|
||||
mel_bins=mel_bins,
|
||||
batch=batch_size,
|
||||
channels=8, # Audio latent channels
|
||||
),
|
||||
device=device,
|
||||
)
|
||||
|
||||
return latent_coords.to(dtype)
|
||||
|
||||
@staticmethod
|
||||
def _create_per_token_timesteps(conditioning_mask: Tensor, sampled_sigma: Tensor) -> Tensor:
|
||||
"""Create per-token timesteps based on conditioning mask.
|
||||
Args:
|
||||
conditioning_mask: Boolean mask of shape (batch_size, sequence_length),
|
||||
where True = conditioning token (timestep=0), False = target token (use sigma)
|
||||
sampled_sigma: Sampled sigma values of shape (batch_size,) or (batch_size, 1, 1)
|
||||
Returns:
|
||||
Timesteps tensor of shape [batch_size, sequence_length]
|
||||
"""
|
||||
# Expand to match conditioning mask shape [B, seq_len]
|
||||
expanded_sigma = sampled_sigma.view(-1, 1).expand_as(conditioning_mask)
|
||||
|
||||
# Conditioning tokens get 0, target tokens get the sampled sigma
|
||||
return torch.where(conditioning_mask, torch.zeros_like(expanded_sigma), expanded_sigma)
|
||||
|
||||
@staticmethod
|
||||
def _create_first_frame_conditioning_mask(
|
||||
batch_size: int,
|
||||
sequence_length: int,
|
||||
height: int,
|
||||
width: int,
|
||||
device: torch.device,
|
||||
first_frame_conditioning_p: float = 0.0,
|
||||
) -> Tensor:
|
||||
"""Create conditioning mask for first frame conditioning.
|
||||
Args:
|
||||
batch_size: Batch size
|
||||
sequence_length: Total sequence length
|
||||
height: Latent height
|
||||
width: Latent width
|
||||
device: Target device
|
||||
first_frame_conditioning_p: Probability of conditioning on the first frame
|
||||
Returns:
|
||||
Boolean mask where True indicates first frame tokens (if conditioning is enabled)
|
||||
"""
|
||||
conditioning_mask = torch.zeros(batch_size, sequence_length, dtype=torch.bool, device=device)
|
||||
|
||||
if first_frame_conditioning_p > 0 and random.random() < first_frame_conditioning_p:
|
||||
first_frame_end_idx = height * width
|
||||
if first_frame_end_idx < sequence_length:
|
||||
conditioning_mask[:, :first_frame_end_idx] = True
|
||||
|
||||
return conditioning_mask
|
||||
@@ -0,0 +1,289 @@
|
||||
"""Text-to-video training strategy.
|
||||
This strategy implements standard text-to-video generation training where:
|
||||
- Only target latents are used (no reference videos)
|
||||
- Standard noise application and loss computation
|
||||
- Supports first frame conditioning
|
||||
- Optionally supports joint audio-video training
|
||||
"""
|
||||
|
||||
from typing import Any, Literal
|
||||
|
||||
import torch
|
||||
from pydantic import Field
|
||||
from torch import Tensor
|
||||
|
||||
from ltx_core.model.transformer.modality import Modality
|
||||
from ltx_trainer import logger
|
||||
from ltx_trainer.timestep_samplers import TimestepSampler
|
||||
from ltx_trainer.training_strategies.base_strategy import (
|
||||
DEFAULT_FPS,
|
||||
ModelInputs,
|
||||
TrainingStrategy,
|
||||
TrainingStrategyConfigBase,
|
||||
)
|
||||
|
||||
|
||||
class TextToVideoConfig(TrainingStrategyConfigBase):
|
||||
"""Configuration for text-to-video training strategy."""
|
||||
|
||||
name: Literal["text_to_video"] = "text_to_video"
|
||||
|
||||
first_frame_conditioning_p: float = Field(
|
||||
default=0.1,
|
||||
description="Probability of conditioning on the first frame during training",
|
||||
ge=0.0,
|
||||
le=1.0,
|
||||
)
|
||||
|
||||
with_audio: bool = Field(
|
||||
default=False,
|
||||
description="Whether to include audio in training (joint audio-video generation)",
|
||||
)
|
||||
|
||||
audio_latents_dir: str = Field(
|
||||
default="audio_latents",
|
||||
description="Directory name for audio latents when with_audio is True",
|
||||
)
|
||||
|
||||
|
||||
class TextToVideoStrategy(TrainingStrategy):
|
||||
"""Text-to-video training strategy.
|
||||
This strategy implements regular video generation training where:
|
||||
- Only target latents are used (no reference videos)
|
||||
- Standard noise application and loss computation
|
||||
- Supports first frame conditioning
|
||||
- Optionally supports joint audio-video training when with_audio=True
|
||||
"""
|
||||
|
||||
config: TextToVideoConfig
|
||||
|
||||
def __init__(self, config: TextToVideoConfig):
|
||||
"""Initialize strategy with configuration.
|
||||
Args:
|
||||
config: Text-to-video configuration
|
||||
"""
|
||||
super().__init__(config)
|
||||
|
||||
@property
|
||||
def requires_audio(self) -> bool:
|
||||
"""Whether this training strategy requires audio components."""
|
||||
return self.config.with_audio
|
||||
|
||||
def get_data_sources(self) -> list[str] | dict[str, str]:
|
||||
"""
|
||||
Text-to-video training requires latents and text conditions.
|
||||
When with_audio is True, also requires audio latents.
|
||||
"""
|
||||
sources = {
|
||||
"latents": "latents",
|
||||
"conditions": "conditions",
|
||||
}
|
||||
|
||||
if self.config.with_audio:
|
||||
sources[self.config.audio_latents_dir] = "audio_latents"
|
||||
|
||||
return sources
|
||||
|
||||
def prepare_training_inputs(
|
||||
self,
|
||||
batch: dict[str, Any],
|
||||
timestep_sampler: TimestepSampler,
|
||||
) -> ModelInputs:
|
||||
"""Prepare inputs for text-to-video training."""
|
||||
# Get pre-encoded latents - dataset provides uniform non-patchified format [B, C, F, H, W]
|
||||
latents = batch["latents"]
|
||||
video_latents = latents["latents"]
|
||||
|
||||
# Get video dimensions (assume same for all batch elements)
|
||||
num_frames = latents["num_frames"][0].item()
|
||||
height = latents["height"][0].item()
|
||||
width = latents["width"][0].item()
|
||||
|
||||
# Patchify latents: [B, C, F, H, W] -> [B, seq_len, C]
|
||||
video_latents = self._video_patchifier.patchify(video_latents)
|
||||
|
||||
# Handle FPS with backward compatibility
|
||||
fps = latents.get("fps", None)
|
||||
if fps is not None and not torch.all(fps == fps[0]):
|
||||
logger.warning(
|
||||
f"Different FPS values found in the batch. Found: {fps.tolist()}, using the first one: {fps[0].item()}"
|
||||
)
|
||||
fps = fps[0].item() if fps is not None else DEFAULT_FPS
|
||||
|
||||
# Get text embeddings (already processed by embedding connectors in trainer)
|
||||
conditions = batch["conditions"]
|
||||
video_prompt_embeds = conditions["video_prompt_embeds"]
|
||||
audio_prompt_embeds = conditions["audio_prompt_embeds"]
|
||||
prompt_attention_mask = conditions["prompt_attention_mask"]
|
||||
|
||||
batch_size = video_latents.shape[0]
|
||||
video_seq_len = video_latents.shape[1]
|
||||
device = video_latents.device
|
||||
dtype = video_latents.dtype
|
||||
|
||||
# Create conditioning mask (first frame conditioning)
|
||||
video_conditioning_mask = self._create_first_frame_conditioning_mask(
|
||||
batch_size=batch_size,
|
||||
sequence_length=video_seq_len,
|
||||
height=height,
|
||||
width=width,
|
||||
device=device,
|
||||
first_frame_conditioning_p=self.config.first_frame_conditioning_p,
|
||||
)
|
||||
|
||||
# Sample noise and sigmas
|
||||
sigmas = timestep_sampler.sample_for(video_latents)
|
||||
video_noise = torch.randn_like(video_latents)
|
||||
|
||||
# Apply noise: noisy = (1 - sigma) * clean + sigma * noise
|
||||
sigmas_expanded = sigmas.view(-1, 1, 1)
|
||||
noisy_video = (1 - sigmas_expanded) * video_latents + sigmas_expanded * video_noise
|
||||
|
||||
# For conditioning tokens, use clean latents
|
||||
conditioning_mask_expanded = video_conditioning_mask.unsqueeze(-1)
|
||||
noisy_video = torch.where(conditioning_mask_expanded, video_latents, noisy_video)
|
||||
|
||||
# Compute video targets (velocity prediction)
|
||||
video_targets = video_noise - video_latents
|
||||
|
||||
# Create per-token timesteps
|
||||
video_timesteps = self._create_per_token_timesteps(video_conditioning_mask, sigmas.squeeze())
|
||||
|
||||
# Generate video positions using ltx_core's native implementation
|
||||
video_positions = self._get_video_positions(
|
||||
num_frames=num_frames,
|
||||
height=height,
|
||||
width=width,
|
||||
batch_size=batch_size,
|
||||
fps=fps,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
# Create video Modality
|
||||
video_modality = Modality(
|
||||
enabled=True,
|
||||
latent=noisy_video,
|
||||
timesteps=video_timesteps,
|
||||
positions=video_positions,
|
||||
context=video_prompt_embeds,
|
||||
context_mask=prompt_attention_mask,
|
||||
)
|
||||
|
||||
# Video loss mask: True for tokens we want to compute loss on (non-conditioning tokens)
|
||||
video_loss_mask = ~video_conditioning_mask
|
||||
|
||||
# Handle audio if enabled
|
||||
audio_modality = None
|
||||
audio_targets = None
|
||||
audio_loss_mask = None
|
||||
|
||||
if self.config.with_audio:
|
||||
audio_modality, audio_targets, audio_loss_mask = self._prepare_audio_inputs(
|
||||
batch=batch,
|
||||
sigmas=sigmas,
|
||||
audio_prompt_embeds=audio_prompt_embeds,
|
||||
prompt_attention_mask=prompt_attention_mask,
|
||||
batch_size=batch_size,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
return ModelInputs(
|
||||
video=video_modality,
|
||||
audio=audio_modality,
|
||||
video_targets=video_targets,
|
||||
audio_targets=audio_targets,
|
||||
video_loss_mask=video_loss_mask,
|
||||
audio_loss_mask=audio_loss_mask,
|
||||
)
|
||||
|
||||
def _prepare_audio_inputs(
|
||||
self,
|
||||
batch: dict[str, Any],
|
||||
sigmas: Tensor,
|
||||
audio_prompt_embeds: Tensor,
|
||||
prompt_attention_mask: Tensor,
|
||||
batch_size: int,
|
||||
device: torch.device,
|
||||
dtype: torch.dtype,
|
||||
) -> tuple[Modality, Tensor, Tensor]:
|
||||
"""Prepare audio inputs for joint audio-video training.
|
||||
Args:
|
||||
batch: Raw batch data containing audio_latents
|
||||
sigmas: Sampled sigma values (same as video)
|
||||
audio_prompt_embeds: Audio context embeddings
|
||||
prompt_attention_mask: Attention mask for context
|
||||
batch_size: Batch size
|
||||
device: Target device
|
||||
dtype: Target dtype
|
||||
Returns:
|
||||
Tuple of (audio_modality, audio_targets, audio_loss_mask)
|
||||
"""
|
||||
# Get audio latents - dataset provides uniform non-patchified format [B, C, T, F]
|
||||
audio_data = batch["audio_latents"]
|
||||
audio_latents = audio_data["latents"]
|
||||
|
||||
# Patchify audio latents: [B, C, T, F] -> [B, T, C*F]
|
||||
audio_latents = self._audio_patchifier.patchify(audio_latents)
|
||||
|
||||
audio_seq_len = audio_latents.shape[1]
|
||||
|
||||
# Sample audio noise
|
||||
audio_noise = torch.randn_like(audio_latents)
|
||||
|
||||
# Apply noise to audio (same sigma as video)
|
||||
sigmas_expanded = sigmas.view(-1, 1, 1)
|
||||
noisy_audio = (1 - sigmas_expanded) * audio_latents + sigmas_expanded * audio_noise
|
||||
|
||||
# Compute audio targets
|
||||
audio_targets = audio_noise - audio_latents
|
||||
|
||||
# Audio timesteps: all tokens use the sampled sigma (no conditioning mask)
|
||||
audio_timesteps = sigmas.view(-1, 1).expand(-1, audio_seq_len)
|
||||
|
||||
# Generate audio positions
|
||||
audio_positions = self._get_audio_positions(
|
||||
num_time_steps=audio_seq_len,
|
||||
batch_size=batch_size,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
# Create audio Modality
|
||||
audio_modality = Modality(
|
||||
enabled=True,
|
||||
latent=noisy_audio,
|
||||
timesteps=audio_timesteps,
|
||||
positions=audio_positions,
|
||||
context=audio_prompt_embeds,
|
||||
context_mask=prompt_attention_mask,
|
||||
)
|
||||
|
||||
# Audio loss mask: all tokens contribute to loss (no conditioning)
|
||||
audio_loss_mask = torch.ones(batch_size, audio_seq_len, dtype=torch.bool, device=device)
|
||||
|
||||
return audio_modality, audio_targets, audio_loss_mask
|
||||
|
||||
def compute_loss(
|
||||
self,
|
||||
video_pred: Tensor,
|
||||
audio_pred: Tensor | None,
|
||||
inputs: ModelInputs,
|
||||
) -> Tensor:
|
||||
"""Compute masked MSE loss for video and optionally audio."""
|
||||
# Video loss
|
||||
video_loss = (video_pred - inputs.video_targets).pow(2)
|
||||
video_loss_mask = inputs.video_loss_mask.unsqueeze(-1).float()
|
||||
video_loss = video_loss.mul(video_loss_mask).div(video_loss_mask.mean())
|
||||
video_loss = video_loss.mean()
|
||||
|
||||
# If no audio, return video loss only
|
||||
if not self.config.with_audio or audio_pred is None or inputs.audio_targets is None:
|
||||
return video_loss
|
||||
|
||||
# Audio loss (no conditioning mask)
|
||||
audio_loss = (audio_pred - inputs.audio_targets).pow(2).mean()
|
||||
|
||||
# Combined loss
|
||||
return video_loss + audio_loss
|
||||
@@ -0,0 +1,223 @@
|
||||
"""Video-to-video training strategy for IC-LoRA.
|
||||
This strategy implements training with reference video conditioning where:
|
||||
- Reference latents (clean) are concatenated with target latents (noised)
|
||||
- Video coordinates handle both reference and target sequences
|
||||
- Loss is computed only on the target portion
|
||||
"""
|
||||
|
||||
from typing import Any, Literal
|
||||
|
||||
import torch
|
||||
from pydantic import Field
|
||||
from torch import Tensor
|
||||
|
||||
from ltx_core.model.transformer.modality import Modality
|
||||
from ltx_trainer import logger
|
||||
from ltx_trainer.timestep_samplers import TimestepSampler
|
||||
from ltx_trainer.training_strategies.base_strategy import (
|
||||
DEFAULT_FPS,
|
||||
ModelInputs,
|
||||
TrainingStrategy,
|
||||
TrainingStrategyConfigBase,
|
||||
)
|
||||
|
||||
|
||||
class VideoToVideoConfig(TrainingStrategyConfigBase):
|
||||
"""Configuration for video-to-video (IC-LoRA) training strategy."""
|
||||
|
||||
name: Literal["video_to_video"] = "video_to_video"
|
||||
|
||||
first_frame_conditioning_p: float = Field(
|
||||
default=0.1,
|
||||
description="Probability of conditioning on the first frame during training",
|
||||
ge=0.0,
|
||||
le=1.0,
|
||||
)
|
||||
|
||||
reference_latents_dir: str = Field(
|
||||
default="reference_latents",
|
||||
description="Directory name for latents of reference videos",
|
||||
)
|
||||
|
||||
|
||||
class VideoToVideoStrategy(TrainingStrategy):
|
||||
"""Video-to-video training strategy for IC-LoRA.
|
||||
This strategy implements training with reference video conditioning where:
|
||||
- Reference latents (clean) are concatenated with target latents (noised)
|
||||
- Video coordinates handle both reference and target sequences
|
||||
- Loss is computed only on the target portion
|
||||
"""
|
||||
|
||||
config: VideoToVideoConfig
|
||||
|
||||
def __init__(self, config: VideoToVideoConfig):
|
||||
"""Initialize strategy with configuration.
|
||||
Args:
|
||||
config: Video-to-video configuration
|
||||
"""
|
||||
super().__init__(config)
|
||||
|
||||
def get_data_sources(self) -> dict[str, str]:
|
||||
"""IC-LoRA training requires latents, conditions, and reference latents."""
|
||||
return {
|
||||
"latents": "latents",
|
||||
"conditions": "conditions",
|
||||
self.config.reference_latents_dir: "ref_latents",
|
||||
}
|
||||
|
||||
def prepare_training_inputs(
|
||||
self,
|
||||
batch: dict[str, Any],
|
||||
timestep_sampler: TimestepSampler,
|
||||
) -> ModelInputs:
|
||||
"""Prepare inputs for IC-LoRA training with reference videos."""
|
||||
# Get pre-encoded latents - dataset provides uniform non-patchified format [B, C, F, H, W]
|
||||
latents = batch["latents"]
|
||||
target_latents = latents["latents"]
|
||||
ref_latents = batch["ref_latents"]["latents"]
|
||||
|
||||
# Get dimensions
|
||||
num_frames = latents["num_frames"][0].item()
|
||||
height = latents["height"][0].item()
|
||||
width = latents["width"][0].item()
|
||||
|
||||
ref_latents_info = batch["ref_latents"]
|
||||
ref_frames = ref_latents_info["num_frames"][0].item()
|
||||
ref_height = ref_latents_info["height"][0].item()
|
||||
ref_width = ref_latents_info["width"][0].item()
|
||||
|
||||
# Patchify latents: [B, C, F, H, W] -> [B, seq_len, C]
|
||||
target_latents = self._video_patchifier.patchify(target_latents)
|
||||
ref_latents = self._video_patchifier.patchify(ref_latents)
|
||||
|
||||
# Handle FPS
|
||||
fps = latents.get("fps", None)
|
||||
if fps is not None and not torch.all(fps == fps[0]):
|
||||
logger.warning(
|
||||
f"Different FPS values found in the batch. Found: {fps.tolist()}, using the first one: {fps[0].item()}"
|
||||
)
|
||||
fps = fps[0].item() if fps is not None else DEFAULT_FPS
|
||||
|
||||
# Get text embeddings (already processed by embedding connectors in trainer)
|
||||
# Video-to-video uses only video embeddings
|
||||
conditions = batch["conditions"]
|
||||
prompt_embeds = conditions["video_prompt_embeds"]
|
||||
prompt_attention_mask = conditions["prompt_attention_mask"]
|
||||
|
||||
batch_size = target_latents.shape[0]
|
||||
ref_seq_len = ref_latents.shape[1]
|
||||
target_seq_len = target_latents.shape[1]
|
||||
device = target_latents.device
|
||||
dtype = target_latents.dtype
|
||||
|
||||
# Create conditioning mask
|
||||
# Reference tokens are always conditioning (timestep=0)
|
||||
ref_conditioning_mask = torch.ones(batch_size, ref_seq_len, dtype=torch.bool, device=device)
|
||||
|
||||
# Target tokens: check for first frame conditioning
|
||||
target_conditioning_mask = self._create_first_frame_conditioning_mask(
|
||||
batch_size=batch_size,
|
||||
sequence_length=target_seq_len,
|
||||
height=height,
|
||||
width=width,
|
||||
device=device,
|
||||
first_frame_conditioning_p=self.config.first_frame_conditioning_p,
|
||||
)
|
||||
|
||||
# Combined conditioning mask
|
||||
conditioning_mask = torch.cat([ref_conditioning_mask, target_conditioning_mask], dim=1)
|
||||
|
||||
# Sample noise and sigmas for target
|
||||
sigmas = timestep_sampler.sample_for(target_latents)
|
||||
noise = torch.randn_like(target_latents)
|
||||
sigmas_expanded = sigmas.view(-1, 1, 1)
|
||||
|
||||
# Apply noise to target
|
||||
noisy_target = (1 - sigmas_expanded) * target_latents + sigmas_expanded * noise
|
||||
|
||||
# For first frame conditioning in target, use clean latents
|
||||
target_conditioning_mask_expanded = target_conditioning_mask.unsqueeze(-1)
|
||||
noisy_target = torch.where(target_conditioning_mask_expanded, target_latents, noisy_target)
|
||||
|
||||
# Targets for loss computation
|
||||
targets = noise - target_latents
|
||||
|
||||
# Concatenate reference (clean) and target (noisy)
|
||||
combined_latents = torch.cat([ref_latents, noisy_target], dim=1)
|
||||
|
||||
# Create per-token timesteps
|
||||
timesteps = self._create_per_token_timesteps(conditioning_mask, sigmas.squeeze())
|
||||
|
||||
# Generate positions for reference and target separately, then concatenate
|
||||
ref_positions = self._get_video_positions(
|
||||
num_frames=ref_frames,
|
||||
height=ref_height,
|
||||
width=ref_width,
|
||||
batch_size=batch_size,
|
||||
fps=fps,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
target_positions = self._get_video_positions(
|
||||
num_frames=num_frames,
|
||||
height=height,
|
||||
width=width,
|
||||
batch_size=batch_size,
|
||||
fps=fps,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
# Concatenate positions along sequence dimension
|
||||
positions = torch.cat([ref_positions, target_positions], dim=2)
|
||||
|
||||
# Create video Modality
|
||||
video_modality = Modality(
|
||||
enabled=True,
|
||||
latent=combined_latents,
|
||||
timesteps=timesteps,
|
||||
positions=positions,
|
||||
context=prompt_embeds,
|
||||
context_mask=prompt_attention_mask,
|
||||
)
|
||||
|
||||
# Loss mask: only compute loss on non-conditioning target tokens
|
||||
# Reference tokens: all False (no loss)
|
||||
# Target tokens: True where not conditioning
|
||||
ref_loss_mask = torch.zeros(batch_size, ref_seq_len, dtype=torch.bool, device=device)
|
||||
target_loss_mask = ~target_conditioning_mask
|
||||
video_loss_mask = torch.cat([ref_loss_mask, target_loss_mask], dim=1)
|
||||
|
||||
return ModelInputs(
|
||||
video=video_modality,
|
||||
audio=None,
|
||||
video_targets=targets,
|
||||
audio_targets=None,
|
||||
video_loss_mask=video_loss_mask,
|
||||
audio_loss_mask=None,
|
||||
ref_seq_len=ref_seq_len,
|
||||
)
|
||||
|
||||
def compute_loss(
|
||||
self,
|
||||
video_pred: Tensor,
|
||||
_audio_pred: Tensor | None,
|
||||
inputs: ModelInputs,
|
||||
) -> Tensor:
|
||||
"""Compute masked loss only on target portion."""
|
||||
# Extract target portion of prediction
|
||||
ref_seq_len = inputs.ref_seq_len
|
||||
target_pred = video_pred[:, ref_seq_len:, :]
|
||||
|
||||
# Get target portion of loss mask
|
||||
target_loss_mask = inputs.video_loss_mask[:, ref_seq_len:]
|
||||
|
||||
# Compute loss
|
||||
loss = (target_pred - inputs.video_targets).pow(2)
|
||||
|
||||
# Apply loss mask
|
||||
loss_mask = target_loss_mask.unsqueeze(-1).float()
|
||||
loss = loss.mul(loss_mask).div(loss_mask.mean())
|
||||
|
||||
return loss.mean()
|
||||
Reference in New Issue
Block a user