text_embeddings_connectors.py

"""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)