Blame for ltx2/modules/text_encoder.py

"""
Self-contained text encoder adapted from the public LTX-2 code.
Gemma-based text encoder.
"""
from typing import NamedTuple
from pathlib import Path

import torch
import torch.nn as nn
from einops import rearrange
from transformers import AutoTokenizer, Gemma3ForConditionalGeneration

from .model import FeedForward, LTXRopeType, precompute_freqs_cis, rms_norm
# Import Attention from model_ltx_2_3 which supports apply_gated_attention (to_gate_logits)
# The LTX 2.0 Attention in model.py lacks this feature, causing silent weight drops for 2.3 checkpoints
try:
    from .model_ltx_2_3 import Attention
except ImportError:
    from .model import Attention

__all__ = ['AVGemmaTextEncoderModel', 'LTXVGemmaTokenizer']


# ============================================================
# Tokenizer
# ============================================================

class LTXVGemmaTokenizer:
    """Tokenizer wrapper."""
    def __init__(self, tokenizer_path, max_length=1024):
        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, local_files_only=True, model_max_length=max_length)
        self.tokenizer.padding_side = "left"
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token
        self.max_length = max_length
    
    def tokenize_with_weights(self, text):
        encoded = self.tokenizer(text.strip(), padding="max_length", max_length=self.max_length, truncation=True, return_tensors="pt")
        return {"gemma": [(t.item(), a.item()) for t, a in zip(encoded.input_ids[0], encoded.attention_mask[0])]}


# ============================================================
# Feature Extractor
# ============================================================

class GemmaFeaturesExtractorProjLinear(nn.Module):
    """Feature extractor: project multi-layer hidden states to the target dimension.

    V1 (19B): aggregate_embed Linear(188160, 3840, bias=False)
    V2 (22B): video_aggregate_embed Linear(188160, 4096, bias=True)
    """
    def __init__(self, out_dim=3840, bias=False, use_video_key=False):
        super().__init__()
        in_dim = 3840 * 49  # gemma hidden_size * (num_hidden_layers + 1)
        if use_video_key:
            self.video_aggregate_embed = nn.Linear(in_dim, out_dim, bias=bias)
        else:
            self.aggregate_embed = nn.Linear(in_dim, out_dim, bias=bias)
        self._use_video_key = use_video_key

    def forward(self, x):
        if self._use_video_key:
            return self.video_aggregate_embed(x)
        return self.aggregate_embed(x)


# ============================================================
# Embeddings Connector
# ============================================================

class _BasicTransformerBlock1D(nn.Module):
    """1D transformer block (matches the official implementation)."""
    def __init__(self, dim, heads, dim_head, rope_type=LTXRopeType.SPLIT, apply_gated_attention=False):
        super().__init__()
        self.attn1 = Attention(
            query_dim=dim, context_dim=None, heads=heads, dim_head=dim_head,
            rope_type=rope_type, apply_gated_attention=apply_gated_attention,
        )
        self.ff = FeedForward(dim, dim)

    def forward(self, x, attention_mask=None, pe=None):
        norm_x = rms_norm(x).squeeze(1)
        x = self.attn1(norm_x, mask=attention_mask, pe=pe) + x
        if x.ndim == 4:
            x = x.squeeze(1)
        x = self.ff(rms_norm(x)) + x
        if x.ndim == 4:
            x = x.squeeze(1)
        return x


