File size: 2,796 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 | # SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0
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" # noqa: E501
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, # Best configuration for wh t3k and galaxy
)
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:]
|