| |
|
|
| |
|
|
| 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, |
| 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 |
|
|
| |
| 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) |
|
|
| |
| |
| 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 |
|
|
| |
| 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] |
|
|
| |
| 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)) |
|
|