class Embeddings1DConnector(nn.Module):
    """1D embedding connector for text embeddings (matches the official implementation)."""
    def __init__(self, attention_head_dim=128, num_attention_heads=32, num_layers=8,
                 positional_embedding_theta=10000.0, positional_embedding_max_pos=None,
                 num_learnable_registers=128, rope_type=LTXRopeType.SPLIT,
                 apply_gated_attention=True):
        super().__init__()
        self.num_attention_heads = num_attention_heads
        self.inner_dim = num_attention_heads * attention_head_dim
        self.positional_embedding_theta = positional_embedding_theta
        self.positional_embedding_max_pos = positional_embedding_max_pos or [1]
        self.rope_type = rope_type
        self.num_learnable_registers = num_learnable_registers

        self.transformer_1d_blocks = nn.ModuleList([
            _BasicTransformerBlock1D(
                self.inner_dim, num_attention_heads, attention_head_dim,
                rope_type, apply_gated_attention=apply_gated_attention,
            )
            for _ in range(num_layers)
        ])
        
        if num_learnable_registers:
            self.learnable_registers = nn.Parameter(torch.rand(num_learnable_registers, self.inner_dim, dtype=torch.bfloat16) * 2 - 1)
    
    def _replace_padded_with_learnable_registers(self, hidden_states, attention_mask):
        n = self.num_learnable_registers
        registers = torch.tile(self.learnable_registers, (hidden_states.shape[1] // n, 1))
        mask_binary = (attention_mask.squeeze(1).squeeze(1).unsqueeze(-1) >= -9000.0).int()
        
        non_zero = hidden_states[:, mask_binary.squeeze().bool(), :]
        pad_len = hidden_states.shape[1] - non_zero.shape[1]
        adjusted = nn.functional.pad(non_zero, (0, 0, 0, pad_len), value=0)
        flipped = torch.flip(mask_binary, dims=[1])
        hidden_states = flipped * adjusted + (1 - flipped) * registers
        attention_mask = torch.zeros_like(attention_mask)
        return hidden_states, attention_mask
    
    def forward(self, hidden_states, attention_mask=None):
        if self.num_learnable_registers:
            hidden_states, attention_mask = self._replace_padded_with_learnable_registers(hidden_states, attention_mask)
        
        indices_grid = torch.arange(hidden_states.shape[1], dtype=torch.float32, device=hidden_states.device)[None, None, :]
        pe = precompute_freqs_cis(indices_grid, self.inner_dim, hidden_states.dtype, self.positional_embedding_theta,
                                  self.positional_embedding_max_pos, False, self.num_attention_heads, self.rope_type)
        
        for block in self.transformer_1d_blocks:
            hidden_states = block(hidden_states, attention_mask, pe)
        
        return rms_norm(hidden_states), attention_mask


# ============================================================
# Encoder Output
# ============================================================

class AVGemmaEncoderOutput(NamedTuple):
    video_encoding: torch.Tensor
    audio_encoding: torch.Tensor
    attention_mask: torch.Tensor


# ============================================================
# AV Gemma Text Encoder
# ============================================================

def _norm_and_concat_padded_batch(encoded_text, sequence_lengths, padding_side="right"):
    """V1 normalization: per-segment mean/range normalization."""
    b, t, d, l = encoded_text.shape
    device = encoded_text.device

    token_indices = torch.arange(t, device=device)[None, :]
    if padding_side == "right":
        mask = token_indices < sequence_lengths[:, None]
    else:
        mask = token_indices >= (t - sequence_lengths[:, None])
    mask = rearrange(mask, "b t -> b t 1 1")

    eps = 1e-6
    masked = encoded_text.masked_fill(~mask, 0.0)
    denom = (sequence_lengths * d).view(b, 1, 1, 1)
    mean = masked.sum(dim=(1, 2), keepdim=True) / (denom + eps)

    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

    normed = 8 * (encoded_text - mean) / (range_ + eps)
    normed = normed.reshape(b, t, -1)

    mask_flat = rearrange(mask, "b t 1 1 -> b t 1").expand(-1, -1, d * l)
    return normed.masked_fill(~mask_flat, 0.0)


def _norm_and_concat_per_token_rms(encoded_text, attention_mask):
    """V2 normalization: per-token RMS normalization.

    Each token is RMS-normalized independently, then all layers are concatenated.

    Args:
        encoded_text: [B, T, D, L] - hidden states from all layers
        attention_mask: [B, T] - binary attention mask
    Returns:
        [B, T, D*L] - normalized and concatenated tensor
    """
    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)
    normed = normed.reshape(B, T, D * L)
    mask_3d = attention_mask.bool().unsqueeze(-1)  # [B, T, 1]
    return torch.where(mask_3d, normed, torch.zeros_like(normed))


def _rescale_norm(x, target_dim, source_dim):
    """Rescale normalization: x * sqrt(target_dim / source_dim)"""
    import math
    return x * math.sqrt(target_dim / source_dim)


class AVGemmaTextEncoderModel(nn.Module):
    """Audio-video text encoder."""
    def __init__(self, feature_extractor_linear, embeddings_connector, audio_embeddings_connector=None,
                 tokenizer=None, model=None, dtype=torch.bfloat16,
                 use_v2_norm=False, gemma_embedding_dim=3840):
        super().__init__()
        self._gemma_root = None
        self.tokenizer = tokenizer
        self.model = model
        self.feature_extractor_linear = feature_extractor_linear.to(dtype=dtype)
        self.embeddings_connector = embeddings_connector.to(dtype=dtype)
        self.audio_embeddings_connector = audio_embeddings_connector.to(dtype=dtype) if audio_embeddings_connector is not None else None
        self.use_v2_norm = use_v2_norm
        self.gemma_embedding_dim = gemma_embedding_dim
    
    def _convert_to_additive_mask(self, attention_mask, dtype):
        return (attention_mask - 1).to(dtype).reshape(attention_mask.shape[0], 1, -1, attention_mask.shape[-1]) * torch.finfo(dtype).max
    
    def _run_feature_extractor(self, hidden_states, attention_mask, padding_side="right"):
        encoded = torch.stack(hidden_states, dim=-1)  # [B, T, D, L]
        if self.use_v2_norm:
            # V2 (22B): per-token RMS normalization + rescale
            normed = _norm_and_concat_per_token_rms(encoded, attention_mask)
            normed = normed.to(encoded.dtype)
            # Rescale: x * sqrt(target_dim / source_dim) to maintain magnitude after projection
            v_dim = self.feature_extractor_linear.video_aggregate_embed.out_features \
                if hasattr(self.feature_extractor_linear, 'video_aggregate_embed') \
                else self.feature_extractor_linear.aggregate_embed.out_features
            normed = _rescale_norm(normed, v_dim, self.gemma_embedding_dim)
        else:
            # V1 (19B): per-segment mean/range normalization
            seq_lengths = attention_mask.sum(dim=-1)
            normed = _norm_and_concat_padded_batch(encoded, seq_lengths, padding_side)
            normed = normed.to(encoded.dtype)
        return self.feature_extractor_linear(normed)
    
    def _preprocess_text(self, text, padding_side="left"):
        token_pairs = self.tokenizer.tokenize_with_weights(text)["gemma"]
        input_ids = torch.tensor([[t[0] for t in token_pairs]], device=self.model.device)
        attention_mask = torch.tensor([[t[1] for t in token_pairs]], device=self.model.device)
        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True)
        projected = self._run_feature_extractor(outputs.hidden_states, attention_mask, padding_side)
        return projected, attention_mask
    
    def _run_connectors(self, encoded_input, attention_mask):
        connector_mask = self._convert_to_additive_mask(attention_mask, encoded_input.dtype)
        encoded, encoded_mask = self.embeddings_connector(encoded_input, connector_mask)
        attention_mask = (encoded_mask < 0.000001).to(torch.int64).reshape(encoded.shape[0], encoded.shape[1], 1)
        encoded = encoded * attention_mask
        if self.audio_embeddings_connector is not None:
            encoded_audio, _ = self.audio_embeddings_connector(encoded_input, connector_mask)
        else:
            encoded_audio = encoded  # video-only mode
        return encoded, encoded_audio, attention_mask.squeeze(-1)
    
    def forward(self, text, padding_side="left"):
        encoded_inputs, attention_mask = self._preprocess_text(text, padding_side)
        video_enc, audio_enc, mask = self._run_connectors(encoded_inputs, attention_mask)
        return AVGemmaEncoderOutput(video_enc, audio_enc, mask)


