Automated PR - 2026-01-29

This commit is contained in:
sync-bot
2026-01-29 18:42:17 +00:00
parent 727c43e998
commit ca1623ad2a
31 changed files with 1723 additions and 663 deletions
+1 -1
View File
@@ -9,7 +9,7 @@ dependencies = [
"torchaudio",
"einops",
"numpy",
"transformers",
"transformers~=4.57.0",
"safetensors",
"accelerate",
"scipy>=1.14",
@@ -1,4 +1,5 @@
from dataclasses import dataclass
import math
from dataclasses import dataclass, field
import torch
@@ -189,6 +190,89 @@ class LegacyStatefulAPGGuider(GuiderProtocol):
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:
batch_size = to_project.shape[0]
positive_flat = to_project.reshape(batch_size, -1)
@@ -2,11 +2,16 @@
from ltx_core.conditioning.exceptions import ConditioningError
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__ = [
"ConditioningError",
"ConditioningItem",
"VideoConditionByKeyframeIndex",
"VideoConditionByLatentIndex",
"VideoConditionByReferenceLatent",
]
@@ -2,8 +2,10 @@
from ltx_core.conditioning.types.keyframe_cond import VideoConditionByKeyframeIndex
from ltx_core.conditioning.types.latent_cond import VideoConditionByLatentIndex
from ltx_core.conditioning.types.reference_video_cond import VideoConditionByReferenceLatent
__all__ = [
"VideoConditionByKeyframeIndex",
"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
config = self.model_config()
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)
lora_strengths = [lora.strength for lora in self.loras]
@@ -140,7 +140,11 @@ class BasicAVTransformerBlock(torch.nn.Module):
audio: TransformerArgs | None,
perturbations: BatchedPerturbationConfig | None = 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:
perturbations = BatchedPerturbationConfig.empty(batch_size)
@@ -211,7 +215,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
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
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)
@@ -226,7 +230,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
* 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
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)
@@ -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
import torch
from transformers import Gemma3Config
from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
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.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 (
Embeddings1DConnector,
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.tokenizer import LTXVGemmaTokenizer
@@ -76,7 +83,11 @@ class AVGemmaTextEncoderModelConfigurator(ModelConfigurator[AVGemmaTextEncoderMo
feature_extractor_linear = GemmaFeaturesExtractorProjLinear.from_config(config)
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(
model=model,
feature_extractor_linear=feature_extractor_linear,
embeddings_connector=embeddings_connector,
audio_embeddings_connector=audio_embeddings_connector,
@@ -85,10 +96,57 @@ class AVGemmaTextEncoderModelConfigurator(ModelConfigurator[AVGemmaTextEncoderMo
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="model.diffusion_model.audio_embeddings_connector.")
.with_matching(prefix="model.diffusion_model.video_embeddings_connector.")
.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_matching(prefix="model.diffusion_model.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.text_encoders.gemma.feature_extractor import GemmaFeaturesExtractorProjLinear
from ltx_core.text_encoders.gemma.tokenizer import LTXVGemmaTokenizer
from ltx_core.utils import find_matching_file
class GemmaTextEncoderModelBase(torch.nn.Module):
@@ -77,7 +78,7 @@ class GemmaTextEncoderModelBase(torch.nn.Module):
messages: list[dict[str, str]],
image: torch.Tensor | None = None,
max_new_tokens: int = 512,
seed: int = 42,
seed: int = 10,
) -> str:
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,
max_new_tokens: int = 512,
system_prompt: str | None = None,
seed: int = 42,
seed: int = 10,
) -> str:
"""Enhance a text prompt for T2V generation."""
@@ -126,7 +127,7 @@ class GemmaTextEncoderModelBase(torch.nn.Module):
image: torch.Tensor,
max_new_tokens: int = 512,
system_prompt: str | None = None,
seed: int = 42,
seed: int = 10,
) -> str:
"""Enhance a text prompt for I2V generation using a reference image."""
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()
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, ...]:
gemma_path = _find_matching_dir(gemma_root, "model*.safetensors")
tokenizer_path = _find_matching_dir(gemma_root, "tokenizer.model")
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
tokenizer_root = str(find_matching_file(gemma_root, "tokenizer.model").parent)
processor_root = str(find_matching_file(gemma_root, "preprocessor_config.json").parent)
def load_tokenizer(module: GemmaTextEncoderModelBase) -> GemmaTextEncoderModelBase:
module.tokenizer = LTXVGemmaTokenizer(tokenizer_path, 1024)
module.tokenizer = LTXVGemmaTokenizer(tokenizer_root, 1024)
return module
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:
raise ValueError("Tokenizer model operation must be performed before processor model operation")
module.processor = Gemma3Processor(image_processor=image_processor, tokenizer=module.tokenizer.tokenizer)
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(
"TokenizerLoad",
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,
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]]:
+11
View File
@@ -1,3 +1,4 @@
from pathlib import Path
from typing import Any
import torch
@@ -49,3 +50,13 @@ def to_denoised(
if isinstance(sigma, torch.Tensor):
sigma = sigma.to(calc_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]