# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE # # Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any # output and results therefrom are provided "AS IS" without any express or implied warranties of # any kind including any warranties of title, merchantability, noninfringement, course of dealing, # usage of trade, or fitness for a particular purpose. You are solely responsible for determining the # appropriateness of using, reproducing, modifying, performing, displaying or distributing any of # the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a # third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise # of rights and permissions under this agreement. # See the License for the specific language governing permissions and limitations under the License. from shared.utils.phase_progress import text_encoding_progress import json import os from copy import deepcopy from dataclasses import dataclass from typing import Optional, Tuple from shared.utils import files_locator as fl import torch import torch.nn as nn from transformers import Qwen2_5_VLForConditionalGeneration from transformers import ( AutoTokenizer, AutoModel, ) from transformers.utils import ModelOutput def use_default(value, default): """Utility: return value if not None, else default.""" return value if value is not None else default # Prompt templates for different models and tasks __all__ = [ "C_SCALE", "PROMPT_TEMPLATE", "MODEL_BASE", ] # =================== Constant Values ===================== # Computation scale factor, 1P = 1_000_000_000_000_000. Tensorboard will display the value in PetaFLOPS to avoid # overflow error when tensorboard logging values. C_SCALE = 1_000_000_000_000_000 PROMPT_TEMPLATE_ENCODE_IMAGE_JSON = [ {"role": "system", "content": "You are a helpful assistant. Describe the image by detailing the following aspects: \ 1. The main content and theme of the image. \ 2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects. \ 3. The background environment, light, style and atmosphere."}, {"role": "user", "content": "{}"} ] PROMPT_TEMPLATE_ENCODE_VIDEO_JSON = [ {"role": "system", "content": "You are a helpful assistant. Describe the video by detailing the following aspects: \ 1. The main content and theme of the video. \ 2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects. \ 3. Actions, events, behaviors temporal relationships, physical movement changes of the objects. \ 4. background environment, light, style and atmosphere. \ 5. camera angles, movements, and transitions used in the video."}, {"role": "user", "content": "{}"} ] PROMPT_TEMPLATE = { "li-dit-encode-image-json": {"template": PROMPT_TEMPLATE_ENCODE_IMAGE_JSON, "crop_start": -1}, # auto-calculate crop_start "li-dit-encode-video-json": {"template": PROMPT_TEMPLATE_ENCODE_VIDEO_JSON, "crop_start": -1}, # auto-calculate crop_start } MODEL_BASE = os.getenv("MODEL_BASE", "") TEXT_ENCODER_PATH = {} TOKENIZER_PATH = {} PRECISION_TO_TYPE = { 'fp32': torch.float32, 'fp16': torch.float16, 'bf16': torch.bfloat16, } def load_text_encoder( text_encoder_type, text_encoder_precision=None, text_encoder_path=None, logger=None, device=None, ): if text_encoder_path is None: if text_encoder_type not in TEXT_ENCODER_PATH: raise ValueError(f"Unsupported text encoder type: {text_encoder_type}") text_encoder_path = TEXT_ENCODER_PATH[text_encoder_type] from mmgp import offload # text_encoder = offload.fast_load_transformers_model(text_encoder_path, forcedConfigPath= os.path.join(os.path.dirname(text_encoder_path), "config.json")) text_encoder = offload.fast_load_transformers_model(text_encoder_path, writable_tensors= True , modelClass=Qwen2_5_VLForConditionalGeneration, defaultConfigPath= fl.locate_file(os.path.join("Qwen2.5-VL-7B-Instruct", "config.json")) ) # text_encoder = AutoModel.from_pretrained(text_encoder_path, low_cpu_mem_usage=True) lm = getattr(text_encoder, "language_model", None) or getattr(getattr(text_encoder, "model", None), "language_model", None) text_encoder.final_layer_norm = lm.norm if lm is not None else text_encoder.norm # from_pretrained will ensure that the model is in eval mode. # if text_encoder_precision is not None: # text_encoder = text_encoder.to(dtype=PRECISION_TO_TYPE[text_encoder_precision]) text_encoder.requires_grad_(False) if device is not None: text_encoder = text_encoder.to(device) return text_encoder, text_encoder_path def load_tokenizer( tokenizer_type, tokenizer_path=None, padding_side="right", logger=None ): processor = None if tokenizer_path is None: if tokenizer_type not in TOKENIZER_PATH: raise ValueError(f"Unsupported tokenizer type: {tokenizer_type}") tokenizer_path = TOKENIZER_PATH[tokenizer_type] tokenizer = AutoTokenizer.from_pretrained( tokenizer_path, padding_side=padding_side ) # If the checkpoint ships a chat template separately, load it when missing. if getattr(tokenizer, "chat_template", None) in (None, ""): tokenizer_root = tokenizer_path if os.path.isdir(tokenizer_path) else os.path.dirname(tokenizer_path) json_path = os.path.join(tokenizer_root, "chat_template.json") jinja_path = os.path.join(tokenizer_root, "chat_template.jinja") template_path = json_path if os.path.exists(json_path) else jinja_path if os.path.exists(template_path): with open(template_path, "r", encoding="utf-8") as f: template_raw = f.read() if template_path.endswith(".json"): try: template_data = json.loads(template_raw) except json.JSONDecodeError: template_data = None tokenizer.chat_template = ( template_data.get("chat_template", template_raw) if isinstance(template_data, dict) else template_raw ) else: tokenizer.chat_template = template_raw return tokenizer, tokenizer_path, processor @dataclass class TextEncoderModelOutput(ModelOutput): """ Base class for model's outputs that also contains a pooling of the last hidden states. Args: hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): Sequence of hidden-states at the output of the last layer of the model. attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Mask to avoid performing attention on padding token indices. Mask values selected in ``[0, 1]``: hidden_states_list (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed): Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. text_outputs (`list`, *optional*, returned when `return_texts=True` is passed): List of decoded texts. """ hidden_state: torch.FloatTensor = None attention_mask: Optional[torch.LongTensor] = None hidden_states_list: Optional[Tuple[torch.FloatTensor, ...]] = None text_outputs: Optional[list] = None image_features: Optional[list] = None class TextEncoder(nn.Module): def __init__( self, text_encoder_type: str, max_length: int, text_encoder_precision: Optional[str] = None, text_encoder_path: Optional[str] = None, tokenizer_type: Optional[str] = None, tokenizer_path: Optional[str] = None, output_key: Optional[str] = None, use_attention_mask: bool = True, prompt_template: Optional[dict] = None, prompt_template_video: Optional[dict] = None, hidden_state_skip_layer: Optional[int] = None, apply_final_norm: bool = False, reproduce: bool = False, logger=None, i2v_mode = None, image_embed_interleave = None, device=None, ): super().__init__() self.text_encoder_type = text_encoder_type self.max_length = max_length self.precision = text_encoder_precision self.model_path = text_encoder_path self.tokenizer_type = ( tokenizer_type if tokenizer_type is not None else text_encoder_type ) self.tokenizer_path = ( tokenizer_path if tokenizer_path is not None else text_encoder_path ) self.use_attention_mask = use_attention_mask if prompt_template_video is not None: assert ( use_attention_mask is True ), "Attention mask is True required when training videos." self.prompt_template = prompt_template self.prompt_template_video = prompt_template_video self.hidden_state_skip_layer = hidden_state_skip_layer self.apply_final_norm = apply_final_norm self.reproduce = reproduce self.logger = logger self.use_template = self.prompt_template is not None if self.use_template: assert ( isinstance(self.prompt_template, dict) and "template" in self.prompt_template ), f"`prompt_template` must be a dictionary with a key 'template', got {self.prompt_template}" assert "{}" in str(self.prompt_template["template"]), ( "`prompt_template['template']` must contain a placeholder `{}` for the input text, " f"got {self.prompt_template['template']}" ) self.use_video_template = self.prompt_template_video is not None if self.use_video_template: if self.prompt_template_video is not None: assert ( isinstance(self.prompt_template_video, dict) and "template" in self.prompt_template_video ), f"`prompt_template_video` must be a dictionary with a key 'template', got {self.prompt_template_video}" assert "{}" in str(self.prompt_template_video["template"]), ( "`prompt_template_video['template']` must contain a placeholder `{}` for the input text, " f"got {self.prompt_template_video['template']}" ) if text_encoder_type != "llm": raise ValueError(f"Unsupported text encoder type: {text_encoder_type}") self.output_key = output_key or "last_hidden_state" self.model, self.model_path = load_text_encoder( text_encoder_type=self.text_encoder_type, text_encoder_precision=self.precision, text_encoder_path=self.model_path, logger=self.logger, device=device, ) self.tokenizer, self.tokenizer_path, self.processor = load_tokenizer( tokenizer_type=self.tokenizer_type, tokenizer_path=self.tokenizer_path, padding_side="right", logger=self.logger, ) # pre-calculate crop_start for image and video if self.use_template and self.prompt_template is not None: self.text2tokens("a photo of a cat", data_type="image") if self.use_video_template and self.prompt_template_video is not None: self.text2tokens("a photo of a cat", data_type="video") @property def dtype(self): return self.model.dtype @property def device(self): return self.model.device def __repr__(self): return f"{self.text_encoder_type} ({self.precision} - {self.model_path})" @staticmethod def apply_text_to_template(text, template, prevent_empty_text=True): """ Apply text to template. Args: text (str): Input text. template (str or list): Template string or list of chat conversation. prevent_empty_text (bool): If Ture, we will prevent the user text from being empty by adding a space. Defaults to True. """ if isinstance(template, str): # Will send string to tokenizer. Used for llm return template.format(text) elif isinstance(template, list): # For JSON list template format (chat conversation) # Create a deep copy to avoid modifying the original template template_copy = deepcopy(template) for item in template_copy: if isinstance(item, dict) and "content" in item: # Replace placeholder with text in the content field item["content"] = item["content"].format(text if text else (" " if prevent_empty_text else "")) return template_copy else: raise TypeError(f"Unsupported template type: {type(template)}") def calculate_crop_start(self, tokenized_input): """ Automatically calculate the crop_start position based on identifying user tokens. Args: tokenized_input: The output from the tokenizer containing input_ids Returns: int: The position where the actual prompt content begins (after user markers) """ input_ids = tokenized_input["input_ids"][0].tolist() # Get the first example's tokens marker = "<|im_start|>user\n" # Tokenize just the marker to get its token IDs marker_tokens = self.tokenizer(marker, add_special_tokens=False)["input_ids"] # Find the end position of the marker in the input sequence for i in range(len(input_ids) - len(marker_tokens) + 1): if input_ids[i:i+len(marker_tokens)] == marker_tokens: # Return the position after the marker return i + len(marker_tokens) # If marker not found, try to find based on special tokens if hasattr(self.tokenizer, 'special_tokens_map'): # Check for user token or any other special token that might indicate user input start for token_name, token_value in self.tokenizer.special_tokens_map.items(): if 'user' in token_name.lower(): user_token_id = self.tokenizer.convert_tokens_to_ids(token_value) if user_token_id in input_ids: return input_ids.index(user_token_id) + 1 # Default fallback: return 0 (no cropping) return 0 def text2tokens(self, text, data_type="image", max_length=300): """ Tokenize the input text. Args: text (str or list): Input text. """ tokenize_input_type = "str" if self.use_template or self.use_video_template: if data_type == "image": prompt_template = self.prompt_template["template"] crop_start = self.prompt_template.get("crop_start", -1) elif data_type == "video": prompt_template = self.prompt_template_video["template"] crop_start = self.prompt_template_video.get("crop_start", -1) else: raise ValueError(f"Unsupported data type: {data_type}") if isinstance(text, (list, tuple)): text = [ self.apply_text_to_template(one_text, prompt_template) for one_text in text ] if isinstance(text[0], list): tokenize_input_type = "list" elif isinstance(text, str): text = self.apply_text_to_template(text, prompt_template) if isinstance(text, list): tokenize_input_type = "list" else: raise TypeError(f"Unsupported text type: {type(text)}") # First pass: tokenize with arbitrary max_length to find crop_start if crop_start == -1: # Use temporary max_length for the first pass (large enough) temp_kwargs = dict( truncation=True, max_length=256, # Temporary large value padding="max_length", return_tensors="pt", ) # First tokenization pass to calculate crop_start if tokenize_input_type == "str": temp_tokenized = self.tokenizer( text, return_length=False, return_overflowing_tokens=False, return_attention_mask=True, **temp_kwargs, ) elif tokenize_input_type == "list": temp_tokenized = self.tokenizer.apply_chat_template( text, add_generation_prompt=True, tokenize=True, return_dict=True, **temp_kwargs, ) # Calculate the crop_start from this first pass crop_start = self.calculate_crop_start(temp_tokenized) # Store the calculated crop_start for future use if data_type == "image": self.prompt_template["crop_start"] = crop_start else: self.prompt_template_video["crop_start"] = crop_start else: crop_start = 0 # Second pass: tokenize with the proper max_length using the found crop_start kwargs = dict( truncation=True, max_length=max_length + (crop_start if crop_start > 0 else 0), padding="max_length", return_tensors="pt", ) if tokenize_input_type == "str": tokenized_output = self.tokenizer( text, return_length=False, return_overflowing_tokens=False, return_attention_mask=True, **kwargs, ) elif tokenize_input_type == "list": tokenized_output = self.tokenizer.apply_chat_template( text, add_generation_prompt=True, tokenize=True, return_dict=True, **kwargs, ) else: raise ValueError(f"Unsupported tokenize_input_type: {tokenize_input_type}") return tokenized_output def encode( self, batch_encoding, use_attention_mask=None, output_hidden_states=False, do_sample=None, hidden_state_skip_layer=None, return_texts=False, data_type="image", device=None, is_uncond=False ): """ Args: batch_encoding (dict): Batch encoding from tokenizer. use_attention_mask (bool): Whether to use attention mask. If None, use self.use_attention_mask. Defaults to None. output_hidden_states (bool): Whether to output hidden states. If False, return the value of self.output_key. If True, return the entire output. If set self.hidden_state_skip_layer, output_hidden_states will be set True. Defaults to False. do_sample (bool): Whether to sample from the model. Used for Decoder-Only LLMs. Defaults to None. When self.produce is False, do_sample is set to True by default. hidden_state_skip_layer (int): Number of hidden states to hidden_state_skip_layer. 0 means the last layer. If None, self.output_key will be used. Defaults to None. return_texts (bool): Whether to return the decoded texts. Defaults to False. """ device = self.model.device if device is None else device use_attention_mask = use_default(use_attention_mask, self.use_attention_mask) hidden_state_skip_layer = use_default( hidden_state_skip_layer, self.hidden_state_skip_layer ) do_sample = use_default(do_sample, not self.reproduce) attention_mask = ( batch_encoding["attention_mask"].to(device) if use_attention_mask else None ) with text_encoding_progress(self.model.model.language_model.layers, prompt_count=batch_encoding["input_ids"].shape[0]): outputs = self.model( input_ids=batch_encoding["input_ids"].to(device), attention_mask=attention_mask, output_hidden_states=output_hidden_states or hidden_state_skip_layer is not None, ) if hidden_state_skip_layer is not None: last_hidden_state = outputs.hidden_states[-(hidden_state_skip_layer + 1)] # Real last hidden state already has layer norm applied. So here we only apply it # for intermediate layers. if hidden_state_skip_layer > 0 and self.apply_final_norm: last_hidden_state = self.model.final_layer_norm(last_hidden_state) else: last_hidden_state = outputs[self.output_key] # Remove hidden states of instruction tokens, only keep prompt tokens. if self.use_template: if data_type == "image": crop_start = self.prompt_template.get("crop_start", 0) elif data_type == "video": crop_start = self.prompt_template_video.get("crop_start", 0) else: raise ValueError(f"Unsupported data type: {data_type}") if crop_start > 0: last_hidden_state = last_hidden_state[:, crop_start:] attention_mask = ( attention_mask[:, crop_start:] if use_attention_mask else None ) if output_hidden_states: return TextEncoderModelOutput( last_hidden_state, attention_mask, outputs.hidden_states ) return TextEncoderModelOutput(last_hidden_state, attention_mask) def forward( self, text, use_attention_mask=None, output_hidden_states=False, do_sample=False, hidden_state_skip_layer=None, return_texts=False, ): batch_encoding = self.text2tokens(text, max_length=self.max_length) return self.encode( batch_encoding, use_attention_mask=use_attention_mask, output_hidden_states=output_hidden_states, do_sample=do_sample, hidden_state_skip_layer=hidden_state_skip_layer, return_texts=return_texts, )