Automated PR - 2026-01-29
This commit is contained in:
@@ -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]]:
|
||||
|
||||
@@ -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]
|
||||
|
||||
Reference in New Issue
Block a user