# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc. # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations from typing import TYPE_CHECKING import torch from loguru import logger from transformers import AutoTokenizer, PreTrainedTokenizerBase, T5EncoderModel import ttnn from ...parallel.config import EncoderParallelConfig from ...parallel.manager import CCLManager from ...utils import cache from ...utils.tracing import Tracer from .model_t5 import T5Config, T5Encoder if TYPE_CHECKING: from collections.abc import Iterable class T5TokenizerEncoderPair: def __init__( self, checkpoint: str, *, sequence_length: int | None, empty_sequence_length: int | None, embedding_dim: int | None, zero_masking: bool, device: ttnn.MeshDevice, ccl_manager: CCLManager, parallel_config: EncoderParallelConfig, use_attention_mask: bool, use_torch: bool, enabled: bool, ) -> None: self._device = device self._ccl_manager = ccl_manager self._parallel_config = parallel_config self._tokenizer = AutoTokenizer.from_pretrained(checkpoint) self._sequence_length = sequence_length if sequence_length is not None else self._tokenizer.model_max_length self._empty_sequence_length = ( empty_sequence_length if empty_sequence_length is not None else self._sequence_length ) self._embedding_dim = embedding_dim self._zero_masking = zero_masking self._use_attention_mask = use_attention_mask self._encoder = self._load_encoder(checkpoint, use_torch=use_torch) if enabled else None self._tracer = ( None if use_torch or not enabled else Tracer(self._encoder.forward, device=device, clone_prep_inputs=False) ) def load_torch_model(self, checkpoint: str) -> T5EncoderModel: return T5EncoderModel.from_pretrained(checkpoint) def _load_encoder(self, checkpoint: str, *, use_torch: bool) -> T5Encoder | T5EncoderModel: torch_model = self.load_torch_model(checkpoint) if use_torch: return torch_model config = T5Config( vocab_size=torch_model.config.vocab_size, embed_dim=torch_model.config.d_model, ff_dim=torch_model.config.d_ff, kv_dim=torch_model.config.d_kv, num_heads=torch_model.config.num_heads, num_hidden_layers=torch_model.config.num_layers, max_prompt_length=self._sequence_length, layer_norm_eps=torch_model.config.layer_norm_epsilon, relative_attention_num_buckets=torch_model.config.relative_attention_num_buckets, relative_attention_max_distance=torch_model.config.relative_attention_max_distance, ) model = T5Encoder( config=config, mesh_device=self._device, ccl_manager=self._ccl_manager, parallel_config=self._parallel_config, ) cache.load_model( tt_model=model, get_torch_state_dict=torch_model.state_dict, model_name=checkpoint, subfolder="", parallel_config=self._parallel_config, mesh_shape=tuple(self._device.shape), mesh_device=self._device, ) return model def encode( self, prompts: Iterable[str], *, num_images_per_prompt: int, enable_tracing: bool = False ) -> torch.Tensor: return _get_t5_prompt_embeds( prompts=prompts, num_images_per_prompt=num_images_per_prompt, tokenizer=self._tokenizer, text_encoder=self._encoder, tracer=self._tracer if enable_tracing else None, sequence_length=self._sequence_length, empty_sequence_length=self._empty_sequence_length, embedding_dim=self._embedding_dim, zero_masking=self._zero_masking, use_attention_mask=self._use_attention_mask, mesh_device=self._device, ) # adapted from https://github.com/huggingface/diffusers/blob/v0.31.0/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py def _get_t5_prompt_embeds( *, prompts: Iterable[str], text_encoder: T5Encoder | T5EncoderModel | None, tracer: Tracer | None = None, tokenizer: PreTrainedTokenizerBase, sequence_length: int, empty_sequence_length: int, num_images_per_prompt: int, mesh_device: ttnn.MeshDevice | None, embedding_dim: int | None, zero_masking: bool, use_attention_mask: bool, ) -> torch.Tensor: prompts = list(prompts) if text_encoder is None: assert embedding_dim is not None return torch.zeros([len(prompts) * num_images_per_prompt, empty_sequence_length, embedding_dim]) tokenizer_out = tokenizer( prompts, return_tensors="pt", padding="max_length", max_length=sequence_length, truncation=True, ) tokens = tokenizer_out.input_ids attention_mask = tokenizer_out.attention_mask if use_attention_mask else None untruncated_tokens = tokenizer( prompts, return_tensors="pt", padding="longest", ).input_ids if untruncated_tokens.shape[-1] >= tokens.shape[-1] and not torch.equal(tokens, untruncated_tokens): logger.warning("T5 input text was truncated") if isinstance(text_encoder, T5Encoder): assert mesh_device is not None tt_tokens = ttnn.from_torch( tokens, layout=ttnn.TILE_LAYOUT, dtype=ttnn.uint32, device=mesh_device, mesh_mapper=ttnn.replicate_tensor_to_mesh_mapper(mesh_device), ) tt_attention_mask = ( ttnn.from_torch( attention_mask[:, None, None, :], layout=ttnn.TILE_LAYOUT, dtype=ttnn.bfloat16, device=mesh_device, mesh_mapper=ttnn.replicate_tensor_to_mesh_mapper(mesh_device), ) if attention_mask is not None else None ) tt_hidden_states = (tracer or text_encoder.forward)(prompt=tt_tokens, attention_mask=tt_attention_mask) tt_prompt_embeds = tt_hidden_states[-1] prompt_embeds = ttnn.to_torch(ttnn.get_device_tensors(tt_prompt_embeds)[0]) else: tokens = tokens.to(device=text_encoder.device) with torch.no_grad(): output = text_encoder.forward(tokens, attention_mask=attention_mask) prompt_embeds = output.last_hidden_state.to("cpu") if zero_masking: prompt_embeds = prompt_embeds * (tokens != tokenizer.pad_token_id).unsqueeze(-1) if embedding_dim is not None: assert prompt_embeds.shape[-1] == embedding_dim return prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0)