github-actions[bot]
719 lines
29 KiB
Python
719 lines
29 KiB
Python
"""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
|