File size: 7,083 Bytes
629b123 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 | # SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
import html
import os
from typing import TYPE_CHECKING
import ftfy
import regex as re
from transformers import AutoTokenizer, UMT5EncoderModel
import ttnn
from models.tt_dit.encoders.umt5.model_umt5 import UMT5Config, UMT5Encoder
from models.tt_dit.parallel.config import DiTParallelConfig, EncoderParallelConfig
from models.tt_dit.parallel.manager import CCLManager
from models.tt_dit.pipelines.events import PipelineEventCallback, SectionEnd, SectionStart, null_callback
from models.tt_dit.utils import cache
if TYPE_CHECKING:
from collections.abc import Sequence
def _basic_clean(text: str) -> str:
text = ftfy.fix_text(text)
text = html.unescape(html.unescape(text))
return text.strip()
def _whitespace_clean(text: str) -> str:
text = re.sub(r"\s+", " ", text)
return text.strip()
def prompt_clean(text: str) -> str:
return _whitespace_clean(_basic_clean(text))
class TextEncoder:
def __init__(
self,
*,
checkpoint_name: str,
device: ttnn.MeshDevice,
ccl_manager: CCLManager,
encoder_parallel_config: EncoderParallelConfig,
dit_parallel_config: DiTParallelConfig,
max_sequence_length: int = 512,
) -> None:
self._device = device
self._ccl_manager = ccl_manager
self._encoder_parallel_config = encoder_parallel_config
self._dit_parallel_config = dit_parallel_config
self._checkpoint_name = checkpoint_name
self._max_sequence_length = max_sequence_length
self._tokenizer = AutoTokenizer.from_pretrained(checkpoint_name, subfolder="tokenizer", trust_remote_code=True)
self._torch_text_encoder = UMT5EncoderModel.from_pretrained(
checkpoint_name, subfolder="text_encoder", trust_remote_code=True
)
umt5_config = UMT5Config(
vocab_size=self._torch_text_encoder.config.vocab_size,
embed_dim=self._torch_text_encoder.config.d_model,
ff_dim=self._torch_text_encoder.config.d_ff,
kv_dim=self._torch_text_encoder.config.d_kv,
num_heads=self._torch_text_encoder.config.num_heads,
num_hidden_layers=self._torch_text_encoder.config.num_layers,
max_prompt_length=512, # TODO: Consider removing
layer_norm_eps=self._torch_text_encoder.config.layer_norm_epsilon,
relative_attention_num_buckets=self._torch_text_encoder.config.relative_attention_num_buckets,
relative_attention_max_distance=self._torch_text_encoder.config.relative_attention_max_distance,
)
self._tt_encoder = UMT5Encoder(
config=umt5_config,
mesh_device=device,
ccl_manager=ccl_manager,
parallel_config=encoder_parallel_config,
)
def prepare(self) -> None:
cache.load_model(
self._tt_encoder,
model_name=os.path.basename(self._checkpoint_name),
subfolder="text_encoder",
parallel_config=self._encoder_parallel_config,
mesh_shape=tuple(self._device.shape),
mesh_device=self._device,
get_torch_state_dict=lambda: self._torch_text_encoder.state_dict(),
)
def encode_cfg(
self,
prompts: Sequence[str],
negative_prompts: Sequence[str],
*,
cfg_enabled: bool,
num_videos_per_prompt: int = 1,
on_event: PipelineEventCallback = null_callback,
) -> tuple[object, object]:
on_event(SectionStart("t5_encoding"))
prompts_list = list(prompts)
batch_size = len(prompts_list)
all_input_prompts: list[str] = list(prompts_list)
pos_prompt_end_idx = batch_size * num_videos_per_prompt
neg_prompt_end_idx = pos_prompt_end_idx
if cfg_enabled:
negative_prompts_list = list(negative_prompts)
assert batch_size == len(
negative_prompts_list
), f"`negative_prompt` has batch size {len(negative_prompts_list)}, but `prompt` has batch size {batch_size}."
all_input_prompts += negative_prompts_list
neg_prompt_end_idx = pos_prompt_end_idx + batch_size * num_videos_per_prompt
# Pad batch list of prompts to ensure proper sharding on batch dimension.
total_prompts = len(all_input_prompts)
num_devices = self._device.shape[1 - self._dit_parallel_config.tensor_parallel.mesh_axis]
all_input_prompts += [" "] * ((num_devices - (total_prompts % num_devices)) % num_devices)
all_prompt_embeds = self._encode(
all_input_prompts,
num_videos_per_prompt=num_videos_per_prompt,
)
prompt_embeds = all_prompt_embeds[:, :pos_prompt_end_idx]
negative_prompt_embeds = all_prompt_embeds[:, pos_prompt_end_idx:neg_prompt_end_idx] if cfg_enabled else None
on_event(SectionEnd("t5_encoding"))
return prompt_embeds, negative_prompt_embeds
def _encode(
self,
prompts: list[str],
*,
num_videos_per_prompt: int,
) -> object:
prompts = [prompt_clean(u) for u in prompts]
batch_size = len(prompts)
# NOTE: while the reference impl does not pad to max_sequence_length, for some reason this
# seems to be necessary for correctness in this pipeline. TODO: investigate
text_inputs = self._tokenizer(
prompts,
padding="max_length",
max_length=self._max_sequence_length,
truncation=True,
add_special_tokens=True,
return_attention_mask=True,
return_tensors="pt",
)
text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask
# Shard on batch dimension. On non TP axis
dims = [None, None]
DP_axis = 1 - self._dit_parallel_config.tensor_parallel.mesh_axis
dims[DP_axis] = 0
mesh_mapper = ttnn.ShardTensor2dMesh(self._device, mesh_shape=tuple(self._device.shape), dims=dims)
tt_prompt = ttnn.from_torch(
text_input_ids,
layout=ttnn.ROW_MAJOR_LAYOUT,
device=self._device,
mesh_mapper=mesh_mapper,
)
tt_mask = ttnn.from_torch(
mask,
dtype=ttnn.bfloat16,
layout=ttnn.TILE_LAYOUT,
device=self._device,
mesh_mapper=mesh_mapper,
)
prompt_embeds = self._tt_encoder(tt_prompt, attention_mask=tt_mask)[-1]
# use the mask to zero out the padding tokens.
prompt_embeds = prompt_embeds * ttnn.unsqueeze(tt_mask, -1)
prompt_embeds = self._ccl_manager.all_gather(prompt_embeds, dim=0, mesh_axis=DP_axis, use_hyperparams=True)
_, seq_len, _ = prompt_embeds.shape
prompt_embeds = ttnn.repeat(prompt_embeds, (1, num_videos_per_prompt, 1))
return ttnn.reshape(prompt_embeds, (1, batch_size * num_videos_per_prompt, seq_len, -1))
|