Merge pull request #97 from Lightricks/pr-2026-01-29

Open a PR to sync - 2026-01-29
This commit is contained in:
Michael Kupchick
2026-01-29 21:30:12 +02:00
committed by GitHub
31 changed files with 1723 additions and 663 deletions
+1 -1
View File
@@ -9,7 +9,7 @@ dependencies = [
"torchaudio", "torchaudio",
"einops", "einops",
"numpy", "numpy",
"transformers", "transformers~=4.57.0",
"safetensors", "safetensors",
"accelerate", "accelerate",
"scipy>=1.14", "scipy>=1.14",
@@ -1,4 +1,5 @@
from dataclasses import dataclass import math
from dataclasses import dataclass, field
import torch import torch
@@ -189,6 +190,89 @@ class LegacyStatefulAPGGuider(GuiderProtocol):
return self.scale != 0.0 return self.scale != 0.0
@dataclass(frozen=True)
class MultiModalGuiderParams:
"""
Parameters for the multi-modal guider.
"""
cfg_scale: float = 1.0
"CFG (Classifier-free guidance) scale controlling how strongly the model adheres to the prompt."
stg_scale: float = 0.0
"STG (Spatio-Temporal Guidance) scale controls how strongly the model reacts to the perturbation of the modality."
stg_blocks: list[int] | None = field(default_factory=list)
"Which transformer blocks to perturb for STG."
rescale_scale: float = 0.0
"Rescale scale controlling how strongly the model rescales the modality after applying other guidance."
modality_scale: float = 1.0
"Modality scale controlling how strongly the model reacts to the perturbation of the modality."
skip_step: int = 0
"Skip step controlling how often the model skips the step."
@dataclass(frozen=True)
class MultiModalGuider:
"""
Multi-modal guider.
"""
params: MultiModalGuiderParams
negative_context: torch.Tensor | None = None
def calculate(
self,
cond: torch.Tensor,
uncond_text: torch.Tensor | float,
uncond_perturbed: torch.Tensor | float,
uncond_modality: torch.Tensor | float,
) -> torch.Tensor:
"""
The guider calculates the guidance delta as (scale - 1) * (cond - uncond) for cfg and modality cfg,
and as scale * (cond - uncond) for stg, steering the denoising process away from the unconditioned
prediction.
"""
pred = (
cond
+ (self.params.cfg_scale - 1) * (cond - uncond_text)
+ self.params.stg_scale * (cond - uncond_perturbed)
+ (self.params.modality_scale - 1) * (cond - uncond_modality)
)
if self.params.rescale_scale != 0:
factor = cond.std() / pred.std()
factor = self.params.rescale_scale * factor + (1 - self.params.rescale_scale)
pred = pred * factor
return pred
def do_unconditional_generation(self) -> bool:
"""
Returns True if the guider is doing unconditional generation.
"""
return not math.isclose(self.params.cfg_scale, 1.0)
def do_perturbed_generation(self) -> bool:
"""
Returns True if the guider is doing perturbed generation.
"""
return not math.isclose(self.params.stg_scale, 0.0)
def do_isolated_modality_generation(self) -> bool:
"""
Returns True if the guider is doing isolated modality generation.
"""
return not math.isclose(self.params.modality_scale, 1.0)
def should_skip_step(self, step: int) -> bool:
"""
Returns True if the guider should skip the step.
"""
if self.params.skip_step == 0:
return False
return step % (self.params.skip_step + 1) != 0
def projection_coef(to_project: torch.Tensor, project_onto: torch.Tensor) -> torch.Tensor: def projection_coef(to_project: torch.Tensor, project_onto: torch.Tensor) -> torch.Tensor:
batch_size = to_project.shape[0] batch_size = to_project.shape[0]
positive_flat = to_project.reshape(batch_size, -1) positive_flat = to_project.reshape(batch_size, -1)
@@ -2,11 +2,16 @@
from ltx_core.conditioning.exceptions import ConditioningError from ltx_core.conditioning.exceptions import ConditioningError
from ltx_core.conditioning.item import ConditioningItem from ltx_core.conditioning.item import ConditioningItem
from ltx_core.conditioning.types import VideoConditionByKeyframeIndex, VideoConditionByLatentIndex from ltx_core.conditioning.types import (
VideoConditionByKeyframeIndex,
VideoConditionByLatentIndex,
VideoConditionByReferenceLatent,
)
__all__ = [ __all__ = [
"ConditioningError", "ConditioningError",
"ConditioningItem", "ConditioningItem",
"VideoConditionByKeyframeIndex", "VideoConditionByKeyframeIndex",
"VideoConditionByLatentIndex", "VideoConditionByLatentIndex",
"VideoConditionByReferenceLatent",
] ]
@@ -2,8 +2,10 @@
from ltx_core.conditioning.types.keyframe_cond import VideoConditionByKeyframeIndex from ltx_core.conditioning.types.keyframe_cond import VideoConditionByKeyframeIndex
from ltx_core.conditioning.types.latent_cond import VideoConditionByLatentIndex from ltx_core.conditioning.types.latent_cond import VideoConditionByLatentIndex
from ltx_core.conditioning.types.reference_video_cond import VideoConditionByReferenceLatent
__all__ = [ __all__ = [
"VideoConditionByKeyframeIndex", "VideoConditionByKeyframeIndex",
"VideoConditionByLatentIndex", "VideoConditionByLatentIndex",
"VideoConditionByReferenceLatent",
] ]
@@ -0,0 +1,78 @@
"""Reference video conditioning for IC-LoRA inference."""
import torch
from ltx_core.components.patchifiers import get_pixel_coords
from ltx_core.conditioning.item import ConditioningItem
from ltx_core.tools import VideoLatentTools
from ltx_core.types import LatentState, VideoLatentShape
class VideoConditionByReferenceLatent(ConditioningItem):
"""
Conditions video generation on a reference video latent for IC-LoRA inference.
IC-LoRAs are trained by concatenating reference (control signal) and target tokens,
learning to attend across both. This class replicates that setup at inference by
appending reference tokens to the latent sequence.
IC-LoRAs can be trained with lower-resolution references than the target (e.g., 384px
reference for 768px output) for efficiency and better generalization. The
`downscale_factor` scales reference positions to match target coordinates, preserving
the learned positional relationships. This must match the factor used during training
(stored in LoRA metadata).
Args:
latent: Reference video latents [B, C, F, H, W]
downscale_factor: Target/reference resolution ratio (e.g., 2 = half-resolution
reference). Spatial positions are scaled by this factor.
strength: Conditioning strength. 1.0 = full (reference kept clean),
0.0 = none (reference denoised). Default 1.0.
"""
def __init__(
self,
latent: torch.Tensor,
downscale_factor: int = 1,
strength: float = 1.0,
):
self.latent = latent
self.downscale_factor = downscale_factor
self.strength = strength
def apply_to(
self,
latent_state: LatentState,
latent_tools: VideoLatentTools,
) -> LatentState:
"""Append reference video tokens with scaled positions."""
tokens = latent_tools.patchifier.patchify(self.latent)
# Compute positions for the reference video's actual dimensions
latent_coords = latent_tools.patchifier.get_patch_grid_bounds(
output_shape=VideoLatentShape.from_torch_shape(self.latent.shape),
device=self.latent.device,
)
positions = get_pixel_coords(
latent_coords=latent_coords,
scale_factors=latent_tools.scale_factors,
causal_fix=latent_tools.causal_fix,
)
positions = positions.to(dtype=torch.float32)
positions[:, 0, ...] /= latent_tools.fps
# Scale spatial positions to match target coordinate space
if self.downscale_factor != 1:
positions[:, 1, ...] *= self.downscale_factor # height axis
positions[:, 2, ...] *= self.downscale_factor # width axis
denoise_mask = torch.full(
size=(*tokens.shape[:2], 1),
fill_value=1.0 - self.strength,
device=self.latent.device,
dtype=self.latent.dtype,
)
return LatentState(
latent=torch.cat([latent_state.latent, tokens], dim=1),
denoise_mask=torch.cat([latent_state.denoise_mask, denoise_mask], dim=1),
positions=torch.cat([latent_state.positions, positions], dim=2),
clean_latent=torch.cat([latent_state.clean_latent, tokens], dim=1),
)
@@ -73,7 +73,7 @@ class SingleGPUModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType],
device = torch.device("cuda") if device is None else device device = torch.device("cuda") if device is None else device
config = self.model_config() config = self.model_config()
meta_model = self.meta_model(config, self.module_ops) meta_model = self.meta_model(config, self.module_ops)
model_paths = self.model_path if isinstance(self.model_path, tuple) else [self.model_path] model_paths = list(self.model_path) if isinstance(self.model_path, tuple) else [self.model_path]
model_state_dict = self.load_sd(model_paths, sd_ops=self.model_sd_ops, registry=self.registry, device=device) model_state_dict = self.load_sd(model_paths, sd_ops=self.model_sd_ops, registry=self.registry, device=device)
lora_strengths = [lora.strength for lora in self.loras] lora_strengths = [lora.strength for lora in self.loras]
@@ -140,7 +140,11 @@ class BasicAVTransformerBlock(torch.nn.Module):
audio: TransformerArgs | None, audio: TransformerArgs | None,
perturbations: BatchedPerturbationConfig | None = None, perturbations: BatchedPerturbationConfig | None = None,
) -> tuple[TransformerArgs | None, TransformerArgs | None]: ) -> tuple[TransformerArgs | None, TransformerArgs | None]:
batch_size = video.x.shape[0] if video is None and audio is None:
raise ValueError("At least one of video or audio must be provided")
batch_size = (video or audio).x.shape[0]
if perturbations is None: if perturbations is None:
perturbations = BatchedPerturbationConfig.empty(batch_size) perturbations = BatchedPerturbationConfig.empty(batch_size)
@@ -211,7 +215,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
video.cross_gate_timestep, video.cross_gate_timestep,
) )
if run_a2v: if run_a2v and not perturbations.all_in_batch(PerturbationType.SKIP_A2V_CROSS_ATTN, self.idx):
vx_scaled = vx_norm3 * (1 + scale_ca_video_hidden_states_a2v) + shift_ca_video_hidden_states_a2v vx_scaled = vx_norm3 * (1 + scale_ca_video_hidden_states_a2v) + shift_ca_video_hidden_states_a2v
ax_scaled = ax_norm3 * (1 + scale_ca_audio_hidden_states_a2v) + shift_ca_audio_hidden_states_a2v ax_scaled = ax_norm3 * (1 + scale_ca_audio_hidden_states_a2v) + shift_ca_audio_hidden_states_a2v
a2v_mask = perturbations.mask_like(PerturbationType.SKIP_A2V_CROSS_ATTN, self.idx, vx) a2v_mask = perturbations.mask_like(PerturbationType.SKIP_A2V_CROSS_ATTN, self.idx, vx)
@@ -226,7 +230,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
* a2v_mask * a2v_mask
) )
if run_v2a: if run_v2a and not perturbations.all_in_batch(PerturbationType.SKIP_V2A_CROSS_ATTN, self.idx):
ax_scaled = ax_norm3 * (1 + scale_ca_audio_hidden_states_v2a) + shift_ca_audio_hidden_states_v2a ax_scaled = ax_norm3 * (1 + scale_ca_audio_hidden_states_v2a) + shift_ca_audio_hidden_states_v2a
vx_scaled = vx_norm3 * (1 + scale_ca_video_hidden_states_v2a) + shift_ca_video_hidden_states_v2a vx_scaled = vx_norm3 * (1 + scale_ca_video_hidden_states_v2a) + shift_ca_video_hidden_states_v2a
v2a_mask = perturbations.mask_like(PerturbationType.SKIP_V2A_CROSS_ATTN, self.idx, ax) v2a_mask = perturbations.mask_like(PerturbationType.SKIP_V2A_CROSS_ATTN, self.idx, ax)
@@ -0,0 +1,75 @@
from dataclasses import asdict, dataclass, field
@dataclass
class Gemma3RopeScaling:
factor: float = 8.0
rope_type: str = "linear"
@dataclass
class Gemma3TextConfig:
attention_bias: bool = False
attention_dropout: float = 0.0
attn_logit_softcapping: float | None = None
cache_implementation: str = "hybrid"
final_logit_softcapping: float | None = None
head_dim: int = 256
hidden_activation: str = "gelu_pytorch_tanh"
hidden_size: int = 3840
initializer_range: float = 0.02
intermediate_size: int = 15360
max_position_embeddings: int = 131072
model_type: str = "gemma3_text"
num_attention_heads: int = 16
num_hidden_layers: int = 48
num_key_value_heads: int = 8
query_pre_attn_scalar: int = 256
rms_norm_eps: float = 1e-06
rope_local_base_freq: int = 10000
rope_scaling: Gemma3RopeScaling = field(default_factory=Gemma3RopeScaling)
rope_theta: int = 1000000
sliding_window: int = 1024
sliding_window_pattern: int = 6
torch_dtype: str = "float32"
use_cache: bool = True
vocab_size: int = 262208
@dataclass
class Gemma3VisionConfig:
attention_dropout: float = 0.0
hidden_act: str = "gelu_pytorch_tanh"
hidden_size: int = 1152
image_size: int = 896
intermediate_size: int = 4304
layer_norm_eps: float = 1e-06
model_type: str = "siglip_vision_model"
num_attention_heads: int = 16
num_channels: int = 3
num_hidden_layers: int = 27
patch_size: int = 14
torch_dtype: str = "float32"
vision_use_head: bool = False
@dataclass
class Gemma3ConfigData:
architectures: list[str] = field(default_factory=lambda: ["Gemma3ForConditionalGeneration"])
boi_token_index: int = 255999
eoi_token_index: int = 256000
eos_token_id: list[int] = field(default_factory=lambda: [1, 106])
image_token_index: int = 262144
initializer_range: float = 0.02
mm_tokens_per_image: int = 256
model_type: str = "gemma3"
text_config: Gemma3TextConfig = field(default_factory=Gemma3TextConfig)
torch_dtype: str = "bfloat16"
transformers_version: str = "4.51.0"
vision_config: Gemma3VisionConfig = field(default_factory=Gemma3VisionConfig)
def to_dict(self) -> dict:
return asdict(self)
GEMMA3_CONFIG_FOR_LTX = Gemma3ConfigData()
@@ -1,15 +1,22 @@
from typing import NamedTuple from typing import NamedTuple
import torch import torch
from transformers import Gemma3Config
from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
from transformers.models.gemma3 import Gemma3ForConditionalGeneration from transformers.models.gemma3 import Gemma3ForConditionalGeneration
from ltx_core.loader import KeyValueOperationResult
from ltx_core.loader.module_ops import ModuleOps
from ltx_core.loader.sd_ops import SDOps from ltx_core.loader.sd_ops import SDOps
from ltx_core.model.model_protocol import ModelConfigurator from ltx_core.model.model_protocol import ModelConfigurator
from ltx_core.text_encoders.gemma.config import GEMMA3_CONFIG_FOR_LTX
from ltx_core.text_encoders.gemma.embeddings_connector import ( from ltx_core.text_encoders.gemma.embeddings_connector import (
Embeddings1DConnector, Embeddings1DConnector,
Embeddings1DConnectorConfigurator, Embeddings1DConnectorConfigurator,
) )
from ltx_core.text_encoders.gemma.encoders.base_encoder import GemmaTextEncoderModelBase from ltx_core.text_encoders.gemma.encoders.base_encoder import (
GemmaTextEncoderModelBase,
)
from ltx_core.text_encoders.gemma.feature_extractor import GemmaFeaturesExtractorProjLinear from ltx_core.text_encoders.gemma.feature_extractor import GemmaFeaturesExtractorProjLinear
from ltx_core.text_encoders.gemma.tokenizer import LTXVGemmaTokenizer from ltx_core.text_encoders.gemma.tokenizer import LTXVGemmaTokenizer
@@ -76,7 +83,11 @@ class AVGemmaTextEncoderModelConfigurator(ModelConfigurator[AVGemmaTextEncoderMo
feature_extractor_linear = GemmaFeaturesExtractorProjLinear.from_config(config) feature_extractor_linear = GemmaFeaturesExtractorProjLinear.from_config(config)
embeddings_connector = Embeddings1DConnectorConfigurator.from_config(config) embeddings_connector = Embeddings1DConnectorConfigurator.from_config(config)
audio_embeddings_connector = Embeddings1DConnectorConfigurator.from_config(config) audio_embeddings_connector = Embeddings1DConnectorConfigurator.from_config(config)
gemma_config = Gemma3Config.from_dict(GEMMA3_CONFIG_FOR_LTX.to_dict())
with torch.device("meta"):
model = Gemma3ForConditionalGeneration(gemma_config)
return AVGemmaTextEncoderModel( return AVGemmaTextEncoderModel(
model=model,
feature_extractor_linear=feature_extractor_linear, feature_extractor_linear=feature_extractor_linear,
embeddings_connector=embeddings_connector, embeddings_connector=embeddings_connector,
audio_embeddings_connector=audio_embeddings_connector, audio_embeddings_connector=audio_embeddings_connector,
@@ -85,10 +96,57 @@ class AVGemmaTextEncoderModelConfigurator(ModelConfigurator[AVGemmaTextEncoderMo
AV_GEMMA_TEXT_ENCODER_KEY_OPS = ( AV_GEMMA_TEXT_ENCODER_KEY_OPS = (
SDOps("AV_GEMMA_TEXT_ENCODER_KEY_OPS") SDOps("AV_GEMMA_TEXT_ENCODER_KEY_OPS")
# 1. Map the feature extractor
.with_matching(prefix="text_embedding_projection.") .with_matching(prefix="text_embedding_projection.")
.with_matching(prefix="model.diffusion_model.audio_embeddings_connector.")
.with_matching(prefix="model.diffusion_model.video_embeddings_connector.")
.with_replacement("text_embedding_projection.", "feature_extractor_linear.") .with_replacement("text_embedding_projection.", "feature_extractor_linear.")
# 2. Map the connectors (fixing the swapped prefixes from before)
.with_matching(prefix="model.diffusion_model.video_embeddings_connector.")
.with_replacement("model.diffusion_model.video_embeddings_connector.", "embeddings_connector.") .with_replacement("model.diffusion_model.video_embeddings_connector.", "embeddings_connector.")
.with_matching(prefix="model.diffusion_model.audio_embeddings_connector.")
.with_replacement("model.diffusion_model.audio_embeddings_connector.", "audio_embeddings_connector.") .with_replacement("model.diffusion_model.audio_embeddings_connector.", "audio_embeddings_connector.")
# 3. Map language model layers (note the double .model prefix)
.with_matching(prefix="language_model.model.")
.with_replacement("language_model.model.", "model.model.language_model.")
# 4. Map the Vision Tower
.with_matching(prefix="vision_tower.")
.with_replacement("vision_tower.", "model.model.vision_tower.")
# 5. Map the Multi-Modal Projector
.with_matching(prefix="multi_modal_projector.")
.with_replacement("multi_modal_projector.", "model.model.multi_modal_projector.")
.with_kv_operation(
operation=lambda key, value: [
KeyValueOperationResult(key, value),
KeyValueOperationResult("model.lm_head.weight", value),
],
key_prefix="model.model.language_model.embed_tokens.weight",
)
)
def create_and_populate(module: AVGemmaTextEncoderModel) -> AVGemmaTextEncoderModel:
model = module.model
v_model = model.model.vision_tower.vision_model
l_model = model.model.language_model
config = model.config.text_config
dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
base = config.rope_local_base_freq
local_rope_freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(dtype=torch.float) / dim))
inv_freqs, _ = ROPE_INIT_FUNCTIONS[config.rope_scaling["rope_type"]](config)
positions_length = len(v_model.embeddings.position_ids[0])
position_ids = torch.arange(positions_length, dtype=torch.long, device="cpu").unsqueeze(0)
v_model.embeddings.register_buffer("position_ids", position_ids)
embed_scale = torch.tensor(model.config.text_config.hidden_size**0.5, device="cpu")
l_model.embed_tokens.register_buffer("embed_scale", embed_scale)
l_model.rotary_emb_local.register_buffer("inv_freq", local_rope_freqs)
l_model.rotary_emb.register_buffer("inv_freq", inv_freqs)
return module
GEMMA_MODEL_OPS = ModuleOps(
name="GemmaModel",
matcher=lambda module: hasattr(module, "model") and isinstance(module.model, Gemma3ForConditionalGeneration),
mutator=create_and_populate,
) )
@@ -8,6 +8,7 @@ from transformers import AutoImageProcessor, Gemma3ForConditionalGeneration, Gem
from ltx_core.loader.module_ops import ModuleOps from ltx_core.loader.module_ops import ModuleOps
from ltx_core.text_encoders.gemma.feature_extractor import GemmaFeaturesExtractorProjLinear from ltx_core.text_encoders.gemma.feature_extractor import GemmaFeaturesExtractorProjLinear
from ltx_core.text_encoders.gemma.tokenizer import LTXVGemmaTokenizer from ltx_core.text_encoders.gemma.tokenizer import LTXVGemmaTokenizer
from ltx_core.utils import find_matching_file
class GemmaTextEncoderModelBase(torch.nn.Module): class GemmaTextEncoderModelBase(torch.nn.Module):
@@ -77,7 +78,7 @@ class GemmaTextEncoderModelBase(torch.nn.Module):
messages: list[dict[str, str]], messages: list[dict[str, str]],
image: torch.Tensor | None = None, image: torch.Tensor | None = None,
max_new_tokens: int = 512, max_new_tokens: int = 512,
seed: int = 42, seed: int = 10,
) -> str: ) -> str:
text = self.processor.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) text = self.processor.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
@@ -107,7 +108,7 @@ class GemmaTextEncoderModelBase(torch.nn.Module):
prompt: str, prompt: str,
max_new_tokens: int = 512, max_new_tokens: int = 512,
system_prompt: str | None = None, system_prompt: str | None = None,
seed: int = 42, seed: int = 10,
) -> str: ) -> str:
"""Enhance a text prompt for T2V generation.""" """Enhance a text prompt for T2V generation."""
@@ -126,7 +127,7 @@ class GemmaTextEncoderModelBase(torch.nn.Module):
image: torch.Tensor, image: torch.Tensor,
max_new_tokens: int = 512, max_new_tokens: int = 512,
system_prompt: str | None = None, system_prompt: str | None = None,
seed: int = 42, seed: int = 10,
) -> str: ) -> str:
"""Enhance a text prompt for I2V generation using a reference image.""" """Enhance a text prompt for I2V generation using a reference image."""
system_prompt = system_prompt or self.default_gemma_i2v_system_prompt system_prompt = system_prompt or self.default_gemma_i2v_system_prompt
@@ -219,44 +220,21 @@ def _load_system_prompt(prompt_name: str) -> str:
return f.read() return f.read()
def _find_matching_dir(root_path: str, pattern: str) -> str:
"""
Recursively search for files matching a glob pattern and return the parent directory of the first match.
"""
matches = list(Path(root_path).rglob(pattern))
if not matches:
raise FileNotFoundError(f"No files matching pattern '{pattern}' found under {root_path}")
return str(matches[0].parent)
def module_ops_from_gemma_root(gemma_root: str) -> tuple[ModuleOps, ...]: def module_ops_from_gemma_root(gemma_root: str) -> tuple[ModuleOps, ...]:
gemma_path = _find_matching_dir(gemma_root, "model*.safetensors") tokenizer_root = str(find_matching_file(gemma_root, "tokenizer.model").parent)
tokenizer_path = _find_matching_dir(gemma_root, "tokenizer.model") processor_root = str(find_matching_file(gemma_root, "preprocessor_config.json").parent)
processor_path = _find_matching_dir(gemma_root, "preprocessor_config.json")
def load_gemma(module: GemmaTextEncoderModelBase) -> GemmaTextEncoderModelBase:
module.model = Gemma3ForConditionalGeneration.from_pretrained(
gemma_path, local_files_only=True, torch_dtype=torch.bfloat16
)
return module
def load_tokenizer(module: GemmaTextEncoderModelBase) -> GemmaTextEncoderModelBase: def load_tokenizer(module: GemmaTextEncoderModelBase) -> GemmaTextEncoderModelBase:
module.tokenizer = LTXVGemmaTokenizer(tokenizer_path, 1024) module.tokenizer = LTXVGemmaTokenizer(tokenizer_root, 1024)
return module return module
def load_processor(module: GemmaTextEncoderModelBase) -> GemmaTextEncoderModelBase: def load_processor(module: GemmaTextEncoderModelBase) -> GemmaTextEncoderModelBase:
image_processor = AutoImageProcessor.from_pretrained(processor_path, local_files_only=True) image_processor = AutoImageProcessor.from_pretrained(processor_root, local_files_only=True)
if not module.tokenizer: if not module.tokenizer:
raise ValueError("Tokenizer model operation must be performed before processor model operation") raise ValueError("Tokenizer model operation must be performed before processor model operation")
module.processor = Gemma3Processor(image_processor=image_processor, tokenizer=module.tokenizer.tokenizer) module.processor = Gemma3Processor(image_processor=image_processor, tokenizer=module.tokenizer.tokenizer)
return module return module
gemma_load_ops = ModuleOps(
"GemmaLoad",
matcher=lambda module: isinstance(module, GemmaTextEncoderModelBase) and module.model is None,
mutator=load_gemma,
)
tokenizer_load_ops = ModuleOps( tokenizer_load_ops = ModuleOps(
"TokenizerLoad", "TokenizerLoad",
matcher=lambda module: isinstance(module, GemmaTextEncoderModelBase) and module.tokenizer is None, matcher=lambda module: isinstance(module, GemmaTextEncoderModelBase) and module.tokenizer is None,
@@ -267,7 +245,7 @@ def module_ops_from_gemma_root(gemma_root: str) -> tuple[ModuleOps, ...]:
matcher=lambda module: isinstance(module, GemmaTextEncoderModelBase) and module.processor is None, matcher=lambda module: isinstance(module, GemmaTextEncoderModelBase) and module.processor is None,
mutator=load_processor, mutator=load_processor,
) )
return (gemma_load_ops, tokenizer_load_ops, processor_load_ops) return (tokenizer_load_ops, processor_load_ops)
def encode_text(text_encoder: GemmaTextEncoderModelBase, prompts: list[str]) -> list[tuple[torch.Tensor, torch.Tensor]]: def encode_text(text_encoder: GemmaTextEncoderModelBase, prompts: list[str]) -> list[tuple[torch.Tensor, torch.Tensor]]:
+11
View File
@@ -1,3 +1,4 @@
from pathlib import Path
from typing import Any from typing import Any
import torch import torch
@@ -49,3 +50,13 @@ def to_denoised(
if isinstance(sigma, torch.Tensor): if isinstance(sigma, torch.Tensor):
sigma = sigma.to(calc_dtype) sigma = sigma.to(calc_dtype)
return (sample.to(calc_dtype) - velocity.to(calc_dtype) * sigma).to(sample.dtype) return (sample.to(calc_dtype) - velocity.to(calc_dtype) * sigma).to(sample.dtype)
def find_matching_file(root_path: str, pattern: str) -> Path:
"""
Recursively search for files matching a glob pattern and return the first match.
"""
matches = list(Path(root_path).rglob(pattern))
if not matches:
raise FileNotFoundError(f"No files matching pattern '{pattern}' found under {root_path}")
return matches[0]
+81 -6
View File
@@ -89,7 +89,7 @@ Do you need to condition on existing images/videos?
### Features Comparison ### Features Comparison
| Pipeline | Stages | CFG | Upsampling | Conditioning | Best For | | Pipeline | Stages | [Multimodal Guidance](#%EF%B8%8F-multimodal-guidance) | Upsampling | Conditioning | Best For |
| -------- | ------ | --- | ---------- | ------------- | -------- | | -------- | ------ | --- | ---------- | ------------- | -------- |
| **TI2VidTwoStagesPipeline** | 2 | ✅ | ✅ | Image | **Production quality** (recommended) | | **TI2VidTwoStagesPipeline** | 2 | ✅ | ✅ | Image | **Production quality** (recommended) |
| **TI2VidOneStagePipeline** | 1 | ✅ | ❌ | Image | Educational, prototyping | | **TI2VidOneStagePipeline** | 1 | ✅ | ❌ | Image | Educational, prototyping |
@@ -107,7 +107,7 @@ Do you need to condition on existing images/videos?
**Source**: [`src/ltx_pipelines/ti2vid_two_stages.py`](src/ltx_pipelines/ti2vid_two_stages.py) **Source**: [`src/ltx_pipelines/ti2vid_two_stages.py`](src/ltx_pipelines/ti2vid_two_stages.py)
Two-stage generation: Stage 1 generates low-resolution video with CFG guidance, Stage 2 upsamples to 2x resolution with distilled LoRA refinement. Supports image conditioning. Highest quality output, slower than one-stage but significantly better quality. Two-stage generation: Stage 1 generates low-resolution video with [multimodal guidance](#%EF%B8%8F-multimodal-guidance), Stage 2 upsamples to 2x resolution with distilled LoRA refinement. Supports image conditioning. Highest quality output, slower than one-stage but significantly better quality.
**Use when:** Production-quality video generation, higher resolution needed, quality over speed, text-to-video with image conditioning. **Use when:** Production-quality video generation, higher resolution needed, quality over speed, text-to-video with image conditioning.
@@ -121,7 +121,7 @@ Two-stage generation: Stage 1 generates low-resolution video with CFG guidance,
> **⚠️ Important:** This pipeline is primarily for educational purposes. For production-quality results, use `TI2VidTwoStagesPipeline` or other two-stage pipelines. > **⚠️ Important:** This pipeline is primarily for educational purposes. For production-quality results, use `TI2VidTwoStagesPipeline` or other two-stage pipelines.
Single-stage generation (no upsampling) with CFG guidance and image conditioning support. Faster inference but lower resolution output (typically 512x768). Single-stage generation (no upsampling) with [multimodal guidance](#%EF%B8%8F-multimodal-guidance) and image conditioning support. Faster inference but lower resolution output (typically 512x768).
**Use when:** Learning how the pipeline works, quick prototyping, testing, or when high resolution is not needed. **Use when:** Learning how the pipeline works, quick prototyping, testing, or when high resolution is not needed.
@@ -133,7 +133,7 @@ Single-stage generation (no upsampling) with CFG guidance and image conditioning
**Source**: [`src/ltx_pipelines/distilled.py`](src/ltx_pipelines/distilled.py) **Source**: [`src/ltx_pipelines/distilled.py`](src/ltx_pipelines/distilled.py)
Two-stage generation with 8 predefined sigmas (8 steps in stage 1, 4 steps in stage 2). No CFG guidance required. Fastest inference among all pipelines. Supports image conditioning. Requires spatial upsampler. Two-stage generation with 8 predefined sigmas (8 steps in stage 1, 4 steps in stage 2). No guidance required. Fastest inference among all pipelines. Supports image conditioning. Requires spatial upsampler.
**Use when:** Fastest inference is critical, batch processing many videos, or when you have a distilled model checkpoint. **Use when:** Fastest inference is critical, batch processing many videos, or when you have a distilled model checkpoint.
@@ -157,7 +157,7 @@ Two-stage generation with IC-LoRA support. Can condition on reference videos (vi
**Source**: [`src/ltx_pipelines/keyframe_interpolation.py`](src/ltx_pipelines/keyframe_interpolation.py) **Source**: [`src/ltx_pipelines/keyframe_interpolation.py`](src/ltx_pipelines/keyframe_interpolation.py)
Two-stage generation with keyframe interpolation. Uses guiding latents (additive conditioning) instead of replacing latents for smoother transitions. CFG guidance in stage 1, upsampling in stage 2. Two-stage generation with keyframe interpolation. Uses guiding latents (additive conditioning) instead of replacing latents for smoother transitions. [Multimodal guidance](#%EF%B8%8F-multimodal-guidance) in stage 1, upsampling in stage 2.
**Use when:** You have keyframe images and want to interpolate between them, creating smooth transitions, or animation/motion interpolation tasks. **Use when:** You have keyframe images and want to interpolate between them, creating smooth transitions, or animation/motion interpolation tasks.
@@ -190,6 +190,61 @@ All pipelines support image conditioning, but with different methods:
--- ---
## 🎛️ Multimodal Guidance
LTX-2 pipelines use **multimodal guidance** to steer the diffusion process for both video and audio modalities. Each modality (video, audio) has its own guider with independent parameters, allowing fine-grained control over generation quality and adherence to prompts.
### Guidance Parameters
The `MultiModalGuiderParams` dataclass controls guidance behavior:
| Parameter | Description |
| --------- | ----------- |
| `cfg_scale` | **Classifier-Free Guidance** scale. Higher values make the output adhere more strongly to the text prompt. Typical values: 2.05.0. Set to **1.0** to disable. |
| `stg_scale` | **Spatio-Temporal Guidance** scale. Controls perturbation-based guidance for improved temporal coherence. Typical values: 0.51.5. Set to **0.0** to disable. |
| `stg_blocks` | Which transformer blocks to perturb for STG (e.g., `[29]` for the last block). Set to **`[]`** to disable STG. |
| `rescale_scale` | Rescales the guided prediction to match the variance of the conditional prediction. Helps prevent over-saturation. Typical values: 0.50.7. Set to **0.0** to disable. |
| `modality_scale` | **Modality CFG** scale. Steers the model away from unsynced video and audio results, improving audio-visual coherence. Set to **1.0** to disable. |
| `skip_step` | Skip guidance every N steps. Can speed up inference with minimal quality loss. Set to **0** to disable (never skip). |
### How It Works
The multimodal guider combines three guidance signals during each denoising step:
1. **CFG (Text Guidance)**: Steers generation toward the text prompt by computing `(cond - uncond_text)`.
2. **STG (Perturbation Guidance)**: Improves structural coherence by perturbing specific transformer blocks and steering away from the perturbed prediction.
3. **Modality CFG**: For joint audio-video generation, steers the model away from unsynced video and audio results.
### Example Configuration
```python
from ltx_core.components.guiders import MultiModalGuiderParams
# Video guider: moderate CFG, STG enabled, modality isolation
video_guider_params = MultiModalGuiderParams(
cfg_scale=3.0,
stg_scale=1.0,
rescale_scale=0.7,
modality_scale=3.0,
stg_blocks=[29],
)
# Audio guider: higher CFG for stronger prompt adherence
audio_guider_params = MultiModalGuiderParams(
cfg_scale=7.0,
stg_scale=1.0,
rescale_scale=0.7,
modality_scale=3.0,
stg_blocks=[29],
)
```
> **Tip:** Start with the default values from [`constants.py`](src/ltx_pipelines/utils/constants.py) and adjust based on your use case. Higher `cfg_scale` = stronger prompt adherence but potentially less natural motion; higher `stg_scale` = better temporal coherence but slower inference (requires extra forward passes).
>
> **Tip:** When generating video with audio, set `modality_scale` > 1.0 (e.g., 3.0) to improve audio-visual sync. If generating video-only, set it to 1.0 to disable.
---
## ⚡ Optimization Tips ## ⚡ Optimization Tips
@@ -276,6 +331,7 @@ This allows you to use **20-30 steps instead of 40** while maintaining quality.
```python ```python
from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
from ltx_core.components.guiders import MultiModalGuiderParams
distilled_lora = [ distilled_lora = [
LoraPathStrengthAndSDOps( LoraPathStrengthAndSDOps(
@@ -293,6 +349,24 @@ pipeline = TI2VidTwoStagesPipeline(
loras=[], loras=[],
) )
video_guider_params = MultiModalGuiderParams(
cfg_scale=3.0,
stg_scale=1.0,
rescale_scale=0.7,
modality_scale=3.0,
skip_step=0,
stg_blocks=[29],
)
audio_guider_params = MultiModalGuiderParams(
cfg_scale=7.0,
stg_scale=1.0,
rescale_scale=0.7,
modality_scale=3.0,
skip_step=0,
stg_blocks=[29],
)
# Generate video from image # Generate video from image
pipeline( pipeline(
prompt="A serene landscape with mountains in the background", prompt="A serene landscape with mountains in the background",
@@ -303,7 +377,8 @@ pipeline(
num_frames=121, num_frames=121,
frame_rate=25.0, frame_rate=25.0,
num_inference_steps=40, num_inference_steps=40,
cfg_guidance_scale=3.0, video_guider_params=video_guider_params,
audio_guider_params=audio_guider_params,
images=[("input_image.jpg", 0, 1.0)], # Image at frame 0, strength 1.0 images=[("input_image.jpg", 0, 1.0)], # Image at frame 0, strength 1.0
) )
``` ```
@@ -2,11 +2,12 @@ import logging
from collections.abc import Iterator from collections.abc import Iterator
import torch import torch
from safetensors import safe_open
from ltx_core.components.diffusion_steps import EulerDiffusionStep from ltx_core.components.diffusion_steps import EulerDiffusionStep
from ltx_core.components.noisers import GaussianNoiser from ltx_core.components.noisers import GaussianNoiser
from ltx_core.components.protocols import DiffusionStepProtocol from ltx_core.components.protocols import DiffusionStepProtocol
from ltx_core.conditioning import ConditioningItem, VideoConditionByKeyframeIndex from ltx_core.conditioning import ConditioningItem, VideoConditionByReferenceLatent
from ltx_core.loader import LoraPathStrengthAndSDOps from ltx_core.loader import LoraPathStrengthAndSDOps
from ltx_core.model.audio_vae import decode_audio as vae_decode_audio from ltx_core.model.audio_vae import decode_audio as vae_decode_audio
from ltx_core.model.upsampler import upsample_video from ltx_core.model.upsampler import upsample_video
@@ -81,6 +82,21 @@ class ICLoraPipeline:
) )
self.device = device self.device = device
# Read reference downscale factor from LoRA metadata.
# IC-LoRAs trained with low-resolution reference videos store this factor
# so inference can resize reference videos to match training conditions.
self.reference_downscale_factor = 1
for lora in loras:
scale = _read_lora_reference_downscale_factor(lora.path)
if scale != 1:
if self.reference_downscale_factor not in (1, scale):
raise ValueError(
f"Conflicting reference_downscale_factor values in LoRAs: "
f"already have {self.reference_downscale_factor}, but {lora.path} "
f"specifies {scale}. Cannot combine LoRAs with different reference scales."
)
self.reference_downscale_factor = scale
@torch.inference_mode() @torch.inference_mode()
def __call__( def __call__(
self, self,
@@ -249,17 +265,34 @@ class ICLoraPipeline:
device=self.device, device=self.device,
) )
# Calculate scaled dimensions for reference video conditioning.
# IC-LoRAs trained with downscaled reference videos expect the same ratio at inference.
scale = self.reference_downscale_factor
if scale != 1 and (height % scale != 0 or width % scale != 0):
raise ValueError(
f"Output dimensions ({height}x{width}) must be divisible by reference_downscale_factor ({scale})"
)
ref_height = height // scale
ref_width = width // scale
for video_path, strength in video_conditioning: for video_path, strength in video_conditioning:
# Load video at scaled-down resolution (if scale > 1)
video = load_video_conditioning( video = load_video_conditioning(
video_path=video_path, video_path=video_path,
height=height, height=ref_height,
width=width, width=ref_width,
frame_cap=num_frames, frame_cap=num_frames,
dtype=self.dtype, dtype=self.dtype,
device=self.device, device=self.device,
) )
encoded_video = video_encoder(video) encoded_video = video_encoder(video)
conditionings.append(VideoConditionByKeyframeIndex(keyframes=encoded_video, frame_idx=0, strength=strength)) conditionings.append(
VideoConditionByReferenceLatent(
latent=encoded_video,
downscale_factor=scale,
strength=strength,
)
)
return conditionings return conditionings
@@ -307,5 +340,26 @@ def main() -> None:
) )
def _read_lora_reference_downscale_factor(lora_path: str) -> int:
"""Read reference_downscale_factor from LoRA safetensors metadata.
Some IC-LoRA models are trained with reference videos at lower resolution than
the target output. This allows for more efficient training and can improve
generalization. The downscale factor indicates the ratio between target and
reference resolutions (e.g., factor=2 means reference is half the resolution).
Args:
lora_path: Path to the LoRA .safetensors file
Returns:
The reference downscale factor (1 if not specified in metadata, meaning
reference and target have the same resolution)
"""
try:
with safe_open(lora_path, framework="pt") as f:
metadata = f.metadata() or {}
return int(metadata.get("reference_downscale_factor", 1))
except Exception as e:
logging.warning(f"Failed to read metadata from LoRA file '{lora_path}': {e}")
return 1
if __name__ == "__main__": if __name__ == "__main__":
main() main()
@@ -4,7 +4,7 @@ from collections.abc import Iterator
import torch import torch
from ltx_core.components.diffusion_steps import EulerDiffusionStep from ltx_core.components.diffusion_steps import EulerDiffusionStep
from ltx_core.components.guiders import CFGGuider from ltx_core.components.guiders import MultiModalGuider, MultiModalGuiderParams
from ltx_core.components.noisers import GaussianNoiser from ltx_core.components.noisers import GaussianNoiser
from ltx_core.components.protocols import DiffusionStepProtocol from ltx_core.components.protocols import DiffusionStepProtocol
from ltx_core.components.schedulers import LTX2Scheduler from ltx_core.components.schedulers import LTX2Scheduler
@@ -28,8 +28,8 @@ from ltx_pipelines.utils.helpers import (
euler_denoising_loop, euler_denoising_loop,
generate_enhanced_prompt, generate_enhanced_prompt,
get_device, get_device,
guider_denoising_func,
image_conditionings_by_adding_guiding_latent, image_conditionings_by_adding_guiding_latent,
multi_modal_guider_denoising_func,
simple_denoising_func, simple_denoising_func,
) )
from ltx_pipelines.utils.media_io import encode_video from ltx_pipelines.utils.media_io import encode_video
@@ -86,7 +86,8 @@ class KeyframeInterpolationPipeline:
num_frames: int, num_frames: int,
frame_rate: float, frame_rate: float,
num_inference_steps: int, num_inference_steps: int,
cfg_guidance_scale: float, video_guider_params: MultiModalGuiderParams,
audio_guider_params: MultiModalGuiderParams,
images: list[tuple[str, int, float]], images: list[tuple[str, int, float]],
tiling_config: TilingConfig | None = None, tiling_config: TilingConfig | None = None,
enhance_prompt: bool = False, enhance_prompt: bool = False,
@@ -96,7 +97,6 @@ class KeyframeInterpolationPipeline:
generator = torch.Generator(device=self.device).manual_seed(seed) generator = torch.Generator(device=self.device).manual_seed(seed)
noiser = GaussianNoiser(generator=generator) noiser = GaussianNoiser(generator=generator)
stepper = EulerDiffusionStep() stepper = EulerDiffusionStep()
cfg_guider = CFGGuider(cfg_guidance_scale)
dtype = torch.bfloat16 dtype = torch.bfloat16
text_encoder = self.stage_1_model_ledger.text_encoder() text_encoder = self.stage_1_model_ledger.text_encoder()
@@ -125,12 +125,17 @@ class KeyframeInterpolationPipeline:
video_state=video_state, video_state=video_state,
audio_state=audio_state, audio_state=audio_state,
stepper=stepper, stepper=stepper,
denoise_fn=guider_denoising_func( denoise_fn=multi_modal_guider_denoising_func(
cfg_guider, video_guider=MultiModalGuider(
v_context_p, params=video_guider_params,
v_context_n, negative_context=v_context_n,
a_context_p, ),
a_context_n, audio_guider=MultiModalGuider(
params=audio_guider_params,
negative_context=a_context_n,
),
v_context=v_context_p,
a_context=a_context_p,
transformer=transformer, # noqa: F821 transformer=transformer, # noqa: F821
), ),
) )
@@ -256,7 +261,22 @@ def main() -> None:
num_frames=args.num_frames, num_frames=args.num_frames,
frame_rate=args.frame_rate, frame_rate=args.frame_rate,
num_inference_steps=args.num_inference_steps, num_inference_steps=args.num_inference_steps,
cfg_guidance_scale=args.cfg_guidance_scale, video_guider_params=MultiModalGuiderParams(
cfg_scale=args.video_cfg_guidance_scale,
stg_scale=args.video_stg_guidance_scale,
rescale_scale=args.video_rescale_scale,
modality_scale=args.a2v_guidance_scale,
skip_step=args.video_skip_step,
stg_blocks=args.video_stg_blocks,
),
audio_guider_params=MultiModalGuiderParams(
cfg_scale=args.audio_cfg_guidance_scale,
stg_scale=args.audio_stg_guidance_scale,
rescale_scale=args.audio_rescale_scale,
modality_scale=args.v2a_guidance_scale,
skip_step=args.audio_skip_step,
stg_blocks=args.audio_stg_blocks,
),
images=args.images, images=args.images,
tiling_config=tiling_config, tiling_config=tiling_config,
) )
@@ -4,7 +4,7 @@ from collections.abc import Iterator
import torch import torch
from ltx_core.components.diffusion_steps import EulerDiffusionStep from ltx_core.components.diffusion_steps import EulerDiffusionStep
from ltx_core.components.guiders import CFGGuider from ltx_core.components.guiders import MultiModalGuider, MultiModalGuiderParams
from ltx_core.components.noisers import GaussianNoiser from ltx_core.components.noisers import GaussianNoiser
from ltx_core.components.protocols import DiffusionStepProtocol from ltx_core.components.protocols import DiffusionStepProtocol
from ltx_core.components.schedulers import LTX2Scheduler from ltx_core.components.schedulers import LTX2Scheduler
@@ -23,8 +23,8 @@ from ltx_pipelines.utils.helpers import (
euler_denoising_loop, euler_denoising_loop,
generate_enhanced_prompt, generate_enhanced_prompt,
get_device, get_device,
guider_denoising_func,
image_conditionings_by_replacing_latent, image_conditionings_by_replacing_latent,
multi_modal_guider_denoising_func,
) )
from ltx_pipelines.utils.media_io import encode_video from ltx_pipelines.utils.media_io import encode_video
from ltx_pipelines.utils.types import PipelineComponents from ltx_pipelines.utils.types import PipelineComponents
@@ -73,7 +73,8 @@ class TI2VidOneStagePipeline:
num_frames: int, num_frames: int,
frame_rate: float, frame_rate: float,
num_inference_steps: int, num_inference_steps: int,
cfg_guidance_scale: float, video_guider_params: MultiModalGuiderParams,
audio_guider_params: MultiModalGuiderParams,
images: list[tuple[str, int, float]], images: list[tuple[str, int, float]],
enhance_prompt: bool = False, enhance_prompt: bool = False,
) -> tuple[Iterator[torch.Tensor], torch.Tensor]: ) -> tuple[Iterator[torch.Tensor], torch.Tensor]:
@@ -82,7 +83,6 @@ class TI2VidOneStagePipeline:
generator = torch.Generator(device=self.device).manual_seed(seed) generator = torch.Generator(device=self.device).manual_seed(seed)
noiser = GaussianNoiser(generator=generator) noiser = GaussianNoiser(generator=generator)
stepper = EulerDiffusionStep() stepper = EulerDiffusionStep()
cfg_guider = CFGGuider(cfg_guidance_scale)
dtype = torch.bfloat16 dtype = torch.bfloat16
text_encoder = self.model_ledger.text_encoder() text_encoder = self.model_ledger.text_encoder()
@@ -111,12 +111,17 @@ class TI2VidOneStagePipeline:
video_state=video_state, video_state=video_state,
audio_state=audio_state, audio_state=audio_state,
stepper=stepper, stepper=stepper,
denoise_fn=guider_denoising_func( denoise_fn=multi_modal_guider_denoising_func(
cfg_guider, video_guider=MultiModalGuider(
v_context_p, params=video_guider_params,
v_context_n, negative_context=v_context_n,
a_context_p, ),
a_context_n, audio_guider=MultiModalGuider(
params=audio_guider_params,
negative_context=a_context_n,
),
v_context=v_context_p,
a_context=a_context_p,
transformer=transformer, # noqa: F821 transformer=transformer, # noqa: F821
), ),
) )
@@ -175,7 +180,22 @@ def main() -> None:
num_frames=args.num_frames, num_frames=args.num_frames,
frame_rate=args.frame_rate, frame_rate=args.frame_rate,
num_inference_steps=args.num_inference_steps, num_inference_steps=args.num_inference_steps,
cfg_guidance_scale=args.cfg_guidance_scale, video_guider_params=MultiModalGuiderParams(
cfg_scale=args.video_cfg_guidance_scale,
stg_scale=args.video_stg_guidance_scale,
rescale_scale=args.video_rescale_scale,
modality_scale=args.a2v_guidance_scale,
skip_step=args.video_skip_step,
stg_blocks=args.video_stg_blocks,
),
audio_guider_params=MultiModalGuiderParams(
cfg_scale=args.audio_cfg_guidance_scale,
stg_scale=args.audio_stg_guidance_scale,
rescale_scale=args.audio_rescale_scale,
modality_scale=args.v2a_guidance_scale,
skip_step=args.audio_skip_step,
stg_blocks=args.audio_stg_blocks,
),
images=args.images, images=args.images,
) )
@@ -4,7 +4,7 @@ from collections.abc import Iterator
import torch import torch
from ltx_core.components.diffusion_steps import EulerDiffusionStep from ltx_core.components.diffusion_steps import EulerDiffusionStep
from ltx_core.components.guiders import CFGGuider from ltx_core.components.guiders import MultiModalGuider, MultiModalGuiderParams
from ltx_core.components.noisers import GaussianNoiser from ltx_core.components.noisers import GaussianNoiser
from ltx_core.components.protocols import DiffusionStepProtocol from ltx_core.components.protocols import DiffusionStepProtocol
from ltx_core.components.schedulers import LTX2Scheduler from ltx_core.components.schedulers import LTX2Scheduler
@@ -28,8 +28,8 @@ from ltx_pipelines.utils.helpers import (
euler_denoising_loop, euler_denoising_loop,
generate_enhanced_prompt, generate_enhanced_prompt,
get_device, get_device,
guider_denoising_func,
image_conditionings_by_replacing_latent, image_conditionings_by_replacing_latent,
multi_modal_guider_denoising_func,
simple_denoising_func, simple_denoising_func,
) )
from ltx_pipelines.utils.media_io import encode_video from ltx_pipelines.utils.media_io import encode_video
@@ -88,7 +88,8 @@ class TI2VidTwoStagesPipeline:
num_frames: int, num_frames: int,
frame_rate: float, frame_rate: float,
num_inference_steps: int, num_inference_steps: int,
cfg_guidance_scale: float, video_guider_params: MultiModalGuiderParams,
audio_guider_params: MultiModalGuiderParams,
images: list[tuple[str, int, float]], images: list[tuple[str, int, float]],
tiling_config: TilingConfig | None = None, tiling_config: TilingConfig | None = None,
enhance_prompt: bool = False, enhance_prompt: bool = False,
@@ -98,7 +99,6 @@ class TI2VidTwoStagesPipeline:
generator = torch.Generator(device=self.device).manual_seed(seed) generator = torch.Generator(device=self.device).manual_seed(seed)
noiser = GaussianNoiser(generator=generator) noiser = GaussianNoiser(generator=generator)
stepper = EulerDiffusionStep() stepper = EulerDiffusionStep()
cfg_guider = CFGGuider(cfg_guidance_scale)
dtype = torch.bfloat16 dtype = torch.bfloat16
text_encoder = self.stage_1_model_ledger.text_encoder() text_encoder = self.stage_1_model_ledger.text_encoder()
@@ -127,12 +127,17 @@ class TI2VidTwoStagesPipeline:
video_state=video_state, video_state=video_state,
audio_state=audio_state, audio_state=audio_state,
stepper=stepper, stepper=stepper,
denoise_fn=guider_denoising_func( denoise_fn=multi_modal_guider_denoising_func(
cfg_guider, video_guider=MultiModalGuider(
v_context_p, params=video_guider_params,
v_context_n, negative_context=v_context_n,
a_context_p, ),
a_context_n, audio_guider=MultiModalGuider(
params=audio_guider_params,
negative_context=a_context_n,
),
v_context=v_context_p,
a_context=a_context_p,
transformer=transformer, # noqa: F821 transformer=transformer, # noqa: F821
), ),
) )
@@ -259,7 +264,22 @@ def main() -> None:
num_frames=args.num_frames, num_frames=args.num_frames,
frame_rate=args.frame_rate, frame_rate=args.frame_rate,
num_inference_steps=args.num_inference_steps, num_inference_steps=args.num_inference_steps,
cfg_guidance_scale=args.cfg_guidance_scale, video_guider_params=MultiModalGuiderParams(
cfg_scale=args.video_cfg_guidance_scale,
stg_scale=args.video_stg_guidance_scale,
rescale_scale=args.video_rescale_scale,
modality_scale=args.a2v_guidance_scale,
skip_step=args.video_skip_step,
stg_blocks=args.video_stg_blocks,
),
audio_guider_params=MultiModalGuiderParams(
cfg_scale=args.audio_cfg_guidance_scale,
stg_scale=args.audio_stg_guidance_scale,
rescale_scale=args.audio_rescale_scale,
modality_scale=args.v2a_guidance_scale,
skip_step=args.audio_skip_step,
stg_blocks=args.audio_stg_blocks,
),
images=args.images, images=args.images,
tiling_config=tiling_config, tiling_config=tiling_config,
) )
@@ -7,13 +7,14 @@ from ltx_pipelines.utils.constants import (
DEFAULT_1_STAGE_WIDTH, DEFAULT_1_STAGE_WIDTH,
DEFAULT_2_STAGE_HEIGHT, DEFAULT_2_STAGE_HEIGHT,
DEFAULT_2_STAGE_WIDTH, DEFAULT_2_STAGE_WIDTH,
DEFAULT_CFG_GUIDANCE_SCALE, DEFAULT_AUDIO_GUIDER_PARAMS,
DEFAULT_FRAME_RATE, DEFAULT_FRAME_RATE,
DEFAULT_LORA_STRENGTH, DEFAULT_LORA_STRENGTH,
DEFAULT_NEGATIVE_PROMPT, DEFAULT_NEGATIVE_PROMPT,
DEFAULT_NUM_FRAMES, DEFAULT_NUM_FRAMES,
DEFAULT_NUM_INFERENCE_STEPS, DEFAULT_NUM_INFERENCE_STEPS,
DEFAULT_SEED, DEFAULT_SEED,
DEFAULT_VIDEO_GUIDER_PARAMS,
) )
@@ -185,16 +186,6 @@ def basic_arg_parser() -> argparse.ArgumentParser:
def default_1_stage_arg_parser() -> argparse.ArgumentParser: def default_1_stage_arg_parser() -> argparse.ArgumentParser:
parser = basic_arg_parser() parser = basic_arg_parser()
parser.add_argument(
"--cfg-guidance-scale",
type=float,
default=DEFAULT_CFG_GUIDANCE_SCALE,
help=(
f"Classifier-free guidance (CFG) scale controlling how strongly "
f"the model adheres to the prompt. Higher values increase prompt "
f"adherence but may reduce diversity (default: {DEFAULT_CFG_GUIDANCE_SCALE})."
),
)
parser.add_argument( parser.add_argument(
"--negative-prompt", "--negative-prompt",
type=str, type=str,
@@ -205,7 +196,126 @@ def default_1_stage_arg_parser() -> argparse.ArgumentParser:
"Default: a comprehensive negative prompt covering common artifacts and quality issues." "Default: a comprehensive negative prompt covering common artifacts and quality issues."
), ),
) )
parser.add_argument(
"--video-cfg-guidance-scale",
type=float,
default=DEFAULT_VIDEO_GUIDER_PARAMS.cfg_scale,
help=(
f"Classifier-free guidance (CFG) scale controlling how strongly "
f"the model adheres to the video prompt. Higher values increase prompt "
"adherence but may reduce diversity. 1.0 means no effect "
f"(default: {DEFAULT_VIDEO_GUIDER_PARAMS.cfg_scale})."
),
)
parser.add_argument(
"--video-stg-guidance-scale",
type=float,
default=DEFAULT_VIDEO_GUIDER_PARAMS.stg_scale,
help=(
f"STG (Spatio-Temporal Guidance) scale controlling how strongly "
f"the model reacts to the perturbation of the video modality. Higher values increase "
f"the effect but may reduce quality. 0.0 means no effect "
f"(default: {DEFAULT_VIDEO_GUIDER_PARAMS.stg_scale})."
),
)
parser.add_argument(
"--video-rescale-scale",
type=float,
default=DEFAULT_VIDEO_GUIDER_PARAMS.rescale_scale,
help=(
f"Rescale scale controlling how strongly "
f"the model rescales the video modality after applying other guidance. Higher values tend to decrease "
f"oversaturation effects. 0.0 means no effect (default: {DEFAULT_VIDEO_GUIDER_PARAMS.rescale_scale})."
),
)
parser.add_argument(
"--video-stg-blocks",
type=int,
nargs="*",
default=DEFAULT_VIDEO_GUIDER_PARAMS.stg_blocks,
help=(f"Which transformer blocks to perturb for STG. Default: {DEFAULT_VIDEO_GUIDER_PARAMS.stg_blocks}."),
)
parser.add_argument(
"--a2v-guidance-scale",
type=float,
default=DEFAULT_VIDEO_GUIDER_PARAMS.modality_scale,
help=(
f"A2V (Audio-to-Video) guidance scale controlling how strongly "
f"the model reacts to the perturbation of the audio-to-video cross-attention. Higher values may increase "
f"lipsync quality. 1.0 means no effect (default: {DEFAULT_VIDEO_GUIDER_PARAMS.modality_scale})."
),
)
parser.add_argument(
"--video-skip-step",
type=int,
default=DEFAULT_VIDEO_GUIDER_PARAMS.skip_step,
help=(
"Video skip step N controls periodic skipping during the video diffusion process: "
"only steps where step_index % (N + 1) == 0 are processed, all others are skipped "
f"(e.g., 0 = no skipping; 1 = skip every other step; 2 = skip 2 of every 3 steps; "
f"default: {DEFAULT_VIDEO_GUIDER_PARAMS.skip_step})."
),
)
parser.add_argument(
"--audio-cfg-guidance-scale",
type=float,
default=DEFAULT_AUDIO_GUIDER_PARAMS.cfg_scale,
help=(
f"Audio CFG (Classifier-free guidance) scale controlling how strongly "
f"the model adheres to the audio prompt. Higher values increase prompt "
f"adherence but may reduce diversity. 1.0 means no effect "
f"(default: {DEFAULT_AUDIO_GUIDER_PARAMS.cfg_scale})."
),
)
parser.add_argument(
"--audio-stg-guidance-scale",
type=float,
default=DEFAULT_AUDIO_GUIDER_PARAMS.stg_scale,
help=(
f"Audio STG (Spatio-Temporal Guidance) scale controlling how strongly "
f"the model reacts to the perturbation of the audio modality. Higher values increase "
f"the effect but may reduce quality. 0.0 means no effect "
f"(default: {DEFAULT_AUDIO_GUIDER_PARAMS.stg_scale})."
),
)
parser.add_argument(
"--audio-rescale-scale",
type=float,
default=DEFAULT_AUDIO_GUIDER_PARAMS.rescale_scale,
help=(
f"Audio rescale scale controlling how strongly "
f"the model rescales the audio modality after applying other guidance. "
f"Experimental. 0.0 means no effect (default: {DEFAULT_AUDIO_GUIDER_PARAMS.rescale_scale})."
),
)
parser.add_argument(
"--audio-stg-blocks",
type=int,
nargs="*",
default=DEFAULT_AUDIO_GUIDER_PARAMS.stg_blocks,
help=(f"Which transformer blocks to perturb for Audio STG. Default: {DEFAULT_AUDIO_GUIDER_PARAMS.stg_blocks}."),
)
parser.add_argument(
"--v2a-guidance-scale",
type=float,
default=DEFAULT_AUDIO_GUIDER_PARAMS.modality_scale,
help=(
f"V2A (Video-to-Audio) guidance scale controlling how strongly "
f"the model reacts to the perturbation of the video-to-audio cross-attention. Higher values may increase "
f"lipsync quality. 1.0 means no effect (default: {DEFAULT_AUDIO_GUIDER_PARAMS.modality_scale})."
),
)
parser.add_argument(
"--audio-skip-step",
type=int,
default=DEFAULT_AUDIO_GUIDER_PARAMS.skip_step,
help=(
"Audio skip step N controls periodic skipping during the audio diffusion process: "
"only steps where step_index % (N + 1) == 0 are processed, all others are skipped "
f"(e.g., 0 = no skipping; 1 = skip every other step; 2 = skip 2 of every 3 steps; "
f"default: {DEFAULT_AUDIO_GUIDER_PARAMS.skip_step})."
),
)
return parser return parser
@@ -4,6 +4,7 @@
# Noise schedule for the distilled pipeline. These sigma values control noise # Noise schedule for the distilled pipeline. These sigma values control noise
# levels at each denoising step and were tuned to match the distillation process. # levels at each denoising step and were tuned to match the distillation process.
from ltx_core.components.guiders import MultiModalGuiderParams
from ltx_core.types import SpatioTemporalScaleFactors from ltx_core.types import SpatioTemporalScaleFactors
DISTILLED_SIGMA_VALUES = [1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875, 0.0] DISTILLED_SIGMA_VALUES = [1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875, 0.0]
@@ -24,13 +25,28 @@ DEFAULT_2_STAGE_WIDTH = DEFAULT_1_STAGE_WIDTH * 2
DEFAULT_NUM_FRAMES = 121 DEFAULT_NUM_FRAMES = 121
DEFAULT_FRAME_RATE = 24.0 DEFAULT_FRAME_RATE = 24.0
DEFAULT_NUM_INFERENCE_STEPS = 40 DEFAULT_NUM_INFERENCE_STEPS = 40
DEFAULT_CFG_GUIDANCE_SCALE = 4.0 DEFAULT_VIDEO_GUIDER_PARAMS = MultiModalGuiderParams(
cfg_scale=3.0,
stg_scale=1.0,
rescale_scale=0.7,
modality_scale=3.0,
skip_step=0,
stg_blocks=[29],
)
# ============================================================================= # =============================================================================
# Audio # Audio
# ============================================================================= # =============================================================================
DEFAULT_AUDIO_GUIDER_PARAMS = MultiModalGuiderParams(
cfg_scale=7.0,
stg_scale=1.0,
rescale_scale=0.7,
modality_scale=3.0,
skip_step=0,
stg_blocks=[29],
)
AUDIO_SAMPLE_RATE = 24000 AUDIO_SAMPLE_RATE = 24000
@@ -5,6 +5,7 @@ from dataclasses import replace
import torch import torch
from tqdm import tqdm from tqdm import tqdm
from ltx_core.components.guiders import MultiModalGuider
from ltx_core.components.noisers import Noiser from ltx_core.components.noisers import Noiser
from ltx_core.components.protocols import DiffusionStepProtocol, GuiderProtocol from ltx_core.components.protocols import DiffusionStepProtocol, GuiderProtocol
from ltx_core.conditioning import ( from ltx_core.conditioning import (
@@ -12,6 +13,12 @@ from ltx_core.conditioning import (
VideoConditionByKeyframeIndex, VideoConditionByKeyframeIndex,
VideoConditionByLatentIndex, VideoConditionByLatentIndex,
) )
from ltx_core.guidance.perturbations import (
BatchedPerturbationConfig,
Perturbation,
PerturbationConfig,
PerturbationType,
)
from ltx_core.model.transformer import Modality, X0Model from ltx_core.model.transformer import Modality, X0Model
from ltx_core.model.video_vae import VideoEncoder from ltx_core.model.video_vae import VideoEncoder
from ltx_core.text_encoders.gemma import GemmaTextEncoderModelBase from ltx_core.text_encoders.gemma import GemmaTextEncoderModelBase
@@ -376,6 +383,113 @@ def guider_denoising_func(
return guider_denoising_step return guider_denoising_step
def multi_modal_guider_denoising_func(
video_guider: MultiModalGuider,
audio_guider: MultiModalGuider,
v_context: torch.Tensor,
a_context: torch.Tensor,
transformer: X0Model,
) -> DenoisingFunc:
last_denoised_video = None
last_denoised_audio = None
def guider_denoising_step(
video_state: LatentState, audio_state: LatentState, sigmas: torch.Tensor, step_index: int
) -> tuple[torch.Tensor, torch.Tensor]:
nonlocal last_denoised_video, last_denoised_audio
if video_guider.should_skip_step(step_index) and audio_guider.should_skip_step(step_index):
return last_denoised_video, last_denoised_audio
sigma = sigmas[step_index]
pos_video_modality = modality_from_latent_state(
video_state, v_context, sigma, enabled=not video_guider.should_skip_step(step_index)
)
pos_audio_modality = modality_from_latent_state(
audio_state, a_context, sigma, enabled=not audio_guider.should_skip_step(step_index)
)
denoised_video, denoised_audio = transformer(
video=pos_video_modality, audio=pos_audio_modality, perturbations=None
)
neg_denoised_video, neg_denoised_audio = 0.0, 0.0
if video_guider.do_unconditional_generation() or audio_guider.do_unconditional_generation():
if video_guider.do_unconditional_generation() and video_guider.negative_context is None:
raise ValueError("Negative context is required for unconditioned denoising")
if audio_guider.do_unconditional_generation() and audio_guider.negative_context is None:
raise ValueError("Negative context is required for unconditioned denoising")
neg_video_modality = modality_from_latent_state(
video_state,
video_guider.negative_context
if video_guider.negative_context is not None
else pos_video_modality.context,
sigma,
)
neg_audio_modality = modality_from_latent_state(
audio_state,
audio_guider.negative_context
if audio_guider.negative_context is not None
else pos_audio_modality.context,
sigma,
)
neg_denoised_video, neg_denoised_audio = transformer(
video=neg_video_modality, audio=neg_audio_modality, perturbations=None
)
ptb_denoised_video, ptb_denoised_audio = 0.0, 0.0
if video_guider.do_perturbed_generation() or audio_guider.do_perturbed_generation():
perturbations = []
if video_guider.do_perturbed_generation():
perturbations.append(
Perturbation(type=PerturbationType.SKIP_VIDEO_SELF_ATTN, blocks=video_guider.params.stg_blocks)
)
if audio_guider.do_perturbed_generation():
perturbations.append(
Perturbation(type=PerturbationType.SKIP_AUDIO_SELF_ATTN, blocks=audio_guider.params.stg_blocks)
)
perturbation_config = PerturbationConfig(perturbations=perturbations)
ptb_denoised_video, ptb_denoised_audio = transformer(
video=pos_video_modality,
audio=pos_audio_modality,
perturbations=BatchedPerturbationConfig(perturbations=[perturbation_config]),
)
mod_denoised_video, mod_denoised_audio = 0.0, 0.0
if video_guider.do_isolated_modality_generation() or audio_guider.do_isolated_modality_generation():
perturbations = [
Perturbation(type=PerturbationType.SKIP_A2V_CROSS_ATTN, blocks=None),
Perturbation(type=PerturbationType.SKIP_V2A_CROSS_ATTN, blocks=None),
]
perturbation_config = PerturbationConfig(perturbations=perturbations)
mod_denoised_video, mod_denoised_audio = transformer(
video=pos_video_modality,
audio=pos_audio_modality,
perturbations=BatchedPerturbationConfig(perturbations=[perturbation_config]),
)
if video_guider.should_skip_step(step_index):
denoised_video = last_denoised_video
else:
denoised_video = video_guider.calculate(
denoised_video, neg_denoised_video, ptb_denoised_video, mod_denoised_video
)
if audio_guider.should_skip_step(step_index):
denoised_audio = last_denoised_audio
else:
denoised_audio = audio_guider.calculate(
denoised_audio, neg_denoised_audio, ptb_denoised_audio, mod_denoised_audio
)
last_denoised_video = denoised_video
last_denoised_audio = denoised_audio
return denoised_video, denoised_audio
return guider_denoising_step
def denoise_audio_video( # noqa: PLR0913 def denoise_audio_video( # noqa: PLR0913
output_shape: VideoPixelShape, output_shape: VideoPixelShape,
conditionings: list[ConditioningItem], conditionings: list[ConditioningItem],
@@ -35,6 +35,8 @@ from ltx_core.text_encoders.gemma import (
AVGemmaTextEncoderModelConfigurator, AVGemmaTextEncoderModelConfigurator,
module_ops_from_gemma_root, module_ops_from_gemma_root,
) )
from ltx_core.text_encoders.gemma.encoders.av_encoder import GEMMA_MODEL_OPS
from ltx_core.utils import find_matching_file
class ModelLedger: class ModelLedger:
@@ -143,12 +145,16 @@ class ModelLedger:
) )
if self.gemma_root_path is not None: if self.gemma_root_path is not None:
module_ops = module_ops_from_gemma_root(self.gemma_root_path)
model_folder = find_matching_file(self.gemma_root_path, "model*.safetensors").parent
weight_paths = [str(p) for p in model_folder.rglob("*.safetensors")]
self.text_encoder_builder = Builder( self.text_encoder_builder = Builder(
model_path=self.checkpoint_path, model_path=(str(self.checkpoint_path), *weight_paths),
model_class_configurator=AVGemmaTextEncoderModelConfigurator, model_class_configurator=AVGemmaTextEncoderModelConfigurator,
model_sd_ops=AV_GEMMA_TEXT_ENCODER_KEY_OPS, model_sd_ops=AV_GEMMA_TEXT_ENCODER_KEY_OPS,
registry=self.registry, registry=self.registry,
module_ops=module_ops_from_gemma_root(self.gemma_root_path), module_ops=(GEMMA_MODEL_OPS, *module_ops),
) )
if self.spatial_upsampler_path is not None: if self.spatial_upsampler_path is not None:
@@ -200,6 +200,12 @@ validation:
- "/path/to/reference_video_1.mp4" - "/path/to/reference_video_1.mp4"
- "/path/to/reference_video_2.mp4" - "/path/to/reference_video_2.mp4"
# Downscale factor for reference videos (for efficient IC-LoRA training)
# When > 1, reference videos are processed at 1/n resolution
# Must match the --reference-downscale-factor used during dataset preprocessing
# Examples: 1 = same resolution, 2 = half resolution (384x384 ref for 768x768 target)
reference_downscale_factor: 1
# Negative prompt to avoid unwanted artifacts # Negative prompt to avoid unwanted artifacts
negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted" negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted"
+73 -17
View File
@@ -111,7 +111,8 @@ IC-LoRA enables a wide range of advanced video-to-video applications, such as:
- **Colorization**: Convert grayscale reference videos into colorized outputs - **Colorization**: Convert grayscale reference videos into colorized outputs
- **Restoration and enhancement**: Denoise, upscale, or restore old or degraded videos - **Restoration and enhancement**: Denoise, upscale, or restore old or degraded videos
By providing paired reference and target videos, IC-LoRA can learn complex transformations that go beyond caption-based conditioning. By providing paired reference and target videos, IC-LoRA can learn complex transformations that go beyond caption-based
conditioning.
IC-LoRA training fundamentally differs from standard LoRA and full fine-tuning: IC-LoRA training fundamentally differs from standard LoRA and full fine-tuning:
@@ -140,8 +141,10 @@ training_strategy:
### Dataset Requirements for IC-LoRA ### Dataset Requirements for IC-LoRA
- Your dataset must contain **paired videos** where each target video has a corresponding reference video - Your dataset must contain **paired videos** where each target video has a corresponding reference video
- Reference and target videos must have **identical resolution and length** - Reference and target videos must have the **same frame count** (length)
- Both reference and target videos should be **preprocessed together** using the same resolution buckets - Reference videos can optionally be at **lower spatial resolution** than target videos (
see [Scaled Reference Conditioning](#scaled-reference-conditioning) below)
- Both reference and target videos should be **preprocessed** before training
**Dataset structure for IC-LoRA training:** **Dataset structure for IC-LoRA training:**
@@ -181,30 +184,82 @@ validation:
reference_videos: reference_videos:
- "/path/to/reference1.mp4" - "/path/to/reference1.mp4"
- "/path/to/reference2.mp4" - "/path/to/reference2.mp4"
reference_downscale_factor: 1 # Set to match preprocessing (e.g., 2 for half resolution)
include_reference_in_output: true # Show reference side-by-side with output include_reference_in_output: true # Show reference side-by-side with output
``` ```
### Scaled Reference Conditioning
For more efficient training and inference, you can use **downscaled reference videos** while keeping target videos at
full resolution. This reduces the number of conditioning tokens, leading to:
- **Faster training** due to shorter sequence lengths
- **Faster inference** with reduced memory usage
- **Same aspect ratio** maintained between reference and target
#### How It Works
When the reference video has resolution `H/n × W/n` and the target video has resolution `H × W`, the trainer
automatically detects this scale factor `n` and adjusts the positional encodings so that the reference positions
map to the correct locations in the target coordinate space.
#### Preprocessing Datasets with Scaled References
Use the `--reference-downscale-factor` option when running `process_dataset.py`:
```bash
# Process dataset with scaled reference videos (half resolution)
uv run python scripts/process_dataset.py dataset.json \
--resolution-buckets 768x768x25 \
--model-path /path/to/ltx2.safetensors \
--text-encoder-path /path/to/gemma \
--reference-column "reference_path" \
--reference-downscale-factor 2
```
This will:
- Process target videos at 768×768 resolution
- Process reference videos at 384×384 resolution (768 / 2)
- The trainer will automatically infer the scale factor from the dimension ratio
**Important**: Set `reference_downscale_factor: 2` in your validation configuration to match the preprocessing:
```yaml
validation:
reference_downscale_factor: 2 # Must match the preprocessing factor
reference_videos:
- "/path/to/reference1.mp4"
- "/path/to/reference2.mp4"
```
> [!NOTE]
> The scale factor must be a positive integer, and all dimensions must be divisible by 32.
> Common scale factors are 1 (no scaling), 2 (half resolution), or 4 (quarter resolution).
## 📊 Training Mode Comparison ## 📊 Training Mode Comparison
| Aspect | LoRA | Audio-Video LoRA | Full Fine-tuning | IC-LoRA | | Aspect | LoRA | Audio-Video LoRA | Full Fine-tuning | IC-LoRA |
|----------------------|------------|------------------|------------------|----------------| |----------------------|--------------------------------|--------------------------------|------------------|--------------------------------|
| **Memory Usage** | Low | Low-Medium | High | Medium | | **Memory Usage** | Low | Low-Medium | High | Medium |
| **Training Speed** | Fast | Fast | Slow | Medium | | **Training Speed** | Fast | Fast | Slow | Medium |
| **Output Size** | 100MB-few GB (depends on rank) | 100MB-few GB (depends on rank) | Tens of GB | 100MB-few GB (depends on rank) | | **Output Size** | 100MB-few GB (depends on rank) | 100MB-few GB (depends on rank) | Tens of GB | 100MB-few GB (depends on rank) |
| **Flexibility** | Medium | Medium | High | Specialized | | **Flexibility** | Medium | Medium | High | Specialized |
| **Audio Support** | Optional | Yes | Optional | No | | **Audio Support** | Optional | Yes | Optional | No |
| **Reference Videos** | No | No | No | Yes (required) | | **Reference Videos** | No | No | No | Yes (required) |
## 🎬 Using Trained Models for Inference ## 🎬 Using Trained Models for Inference
After training, use the [`ltx-pipelines`](../../ltx-pipelines/) package for production inference with your trained LoRAs: After training, use the [`ltx-pipelines`](../../ltx-pipelines/) package for production inference with your trained
LoRAs:
| Training Mode | Recommended Pipeline | | Training Mode | Recommended Pipeline |
|---------------|---------------------| |-------------------------|-------------------------------------------------------|
| LoRA / Audio-Video LoRA | `TI2VidOneStagePipeline` or `TI2VidTwoStagesPipeline` | | LoRA / Audio-Video LoRA | `TI2VidOneStagePipeline` or `TI2VidTwoStagesPipeline` |
| IC-LoRA | `ICLoraPipeline` | | IC-LoRA | `ICLoraPipeline` |
All pipelines support loading custom LoRAs via the `loras` parameter. See the [`ltx-pipelines`](../../ltx-pipelines/) package All pipelines support loading custom LoRAs via the `loras` parameter. See the [`ltx-pipelines`](../../ltx-pipelines/)
package
documentation for detailed usage instructions. documentation for detailed usage instructions.
## 🚀 Next Steps ## 🚀 Next Steps
@@ -216,6 +271,7 @@ Once you've chosen your training mode:
- Start training with the [Training Guide](training-guide.md) - Start training with the [Training Guide](training-guide.md)
> [!TIP] > [!TIP]
> Need a training mode that's not covered here? See [Implementing Custom Training Strategies](custom-training-strategies.md) > Need a training mode that's not covered here?
> See [Implementing Custom Training Strategies](custom-training-strategies.md)
> to learn how to create your own strategy for specialized use cases like video inpainting, audio-only training, or > to learn how to create your own strategy for specialized use cases like video inpainting, audio-only training, or
> custom conditioning. > custom conditioning.
@@ -16,13 +16,14 @@ from pathlib import Path
import typer import typer
from decode_latents import LatentsDecoder from decode_latents import LatentsDecoder
from process_captions import compute_captions_embeddings from process_captions import compute_captions_embeddings
from process_videos import compute_latents, parse_resolution_buckets from process_videos import compute_latents, compute_scaled_resolution_buckets, parse_resolution_buckets
from rich.console import Console from rich.console import Console
from ltx_trainer import logger from ltx_trainer import logger
from ltx_trainer.gpu_utils import free_gpu_memory_context from ltx_trainer.gpu_utils import free_gpu_memory_context
console = Console() console = Console()
app = typer.Typer( app = typer.Typer(
pretty_exceptions_enable=False, pretty_exceptions_enable=False,
no_args_is_help=True, no_args_is_help=True,
@@ -46,6 +47,7 @@ def preprocess_dataset( # noqa: PLR0913
device: str, device: str,
remove_llm_prefixes: bool = False, remove_llm_prefixes: bool = False,
reference_column: str | None = None, reference_column: str | None = None,
reference_downscale_factor: int = 1,
with_audio: bool = False, with_audio: bool = False,
load_text_encoder_in_8bit: bool = False, load_text_encoder_in_8bit: bool = False,
) -> None: ) -> None:
@@ -99,14 +101,33 @@ def preprocess_dataset( # noqa: PLR0913
# Process reference videos if reference_column is provided # Process reference videos if reference_column is provided
if reference_column: if reference_column:
logger.info("Processing reference videos for IC-LoRA training...") # Validate: scaled references with multiple buckets can cause ambiguous bucket matching
if reference_downscale_factor > 1 and len(resolution_buckets) > 1:
raise ValueError(
"When using --reference-downscale-factor > 1, only a single resolution bucket is supported. "
"Using multiple buckets with scaled references can cause ambiguous bucket matching "
"(e.g., a 512x256 reference could match either the scaled-down 1024x512 bucket or the 512x256 "
"bucket). Please use a single resolution bucket or set --reference-downscale-factor to 1."
)
# Calculate and validate scaled resolution buckets for reference videos
reference_buckets = compute_scaled_resolution_buckets(resolution_buckets, reference_downscale_factor)
if reference_downscale_factor > 1:
logger.info(
f"Processing reference videos for IC-LoRA training at 1/{reference_downscale_factor} resolution..."
)
logger.info(f"Reference resolution buckets: {reference_buckets}")
else:
logger.info("Processing reference videos for IC-LoRA training...")
reference_latents_dir = output_base / "reference_latents" reference_latents_dir = output_base / "reference_latents"
compute_latents( compute_latents(
dataset_file=dataset_file, dataset_file=dataset_file,
main_media_column=video_column, main_media_column=video_column,
video_column=reference_column, video_column=reference_column,
resolution_buckets=resolution_buckets, resolution_buckets=reference_buckets,
output_dir=str(reference_latents_dir), output_dir=str(reference_latents_dir),
model_path=model_path, model_path=model_path,
batch_size=batch_size, batch_size=batch_size,
@@ -226,6 +247,11 @@ def main( # noqa: PLR0913
default=False, default=False,
help="Load the Gemma text encoder in 8-bit precision to save GPU memory (requires bitsandbytes)", help="Load the Gemma text encoder in 8-bit precision to save GPU memory (requires bitsandbytes)",
), ),
reference_downscale_factor: int = typer.Option(
default=1,
help="Downscale factor for reference video resolution. When > 1, reference videos are processed at "
"1/n resolution (e.g., 2 means half resolution). Used for efficient IC-LoRA training.",
),
) -> None: ) -> None:
"""Preprocess a video dataset by computing and saving latents and text embeddings. """Preprocess a video dataset by computing and saving latents and text embeddings.
The dataset must be a CSV, JSON, or JSONL file with columns for captions and video paths. The dataset must be a CSV, JSON, or JSONL file with columns for captions and video paths.
@@ -242,6 +268,10 @@ def main( # noqa: PLR0913
python scripts/process_dataset.py dataset.json --resolution-buckets 768x768x25 \\ python scripts/process_dataset.py dataset.json --resolution-buckets 768x768x25 \\
--model-path /path/to/ltx2.safetensors --text-encoder-path /path/to/gemma \\ --model-path /path/to/ltx2.safetensors --text-encoder-path /path/to/gemma \\
--reference-column "reference_path" --reference-column "reference_path"
# Process dataset with scaled reference videos (half resolution) for efficient IC-LoRA
python scripts/process_dataset.py dataset.json --resolution-buckets 768x768x25 \\
--model-path /path/to/ltx2.safetensors --text-encoder-path /path/to/gemma \\
--reference-column "reference_path" --reference-downscale-factor 2
# Process dataset with audio for audio-video training # Process dataset with audio for audio-video training
python scripts/process_dataset.py dataset.json --resolution-buckets 768x512x97 \\ python scripts/process_dataset.py dataset.json --resolution-buckets 768x512x97 \\
--model-path /path/to/ltx2.safetensors --text-encoder-path /path/to/gemma \\ --model-path /path/to/ltx2.safetensors --text-encoder-path /path/to/gemma \\
@@ -255,6 +285,13 @@ def main( # noqa: PLR0913
"When training with multiple resolution buckets, you must use a batch size of 1." "When training with multiple resolution buckets, you must use a batch size of 1."
) )
# Validate reference_downscale_factor
if reference_downscale_factor < 1:
raise typer.BadParameter("--reference-downscale-factor must be >= 1")
if reference_downscale_factor > 1 and not reference_column:
logger.warning("--reference-downscale-factor specified but no --reference-column provided. Ignoring.")
preprocess_dataset( preprocess_dataset(
dataset_file=dataset_path, dataset_file=dataset_path,
caption_column=caption_column, caption_column=caption_column,
@@ -270,6 +307,7 @@ def main( # noqa: PLR0913
device=device, device=device,
remove_llm_prefixes=remove_llm_prefixes, remove_llm_prefixes=remove_llm_prefixes,
reference_column=reference_column, reference_column=reference_column,
reference_downscale_factor=reference_downscale_factor,
with_audio=with_audio, with_audio=with_audio,
load_text_encoder_in_8bit=load_text_encoder_in_8bit, load_text_encoder_in_8bit=load_text_encoder_in_8bit,
) )
@@ -891,6 +891,50 @@ def parse_resolution_buckets(resolution_buckets_str: str) -> list[tuple[int, int
return resolution_buckets return resolution_buckets
def compute_scaled_resolution_buckets(
resolution_buckets: list[tuple[int, int, int]],
scale_factor: int,
) -> list[tuple[int, int, int]]:
"""Compute scaled resolution buckets and validate the results."""
if scale_factor == 1:
return resolution_buckets
scaled_buckets = []
for frames, height, width in resolution_buckets:
# Validate that scale factor evenly divides the dimensions
if height % scale_factor != 0:
raise ValueError(
f"Height {height} is not evenly divisible by scale factor {scale_factor}. "
f"Choose a scale factor that divides {height} evenly."
)
if width % scale_factor != 0:
raise ValueError(
f"Width {width} is not evenly divisible by scale factor {scale_factor}. "
f"Choose a scale factor that divides {width} evenly."
)
scaled_height = height // scale_factor
scaled_width = width // scale_factor
# Validate scaled dimensions are divisible by VAE spatial factor
if scaled_height % VAE_SPATIAL_FACTOR != 0:
raise ValueError(
f"Scaled height {scaled_height} (from {height} / {scale_factor}) "
f"is not divisible by {VAE_SPATIAL_FACTOR}. "
f"Choose a different scale factor or adjust your resolution buckets."
)
if scaled_width % VAE_SPATIAL_FACTOR != 0:
raise ValueError(
f"Scaled width {scaled_width} (from {width} / {scale_factor}) "
f"is not divisible by {VAE_SPATIAL_FACTOR}. "
f"Choose a different scale factor or adjust your resolution buckets."
)
scaled_buckets.append((frames, scaled_height, scaled_width))
return scaled_buckets
@app.command() @app.command()
def main( # noqa: PLR0913 def main( # noqa: PLR0913
dataset_file: str = typer.Argument( dataset_file: str = typer.Argument(
@@ -207,6 +207,14 @@ class ValidationConfig(ConfigBaseModel):
"One video path must be provided for each validation prompt", "One video path must be provided for each validation prompt",
) )
reference_downscale_factor: int = Field(
default=1,
description="Downscale factor for reference videos in IC-LoRA validation. "
"When > 1, reference videos are processed at 1/n resolution (e.g., 2 means half resolution). "
"Must match the factor used during dataset preprocessing.",
ge=1,
)
video_dims: tuple[int, int, int] = Field( video_dims: tuple[int, int, int] = Field(
default=(960, 544, 97), default=(960, 544, 97),
description="Dimensions of validation videos (width, height, frames). " description="Dimensions of validation videos (width, height, frames). "
@@ -334,6 +342,41 @@ class ValidationConfig(ConfigBaseModel):
return v return v
@model_validator(mode="after")
def validate_scaled_reference_dimensions(self) -> "ValidationConfig":
"""Validate that scaled reference dimensions are valid when reference_downscale_factor > 1."""
if self.reference_downscale_factor > 1:
width, height, _frames = self.video_dims
# Validate that downscale factor evenly divides the target dimensions
if width % self.reference_downscale_factor != 0:
raise ValueError(
f"Width {width} is not evenly divisible by reference_downscale_factor "
f"{self.reference_downscale_factor}. Choose a downscale factor that divides {width} evenly."
)
if height % self.reference_downscale_factor != 0:
raise ValueError(
f"Height {height} is not evenly divisible by reference_downscale_factor "
f"{self.reference_downscale_factor}. Choose a downscale factor that divides {height} evenly."
)
scaled_width = width // self.reference_downscale_factor
scaled_height = height // self.reference_downscale_factor
# Validate scaled dimensions are divisible by 32
if scaled_width % 32 != 0:
raise ValueError(
f"Scaled reference width {scaled_width} (from {width} / {self.reference_downscale_factor}) "
f"is not divisible by 32. Choose a different downscale factor or adjust video_dims."
)
if scaled_height % 32 != 0:
raise ValueError(
f"Scaled reference height {scaled_height} (from {height} / {self.reference_downscale_factor}) "
f"is not divisible by 32. Choose a different downscale factor or adjust video_dims."
)
return self
class CheckpointsConfig(ConfigBaseModel): class CheckpointsConfig(ConfigBaseModel):
"""Configuration for model checkpointing during training""" """Configuration for model checkpointing during training"""
@@ -218,16 +218,22 @@ def load_text_encoder(
from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder
from ltx_core.text_encoders.gemma.encoders.av_encoder import ( from ltx_core.text_encoders.gemma.encoders.av_encoder import (
AV_GEMMA_TEXT_ENCODER_KEY_OPS, AV_GEMMA_TEXT_ENCODER_KEY_OPS,
GEMMA_MODEL_OPS,
AVGemmaTextEncoderModelConfigurator, AVGemmaTextEncoderModelConfigurator,
) )
from ltx_core.text_encoders.gemma.encoders.base_encoder import module_ops_from_gemma_root from ltx_core.text_encoders.gemma.encoders.base_encoder import module_ops_from_gemma_root
from ltx_core.utils import find_matching_file
torch_device = _to_torch_device(device) torch_device = _to_torch_device(device)
gemma_model_folder = find_matching_file(str(gemma_model_path), "model*.safetensors").parent
gemma_weight_paths = [str(p) for p in gemma_model_folder.rglob("*.safetensors")]
text_encoder = SingleGPUModelBuilder( text_encoder = SingleGPUModelBuilder(
model_path=str(checkpoint_path), model_path=(str(checkpoint_path), *gemma_weight_paths),
model_class_configurator=AVGemmaTextEncoderModelConfigurator, model_class_configurator=AVGemmaTextEncoderModelConfigurator,
model_sd_ops=AV_GEMMA_TEXT_ENCODER_KEY_OPS, model_sd_ops=AV_GEMMA_TEXT_ENCODER_KEY_OPS,
module_ops=module_ops_from_gemma_root(str(gemma_model_path)), module_ops=(GEMMA_MODEL_OPS, *module_ops_from_gemma_root(str(gemma_model_path))),
).build(device=torch_device, dtype=dtype) ).build(device=torch_device, dtype=dtype)
return text_encoder return text_encoder
@@ -795,6 +795,7 @@ class LtxvTrainer:
seed=self._config.validation.seed, seed=self._config.validation.seed,
condition_image=condition_image, condition_image=condition_image,
reference_video=reference_video, reference_video=reference_video,
reference_downscale_factor=self._config.validation.reference_downscale_factor,
generate_audio=generate_audio, generate_audio=generate_audio,
include_reference_in_output=self._config.validation.include_reference_in_output, include_reference_in_output=self._config.validation.include_reference_in_output,
cached_embeddings=cached_embeddings, cached_embeddings=cached_embeddings,
@@ -885,8 +886,11 @@ class LtxvTrainer:
# Cast to configured precision # Cast to configured precision
state_dict = {k: v.to(save_dtype) if isinstance(v, Tensor) else v for k, v in state_dict.items()} state_dict = {k: v.to(save_dtype) if isinstance(v, Tensor) else v for k, v in state_dict.items()}
# Save to disk # Build metadata for safetensors file
save_file(state_dict, saved_weights_path) metadata = self._build_checkpoint_metadata()
# Save to disk with metadata
save_file(state_dict, saved_weights_path, metadata=metadata)
else: else:
# Cast to configured precision # Cast to configured precision
full_state_dict = {k: v.to(save_dtype) if isinstance(v, Tensor) else v for k, v in full_state_dict.items()} full_state_dict = {k: v.to(save_dtype) if isinstance(v, Tensor) else v for k, v in full_state_dict.items()}
@@ -913,6 +917,21 @@ class LtxvTrainer:
# Update the list to only contain kept checkpoints # Update the list to only contain kept checkpoints
self._checkpoint_paths = self._checkpoint_paths[-self._config.checkpoints.keep_last_n :] self._checkpoint_paths = self._checkpoint_paths[-self._config.checkpoints.keep_last_n :]
def _build_checkpoint_metadata(self) -> dict[str, str]:
"""Build metadata dictionary for safetensors checkpoint.
Delegates to the training strategy to get strategy-specific metadata
that downstream inference pipelines may need.
Returns:
Dictionary of string key-value pairs for safetensors metadata.
Values are converted to strings for safetensors compatibility.
"""
raw_metadata = self._training_strategy.get_checkpoint_metadata()
# Convert all values to strings for safetensors compatibility
metadata = {k: str(v) for k, v in raw_metadata.items()}
if metadata:
logger.info(f"Saving checkpoint metadata: {metadata}")
return metadata
def _save_config(self) -> None: def _save_config(self) -> None:
"""Save the training configuration as a YAML file in the output directory.""" """Save the training configuration as a YAML file in the output directory."""
if not IS_MAIN_PROCESS: if not IS_MAIN_PROCESS:
@@ -128,6 +128,15 @@ class TrainingStrategy(ABC):
Scalar loss tensor 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( def _get_video_positions(
self, self,
num_frames: int, num_frames: int,
@@ -46,9 +46,13 @@ class VideoToVideoStrategy(TrainingStrategy):
- Reference latents (clean) are concatenated with target latents (noised) - Reference latents (clean) are concatenated with target latents (noised)
- Video coordinates handle both reference and target sequences - Video coordinates handle both reference and target sequences
- Loss is computed only on the target portion - Loss is computed only on the target portion
Attributes:
reference_downscale_factor: The inferred downscale factor of reference videos.
This is computed from the first batch and cached for metadata export.
""" """
config: VideoToVideoConfig config: VideoToVideoConfig
reference_downscale_factor: int | None
def __init__(self, config: VideoToVideoConfig): def __init__(self, config: VideoToVideoConfig):
"""Initialize strategy with configuration. """Initialize strategy with configuration.
@@ -56,6 +60,7 @@ class VideoToVideoStrategy(TrainingStrategy):
config: Video-to-video configuration config: Video-to-video configuration
""" """
super().__init__(config) super().__init__(config)
self.reference_downscale_factor = None # Will be inferred from first batch
def get_data_sources(self) -> dict[str, str]: def get_data_sources(self) -> dict[str, str]:
"""IC-LoRA training requires latents, conditions, and reference latents.""" """IC-LoRA training requires latents, conditions, and reference latents."""
@@ -65,7 +70,7 @@ class VideoToVideoStrategy(TrainingStrategy):
self.config.reference_latents_dir: "ref_latents", self.config.reference_latents_dir: "ref_latents",
} }
def prepare_training_inputs( def prepare_training_inputs( # noqa: PLR0915
self, self,
batch: dict[str, Any], batch: dict[str, Any],
timestep_sampler: TimestepSampler, timestep_sampler: TimestepSampler,
@@ -86,6 +91,26 @@ class VideoToVideoStrategy(TrainingStrategy):
ref_height = ref_latents_info["height"][0].item() ref_height = ref_latents_info["height"][0].item()
ref_width = ref_latents_info["width"][0].item() ref_width = ref_latents_info["width"][0].item()
# Infer reference downscale factor from dimension ratios
# This allows training with downscaled reference videos for efficiency
reference_downscale_factor = self._infer_reference_downscale_factor(
target_height=height,
target_width=width,
ref_height=ref_height,
ref_width=ref_width,
)
# Cache the scale factor for metadata export (only on first batch)
if self.reference_downscale_factor is None:
self.reference_downscale_factor = reference_downscale_factor
elif self.reference_downscale_factor != reference_downscale_factor:
raise ValueError(
f"Inconsistent reference downscale factor across batches. "
f"First batch had factor={self.reference_downscale_factor}, "
f"but current batch has factor={reference_downscale_factor}. "
f"All training samples must use the same reference/target resolution ratio."
)
# Patchify latents: [B, C, F, H, W] -> [B, seq_len, C] # Patchify latents: [B, C, F, H, W] -> [B, seq_len, C]
target_latents = self._video_patchifier.patchify(target_latents) target_latents = self._video_patchifier.patchify(target_latents)
ref_latents = self._video_patchifier.patchify(ref_latents) ref_latents = self._video_patchifier.patchify(ref_latents)
@@ -159,6 +184,15 @@ class VideoToVideoStrategy(TrainingStrategy):
dtype=dtype, dtype=dtype,
) )
# Scale reference positions to match target coordinate space
# This maps ref positions from (0, ref_H, ref_W) to (0, target_H, target_W)
# Position tensor shape: [B, 3, seq_len, 2] where dim 1 is (time, height, width)
if reference_downscale_factor != 1:
ref_positions = ref_positions.clone()
ref_positions[:, 1, ...] *= reference_downscale_factor # height axis
ref_positions[:, 2, ...] *= reference_downscale_factor # width axis
# Time axis (index 0) remains unchanged
target_positions = self._get_video_positions( target_positions = self._get_video_positions(
num_frames=num_frames, num_frames=num_frames,
height=height, height=height,
@@ -221,3 +255,48 @@ class VideoToVideoStrategy(TrainingStrategy):
loss = loss.mul(loss_mask).div(loss_mask.mean()) loss = loss.mul(loss_mask).div(loss_mask.mean())
return loss.mean() return loss.mean()
def get_checkpoint_metadata(self) -> dict[str, Any]:
"""Get metadata for checkpoint files."""
metadata: dict[str, Any] = {}
# Always include reference_downscale_factor for IC-LoRAs so inference
# pipelines know the expected scale factor for reference videos.
if self.reference_downscale_factor is not None:
metadata["reference_downscale_factor"] = self.reference_downscale_factor
return metadata
@staticmethod
def _infer_reference_downscale_factor(
target_height: int,
target_width: int,
ref_height: int,
ref_width: int,
) -> int:
"""Infer the reference downscale factor from target and reference dimensions."""
# If dimensions match, no scaling needed
if target_height == ref_height and target_width == ref_width:
return 1
# Calculate scale factors for each dimension
if target_height % ref_height != 0 or target_width % ref_width != 0:
raise ValueError(
f"Target dimensions ({target_height}x{target_width}) must be exact multiples "
f"of reference dimensions ({ref_height}x{ref_width})"
)
scale_h = target_height // ref_height
scale_w = target_width // ref_width
if scale_h != scale_w:
raise ValueError(
f"Reference scale must be uniform. Got height scale {scale_h} and width scale {scale_w}. "
f"Target: {target_height}x{target_width}, Reference: {ref_height}x{ref_width}"
)
if scale_h < 1:
raise ValueError(
f"Reference dimensions ({ref_height}x{ref_width}) cannot be larger than "
f"target dimensions ({target_height}x{target_width})"
)
return scale_h
@@ -85,6 +85,7 @@ class GenerationConfig:
seed: int = 42 # Random seed for reproducibility seed: int = 42 # Random seed for reproducibility
condition_image: Tensor | None = None # Optional first frame image for image-to-video condition_image: Tensor | None = None # Optional first frame image for image-to-video
reference_video: Tensor | None = None # For IC-LoRA: [F, C, H, W] in [0, 1] reference_video: Tensor | None = None # For IC-LoRA: [F, C, H, W] in [0, 1]
reference_downscale_factor: int = 1 # For IC-LoRA: downscale factor (1 = same resolution, 2 = half resolution)
generate_audio: bool = True # Whether to generate audio alongside video 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 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) cached_embeddings: CachedPromptEmbeddings | None = None # Pre-computed text embeddings (avoids loading Gemma)
@@ -251,6 +252,14 @@ class ValidationSampler:
ref_latent, ref_positions = self._encode_video(ref_video_preprocessed, config.frame_rate, device) ref_latent, ref_positions = self._encode_video(ref_video_preprocessed, config.frame_rate, device)
ref_seq_len = ref_latent.shape[1] ref_seq_len = ref_latent.shape[1]
# Scale reference positions to match target coordinate space
# Position tensor shape: [B, 3, seq_len, 2] where dim 1 is (time, height, width)
if config.reference_downscale_factor != 1:
ref_positions = ref_positions.clone()
ref_positions[:, 1, ...] *= config.reference_downscale_factor # height axis
ref_positions[:, 2, ...] *= config.reference_downscale_factor # width axis
# Time axis (index 0) remains unchanged
# Create target video state # Create target video state
video_tools = self._create_video_latent_tools(config) video_tools = self._create_video_latent_tools(config)
target_clean_state = video_tools.create_initial_state(device=device, dtype=torch.bfloat16) target_clean_state = video_tools.create_initial_state(device=device, dtype=torch.bfloat16)
@@ -375,13 +384,28 @@ class ValidationSampler:
@staticmethod @staticmethod
def _preprocess_reference_video(config: GenerationConfig) -> Tensor: def _preprocess_reference_video(config: GenerationConfig) -> Tensor:
"""Preprocess reference video: resize, crop, and convert to model input format. """Preprocess reference video: resize, crop, and convert to model input format.
When reference_downscale_factor > 1, the reference video is downscaled to a smaller
resolution for more efficient inference. The positions will be scaled up later
to match the target coordinate space.
Args: Args:
config: Generation configuration with reference_video config: Generation configuration
Returns: Returns:
Preprocessed video tensor [B, C, F, H, W] in [-1, 1] range Preprocessed video tensor [B, C, F, H, W] in [-1, 1] range
""" """
ref_video = config.reference_video # [F, C, H, W] in [0, 1] ref_video = config.reference_video # [F, C, H, W] in [0, 1]
target_height, target_width = config.height, config.width scale_factor = config.reference_downscale_factor
# Target dimensions for reference (scaled down if scale_factor > 1)
target_height = config.height // scale_factor
target_width = config.width // scale_factor
# Validate scaled dimensions
if target_height % 32 != 0 or target_width % 32 != 0:
raise ValueError(
f"Scaled reference dimensions ({target_height}x{target_width}) must be divisible by 32. "
f"Original: {config.height}x{config.width}, scale_factor: {scale_factor}"
)
current_height, current_width = ref_video.shape[2:] current_height, current_width = ref_video.shape[2:]
# Resize maintaining aspect ratio and center crop if needed # Resize maintaining aspect ratio and center crop if needed
@@ -745,11 +769,28 @@ class ValidationSampler:
If the videos have different frame counts, the shorter one is padded with If the videos have different frame counts, the shorter one is padded with
its last frame repeated. its last frame repeated.
Args: Args:
left_video: Left video tensor [C, F1, H, W] in [0, 1] left_video: Left video tensor [C, F1, H1, W1] in [0, 1]
right_video: Right video tensor [C, F2, H, W] in [0, 1] right_video: Right video tensor [C, F2, H2, W2] in [0, 1]
Returns: Returns:
Concatenated video tensor [C, max(F1,F2), H, W*2] in [0, 1] Concatenated video tensor [C, max(F1,F2), H2, W1_scaled+W2] in [0, 1]
""" """
left_height, left_width = left_video.shape[2], left_video.shape[3]
right_height = right_video.shape[2]
# Resize left video to match right video's height if needed
if left_height != right_height:
# Scale width proportionally to maintain aspect ratio
scale = right_height / left_height
new_width = int(left_width * scale)
# Interpolate expects [N, C, H, W], we have [C, F, H, W]
# Reshape to [C*F, 1, H, W] -> interpolate -> reshape back
c, f, h, w = left_video.shape
left_video = left_video.reshape(c * f, 1, h, w)
left_video = torch.nn.functional.interpolate(
left_video, size=(right_height, new_width), mode="bilinear", align_corners=False
)
left_video = left_video.reshape(c, f, right_height, new_width)
left_frames = left_video.shape[1] left_frames = left_video.shape[1]
right_frames = right_video.shape[1] right_frames = right_video.shape[1]
Generated
+523 -534
View File
File diff suppressed because it is too large Load Diff