| |
|
|
| |
|
|
| """On-device Gemma-3 text encoder + embeddings connectors for LTX-2, following the |
| tt_dit encoder_pair convention (cf. ``encoders/t5/encoder_pair.py``).""" |
|
|
| from __future__ import annotations |
|
|
| import glob |
| import os |
| import time |
| from typing import Callable |
|
|
| import torch |
| from huggingface_hub import snapshot_download |
| from loguru import logger |
| from safetensors import safe_open |
| from safetensors.torch import load_file |
| from transformers import AutoTokenizer |
|
|
| import ttnn |
|
|
| from ...parallel.manager import CCLManager |
| from ...utils import cache as cache_module |
| from ...utils.mochi import get_rot_transformation_mat |
| from ...utils.tensor import bf16_tensor |
| from ...utils.tracing import traced_function |
| from .embeddings_connector import EmbeddingsConnector |
| from .feature_extractor import GemmaFeatureExtractor |
| from .model_gemma import GemmaConfig, GemmaEncoder |
|
|
| GEMMA_SEQUENCE_LENGTH = 1024 |
| VIDEO_EMBED_DIM = 4096 |
| AUDIO_EMBED_DIM = 2048 |
|
|
| _CONNECTOR_PREFIXES = ( |
| "text_embedding_projection.video_aggregate_embed.", |
| "text_embedding_projection.audio_aggregate_embed.", |
| "model.diffusion_model.video_embeddings_connector.", |
| "model.diffusion_model.audio_embeddings_connector.", |
| ) |
|
|
|
|
| def _resolve_gemma_dir(gemma: str) -> str: |
| """Resolve a Gemma reference to a local directory: a local dir is returned as-is, a |
| HuggingFace repo id is snapshot-downloaded (cached). The original reference — not this |
| resolved snapshot dir, whose basename is an opaque revision hash — names the cache.""" |
| if os.path.isdir(gemma): |
| return gemma |
| logger.info(f"Resolving HuggingFace Gemma repo {gemma} (auto-download if missing)") |
| return snapshot_download(repo_id=gemma) |
|
|
|
|
| |
| def _gemma_state_dict(gemma_path: str) -> dict[str, torch.Tensor]: |
| """All Gemma encoder weights, from the HF safetensors shards.""" |
| weight_files = sorted(glob.glob(f"{gemma_path}/model-*.safetensors")) or sorted( |
| glob.glob(f"{gemma_path}/*.safetensors") |
| ) |
| sd: dict[str, torch.Tensor] = {} |
| for f in weight_files: |
| sd.update(load_file(f)) |
| return sd |
|
|
|
|
| def _read_connector_checkpoint(checkpoint_name: str) -> dict[str, torch.Tensor]: |
| """The connector + aggregate_embed tensors from the LTX checkpoint.""" |
| state: dict[str, torch.Tensor] = {} |
| with safe_open(checkpoint_name, "pt") as f: |
| for k in f.keys(): |
| if k.startswith(_CONNECTOR_PREFIXES): |
| state[k] = f.get_tensor(k) |
| return state |
|
|
|
|
| def _feature_extractor_state_dict(ckpt, *, mode: str, gemma_hidden_size: int, gemma_num_layers: int) -> dict: |
| """video/audio aggregate_embed weights, permuted D-major→layer-major to match the |
| on-device layer-major concat (see GemmaFeatureExtractor._weight_to_layer_major).""" |
| sd = {} |
| for axis in ("video", "audio") if mode == "av" else ("video",): |
| prefix = f"text_embedding_projection.{axis}_aggregate_embed." |
| for k, v in ckpt.items(): |
| if k.startswith(prefix): |
| sub = k[len(prefix) :] |
| if sub == "weight": |
| v = GemmaFeatureExtractor._weight_to_layer_major(v, gemma_hidden_size, gemma_num_layers) |
| sd[f"{axis}_aggregate_embed.{sub}"] = v |
| return sd |
|
|
|
|
| def _connector_state_dict(ckpt, axis: str, num_blocks: int) -> dict: |
| """One connector's transformer blocks + norm, dropping blocks beyond ``num_blocks``.""" |
| prefix = f"model.diffusion_model.{axis}_embeddings_connector." |
| sd = {} |
| for k, v in ckpt.items(): |
| if not k.startswith(prefix): |
| continue |
| sub = k[len(prefix) :] |
| if sub.startswith("transformer_1d_blocks.") and int(sub.split(".")[1]) >= num_blocks: |
| continue |
| sd[sub] = v |
| return sd |
|
|
|
|
| class GemmaTokenizerEncoderPair: |
| """Tokenizer + on-device Gemma-3 encoder + embeddings connectors for LTX-2. |
| |
| ``encode(prompts)`` returns ``[(video_embeds, audio_embeds), ...]``. Gemma weights load from |
| ``gemma_path``; connector/feature-extractor weights from the LTX ``checkpoint_name``. |
| """ |
|
|
| def __init__( |
| self, |
| gemma_path: str | None, |
| *, |
| mesh_device: ttnn.MeshDevice, |
| ccl_manager: CCLManager, |
| parallel_config, |
| checkpoint_name: str | None = None, |
| mode: str = "av", |
| dynamic_load: bool = False, |
| sequence_length: int = GEMMA_SEQUENCE_LENGTH, |
| video_dim: int = VIDEO_EMBED_DIM, |
| audio_dim: int = AUDIO_EMBED_DIM, |
| num_layers: int = 48, |
| hidden_layer_index: int = -1, |
| ) -> None: |
| |
| |
| self.gemma_ref = gemma_path |
| self.gemma_path = _resolve_gemma_dir(gemma_path) if gemma_path else None |
| self.mesh_device = mesh_device |
| self.ccl_manager = ccl_manager |
| self.parallel_config = parallel_config |
| self.checkpoint_name = checkpoint_name |
| self.mode = mode |
| self.dynamic_load = dynamic_load |
| self._num_layers = num_layers |
| self._hidden_layer_index = hidden_layer_index |
| self._sequence_length = sequence_length |
| self._video_dim = video_dim |
| self._audio_dim = audio_dim |
|
|
| self.gemma_encoder = None |
| self.tokenizer = None |
| self.feature_extractor = None |
| self.video_connector = None |
| self.audio_connector = None |
| self._coresident_peers: list = [] |
| self._cached_trans_mat = None |
| |
| |
| |
| |
| self._encoder_trace = not dynamic_load |
| |
| |
| |
| |
| self._trace_gate_open = True |
|
|
| |
| @property |
| def sequence_length(self) -> int: |
| return self._sequence_length |
|
|
| @property |
| def video_dim(self) -> int: |
| return self._video_dim |
|
|
| @property |
| def audio_dim(self) -> int: |
| return self._audio_dim |
|
|
| def defer_trace_capture(self) -> None: |
| """Hold the encode trace back until ``open_trace_gate`` (see ``_trace_gate_open``).""" |
| self._trace_gate_open = False |
|
|
| def open_trace_gate(self) -> None: |
| """Allow the next encode to capture/replay its trace (see ``_trace_gate_open``).""" |
| self._trace_gate_open = True |
|
|
| def register_coresident_peers(self, peers: list) -> None: |
| """Store the DiT/VAE peers the encoder modules must not be L1-coresident with. |
| Exclusions are wired at each module's first build, so peers must be registered before |
| the first ``ensure_loaded``/``load_*``.""" |
| self._coresident_peers = list(peers) |
|
|
| def _register_exclusions(self, module) -> None: |
| """Bidirectionally exclude a freshly-built encoder module against the stored peers, |
| before its first load — so the encoder load auto-evicts the DiT (and the later DiT |
| reload auto-evicts the encoder). No-op unless dynamic_load.""" |
| if not self.dynamic_load or not self._coresident_peers: |
| return |
| module.register_coresident_exclusions(*self._coresident_peers) |
| for p in self._coresident_peers: |
| p.register_coresident_exclusions(module) |
|
|
| def is_loaded(self) -> bool: |
| return self.gemma_encoder is not None and self.gemma_encoder.is_loaded() |
|
|
| def ensure_loaded(self, connector_state: dict | Callable[[], dict] | None = None) -> None: |
| """Build/reload the encoder + connector weights (evicted by dynamic_load).""" |
| if self.is_loaded(): |
| return |
| self.load_gemma_encoder() |
| self.load_embeddings_connectors( |
| connector_state |
| if connector_state is not None |
| else (lambda: _read_connector_checkpoint(self.checkpoint_name)) |
| ) |
|
|
| def load_gemma_encoder(self, gemma_path: str | None = None) -> None: |
| """Load the TTNN Gemma-3 encoder. Built once and reused across reloads; weights reload |
| from the shared cache (no re-tilizing 12B params). The CCLManager is supplied by the |
| pipeline (shared with the connectors), so reloads never rebuild it.""" |
| gemma_path = gemma_path or self.gemma_path |
| if self.gemma_encoder is None: |
| config = GemmaConfig( |
| num_hidden_layers=self._num_layers, |
| hidden_layer_index=self._hidden_layer_index, |
| max_position_embeddings=self._sequence_length, |
| ) |
| self.gemma_encoder = GemmaEncoder(config, self.mesh_device, self.ccl_manager, self.parallel_config) |
| self._register_exclusions(self.gemma_encoder) |
| |
| |
| self.tokenizer = AutoTokenizer.from_pretrained(gemma_path) |
|
|
| t0 = time.time() |
| cache_module.load_model( |
| self.gemma_encoder, |
| model_name=os.path.basename(os.path.normpath(self.gemma_ref)), |
| subfolder="text_encoder", |
| parallel_config=self.parallel_config, |
| mesh_shape=tuple(self.mesh_device.shape), |
| mesh_device=self.mesh_device, |
| get_torch_state_dict=lambda: _gemma_state_dict(gemma_path), |
| ) |
| logger.info(f"Loaded TTNN Gemma encoder ({self._num_layers}L) in {time.time()-t0:.0f}s") |
|
|
| def load_embeddings_connectors( |
| self, |
| checkpoint_state: dict[str, torch.Tensor] | Callable[[], dict[str, torch.Tensor]], |
| *, |
| gemma_hidden_size: int = 3840, |
| gemma_num_layers: int = 49, |
| video_num_blocks: int = 8, |
| audio_num_blocks: int = 8, |
| num_heads: int = 32, |
| ) -> None: |
| """Load the feature extractor + video/audio connectors from the LTX checkpoint through |
| the shared cache. Modules are built once and reused across reloads; the CCLManager is |
| supplied by the pipeline (shared with the Gemma encoder). |
| |
| ``checkpoint_state`` is a dict or a zero-arg callable; resolved once and only on a cache |
| miss, so an all-cache-hit reload reads neither the checkpoint nor the tilizer. |
| """ |
| ckpt = _memoize(checkpoint_state) |
| ckpt_name = ( |
| os.path.basename(self.checkpoint_name).removesuffix(".safetensors") |
| if self.checkpoint_name |
| else "ltx-connectors" |
| ) |
|
|
| if self.feature_extractor is None: |
| self.feature_extractor = GemmaFeatureExtractor( |
| input_dim=gemma_hidden_size * gemma_num_layers, |
| embedding_dim=gemma_hidden_size, |
| video_dim=self._video_dim, |
| audio_dim=self._audio_dim if self.mode == "av" else None, |
| mesh_device=self.mesh_device, |
| ccl_manager=self.ccl_manager, |
| parallel_config=self.parallel_config, |
| ) |
| self._register_exclusions(self.feature_extractor) |
| cache_module.load_model( |
| self.feature_extractor, |
| model_name=ckpt_name, |
| subfolder="feature_extractor", |
| parallel_config=self.parallel_config, |
| mesh_shape=tuple(self.mesh_device.shape), |
| mesh_device=self.mesh_device, |
| get_torch_state_dict=lambda: _feature_extractor_state_dict( |
| ckpt(), mode=self.mode, gemma_hidden_size=gemma_hidden_size, gemma_num_layers=gemma_num_layers |
| ), |
| ) |
|
|
| self.video_connector = self._load_connector( |
| "video", self._video_dim, video_num_blocks, num_heads, ckpt_name, ckpt |
| ) |
| self.audio_connector = ( |
| self._load_connector("audio", self._audio_dim, audio_num_blocks, num_heads, ckpt_name, ckpt) |
| if self.mode == "av" |
| else None |
| ) |
|
|
| def _load_connector(self, axis, output_dim, num_blocks, num_heads, ckpt_name, ckpt) -> EmbeddingsConnector: |
| connector = getattr(self, f"{axis}_connector") |
| if connector is None: |
| connector = EmbeddingsConnector( |
| output_dim=output_dim, |
| num_blocks=num_blocks, |
| num_heads=num_heads, |
| mesh_device=self.mesh_device, |
| ccl_manager=self.ccl_manager, |
| parallel_config=self.parallel_config, |
| ) |
| self._register_exclusions(connector) |
| cache_module.load_model( |
| connector, |
| model_name=ckpt_name, |
| subfolder=f"{axis}_connector", |
| parallel_config=self.parallel_config, |
| mesh_shape=tuple(self.mesh_device.shape), |
| mesh_device=self.mesh_device, |
| dtype="float32", |
| get_torch_state_dict=lambda: _connector_state_dict(ckpt(), axis, num_blocks), |
| ) |
| logger.info(f"Loaded {axis} embeddings connector ({num_blocks} blocks, dim={output_dim})") |
| return connector |
|
|
| @traced_function(device=lambda self: self.mesh_device, clone_prep_inputs=False, prep_run=True) |
| def _encode_device(self, tt_ids, tt_gemma_mask, fe_mask, src_idx, keep_mask): |
| """Whole-encode device graph (gemma → feature extractor → connectors), captured as one |
| ttnn trace and replayed per prompt. Per-prompt masks/indices are device inputs (built in |
| encode); rope/trans_mat/registers are captured constants (prep_run allocates them before |
| capture). Returns DEVICE (video, audio) embeds — tokenizer + final to_torch stay in encode.""" |
| trans_mat = self._prepare_trans_mat() |
| all_hidden = self.gemma_encoder(tt_ids, tt_attn_mask=tt_gemma_mask) |
| |
| |
| hs_list = list(all_hidden[:-2]) + [all_hidden[-1]] |
| video_feats, audio_feats = self.feature_extractor(hs_list, fe_mask) |
| video = self.video_connector(video_feats, src_idx, keep_mask, trans_mat=trans_mat) |
| audio = ( |
| self.audio_connector(audio_feats, src_idx, keep_mask, trans_mat=trans_mat) |
| if self.audio_connector is not None |
| else None |
| ) |
| return video, audio |
|
|
| def encode(self, prompts: list[str]) -> list[tuple[torch.Tensor, torch.Tensor | None]]: |
| """Tokenize → traced whole-encode device graph → host embeds, one |
| ``(video_embeds, audio_embeds)`` per prompt. Pure compute; the disk cache lives in the |
| pipeline.""" |
| assert self.gemma_encoder is not None, "Call ensure_loaded() first" |
|
|
| results = [] |
| for prompt in prompts: |
| tokens = self.tokenizer( |
| prompt, return_tensors="pt", padding="max_length", max_length=self._sequence_length, truncation=True |
| ) |
| tt_ids = ttnn.from_torch( |
| tokens.input_ids, device=self.mesh_device, dtype=ttnn.uint32, layout=ttnn.ROW_MAJOR_LAYOUT |
| ) |
| seq = tt_ids.shape[-1] |
| |
| tt_gemma_mask = self.gemma_encoder.build_attn_mask(tokens.attention_mask, seq) |
| fe_mask = self.feature_extractor.build_mask(tokens.attention_mask) |
| |
| src_idx, keep_mask = self.video_connector.build_indices(tokens.attention_mask, seq) |
|
|
| video_dev, audio_dev = self._encode_device( |
| tt_ids, tt_gemma_mask, fe_mask, src_idx, keep_mask, traced=self._encoder_trace and self._trace_gate_open |
| ) |
| video_embeds = ttnn.to_torch(ttnn.get_device_tensors(video_dev)[0]).float() |
| audio_embeds = ( |
| ttnn.to_torch(ttnn.get_device_tensors(audio_dev)[0]).float() if audio_dev is not None else None |
| ) |
| results.append((video_embeds, audio_embeds)) |
| return results |
|
|
| def _prepare_trans_mat(self) -> ttnn.Tensor: |
| """Cached per-tile rotation matrix for rotary_embedding_llama.""" |
| if self._cached_trans_mat is None: |
| self._cached_trans_mat = bf16_tensor(get_rot_transformation_mat(), device=self.mesh_device) |
| return self._cached_trans_mat |
|
|
|
|
| def _memoize(state: dict | Callable[[], dict]) -> Callable[[], dict]: |
| """Wrap a dict-or-callable as a zero-arg callable that resolves at most once.""" |
| cached: dict = {} |
|
|
| def get() -> dict: |
| if "sd" not in cached: |
| cached["sd"] = state() if callable(state) else state |
| return cached["sd"] |
|
|
| return get |
|
|
