377 lines
14 KiB
Python
377 lines
14 KiB
Python
"""Text Embeddings Pipeline: Feature Extractors + Embeddings Processors.
|
|
|
|
Provides the 3-block Gemma text encoder pipeline:
|
|
1. Gemma Model (external) -- runs the LLM, gets hidden states
|
|
2. Feature Extractor (V1/V2) -- normalization + projection
|
|
3. Embeddings Processor (Video / AV) -- wraps Embeddings1DConnector(s)
|
|
"""
|
|
|
|
import json
|
|
import math
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
from comfy.utils import load_torch_file
|
|
from einops import rearrange
|
|
from torch import nn
|
|
|
|
from .embeddings_connector import (
|
|
Embeddings1DConnector,
|
|
load_audio_embeddings_connector,
|
|
load_video_embeddings_connector,
|
|
)
|
|
|
|
_PREFIX_BASE = "model.diffusion_model."
|
|
_PREFIX_TEXT_PROJ = "text_embedding_projection."
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _to_binary_mask(encoded, encoded_mask):
|
|
"""Convert additive mask to binary [B, seq, 1] and zero out padding."""
|
|
binary_mask = (encoded_mask < 0.000001).to(torch.int64)
|
|
binary_mask = binary_mask.reshape(encoded.shape[0], encoded.shape[1], 1)
|
|
return encoded * binary_mask, binary_mask
|
|
|
|
|
|
def _rescale_norm(x: torch.Tensor, target_dim: int, source_dim: int) -> torch.Tensor:
|
|
"""Rescale normalization: x * sqrt(target_dim / source_dim)."""
|
|
return x * math.sqrt(target_dim / source_dim)
|
|
|
|
|
|
def _filter_sd(sd: dict, prefix: str) -> dict:
|
|
"""Extract keys with *prefix* and strip the prefix."""
|
|
return {k[len(prefix) :]: v for k, v in sd.items() if k.startswith(prefix)}
|
|
|
|
|
|
def _load_aggregate_embed(sd: dict, modality: str, dtype) -> nn.Linear:
|
|
"""Load an aggregate_embed Linear from the state dict.
|
|
|
|
Args:
|
|
sd: Full checkpoint state dict.
|
|
modality: ``"video"`` or ``"audio"``.
|
|
dtype: Target dtype.
|
|
|
|
Returns ``linear`` module.
|
|
"""
|
|
name = f"{modality}_aggregate_embed"
|
|
weight_key = f"{_PREFIX_TEXT_PROJ}{name}.weight"
|
|
bias_key = f"{_PREFIX_TEXT_PROJ}{name}.bias"
|
|
weight = sd[weight_key]
|
|
out_features, embedding_dim = weight.shape
|
|
linear = nn.Linear(embedding_dim, out_features, bias=bias_key in sd)
|
|
linear.load_state_dict(_filter_sd(sd, f"{_PREFIX_TEXT_PROJ}{name}."))
|
|
return linear.to(dtype=dtype)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Normalization functions (moved from gemma_encoder.py)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def norm_and_concat_padded_batch(
|
|
encoded_text: torch.Tensor,
|
|
sequence_lengths: torch.Tensor,
|
|
padding_side: str = "right",
|
|
) -> torch.Tensor:
|
|
"""
|
|
Normalize a 4D tensor [B, T, D, L] per sample and per layer, using sequence_lengths to mask.
|
|
Returns [B, T, D * L] tensor with original padding preserved.
|
|
|
|
Args:
|
|
encoded_text: 4D tensor [B, T, D, L]
|
|
sequence_lengths: 1D tensor [B] with actual sequence lengths
|
|
padding_side: "left" or "right" to indicate which side has padding
|
|
"""
|
|
B, T, D, L = encoded_text.shape
|
|
device = encoded_text.device
|
|
|
|
# Build mask: [B, T, 1, 1]
|
|
token_indices = torch.arange(T, device=device)[None, :] # [1, T]
|
|
|
|
if padding_side == "right":
|
|
# For right padding, valid tokens are from 0 to sequence_length-1
|
|
mask = token_indices < sequence_lengths[:, None] # [B, T]
|
|
elif padding_side == "left":
|
|
# For left padding, valid tokens are from (T - sequence_length) to T-1
|
|
start_indices = T - sequence_lengths[:, None] # [B, 1]
|
|
mask = token_indices >= start_indices # [B, T]
|
|
else:
|
|
raise ValueError(f"padding_side must be 'left' or 'right', got {padding_side}")
|
|
|
|
mask = rearrange(mask, "b t -> b t 1 1")
|
|
|
|
# Compute masked mean: [B, 1, 1, L]
|
|
masked = encoded_text.masked_fill(~mask, 0.0)
|
|
# denom = mask.sum(dim=(1, 2), keepdim=True) # [B, 1, 1, 1]
|
|
denom = (sequence_lengths * D).view(B, 1, 1, 1)
|
|
mean = masked.sum(dim=(1, 2), keepdim=True) / (denom + 1e-6)
|
|
|
|
# Compute masked min/max: [B, 1, 1, L]
|
|
x_min = encoded_text.masked_fill(~mask, float("inf")).amin(dim=(1, 2), keepdim=True)
|
|
x_max = encoded_text.masked_fill(~mask, float("-inf")).amax(
|
|
dim=(1, 2), keepdim=True
|
|
)
|
|
range_ = x_max - x_min
|
|
|
|
# Normalize only the valid tokens
|
|
normed = 8 * (encoded_text - mean) / (range_ + 1e-6)
|
|
|
|
# concat to be [Batch, T, D * L] - this preserves the original structure
|
|
normed = normed.reshape(B, T, -1) # [B, T, D * L]
|
|
|
|
# Apply mask to preserve original padding (set padded positions to 0)
|
|
mask_flattened = rearrange(mask, "b t 1 1 -> b t 1").expand(-1, -1, D * L)
|
|
normed = normed.masked_fill(~mask_flattened, 0.0)
|
|
|
|
return normed
|
|
|
|
|
|
def norm_and_concat_per_token_rms(
|
|
encoded_text: torch.Tensor,
|
|
attention_mask: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
"""Per-token RMSNorm normalization.
|
|
|
|
Each token is normalized independently using RMSNorm over its D
|
|
dimension. This is naturally causal and supports packing since each
|
|
token's normalization is self-contained.
|
|
|
|
Args:
|
|
encoded_text: 4D tensor [B, T, D, L]
|
|
attention_mask: [B, T] binary mask (1=real, 0=pad)
|
|
|
|
Returns:
|
|
[B, T, D*L] normalized and flattened tensor with padding zeroed out.
|
|
"""
|
|
B, T, D, L = encoded_text.shape
|
|
variance = torch.mean(encoded_text**2, dim=2, keepdim=True) # [B,T,1,L]
|
|
normed = encoded_text * torch.rsqrt(variance + 1e-6) # [B,T,D,L]
|
|
normed = normed.reshape(B, T, D * L)
|
|
# Zero out padding
|
|
mask_3d = attention_mask.bool().unsqueeze(-1) # [B, T, 1]
|
|
normed = torch.where(mask_3d, normed, torch.zeros_like(normed))
|
|
return normed
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Feature Extractors (Block 2)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class FeatureExtractorV1(nn.Module):
|
|
"""19B: per-segment norm -> aggregate_embed -> 3840"""
|
|
|
|
def __init__(self, aggregate_embed: nn.Module, is_av: bool = False):
|
|
super().__init__()
|
|
self.aggregate_embed = aggregate_embed
|
|
self.is_av = is_av
|
|
|
|
def forward(self, all_layer_hiddens, attention_mask, padding_side="left"):
|
|
sequence_lengths = attention_mask.sum(dim=-1)
|
|
normed = norm_and_concat_padded_batch(
|
|
all_layer_hiddens, sequence_lengths, padding_side
|
|
)
|
|
normed = normed.to(all_layer_hiddens.dtype)
|
|
features = self.aggregate_embed(normed)
|
|
if self.is_av:
|
|
return {"video": features, "audio": features}
|
|
return {"video": features}
|
|
|
|
|
|
class FeatureExtractorV2(nn.Module):
|
|
"""22B: per-token RMS norm -> rescale -> aggregate_embed(s)"""
|
|
|
|
def __init__(
|
|
self,
|
|
video_aggregate_embed: nn.Linear,
|
|
embedding_dim: int,
|
|
audio_aggregate_embed: nn.Linear | None = None,
|
|
):
|
|
super().__init__()
|
|
self.video_aggregate_embed = video_aggregate_embed
|
|
self.audio_aggregate_embed = audio_aggregate_embed
|
|
self.embedding_dim = embedding_dim
|
|
|
|
def forward(self, all_layer_hiddens, attention_mask, padding_side="left"):
|
|
normed = norm_and_concat_per_token_rms(all_layer_hiddens, attention_mask)
|
|
normed = normed.to(all_layer_hiddens.dtype)
|
|
v_dim = self.video_aggregate_embed.out_features
|
|
result = {
|
|
"video": self.video_aggregate_embed(
|
|
_rescale_norm(normed, v_dim, self.embedding_dim)
|
|
)
|
|
}
|
|
if self.audio_aggregate_embed is not None:
|
|
a_dim = self.audio_aggregate_embed.out_features
|
|
result["audio"] = self.audio_aggregate_embed(
|
|
_rescale_norm(normed, a_dim, self.embedding_dim)
|
|
)
|
|
return result
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Embeddings Processors (Block 3)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class VideoEmbeddingsProcessor(nn.Module):
|
|
"""Video-only embeddings processor: single connector."""
|
|
|
|
def __init__(self, video_connector: Embeddings1DConnector):
|
|
super().__init__()
|
|
self.video_connector = video_connector
|
|
|
|
def create_embeddings(self, features, attention_mask):
|
|
encoded, mask = self.video_connector(features["video"], attention_mask)
|
|
encoded, binary_mask = _to_binary_mask(encoded, mask)
|
|
return encoded, binary_mask.squeeze(-1)
|
|
|
|
|
|
class AVEmbeddingsProcessor(nn.Module):
|
|
"""Audio-video embeddings processor: dual connectors + concat."""
|
|
|
|
def __init__(
|
|
self,
|
|
video_connector: Embeddings1DConnector,
|
|
audio_connector: Embeddings1DConnector,
|
|
):
|
|
super().__init__()
|
|
self.video_connector = video_connector
|
|
self.audio_connector = audio_connector
|
|
|
|
def create_embeddings(self, features, attention_mask):
|
|
encoded, mask = self.video_connector(features["video"], attention_mask)
|
|
encoded, binary_mask = _to_binary_mask(encoded, mask)
|
|
audio_encoded, _ = self.audio_connector(features["audio"], attention_mask)
|
|
encoded = torch.cat([encoded, audio_encoded], dim=-1)
|
|
return encoded, binary_mask.squeeze(-1)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Loading helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _load_single_aggregate_embed(sd, dtype):
|
|
"""Load the single aggregate_embed (V1 models).
|
|
|
|
Returns:
|
|
nn.Linear or None if key not found.
|
|
"""
|
|
key = f"{_PREFIX_TEXT_PROJ}aggregate_embed.weight"
|
|
if key not in sd:
|
|
return None
|
|
weight = sd[key]
|
|
in_features, out_features = weight.shape[1], weight.shape[0]
|
|
# Check for bias
|
|
bias_key = f"{_PREFIX_TEXT_PROJ}aggregate_embed.bias"
|
|
has_bias = bias_key in sd
|
|
linear = nn.Linear(in_features, out_features, bias=has_bias)
|
|
linear.load_state_dict(_filter_sd(sd, f"{_PREFIX_TEXT_PROJ}aggregate_embed."))
|
|
return linear.to(dtype=dtype)
|
|
|
|
|
|
def _load_single_aggregate_embed_from_file(path, dtype):
|
|
"""Load aggregate_embed from a standalone checkpoint file (legacy fallback).
|
|
|
|
Args:
|
|
path: Path to safetensors file containing aggregate_embed weights.
|
|
dtype: Target dtype.
|
|
|
|
Returns:
|
|
nn.Linear or None if file does not exist.
|
|
"""
|
|
path = Path(path)
|
|
if not path.exists():
|
|
return None
|
|
loaded_sd = load_torch_file(str(path), return_metadata=False)
|
|
if "aggregate_embed.weight" not in loaded_sd:
|
|
raise ValueError(
|
|
f"Checkpoint {path} does not contain 'aggregate_embed.weight'."
|
|
)
|
|
weight = loaded_sd["aggregate_embed.weight"]
|
|
in_features, out_features = weight.shape[1], weight.shape[0]
|
|
bias_key = "aggregate_embed.bias"
|
|
has_bias = bias_key in loaded_sd
|
|
linear = nn.Linear(in_features, out_features, bias=has_bias)
|
|
linear.load_state_dict(
|
|
{
|
|
k.removeprefix("aggregate_embed."): v
|
|
for k, v in loaded_sd.items()
|
|
if k.startswith("aggregate_embed.")
|
|
}
|
|
)
|
|
return linear.to(dtype=dtype)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Pipeline loader
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def load_text_embeddings_pipeline(
|
|
ltxv_path, dtype=torch.bfloat16, fallback_proj_path=None
|
|
):
|
|
"""Load feature extractor + embeddings processor from LTX-V checkpoint.
|
|
|
|
Auto-detects model variant (19B/22B, video-only/AV).
|
|
|
|
Args:
|
|
ltxv_path: Path to the LTX-V checkpoint.
|
|
dtype: Target dtype for loaded modules.
|
|
fallback_proj_path: Optional path to a standalone ``proj_linear.safetensors``
|
|
file. Used as a legacy fallback when the aggregate_embed is not
|
|
stored inside the LTX-V checkpoint (V1 models).
|
|
|
|
Returns:
|
|
(feature_extractor, embeddings_processor)
|
|
"""
|
|
sd, metadata = load_torch_file(str(ltxv_path), return_metadata=True)
|
|
config = json.loads(metadata.get("config", "{}"))
|
|
transformer_config = config.get("transformer", {})
|
|
|
|
is_av = f"{_PREFIX_BASE}audio_adaln_single.linear.weight" in sd
|
|
has_dual_aggregate = f"{_PREFIX_TEXT_PROJ}video_aggregate_embed.weight" in sd
|
|
|
|
# Load connectors (always needed)
|
|
video_connector = load_video_embeddings_connector(sd, transformer_config, dtype)
|
|
|
|
# Build embeddings processor (Block 3)
|
|
if is_av:
|
|
audio_connector = load_audio_embeddings_connector(sd, transformer_config, dtype)
|
|
processor = AVEmbeddingsProcessor(video_connector, audio_connector)
|
|
else:
|
|
processor = VideoEmbeddingsProcessor(video_connector)
|
|
|
|
# Build feature extractor (Block 2)
|
|
if has_dual_aggregate:
|
|
# V2 (22B): validate config matches expected settings
|
|
_expected = {
|
|
"caption_projection_first_linear": False,
|
|
"caption_proj_input_norm": False,
|
|
"caption_projection_second_linear": False,
|
|
"caption_proj_before_connector": True,
|
|
"text_encoder_norm_type": "per_token_rms",
|
|
}
|
|
for key, expected_val in _expected.items():
|
|
actual = transformer_config.get(key)
|
|
assert actual == expected_val, (
|
|
f"Unexpected config for dual-aggregate model: "
|
|
f"{key}={actual!r}, expected {expected_val!r}"
|
|
)
|
|
video_agg = _load_aggregate_embed(sd, "video", dtype)
|
|
audio_agg = _load_aggregate_embed(sd, "audio", dtype) if is_av else None
|
|
embedding_dim = transformer_config.get("prompt_embedding_dim", 3840)
|
|
return (FeatureExtractorV2(video_agg, embedding_dim, audio_agg), processor)
|
|
# V1 (19B)
|
|
aggregate_embed = _load_single_aggregate_embed(sd, dtype)
|
|
if aggregate_embed is None and fallback_proj_path is not None:
|
|
aggregate_embed = _load_single_aggregate_embed_from_file(
|
|
fallback_proj_path, dtype
|
|
)
|
|
return (FeatureExtractorV1(aggregate_embed, is_av=is_av), processor)
|