"""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_checkpoint_metadata(self) -> dict[str, Any]: """Get strategy-specific metadata to include in checkpoint files. Override this method in subclasses to add custom metadata, e.g. any parameters that a downstream inference pipeline may need. Returns: Dictionary of metadata key-value pairs (values must be JSON-serializable) """ return {} 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