| |
|
|
| |
|
|
| from __future__ import annotations |
|
|
| from typing import TYPE_CHECKING |
|
|
| import torch |
|
|
| import ttnn |
| from models.tt_dit.encoders.qwen25vl.encoder_pair import Qwen25VlTokenizerEncoderPair |
| from models.tt_dit.parallel.config import EncoderParallelConfig |
| from models.tt_dit.parallel.manager import CCLManager |
| from models.tt_dit.pipelines.events import PipelineEventCallback, SectionEnd, SectionStart, null_callback |
|
|
| if TYPE_CHECKING: |
| from collections.abc import Sequence |
|
|
|
|
| PROMPT_TEMPLATE = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n" |
| PROMPT_DROP_IDX = 34 |
|
|
|
|
| class TextEncoder: |
| def __init__( |
| self, |
| *, |
| checkpoint_name: str, |
| device: ttnn.MeshDevice, |
| ccl_manager: CCLManager, |
| parallel_config: EncoderParallelConfig, |
| use_torch: bool, |
| ) -> None: |
| self._encoder = Qwen25VlTokenizerEncoderPair( |
| checkpoint_name, |
| tokenizer_subfolder="tokenizer", |
| encoder_subfolder="text_encoder", |
| device=device, |
| ccl_manager=ccl_manager, |
| parallel_config=parallel_config, |
| use_torch=use_torch, |
| is_fsdp=True, |
| ) |
|
|
| def encoder_loaded(self) -> bool: |
| return self._encoder.encoder_loaded() |
|
|
| def reload_encoder_weights(self) -> None: |
| self._encoder.reload_encoder_weights() |
|
|
| def deallocate_encoder_weights(self) -> None: |
| self._encoder.deallocate_encoder_weights() |
|
|
| @torch.no_grad() |
| def encode_cfg( |
| self, |
| prompts: Sequence[str], |
| negative_prompts: Sequence[str], |
| *, |
| num_images_per_prompt: int, |
| cfg_enabled: bool, |
| traced: bool, |
| on_event: PipelineEventCallback = null_callback, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| assert len(prompts) == len(negative_prompts), "prompts and negative_prompts must have the same length" |
|
|
| all_prompts = [*negative_prompts, *prompts] if cfg_enabled else list(prompts) |
| all_prompts = [PROMPT_TEMPLATE.format(e) for e in all_prompts] |
|
|
| on_event(SectionStart("qwen_encoding")) |
| embeds, mask = self._encoder.encode( |
| all_prompts, |
| num_images_per_prompt=num_images_per_prompt, |
| sequence_length=512 + PROMPT_DROP_IDX, |
| enable_tracing=traced, |
| ) |
| on_event(SectionEnd("qwen_encoding")) |
|
|
| embeds[torch.logical_not(mask)] = 0.0 |
|
|
| return embeds[:, PROMPT_DROP_IDX:], mask[:, PROMPT_DROP_IDX:] |
|
|