# ============================================================
# ============================================================

def create_av_text_encoder(config):
    """Build the text encoder from a config."""
    cfg = config.get("transformer", {})
    rope_type = LTXRopeType(cfg.get("rope_type", "interleaved"))
    pe_max_pos = cfg.get("connector_positional_embedding_max_pos", [1])
    
    feature_extractor = GemmaFeaturesExtractorProjLinear()
    embeddings_connector = Embeddings1DConnector(num_attention_heads=30, attention_head_dim=128, positional_embedding_max_pos=pe_max_pos, rope_type=rope_type)
    audio_connector = Embeddings1DConnector(num_attention_heads=30, attention_head_dim=128, positional_embedding_max_pos=pe_max_pos, rope_type=rope_type)
    
    return AVGemmaTextEncoderModel(feature_extractor, embeddings_connector, audio_connector)


def find_file(root, pattern):
    """Locate a file."""
    matches = list(Path(root).rglob(pattern))
    if not matches:
        raise FileNotFoundError(f"Not found: {pattern} in {root}")
    return str(matches[0].parent)


# ============================================================
# ============================================================

class LTX2TextEncoder(nn.Module):
    """High-level text-encoder wrapper for the inference pipeline."""
    
    def __init__(self, device=None, dtype=torch.bfloat16):
        super().__init__()
        self._device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.dtype = dtype
        self._encoder = None
        self._encode_fn = None
    
    @property
    def device(self):
        return self._device
    
    def set_encoder(self, encoder, encode_fn=None):
        """Set up the encoder."""
        self._encoder = encoder
        self._encode_fn = encode_fn
    
    def forward(self, text_prompts):
        """Encode text prompts."""
        if self._encoder is None:
            raise RuntimeError("Encoder not set")
        
        results = []
        for prompt in text_prompts:
            output = self._encoder(prompt)
            # output.video_encoding is [1, seq_len, dim], squeeze batch dim
            results.append({
                "prompt_embeds": output.video_encoding.squeeze(0),  # [seq_len, dim]
                "audio_embeds": output.audio_encoding.squeeze(0) if output.audio_encoding is not None else None,
                "attention_mask": output.attention_mask.squeeze(0) if output.attention_mask is not None else None,
            })
        
        # Stack batch: [B, seq_len, dim]
        batch = {
            "prompt_embeds": torch.stack([r["prompt_embeds"] for r in results]),
        }
        if results[0]["audio_embeds"] is not None:
            batch["audio_embeds"] = torch.stack([r["audio_embeds"] for r in results])
        return batch