bjf-frzandHongsheng Liu
c33ef34654
Signed-off-by: bjf-frz <frz123db@gmail.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
1473 lines
65 KiB
Python
1473 lines
65 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import logging
|
|
from collections.abc import Iterable
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
from diffusers.schedulers.scheduling_flow_match_euler_discrete import (
|
|
FlowMatchEulerDiscreteScheduler,
|
|
)
|
|
from PIL import Image as PILImage
|
|
from transformers.generation.configuration_utils import GenerationConfig
|
|
from transformers.generation.utils import ALL_CACHE_NAMES, GenerationMixin
|
|
from transformers.utils.generic import ModelOutput
|
|
from vllm.config.vllm import get_current_vllm_config
|
|
from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper
|
|
from vllm.transformers_utils.config import get_config
|
|
|
|
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
|
|
from vllm_omni.diffusion.distributed.utils import get_local_device
|
|
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
|
from vllm_omni.diffusion.models.interface import SupportImageInput
|
|
from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import (
|
|
DiffusionPipelineProfilerMixin,
|
|
)
|
|
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
|
from vllm_omni.inputs.data import OmniTextPrompt
|
|
from vllm_omni.model_executor.models.hunyuan_image3.siglip2 import Siglip2VisionTransformer
|
|
|
|
from .hunyuan_image3_tokenizer import TokenizerWrapper
|
|
from .hunyuan_image3_transformer import (
|
|
CausalMMOutputWithPast,
|
|
HunyuanImage3ImageProcessor,
|
|
HunyuanImage3Model,
|
|
HunyuanImage3PreTrainedModel,
|
|
HunyuanImage3Text2ImagePipeline,
|
|
ImageInfo,
|
|
JointImageInfo,
|
|
LightProjector,
|
|
TimestepEmbedder,
|
|
UNetDown,
|
|
UNetUp,
|
|
build_batch_2d_rope,
|
|
real_batched_index_select,
|
|
)
|
|
from .system_prompt import get_system_prompt
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
BatchRaggedImages = torch.Tensor | list[torch.Tensor | list[torch.Tensor]]
|
|
BatchRaggedTensor = torch.Tensor | list[torch.Tensor]
|
|
|
|
|
|
def default(val, d):
|
|
return val if val is not None else d
|
|
|
|
|
|
def to_device(data, device):
|
|
if device is None:
|
|
return data
|
|
if isinstance(data, torch.Tensor):
|
|
return data.to(device)
|
|
elif isinstance(data, list):
|
|
return [to_device(x, device) for x in data]
|
|
else:
|
|
return data
|
|
|
|
|
|
def _to_pil_image(image: Any) -> PILImage.Image:
|
|
if isinstance(image, PILImage.Image):
|
|
return image
|
|
if isinstance(image, str):
|
|
return PILImage.open(image)
|
|
if isinstance(image, np.ndarray):
|
|
array = image
|
|
if array.dtype != np.uint8:
|
|
if np.issubdtype(array.dtype, np.floating):
|
|
if float(np.min(array)) < 0.0:
|
|
array = (np.clip(array, -1.0, 1.0) + 1.0) / 2.0
|
|
if float(np.max(array)) <= 1.0:
|
|
array = array * 255.0
|
|
array = np.clip(array, 0, 255).astype(np.uint8)
|
|
if array.ndim == 3 and array.shape[0] in (1, 3, 4):
|
|
array = np.transpose(array, (1, 2, 0))
|
|
return PILImage.fromarray(array)
|
|
if isinstance(image, torch.Tensor):
|
|
tensor = image.detach().cpu()
|
|
if tensor.ndim == 4:
|
|
if tensor.shape[0] != 1:
|
|
raise ValueError(f"Only a single image tensor is supported, but got shape {tuple(tensor.shape)}.")
|
|
tensor = tensor.squeeze(0)
|
|
if tensor.ndim == 3 and tensor.shape[0] in (1, 3, 4):
|
|
tensor = tensor.permute(1, 2, 0)
|
|
if tensor.dtype.is_floating_point:
|
|
if float(tensor.min()) < 0.0:
|
|
tensor = (tensor.clamp(-1.0, 1.0) + 1.0) / 2.0
|
|
if float(tensor.max()) > 1.0:
|
|
tensor = tensor / 255.0
|
|
tensor = (tensor.clamp(0.0, 1.0) * 255.0).to(torch.uint8)
|
|
else:
|
|
tensor = tensor.to(torch.uint8)
|
|
return PILImage.fromarray(tensor.numpy())
|
|
raise TypeError(f"Unsupported image input type: {type(image)}")
|
|
|
|
|
|
def _resize_and_crop_center(image: PILImage.Image, target_width: int, target_height: int) -> PILImage.Image:
|
|
# Mirrors HunyuanImage3Processor._resize_and_crop in
|
|
# vllm_omni.model_executor.models.hunyuan_image3.hunyuan_image3 so the AR
|
|
# and DiT stages preprocess condition images identically.
|
|
tw, th = target_width, target_height
|
|
w, h = image.size
|
|
tr = th / tw
|
|
r = h / w
|
|
if r < tr:
|
|
resize_height = th
|
|
resize_width = int(round(th / h * w))
|
|
else:
|
|
resize_width = tw
|
|
resize_height = int(round(tw / w * h))
|
|
resized = image.resize((resize_width, resize_height), PILImage.Resampling.LANCZOS)
|
|
crop_top = int(round((resize_height - th) / 2.0))
|
|
crop_left = int(round((resize_width - tw) / 2.0))
|
|
return resized.crop((crop_left, crop_top, crop_left + tw, crop_top + th))
|
|
|
|
|
|
def _to_python_scalar(value: Any) -> Any:
|
|
if isinstance(value, np.generic):
|
|
return value.item()
|
|
return value
|
|
|
|
|
|
def _image_info_to_payload(image_info: ImageInfo) -> dict[str, Any]:
|
|
return {
|
|
"image_type": image_info.image_type,
|
|
"image_tensor": image_info.image_tensor,
|
|
"image_width": _to_python_scalar(image_info.image_width),
|
|
"image_height": _to_python_scalar(image_info.image_height),
|
|
"token_width": _to_python_scalar(image_info.token_width),
|
|
"token_height": _to_python_scalar(image_info.token_height),
|
|
"image_token_length": _to_python_scalar(image_info.image_token_length),
|
|
"base_size": _to_python_scalar(image_info.base_size),
|
|
"ratio_index": _to_python_scalar(image_info.ratio_index),
|
|
"add_timestep_token": image_info.add_timestep_token,
|
|
"add_guidance_token": image_info.add_guidance_token,
|
|
"use_front_boi_token": image_info.use_front_boi_token,
|
|
"add_image_shape_token": image_info.add_image_shape_token,
|
|
}
|
|
|
|
|
|
def _to_tensor_if_needed(value: Any) -> Any:
|
|
if isinstance(value, np.generic):
|
|
return value.item()
|
|
if isinstance(value, list):
|
|
return torch.tensor(value)
|
|
return value
|
|
|
|
|
|
def _image_info_from_payload(payload: dict[str, Any]) -> ImageInfo:
|
|
return ImageInfo(
|
|
image_type=payload.get("image_type"),
|
|
image_tensor=_to_tensor_if_needed(payload.get("image_tensor")),
|
|
image_width=payload.get("image_width"),
|
|
image_height=payload.get("image_height"),
|
|
token_width=payload.get("token_width"),
|
|
token_height=payload.get("token_height"),
|
|
image_token_length=payload.get("image_token_length"),
|
|
base_size=payload.get("base_size"),
|
|
ratio_index=payload.get("ratio_index"),
|
|
add_timestep_token=payload.get("add_timestep_token", True),
|
|
add_guidance_token=payload.get("add_guidance_token", False),
|
|
use_front_boi_token=payload.get("use_front_boi_token", True),
|
|
add_image_shape_token=payload.get("add_image_shape_token", True),
|
|
)
|
|
|
|
|
|
def _joint_image_info_to_payload(joint_image_info: JointImageInfo) -> dict[str, Any]:
|
|
return {
|
|
"type": "joint_image_info",
|
|
"vae_image_info": _image_info_to_payload(joint_image_info.vae_image_info),
|
|
"vision_image_info": _image_info_to_payload(joint_image_info.vision_image_info),
|
|
"vision_encoder_kwargs": joint_image_info.vision_encoder_kwargs,
|
|
}
|
|
|
|
|
|
def _joint_image_info_from_payload(payload: Any) -> JointImageInfo:
|
|
if isinstance(payload, JointImageInfo):
|
|
return payload
|
|
if not isinstance(payload, dict):
|
|
raise TypeError(f"Expected dict or JointImageInfo for conditional image payload, got {type(payload)}.")
|
|
|
|
vae_image_info = _image_info_from_payload(payload["vae_image_info"])
|
|
vision_image_info = _image_info_from_payload(payload["vision_image_info"])
|
|
vision_encoder_kwargs = payload.get("vision_encoder_kwargs") or {}
|
|
if isinstance(vision_encoder_kwargs, dict):
|
|
vision_encoder_kwargs = {key: _to_tensor_if_needed(value) for key, value in vision_encoder_kwargs.items()}
|
|
return JointImageInfo(
|
|
vae_image_info=vae_image_info,
|
|
vision_image_info=vision_image_info,
|
|
vision_encoder_kwargs=vision_encoder_kwargs,
|
|
)
|
|
|
|
|
|
def get_hunyuan_image_3_pre_process_func(od_config: OmniDiffusionConfig):
|
|
hf_config = get_config(od_config.model, trust_remote_code=True)
|
|
image_processor = HunyuanImage3ImageProcessor(hf_config)
|
|
vae_h_factor = hf_config.vae_downsample_factor[0] * hf_config.patch_size
|
|
vae_w_factor = hf_config.vae_downsample_factor[1] * hf_config.patch_size
|
|
vit_patch_size = getattr(image_processor.vision_encoder_processor, "patch_size", 1)
|
|
if isinstance(vit_patch_size, tuple | list):
|
|
vit_patch_size = int(vit_patch_size[0])
|
|
|
|
def _build_cond_joint_image(raw_image: Any) -> dict[str, Any]:
|
|
pil_image = _to_pil_image(raw_image).convert("RGB")
|
|
orig_width, orig_height = pil_image.size
|
|
|
|
target_width, target_height = image_processor.reso_group.get_target_size(orig_width, orig_height)
|
|
target_width = int(target_width)
|
|
target_height = int(target_height)
|
|
vae_input = _resize_and_crop_center(pil_image, target_width, target_height)
|
|
vae_tensor = image_processor.vae_processor(vae_input)
|
|
base_size, ratio_idx = image_processor.reso_group.get_base_size_and_ratio_index(orig_width, orig_height)
|
|
base_size = int(base_size)
|
|
ratio_idx = int(ratio_idx)
|
|
|
|
vae_info = ImageInfo(
|
|
image_type="vae",
|
|
image_tensor=vae_tensor,
|
|
image_width=target_width,
|
|
image_height=target_height,
|
|
token_width=target_width // vae_w_factor,
|
|
token_height=target_height // vae_h_factor,
|
|
base_size=base_size,
|
|
ratio_index=ratio_idx,
|
|
)
|
|
|
|
vit_inputs = image_processor.vision_encoder_processor(pil_image, return_tensors="pt")
|
|
vit_tensor = vit_inputs["pixel_values"]
|
|
spatial_shapes = vit_inputs["spatial_shapes"].squeeze(0)
|
|
pixel_attention_mask = vit_inputs["pixel_attention_mask"].squeeze(0)
|
|
vit_token_h = int(spatial_shapes[0].item())
|
|
vit_token_w = int(spatial_shapes[1].item())
|
|
|
|
vit_info = ImageInfo(
|
|
image_type="siglip2",
|
|
image_tensor=vit_tensor,
|
|
image_width=vit_token_w * vit_patch_size,
|
|
image_height=vit_token_h * vit_patch_size,
|
|
token_width=vit_token_w,
|
|
token_height=vit_token_h,
|
|
image_token_length=int(vit_tensor.shape[1]),
|
|
)
|
|
|
|
return _joint_image_info_to_payload(
|
|
JointImageInfo(
|
|
vae_image_info=vae_info,
|
|
vision_image_info=vit_info,
|
|
vision_encoder_kwargs={
|
|
"spatial_shapes": spatial_shapes,
|
|
"pixel_attention_mask": pixel_attention_mask,
|
|
},
|
|
)
|
|
)
|
|
|
|
def pre_process_func(request: OmniDiffusionRequest):
|
|
for i, prompt in enumerate(request.prompts):
|
|
if isinstance(prompt, str):
|
|
prompt = OmniTextPrompt(prompt=prompt)
|
|
|
|
if "additional_information" not in prompt:
|
|
prompt["additional_information"] = {}
|
|
|
|
multi_modal_data = prompt.get("multi_modal_data") or {}
|
|
raw_images = multi_modal_data.get("image")
|
|
if raw_images is None:
|
|
raw_images = prompt.get("pil_image")
|
|
has_images = raw_images is not None and (not isinstance(raw_images, list) or len(raw_images) > 0)
|
|
if has_images:
|
|
image_list = raw_images if isinstance(raw_images, list) else [raw_images]
|
|
cond_image_infos = [_build_cond_joint_image(image) for image in image_list]
|
|
prompt["additional_information"]["batch_cond_image_info"] = cond_image_infos
|
|
|
|
bridge_h = prompt.get("height") if isinstance(prompt, dict) else None
|
|
bridge_w = prompt.get("width") if isinstance(prompt, dict) else None
|
|
first_image_w, first_image_h = _to_pil_image(image_list[0]).size
|
|
if request.sampling_params.width is None:
|
|
request.sampling_params.width = int(bridge_w or first_image_w)
|
|
if request.sampling_params.height is None:
|
|
request.sampling_params.height = int(bridge_h or first_image_h)
|
|
|
|
request.prompts[i] = prompt
|
|
|
|
return request
|
|
|
|
return pre_process_func
|
|
|
|
|
|
class HunyuanImage3Pipeline(
|
|
HunyuanImage3PreTrainedModel,
|
|
GenerationMixin,
|
|
SupportImageInput,
|
|
DiffusionPipelineProfilerMixin,
|
|
):
|
|
support_image_input = True
|
|
hf_to_vllm_mapper = WeightsMapper(
|
|
orig_to_new_prefix={
|
|
"model.": "",
|
|
},
|
|
orig_to_new_substr={
|
|
"mlp.gate.wg.": "mlp.gate.",
|
|
"gate_and_up_proj.": "gate_up_proj.",
|
|
},
|
|
)
|
|
_PROFILER_TARGETS = [
|
|
"model.forward",
|
|
"model.layers[0].forward",
|
|
"model.layers[0].self_attn.forward",
|
|
"model.layers[0].mlp.forward",
|
|
"vae.encode",
|
|
"vae.decode",
|
|
"patch_embed.forward",
|
|
"final_layer.forward",
|
|
]
|
|
|
|
def __init__(self, od_config: OmniDiffusionConfig) -> None:
|
|
self.hf_config = get_config(od_config.model, trust_remote_code=True)
|
|
super().__init__(self.hf_config)
|
|
# update diffusion config
|
|
self.generation_config = GenerationConfig.from_pretrained(od_config.model)
|
|
self.od_config = od_config
|
|
self.weights_sources = [
|
|
DiffusersPipelineLoader.ComponentSource(
|
|
model_or_path=od_config.model,
|
|
subfolder=None,
|
|
revision=od_config.revision,
|
|
prefix="",
|
|
fall_back_to_pt=True,
|
|
)
|
|
]
|
|
quant_config = od_config.quantization_config
|
|
self.model = HunyuanImage3Model(self.hf_config, quant_config=quant_config)
|
|
self.transformer = self.model
|
|
# Lazy import to break circular dependency:
|
|
# autoencoder_kl_hunyuan -> hunyuan_image3/__init__ -> pipeline_hunyuan_image3 -> autoencoder_kl_hunyuan
|
|
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_hunyuan import ( # noqa: PLC0415
|
|
DistributedAutoencoderKLHunyuan,
|
|
)
|
|
|
|
self.vae = DistributedAutoencoderKLHunyuan.from_config(self.hf_config.vae)
|
|
self.vae.use_spatial_tiling = self.od_config.vae_use_tiling
|
|
self._pipeline = None
|
|
self._tkwrapper = TokenizerWrapper(od_config.model)
|
|
self.image_processor = HunyuanImage3ImageProcessor(self.hf_config)
|
|
self.vision_model = Siglip2VisionTransformer(self.hf_config.vit)
|
|
# self.vision_model = vision_model.vision_model
|
|
self.vision_aligner = LightProjector(self.hf_config.vit_aligner)
|
|
self.timestep_emb = TimestepEmbedder(hidden_size=self.hf_config.hidden_size)
|
|
if self.hf_config.img_proj_type != "unet":
|
|
raise ValueError(f"Unknown img_proj_type: {self.hf_config.img_proj_type}")
|
|
|
|
self.patch_embed = UNetDown(
|
|
patch_size=self.hf_config.patch_size,
|
|
emb_channels=self.hf_config.hidden_size,
|
|
in_channels=self.hf_config.vae["latent_channels"],
|
|
hidden_channels=self.hf_config.patch_embed_hidden_dim,
|
|
out_channels=self.hf_config.hidden_size,
|
|
)
|
|
self.time_embed = TimestepEmbedder(hidden_size=self.hf_config.hidden_size)
|
|
self.final_layer = UNetUp(
|
|
patch_size=self.hf_config.patch_size,
|
|
emb_channels=self.hf_config.hidden_size,
|
|
in_channels=self.hf_config.hidden_size,
|
|
hidden_channels=self.hf_config.patch_embed_hidden_dim,
|
|
out_channels=self.hf_config.vae["latent_channels"],
|
|
out_norm=True,
|
|
)
|
|
self.time_embed_2 = TimestepEmbedder(hidden_size=self.hf_config.hidden_size)
|
|
self.lm_head = nn.Linear(self.hf_config.hidden_size, self.hf_config.vocab_size, bias=False)
|
|
self.vllm_config = get_current_vllm_config()
|
|
self.post_init()
|
|
self.setup_diffusion_pipeline_profiler(
|
|
enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler,
|
|
)
|
|
|
|
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
|
skip_prefixes = ["lm_head."] if self.hf_config.tie_word_embeddings else []
|
|
# List of unexpected keywords in weight names
|
|
non_model_layer_prefixes = [
|
|
"vae",
|
|
"vision_model",
|
|
"vision_aligner",
|
|
"lm_head",
|
|
"patch_embed",
|
|
"timestep_emb",
|
|
"model.wte",
|
|
"model.ln_f",
|
|
"time_embed",
|
|
"time_embed_2",
|
|
"final_layer.model",
|
|
]
|
|
device = get_local_device()
|
|
named_modules = dict(self.named_modules())
|
|
for prefix in non_model_layer_prefixes:
|
|
mod = named_modules.get(prefix)
|
|
if mod:
|
|
mod.to(device)
|
|
|
|
unexpected_keywords = [
|
|
"guidance_emb",
|
|
"timestep_r_emb",
|
|
]
|
|
skip_prefixes.extend(unexpected_keywords)
|
|
loader = AutoWeightsLoader(
|
|
self,
|
|
skip_prefixes=skip_prefixes,
|
|
)
|
|
return loader.load_weights(weights)
|
|
|
|
def prepare_seed(self, seed=None, batch_size=1):
|
|
# random seed
|
|
if seed is not None:
|
|
return [seed + i for i in range(batch_size)]
|
|
else:
|
|
import random
|
|
|
|
return [random.randint(0, 2**32 - 1) for _ in range(batch_size)]
|
|
|
|
@property
|
|
def pipeline(self):
|
|
if self._pipeline is None:
|
|
# shift hard code
|
|
self.scheduler = FlowMatchEulerDiscreteScheduler(
|
|
num_train_timesteps=1000,
|
|
shift=self.generation_config.flow_shift,
|
|
use_dynamic_shifting=False,
|
|
base_shift=0.5,
|
|
max_shift=1.15,
|
|
time_shift_type="exponential",
|
|
stochastic_sampling=False,
|
|
)
|
|
self._pipeline = HunyuanImage3Text2ImagePipeline(model=self, scheduler=self.scheduler, vae=self.vae)
|
|
return self._pipeline
|
|
|
|
@staticmethod
|
|
def get_pos_emb(custom_pos_emb, position_ids):
|
|
cos, sin = custom_pos_emb
|
|
cos = real_batched_index_select(cos, dim=1, idx=position_ids)
|
|
sin = real_batched_index_select(sin, dim=1, idx=position_ids)
|
|
return cos, sin
|
|
|
|
def instantiate_vae_image_tokens(
|
|
self,
|
|
x: torch.Tensor,
|
|
images: BatchRaggedImages,
|
|
ts: BatchRaggedTensor,
|
|
image_mask: torch.Tensor,
|
|
):
|
|
r"""
|
|
Instantiate the VAE image embeddings into the input embedding sequence.
|
|
Args:
|
|
x (`torch.Tensor`):
|
|
Input sequence tensor with shape `(batch_size, seq_len, n_embd)`.
|
|
images (`BatchRaggedImages`):
|
|
Batch of images to embed. Can be:
|
|
- A 4-D tensor (batch, channels, height, width)
|
|
- A list of 4-D tensors (variable number of images per batch)
|
|
- A list of lists of 3-D tensors (ragged batch structure)
|
|
ts (`BatchRaggedTensor`, *optional*):
|
|
Timestep tensor(s) for conditioning. Can be:
|
|
- A 1-D tensor (single timestep per batch)
|
|
- A list of 1-D tensors (variable timesteps per batch)
|
|
image_mask (`torch.Tensor`, *optional*):
|
|
Boolean mask tensor with shape `(batch_size, seq_len)` indicating which positions
|
|
should be replaced with image embeddings.
|
|
"""
|
|
batch_size, seq_len, n_embd = x.shape
|
|
|
|
if isinstance(images, list):
|
|
index = torch.arange(seq_len, device=x.device).unsqueeze(0).repeat(batch_size, 1)
|
|
t_emb = []
|
|
for i, (image_i, t_i) in enumerate(zip(images, ts)):
|
|
if isinstance(image_i, torch.Tensor):
|
|
# time_embed needs a 1-D tensor as input
|
|
t_i_emb = self.time_embed(t_i)
|
|
# n_{i} x one_image_seq_len x n_embd
|
|
image_i_seq, _, _ = self.patch_embed(image_i, t_i_emb)
|
|
# 1 x (n_{i} * one_image_seq_len)
|
|
image_i_scatter_index = index[i : i + 1].masked_select(image_mask[i : i + 1].bool()).reshape(1, -1)
|
|
x[i : i + 1].scatter_(
|
|
dim=1,
|
|
index=image_i_scatter_index.unsqueeze(-1).repeat(1, 1, n_embd),
|
|
# 1 x (n_{i} * one_image_seq_len) x n_embd
|
|
src=image_i_seq.reshape(1, -1, n_embd), # 1 x (n_{i} * one_image_seq_len) x n_embd
|
|
)
|
|
t_emb.append(t_i_emb)
|
|
elif isinstance(image_i, list):
|
|
# time_embed needs a 1-D tensor as input
|
|
t_i_emb = self.time_embed(t_i) # n_{i} x d
|
|
image_i_seq_list = []
|
|
for j in range(len(image_i)):
|
|
image_ij = image_i[j]
|
|
if image_ij.dim() == 4:
|
|
assert image_i[j].shape[0] == 1, "image_i[j] should have a batch dimension of 1"
|
|
elif image_ij.dim() == 3:
|
|
image_ij = image_ij.unsqueeze(0)
|
|
else:
|
|
raise ValueError(f"image_i[j] should have 3 or 4 dimensions, got {image_ij.dim()}")
|
|
# 1 x one_image_seq_len_{j} x n_embd
|
|
image_i_seq_j, _, _ = self.patch_embed(image_ij, t_i_emb[j : j + 1])
|
|
image_i_seq_list.append(image_i_seq_j)
|
|
# 1 x sum_{j}(one_image_seq_len_{j}) x n_embd
|
|
image_i_seq = torch.cat(image_i_seq_list, dim=1)
|
|
# 1 x sum_{j}(one_image_seq_len_{j})
|
|
image_i_scatter_index = index[i : i + 1].masked_select(image_mask[i : i + 1].bool()).reshape(1, -1)
|
|
x[i : i + 1].scatter_(
|
|
dim=1,
|
|
index=image_i_scatter_index.unsqueeze(-1).repeat(1, 1, n_embd),
|
|
# 1 x sum_{j}(one_image_seq_len_{j}) x n_embd
|
|
src=image_i_seq.reshape(1, -1, n_embd), # 1 x sum_{j}(one_image_seq_len_{j}) x n_embd
|
|
)
|
|
t_emb.append(t_i_emb)
|
|
else:
|
|
raise TypeError(f"image_i should be a torch.Tensor or a list, got {type(image_i)}")
|
|
token_h, token_w = None, None
|
|
else:
|
|
# images is a 4-D tensor
|
|
batch_size, seq_len, n_embd = x.shape
|
|
index = torch.arange(seq_len, device=x.device).unsqueeze(0).repeat(batch_size, 1)
|
|
t_emb = self.time_embed(ts)
|
|
image_seq, token_h, token_w = self.patch_embed(images, t_emb)
|
|
image_scatter_index = index.masked_select(image_mask.bool()).reshape(batch_size, -1)
|
|
x.scatter_(
|
|
dim=1,
|
|
index=image_scatter_index.unsqueeze(-1).repeat(1, 1, n_embd),
|
|
src=image_seq,
|
|
)
|
|
|
|
return x, token_h, token_w
|
|
|
|
def instantiate_timestep_tokens(
|
|
self,
|
|
x: torch.Tensor,
|
|
t: BatchRaggedTensor,
|
|
timestep_scatter_index: BatchRaggedTensor,
|
|
):
|
|
batch_size, seq_len, n_embd = x.shape
|
|
# `_encode_cond_image` returns `t` as list[Tensor] for the
|
|
# multi-image branch (outer length = batch_size, currently fixed
|
|
# at 1 by the stage runtime `max_batch_size`); flatten to a Tensor
|
|
# before reshape.
|
|
if isinstance(t, list):
|
|
t = torch.cat([ti.reshape(-1) for ti in t], dim=0)
|
|
timestep_scatter_src = self.timestep_emb(t.reshape(-1)).reshape(batch_size, -1, n_embd)
|
|
x.scatter_(
|
|
dim=1,
|
|
index=timestep_scatter_index.unsqueeze(-1).repeat(1, 1, n_embd),
|
|
src=timestep_scatter_src,
|
|
)
|
|
|
|
return x
|
|
|
|
def instantiate_vit_image_tokens(
|
|
self,
|
|
x: torch.Tensor,
|
|
cond_vit_images: torch.Tensor | list[torch.Tensor],
|
|
cond_vit_image_mask: torch.Tensor,
|
|
vit_kwargs: dict[str, Any],
|
|
):
|
|
# 1. Forward the vit encoder and vit aligner to get the vit image embeddings and align them to the
|
|
# transformer hidden size
|
|
cond_vit_image_embeds = []
|
|
for batch_idx, image in enumerate(cond_vit_images):
|
|
cur_kwargs = {k: v[batch_idx] for k, v in vit_kwargs.items()}
|
|
image_embed = self.vision_model(image, **cur_kwargs).last_hidden_state
|
|
image_embed = self.vision_aligner(image_embed)
|
|
n, seq_len, dim = image_embed.shape
|
|
image_embed = image_embed.reshape(n * seq_len, dim)
|
|
cond_vit_image_embeds.append(image_embed)
|
|
|
|
# 2. Instantiate the vit image embeddings into the input sequence
|
|
batch_size, seq_len, n_embd = x.shape
|
|
index = torch.arange(seq_len, device=x.device).unsqueeze(0).repeat(batch_size, 1)
|
|
|
|
for i, (image_embed, mask) in enumerate(zip(cond_vit_image_embeds, cond_vit_image_mask)):
|
|
image_scatter_index = index[i : i + 1].masked_select(mask.bool()).reshape(1, -1)
|
|
x[i : i + 1].scatter_(
|
|
dim=1,
|
|
index=image_scatter_index.unsqueeze(-1).repeat(1, 1, n_embd),
|
|
src=image_embed.reshape(1, -1, n_embd),
|
|
)
|
|
|
|
return x
|
|
|
|
def ragged_final_layer(self, x, image_mask, timestep, token_h, token_w, first_step):
|
|
bsz, seq_len, n_embd = x.shape
|
|
if first_step:
|
|
image_output = x.masked_select(image_mask.unsqueeze(-1).bool()).reshape(bsz, -1, n_embd)
|
|
else:
|
|
image_output = x[:, 1:, :]
|
|
timestep_emb = self.time_embed_2(timestep)
|
|
pred = self.final_layer(image_output, timestep_emb, token_h, token_w)
|
|
return pred
|
|
|
|
@staticmethod
|
|
def build_batch_rope_image_info(output, sections):
|
|
rope_image_info = []
|
|
for image_slices, sections_i in zip(output.all_image_slices, sections):
|
|
image_shapes = []
|
|
for section in sections_i:
|
|
if "image" in section["type"]:
|
|
if isinstance(section["token_height"], list):
|
|
assert len(section["token_height"]) == len(section["token_width"]), (
|
|
f"token_height and token_width should have the same length, "
|
|
f"but got {len(section['token_height'])} and {len(section['token_width'])}"
|
|
)
|
|
image_shapes.extend(list(zip(section["token_height"], section["token_width"])))
|
|
else:
|
|
image_shapes.append((section["token_height"], section["token_width"]))
|
|
assert len(image_slices) == len(image_shapes), (
|
|
f"Size miss matching: Image slices({len(image_slices)}) != image shapes({len(image_shapes)})"
|
|
)
|
|
rope_image_info.append(list(zip(image_slices, image_shapes)))
|
|
return rope_image_info
|
|
|
|
def vae_encode(self, image, cfg_factor=1):
|
|
config = self.vae.config
|
|
|
|
if image.ndim == 3:
|
|
image = image.unsqueeze(0)
|
|
if image.ndim == 4:
|
|
image = image.unsqueeze(2)
|
|
if image.ndim != 5:
|
|
raise ValueError(f"Expected image tensor with 3/4/5 dims, got shape {tuple(image.shape)}.")
|
|
|
|
with torch.autocast(device_type=self.model.device.type, dtype=torch.float16, enabled=True):
|
|
vae_encode_result = self.vae.encode(image)
|
|
if isinstance(vae_encode_result, torch.Tensor):
|
|
latents = vae_encode_result
|
|
else:
|
|
# Match HunyuanImage-3's cond encode path: sample the
|
|
# posterior, but use a fixed generator so repeated online
|
|
# requests are deterministic.
|
|
_cond_vae_gen = torch.Generator(device=image.device).manual_seed(0)
|
|
latents = vae_encode_result.latent_dist.sample(_cond_vae_gen)
|
|
if hasattr(config, "shift_factor") and config.shift_factor:
|
|
latents.sub_(config.shift_factor)
|
|
if hasattr(config, "scaling_factor") and config.scaling_factor:
|
|
latents.mul_(config.scaling_factor)
|
|
|
|
if hasattr(self.vae, "ffactor_temporal"):
|
|
assert latents.shape[2] == 1, "latents should have shape [B, C, T, H, W] and T should be 1"
|
|
latents = latents.squeeze(2)
|
|
|
|
# Here we always use t=0 to declare it is a clean conditional image
|
|
t = torch.zeros((latents.shape[0],))
|
|
|
|
if cfg_factor > 1:
|
|
t = t.repeat(cfg_factor)
|
|
latents = latents.repeat(cfg_factor, 1, 1, 1)
|
|
|
|
return t, latents
|
|
|
|
def _encode_cond_image(
|
|
self,
|
|
batch_cond_image_info_list: list[list[JointImageInfo]],
|
|
cfg_factor: int = 1,
|
|
):
|
|
# VAE encode one by one, as we assume cond images have different sizes
|
|
batch_cond_vae_images, batch_cond_t, batch_cond_vit_images = [], [], []
|
|
for cond_image_info_list in batch_cond_image_info_list:
|
|
cond_vae_image_list, cond_t_list, cond_vit_image_list = [], [], []
|
|
for image_info in cond_image_info_list:
|
|
cond_t_, cond_vae_image_ = self.vae_encode(
|
|
image_info.vae_image_info.image_tensor.to(self.device),
|
|
)
|
|
cond_vit_image_list.append(image_info.vision_image_info.image_tensor)
|
|
cond_vae_image_list.append(cond_vae_image_.squeeze(0))
|
|
cond_t_list.append(cond_t_)
|
|
batch_cond_vae_images.append(cond_vae_image_list)
|
|
batch_cond_t.append(cond_t_list)
|
|
batch_cond_vit_images.append(torch.cat(cond_vit_image_list, dim=0))
|
|
|
|
# If only one cond image for each sample and all have the same size, we can batch them together
|
|
# In this case, cond_vae_images is a 4-D tensor.
|
|
if all([len(items) == 1 for items in batch_cond_vae_images]) and all(
|
|
items[0].shape == batch_cond_vae_images[0][0].shape for items in batch_cond_vae_images
|
|
):
|
|
cond_vae_images = torch.stack([items[0] for items in batch_cond_vae_images], dim=0)
|
|
cond_t = torch.cat([items[0] for items in batch_cond_t], dim=0)
|
|
if cfg_factor > 1:
|
|
cond_t = cond_t.repeat(cfg_factor)
|
|
cond_vae_images = cond_vae_images.repeat(cfg_factor, 1, 1, 1)
|
|
else:
|
|
# In this case, cond_vae_images is a list of 4-D tensors or a list of lists of 3-D tensors.
|
|
cond_t = [torch.cat(item, dim=0) for item in batch_cond_t]
|
|
cond_vae_images = []
|
|
for items in batch_cond_vae_images:
|
|
if all(items[0].shape == item.shape for item in items):
|
|
cond_vae_images.append(torch.stack(items, dim=0))
|
|
else:
|
|
cond_vae_images.append(items)
|
|
if cfg_factor > 1:
|
|
cond_t = cond_t * cfg_factor
|
|
cond_vae_images = cond_vae_images * cfg_factor
|
|
|
|
if cfg_factor > 1:
|
|
batch_cond_vit_images = batch_cond_vit_images * cfg_factor
|
|
|
|
return cond_vae_images, cond_t, batch_cond_vit_images
|
|
|
|
@staticmethod
|
|
def check_inputs(prompt=None, message_list=None):
|
|
if prompt is None and message_list is None:
|
|
raise ValueError("Either `prompt` or `message_list` should be provided.")
|
|
if prompt is not None and message_list is not None:
|
|
raise ValueError("Only one of `prompt` or `message_list` should be provided.")
|
|
if prompt is not None:
|
|
assert isinstance(prompt, str) or isinstance(prompt, list), (
|
|
f"`prompt` should be a string or a list of strings, but got {type(prompt)}."
|
|
)
|
|
if isinstance(prompt, list):
|
|
assert len(prompt) > 0 and all(isinstance(p, str) for p in prompt), (
|
|
"`prompt` should be a non-empty list of strings."
|
|
)
|
|
if message_list is not None:
|
|
if not isinstance(message_list, list):
|
|
raise ValueError(f"`message_list` should be a list of messages, but got {type(message_list)}.")
|
|
assert len(message_list) > 0, "`message_list` should be a non-empty list."
|
|
for message in message_list:
|
|
assert isinstance(message, list) or isinstance(message, dict), (
|
|
f"Each message should be a list of dicts or a dict, but got {type(message)}."
|
|
)
|
|
|
|
@staticmethod
|
|
def _normalize_cot_text(cot: str | None) -> str | None:
|
|
"""
|
|
Ensure cot_text starts with the appropriate opening tag.
|
|
|
|
AR-generated text may omit the leading <think> or <recaption> tag
|
|
(since it was used as a generation trigger). This normalizes the text
|
|
so downstream parsing in get_cot_sections works correctly.
|
|
"""
|
|
if not cot:
|
|
return cot
|
|
|
|
if "</think>" in cot and not cot.startswith("<think>"):
|
|
cot = "<think>" + cot
|
|
return cot
|
|
if "</recaption>" in cot and not cot.startswith("<recaption>"):
|
|
cot = "<recaption>" + cot
|
|
return cot
|
|
return cot
|
|
|
|
def prepare_model_inputs(
|
|
self,
|
|
prompt=None,
|
|
mode="gen_image",
|
|
system_prompt=None,
|
|
cot_text=None,
|
|
num_inference_steps=50,
|
|
guidance_scale=5.0,
|
|
image_size="auto",
|
|
message_list=None,
|
|
device=None,
|
|
max_new_tokens=None,
|
|
**kwargs,
|
|
):
|
|
# 1. Sanity check
|
|
self.check_inputs(prompt, message_list)
|
|
device = default(device, self.device)
|
|
|
|
# 2. Format inputs
|
|
batch_message_list = message_list
|
|
batch_prompt = prompt
|
|
batch_cot_text = cot_text
|
|
batch_system_prompt = system_prompt
|
|
batch_gen_image_info = None
|
|
batch_cond_image_info = kwargs.pop("batch_cond_image_info", None)
|
|
|
|
# -- 2.1 message_list
|
|
if batch_message_list is not None:
|
|
if isinstance(batch_message_list[0], dict):
|
|
batch_message_list = [batch_message_list]
|
|
batch_size = len(batch_message_list)
|
|
|
|
batch_gen_image_info = [
|
|
[message["content"] for message in message_list_ if message["type"] == "gen_image"]
|
|
for message_list_ in batch_message_list
|
|
]
|
|
# At most one gen_image is allowed for each message_list
|
|
batch_gen_image_info = [info[-1] if len(info) > 0 else None for info in batch_gen_image_info]
|
|
# Multiple cond images are allowed.
|
|
batch_cond_image_info = [
|
|
[message["content"] for message in message_list_ if message["type"] == "joint_image"]
|
|
for message_list_ in batch_message_list
|
|
]
|
|
|
|
# -- 2.2 Prompt, cot text, system prompt
|
|
else:
|
|
if isinstance(batch_prompt, str):
|
|
batch_prompt = [batch_prompt]
|
|
batch_size = len(batch_prompt)
|
|
|
|
if batch_cot_text is not None:
|
|
if isinstance(batch_cot_text, str):
|
|
batch_cot_text = [batch_cot_text]
|
|
else:
|
|
assert isinstance(batch_cot_text, list) and len(batch_cot_text) == batch_size, (
|
|
"`cot_text` should be a string or a list of strings with the same length as `prompt`."
|
|
)
|
|
|
|
if batch_system_prompt is not None:
|
|
if isinstance(batch_system_prompt, str):
|
|
batch_system_prompt = [batch_system_prompt]
|
|
else:
|
|
assert isinstance(batch_system_prompt, list) and len(batch_system_prompt) == batch_size, (
|
|
"`system_prompts` should be a string or a list of strings with the same length as `prompt`."
|
|
)
|
|
|
|
if mode == "gen_image":
|
|
batch_gen_image_info = [self.image_processor.build_image_info(image_size) for _ in range(batch_size)]
|
|
|
|
if batch_cond_image_info is not None:
|
|
assert isinstance(batch_cond_image_info, list) and len(batch_cond_image_info) == batch_size, (
|
|
"`batch_cond_image_info` should be a list with the same batch size as `prompt`."
|
|
)
|
|
batch_cond_image_info = [cond if isinstance(cond, list) else [cond] for cond in batch_cond_image_info]
|
|
|
|
# -- 2.3 seed
|
|
generator = kwargs.get("generator", None)
|
|
if generator is None:
|
|
seeds = self.prepare_seed(seed=kwargs.get("seed"), batch_size=batch_size)
|
|
generator = [torch.Generator(self.device).manual_seed(seed) for seed in seeds]
|
|
|
|
# 3. apply chat template
|
|
cfg_factor = {"gen_text": 1, "gen_image": 1 + int(guidance_scale > 1.0)}
|
|
bot_task = kwargs.pop("bot_task", "auto")
|
|
# If `drop_think` enabled, always drop <think> parts in the context.
|
|
drop_think = kwargs.get("drop_think", self.generation_config.drop_think)
|
|
# Pull sequence_template from the model's generation_config so the DiT
|
|
# text prefix matches how the model was trained (Instruct for the
|
|
# HunyuanImage-3.0-Instruct checkpoint). Falling back to "pretrain"
|
|
# only if the config does not specify it.
|
|
sequence_template = getattr(self.generation_config, "sequence_template", "pretrain")
|
|
# Apply batched prompt or batched message_list to build input sequence with associated info.
|
|
out = self._tkwrapper.apply_chat_template(
|
|
batch_prompt=batch_prompt,
|
|
batch_message_list=batch_message_list,
|
|
mode=mode,
|
|
batch_gen_image_info=batch_gen_image_info,
|
|
batch_cond_image_info=batch_cond_image_info,
|
|
batch_system_prompt=batch_system_prompt,
|
|
batch_cot_text=batch_cot_text,
|
|
max_length=kwargs.get("max_length"),
|
|
bot_task=bot_task,
|
|
image_base_size=self.config.image_base_size,
|
|
sequence_template=sequence_template,
|
|
cfg_factor=cfg_factor[mode],
|
|
drop_think=drop_think,
|
|
)
|
|
output, sections = out["output"], out["sections"]
|
|
|
|
# 4. Encode conditional images
|
|
if batch_cond_image_info is not None and len(batch_cond_image_info[0]) > 0:
|
|
cond_vae_images, cond_timestep, cond_vit_images = self._encode_cond_image(
|
|
batch_cond_image_info, cfg_factor[mode]
|
|
)
|
|
vit_kwargs = {"spatial_shapes": [], "attention_mask": []}
|
|
for cond_image_info in batch_cond_image_info:
|
|
vit_kwargs["spatial_shapes"].append(
|
|
torch.stack([item.vision_encoder_kwargs["spatial_shapes"] for item in cond_image_info])
|
|
)
|
|
vit_kwargs["attention_mask"].append(
|
|
torch.stack([item.vision_encoder_kwargs["pixel_attention_mask"] for item in cond_image_info])
|
|
)
|
|
if cfg_factor[mode] > 1:
|
|
vit_kwargs["spatial_shapes"] = vit_kwargs["spatial_shapes"] * cfg_factor[mode]
|
|
vit_kwargs["attention_mask"] = vit_kwargs["attention_mask"] * cfg_factor[mode]
|
|
else:
|
|
cond_vae_images, cond_timestep, cond_vit_images = None, None, None
|
|
vit_kwargs = None
|
|
|
|
# 5. Build position embeddings
|
|
rope_image_info = self.build_batch_rope_image_info(output, sections)
|
|
if mode == "gen_text":
|
|
seq_len = self.generation_config.max_length
|
|
else:
|
|
seq_len = output.tokens.shape[1]
|
|
cos, sin = build_batch_2d_rope(
|
|
image_infos=rope_image_info,
|
|
seq_len=seq_len,
|
|
n_elem=self.config.attention_head_dim,
|
|
device=device,
|
|
base=self.config.rope_theta,
|
|
)
|
|
|
|
# 6. Build kv cache
|
|
if bot_task == "img_ratio":
|
|
max_new_tokens = 1
|
|
|
|
# 7. Build position ids
|
|
batch_input_pos = torch.arange(0, output.tokens.shape[1], dtype=torch.long, device=device)[None].expand(
|
|
batch_size * cfg_factor[mode], -1
|
|
) # use expand to share indices to save memory
|
|
|
|
# 8. Build model input kwargs
|
|
tkw = self._tkwrapper
|
|
if image_size == "auto":
|
|
extra_auto_stops = [tkw.special_token_map[f"<img_ratio_{i}>"] for i in range(33)]
|
|
else:
|
|
extra_auto_stops = [tkw.boi_token_id]
|
|
stop_token_id = dict(
|
|
auto=[tkw.eos_token_id] + extra_auto_stops,
|
|
image=[tkw.eos_token_id],
|
|
recaption=[
|
|
tkw.end_recaption_token_id,
|
|
tkw.end_answer_token_id,
|
|
tkw.eos_token_id,
|
|
],
|
|
think=[
|
|
tkw.end_recaption_token_id,
|
|
tkw.end_answer_token_id,
|
|
tkw.eos_token_id,
|
|
],
|
|
img_ratio=extra_auto_stops,
|
|
)
|
|
model_input_kwargs = dict(
|
|
input_ids=output.tokens.to(device),
|
|
position_ids=batch_input_pos,
|
|
past_key_values=None,
|
|
custom_pos_emb=(cos, sin),
|
|
mode=mode,
|
|
num_inference_steps=num_inference_steps,
|
|
guidance_scale=guidance_scale,
|
|
image_mask=to_device(output.gen_image_mask, device),
|
|
gen_timestep_scatter_index=to_device(output.gen_timestep_scatter_index, device),
|
|
cond_vae_images=to_device(cond_vae_images, device),
|
|
cond_timestep=to_device(cond_timestep, device),
|
|
cond_vae_image_mask=to_device(output.cond_vae_image_mask, device),
|
|
cond_vit_images=to_device(cond_vit_images, device),
|
|
cond_vit_image_mask=to_device(output.cond_vit_image_mask, device),
|
|
vit_kwargs={k: to_device(v, self.device) for k, v in vit_kwargs.items()}
|
|
if vit_kwargs is not None
|
|
else None,
|
|
cond_timestep_scatter_index=to_device(output.cond_timestep_scatter_index, device),
|
|
# for inner usage
|
|
tokenizer_output=output,
|
|
batch_gen_image_info=batch_gen_image_info,
|
|
generator=generator,
|
|
# generation config
|
|
eos_token_id=stop_token_id[bot_task],
|
|
max_new_tokens=max_new_tokens,
|
|
)
|
|
|
|
return model_input_kwargs
|
|
|
|
def _prepare_attention_mask_for_generation(
|
|
self,
|
|
inputs_tensor: torch.Tensor,
|
|
generation_config: GenerationConfig,
|
|
model_kwargs: dict[str, Any],
|
|
) -> torch.Tensor:
|
|
# create `4d` bool attention mask (b, 1, seqlen, seqlen) using this implementation to bypass the 2d requirement
|
|
# in the `transformers.generation_utils.GenerationMixin.generate`.
|
|
# This implementation can handle sequences with text and image modalities, where text tokens use causal
|
|
# attention and image tokens use full attention.
|
|
bsz, seq_len = inputs_tensor.shape
|
|
tokenizer_output = model_kwargs["tokenizer_output"]
|
|
batch_image_slices = [
|
|
tokenizer_output.joint_image_slices[i] + tokenizer_output.gen_image_slices[i] for i in range(bsz)
|
|
]
|
|
attention_mask = torch.ones(seq_len, seq_len, dtype=torch.bool).tril(diagonal=0).repeat(bsz, 1, 1)
|
|
full_attn_spans: list[list[tuple[int, int]]] = [[] for _ in range(bsz)]
|
|
for i in range(bsz):
|
|
for j, image_slice in enumerate(batch_image_slices[i]):
|
|
attention_mask[i, image_slice, image_slice] = True
|
|
start = image_slice.start if image_slice.start is not None else 0
|
|
stop = image_slice.stop if image_slice.stop is not None else seq_len
|
|
assert start < stop, f"Invalid image slice: {image_slice}"
|
|
full_attn_spans[i].append((int(start), int(stop)))
|
|
if full_attn_spans[i]:
|
|
full_attn_spans[i].sort(key=lambda x: x[0])
|
|
attention_mask = attention_mask.unsqueeze(1)
|
|
model_kwargs["full_attn_spans"] = full_attn_spans
|
|
return attention_mask
|
|
|
|
def prepare_inputs_for_generation(
|
|
self,
|
|
input_ids,
|
|
past_key_values=None,
|
|
attention_mask=None,
|
|
inputs_embeds=None,
|
|
tokenizer_output=None,
|
|
batch_gen_image_info=None,
|
|
generator=None,
|
|
**kwargs,
|
|
):
|
|
# if `inputs_embeds` are passed, we only want to use them in the 1st generation step
|
|
if inputs_embeds is not None and past_key_values is None:
|
|
model_inputs = {"inputs_embeds": inputs_embeds}
|
|
else:
|
|
if input_ids is not None and input_ids.shape[1] != kwargs["position_ids"].shape[1]: # in decode steps
|
|
input_ids = torch.gather(input_ids, dim=1, index=kwargs["position_ids"])
|
|
model_inputs = {"input_ids": input_ids}
|
|
|
|
model_inputs.update(
|
|
{
|
|
"attention_mask": attention_mask,
|
|
"position_ids": kwargs["position_ids"],
|
|
"past_key_values": past_key_values,
|
|
"use_cache": kwargs.get("use_cache"),
|
|
"custom_pos_emb": kwargs["custom_pos_emb"],
|
|
"mode": kwargs["mode"],
|
|
"images": kwargs.get("images"),
|
|
"image_mask": kwargs.get("image_mask"),
|
|
"timestep": kwargs.get("timestep"),
|
|
"gen_timestep_scatter_index": kwargs.get("gen_timestep_scatter_index"),
|
|
"cond_vae_images": kwargs.get("cond_vae_images"),
|
|
"cond_timestep": kwargs.get("cond_timestep"),
|
|
"cond_vae_image_mask": kwargs.get("cond_vae_image_mask"),
|
|
"cond_vit_images": kwargs.get("cond_vit_images"),
|
|
"cond_vit_image_mask": kwargs.get("cond_vit_image_mask"),
|
|
"vit_kwargs": kwargs.get("vit_kwargs"),
|
|
"cond_timestep_scatter_index": kwargs.get("cond_timestep_scatter_index"),
|
|
"query_lens": kwargs.get("query_lens"),
|
|
"seq_lens": kwargs.get("seq_lens"),
|
|
"num_image_tokens": kwargs.get("num_image_tokens"),
|
|
"ar_kv_reuse_len": kwargs.get("ar_kv_reuse_len", 0),
|
|
"full_attn_spans": kwargs.get("full_attn_spans"),
|
|
}
|
|
)
|
|
return model_inputs
|
|
|
|
def _update_model_kwargs_for_generation(
|
|
self,
|
|
outputs: ModelOutput,
|
|
model_kwargs: dict[str, Any],
|
|
is_encoder_decoder: bool = False,
|
|
num_new_tokens: int = 1,
|
|
) -> dict[str, Any]:
|
|
mode = model_kwargs["mode"]
|
|
|
|
updated_model_kwargs = {
|
|
"mode": mode,
|
|
"custom_pos_emb": model_kwargs["custom_pos_emb"],
|
|
"num_image_tokens": model_kwargs["num_image_tokens"],
|
|
}
|
|
if "full_attn_spans" in model_kwargs:
|
|
updated_model_kwargs["full_attn_spans"] = model_kwargs["full_attn_spans"]
|
|
|
|
# update past_key_values keeping its naming used in model code
|
|
for possible_cache_name in ALL_CACHE_NAMES:
|
|
if possible_cache_name in outputs:
|
|
# TODO (joao): remove output/input mismatch when these old models (xlnet, reformer) are deprecated
|
|
if possible_cache_name in ("past_buckets_states", "mems"):
|
|
cache_name = "past_key_values"
|
|
else:
|
|
cache_name = possible_cache_name
|
|
updated_model_kwargs[cache_name] = getattr(outputs, possible_cache_name)
|
|
break
|
|
|
|
if "tokenizer_output" in model_kwargs:
|
|
if mode == "gen_text":
|
|
# When enable batching, we use right padding, which requires a real_pos to index the valid
|
|
# end position of the sequence. If tokenizer_output in model_kwargs, it means we are in the
|
|
# prefill step of generation.
|
|
real_pos = to_device(model_kwargs["tokenizer_output"].real_pos, self.device)
|
|
updated_model_kwargs["position_ids"] = real_pos
|
|
else:
|
|
# position ids
|
|
image_mask = model_kwargs["image_mask"]
|
|
bsz, seq_len = image_mask.shape
|
|
offset = model_kwargs.get("ar_kv_reuse_offset", 0) # should be an absolute position.
|
|
index = torch.arange(offset, offset + seq_len, device=image_mask.device).unsqueeze(0).repeat(bsz, 1)
|
|
position_ids = index.masked_select(image_mask.bool()).reshape(bsz, -1)
|
|
timestep_position_ids = index[
|
|
torch.arange(bsz), model_kwargs["gen_timestep_scatter_index"][:, -1]
|
|
].unsqueeze(-1)
|
|
updated_model_kwargs["position_ids"] = torch.cat([timestep_position_ids, position_ids], dim=1)
|
|
|
|
# attention mask
|
|
mask_list = []
|
|
current_starts = timestep_position_ids.reshape(-1)
|
|
max_current_start = int(current_starts.max().item())
|
|
for attention_mask_i, position_ids_i, current_start_i in zip(
|
|
model_kwargs["attention_mask"], updated_model_kwargs["position_ids"], current_starts
|
|
):
|
|
query_mask = torch.index_select(attention_mask_i, dim=1, index=position_ids_i.reshape(-1) - offset)
|
|
current_start = int(current_start_i.item())
|
|
prefix_mask = torch.index_select(
|
|
query_mask,
|
|
dim=2,
|
|
index=torch.arange(current_start, device=image_mask.device),
|
|
)
|
|
if current_start < max_current_start:
|
|
prefix_pad = query_mask.new_zeros(
|
|
query_mask.shape[0],
|
|
query_mask.shape[1],
|
|
max_current_start - current_start,
|
|
)
|
|
prefix_mask = torch.cat([prefix_mask, prefix_pad], dim=2)
|
|
current_mask = torch.index_select(query_mask, dim=2, index=position_ids_i.reshape(-1))
|
|
mask_list.append(torch.cat([prefix_mask, current_mask], dim=2))
|
|
attention_mask = torch.stack(mask_list, dim=0)
|
|
updated_model_kwargs["attention_mask"] = attention_mask
|
|
updated_model_kwargs["gen_timestep_scatter_index"] = model_kwargs["gen_timestep_scatter_index"]
|
|
|
|
else:
|
|
if mode == "gen_text":
|
|
# Now we are in the decode steps.
|
|
updated_model_kwargs["position_ids"] = model_kwargs["position_ids"] + 1
|
|
else:
|
|
updated_model_kwargs["position_ids"] = model_kwargs["position_ids"]
|
|
updated_model_kwargs["attention_mask"] = model_kwargs["attention_mask"]
|
|
updated_model_kwargs["gen_timestep_scatter_index"] = model_kwargs["gen_timestep_scatter_index"]
|
|
|
|
return updated_model_kwargs
|
|
|
|
def _generate(
|
|
self,
|
|
generator: list[torch.Generator] | None = None,
|
|
**kwargs,
|
|
):
|
|
mode = kwargs.get("mode", "gen_text")
|
|
# verbose > 1 not support
|
|
if mode == "gen_text":
|
|
raise NotImplementedError("Not support gen text for hunyuan image")
|
|
|
|
elif mode == "gen_image":
|
|
batch_gen_image_info: list[ImageInfo] = kwargs.get("batch_gen_image_info")
|
|
if batch_gen_image_info is None:
|
|
raise ValueError("`batch_gen_image_info` should be provided when `mode` is `gen_image`.")
|
|
|
|
image_info: ImageInfo = batch_gen_image_info[0]
|
|
num_image_tokens = (
|
|
image_info.image_token_length
|
|
+ (1 if image_info.add_timestep_token else 0)
|
|
+ (1 if image_info.add_guidance_token else 0)
|
|
)
|
|
kwargs["num_image_tokens"] = num_image_tokens
|
|
# 50 and 5.0 hard code
|
|
results = self.pipeline(
|
|
batch_size=len(batch_gen_image_info),
|
|
image_size=[
|
|
batch_gen_image_info[0].image_height,
|
|
batch_gen_image_info[0].image_width,
|
|
],
|
|
num_inference_steps=kwargs.get("num_inference_steps", 50),
|
|
guidance_scale=kwargs.get("guidance_scale", 5.0),
|
|
generator=generator,
|
|
model_kwargs=kwargs,
|
|
)
|
|
samples = results[0]
|
|
return samples
|
|
|
|
else:
|
|
raise ValueError(f"Unknown mode {mode}, only `gen_text` and `gen_image` are supported.")
|
|
|
|
@staticmethod
|
|
def _check_inputs(cond, target, check_list):
|
|
if cond:
|
|
for name, item in check_list:
|
|
assert item is not None, f"`{name}` should be provided when `{target}`."
|
|
|
|
def forward_call(
|
|
self,
|
|
input_ids: torch.LongTensor | None = None,
|
|
attention_mask: torch.Tensor | None = None,
|
|
position_ids: torch.LongTensor | None = None,
|
|
past_key_values: list[torch.FloatTensor] | None = None,
|
|
use_cache: bool | None = None,
|
|
output_attentions: bool | None = None,
|
|
output_hidden_states: bool | None = None,
|
|
return_dict: bool | None = None,
|
|
custom_pos_emb: tuple[torch.FloatTensor] | None = None,
|
|
mode: str = "gen_text",
|
|
first_step: bool | None = None,
|
|
# for gen image
|
|
images: BatchRaggedImages | None = None,
|
|
image_mask: torch.Tensor | None = None,
|
|
timestep: BatchRaggedTensor | None = None,
|
|
gen_timestep_scatter_index: torch.Tensor | None = None,
|
|
# for cond image
|
|
cond_vae_images: BatchRaggedImages | None = None,
|
|
cond_timestep: BatchRaggedTensor | None = None,
|
|
cond_vae_image_mask: torch.Tensor | None = None,
|
|
cond_vit_images: BatchRaggedImages | None = None,
|
|
cond_vit_image_mask: torch.Tensor | None = None,
|
|
vit_kwargs: dict[str, Any] | None = None,
|
|
cond_timestep_scatter_index: torch.Tensor | None = None,
|
|
query_lens: list[int] | None = None,
|
|
seq_lens: list[int] | None = None,
|
|
num_image_tokens: int | None = None,
|
|
uncond_cfg_prefill: bool = False,
|
|
ar_kv_reuse_len: int = 0,
|
|
full_attn_spans: list[list[tuple[int, int]]] | None = None,
|
|
) -> tuple | CausalMMOutputWithPast:
|
|
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
|
# Sanity Check of Inputs
|
|
self._check_inputs(
|
|
mode == "gen_image" and not uncond_cfg_prefill,
|
|
"in `gen_image` mode",
|
|
[
|
|
("images", images),
|
|
("timestep", timestep),
|
|
("gen_timestep_scatter_index", gen_timestep_scatter_index),
|
|
],
|
|
)
|
|
self._check_inputs(
|
|
mode == "gen_image" and first_step and not uncond_cfg_prefill,
|
|
"in `gen_image` mode at the first step",
|
|
[
|
|
("image_mask", image_mask),
|
|
],
|
|
)
|
|
self._check_inputs(
|
|
cond_vae_images is not None,
|
|
"`cond_vae_images` is provided",
|
|
[
|
|
("cond_timestep", cond_timestep),
|
|
("cond_vae_image_mask", cond_vae_image_mask),
|
|
("cond_timestep_scatter_index", cond_timestep_scatter_index),
|
|
],
|
|
)
|
|
self._check_inputs(
|
|
cond_vit_images is not None,
|
|
"`cond_vit_images` is provided",
|
|
[
|
|
("cond_vit_image_mask", cond_vit_image_mask),
|
|
("vit_kwargs", vit_kwargs),
|
|
],
|
|
)
|
|
custom_pos_emb = self.get_pos_emb(custom_pos_emb, position_ids)
|
|
|
|
if input_ids is not None:
|
|
inputs_embeds = self.model.embed_tokens(input_ids)
|
|
bsz, seq_len, n_embd = inputs_embeds.shape
|
|
else:
|
|
inputs_embeds = None
|
|
bsz = images.shape[0] if isinstance(images, torch.Tensor) else len(images)
|
|
seq_len = 0
|
|
n_embd = self.config.hidden_size
|
|
|
|
# Instantiate placeholder tokens: <timestep>, <img> for the gen image
|
|
if mode == "gen_text":
|
|
# For gen_text, make sure gen_timestep_scatter_index is None
|
|
gen_timestep_scatter_index = None
|
|
token_h, token_w = None, None
|
|
elif uncond_cfg_prefill:
|
|
token_h, token_w = None, None
|
|
else:
|
|
if first_step:
|
|
assert inputs_embeds is not None
|
|
inputs_embeds, token_h, token_w = self.instantiate_vae_image_tokens(
|
|
inputs_embeds, images, timestep, image_mask
|
|
)
|
|
inputs_embeds = self.instantiate_timestep_tokens(inputs_embeds, timestep, gen_timestep_scatter_index)
|
|
else:
|
|
t_emb = self.time_embed(timestep)
|
|
image_emb, token_h, token_w = self.patch_embed(images, t_emb)
|
|
timestep_emb = self.timestep_emb(timestep).reshape(bsz, -1, n_embd)
|
|
inputs_embeds = torch.cat([timestep_emb, image_emb], dim=1)
|
|
|
|
# Instantiate placeholder tokens: <timestep>, <img> for cond images
|
|
# Should only run once with kv-cache enabled.
|
|
if cond_vae_images is not None:
|
|
inputs_embeds, _, _ = self.instantiate_vae_image_tokens(
|
|
inputs_embeds, cond_vae_images, cond_timestep, cond_vae_image_mask
|
|
)
|
|
inputs_embeds = self.instantiate_timestep_tokens(inputs_embeds, cond_timestep, cond_timestep_scatter_index)
|
|
if cond_vit_images is not None:
|
|
inputs_embeds = self.instantiate_vit_image_tokens(
|
|
inputs_embeds, cond_vit_images, cond_vit_image_mask, vit_kwargs
|
|
)
|
|
assert inputs_embeds is not None
|
|
bsz, seq_len, n_embd = inputs_embeds.shape
|
|
|
|
# decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
|
|
from vllm.forward_context import set_forward_context
|
|
|
|
with set_forward_context(None, self.vllm_config):
|
|
outputs = self.model(
|
|
input_ids=input_ids,
|
|
attention_mask=attention_mask,
|
|
position_ids=position_ids,
|
|
past_key_values=past_key_values,
|
|
inputs_embeds=inputs_embeds,
|
|
use_cache=use_cache,
|
|
output_attentions=output_attentions,
|
|
output_hidden_states=output_hidden_states,
|
|
return_dict=return_dict,
|
|
custom_pos_emb=custom_pos_emb,
|
|
mode=mode,
|
|
first_step=first_step,
|
|
query_lens=query_lens,
|
|
seq_lens=seq_lens,
|
|
num_image_tokens=num_image_tokens,
|
|
gen_timestep_scatter_index=gen_timestep_scatter_index,
|
|
uncond_cfg_prefill=uncond_cfg_prefill,
|
|
ar_kv_reuse_len=ar_kv_reuse_len,
|
|
full_attn_spans=full_attn_spans,
|
|
)
|
|
hidden_states = outputs[0]
|
|
|
|
if mode == "gen_text":
|
|
hidden_states = self.model.ln_f(hidden_states)
|
|
logits = self.lm_head(hidden_states)
|
|
logits = logits.float()
|
|
diffusion_prediction = None
|
|
elif uncond_cfg_prefill:
|
|
logits = None
|
|
diffusion_prediction = None
|
|
else:
|
|
logits = None
|
|
hidden_states = hidden_states.to(inputs_embeds.device)
|
|
assert hidden_states.numel() == bsz * seq_len * n_embd, (
|
|
f"Shape mismatch: {hidden_states.shape} cannot reshape to ({bsz}, {seq_len}, {n_embd})"
|
|
)
|
|
hidden_states = hidden_states.reshape(bsz, seq_len, n_embd)
|
|
diffusion_prediction = self.ragged_final_layer(
|
|
hidden_states, image_mask, timestep, token_h, token_w, first_step
|
|
)
|
|
|
|
if not return_dict:
|
|
output = (logits,) + outputs[1:] + (diffusion_prediction,)
|
|
return output
|
|
|
|
output = CausalMMOutputWithPast(
|
|
logits=logits,
|
|
past_key_values=outputs.past_key_values,
|
|
hidden_states=outputs.hidden_states,
|
|
attentions=outputs.attentions,
|
|
diffusion_prediction=diffusion_prediction,
|
|
)
|
|
|
|
return output
|
|
|
|
def inject_ar_kv_into_layers(
|
|
self,
|
|
ar_kv_data: dict[int, dict[str, torch.Tensor]],
|
|
positive_reuse_len: int,
|
|
) -> None:
|
|
"""Inject AR-stage KV cache into each layer's ImageKVCacheManager.
|
|
|
|
Truncates to positive_reuse_len and sets image_kv_cache_map directly.
|
|
"""
|
|
for layer in self.model.layers:
|
|
layer_idx = layer.layer_idx
|
|
if layer_idx not in ar_kv_data:
|
|
continue
|
|
kv = ar_kv_data[layer_idx]
|
|
cache_mgr = layer.self_attn.image_attn
|
|
k, v = kv["key"], kv["value"]
|
|
cache_mgr._injected_ar_kv = [(k[:positive_reuse_len], v[:positive_reuse_len])]
|
|
|
|
def _extract_ar_kv_from_request(self, req) -> dict[str, Any]:
|
|
kv = getattr(getattr(req, "sampling_params", None), "past_key_values", None)
|
|
if kv is None:
|
|
return {}
|
|
key_cache = getattr(kv, "key_cache", None)
|
|
value_cache = getattr(kv, "value_cache", None)
|
|
if not key_cache or not value_cache:
|
|
return {}
|
|
ar_kv_data = {
|
|
i: {"key": k, "value": v}
|
|
for i, (k, v) in enumerate(zip(key_cache, value_cache))
|
|
if k is not None and v is not None
|
|
}
|
|
if not ar_kv_data:
|
|
return {}
|
|
logger.info(
|
|
f"[AR KV Reuse] Extracted {len(ar_kv_data)} layers of AR KV, "
|
|
f"each with length: {next(iter(ar_kv_data.values()))['key'].shape}"
|
|
)
|
|
return {"ar_kv_data": ar_kv_data}
|
|
|
|
def forward(
|
|
self,
|
|
req: OmniDiffusionRequest,
|
|
prompt: str | list[str] = "",
|
|
image_size="auto",
|
|
height: int = 1024,
|
|
width: int = 1024,
|
|
num_inference_steps: int = 50,
|
|
guidance_scale: float = 5.0,
|
|
generator: torch.Generator | list[torch.Generator] | None = None,
|
|
**kwargs,
|
|
) -> DiffusionOutput:
|
|
extra_args = getattr(getattr(req, "sampling_params", None), "extra_args", {}) or {}
|
|
use_system_prompt = extra_args.get("use_system_prompt")
|
|
system_prompt = extra_args.get("system_prompt")
|
|
# Fall back to per-prompt use_system_prompt forwarded by ar2diffusion
|
|
if req.prompts:
|
|
first_prompt = req.prompts[0]
|
|
if isinstance(first_prompt, dict):
|
|
if use_system_prompt is None:
|
|
use_system_prompt = first_prompt.get("use_system_prompt")
|
|
if system_prompt is None:
|
|
system_prompt = first_prompt.get("system_prompt")
|
|
if use_system_prompt is not None:
|
|
system_prompt = get_system_prompt(use_system_prompt, "image", system_prompt)
|
|
system_prompt = system_prompt.strip() if system_prompt is not None else ""
|
|
prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
|
|
|
|
cot_text_list = [
|
|
(p.get("extra", {}).get("ar_generated_text") if isinstance(p, dict) else None) or None for p in req.prompts
|
|
]
|
|
cot_text = (
|
|
[self._normalize_cot_text(t) for t in cot_text_list] if any(t is not None for t in cot_text_list) else None
|
|
)
|
|
|
|
batch_cond_image_info: list[list[JointImageInfo]] | None = None
|
|
if any(not isinstance(p, str) for p in req.prompts):
|
|
batch_cond_image_info = []
|
|
for prompt_item in req.prompts:
|
|
if isinstance(prompt_item, str):
|
|
batch_cond_image_info.append([])
|
|
continue
|
|
prompt_additional_information = prompt_item.get("additional_information") or {}
|
|
prompt_cond_infos = prompt_additional_information.get("batch_cond_image_info", [])
|
|
if isinstance(prompt_cond_infos, JointImageInfo | dict):
|
|
prompt_cond_infos = [prompt_cond_infos]
|
|
if prompt_cond_infos is None:
|
|
prompt_cond_infos = []
|
|
batch_cond_image_info.append([_joint_image_info_from_payload(item) for item in prompt_cond_infos])
|
|
|
|
has_cond_image = [len(cond_infos) > 0 for cond_infos in batch_cond_image_info]
|
|
if any(has_cond_image) and not all(has_cond_image):
|
|
raise ValueError(
|
|
"When batching Hunyuan image editing requests, every prompt must include input image(s)."
|
|
)
|
|
if not any(has_cond_image):
|
|
batch_cond_image_info = None
|
|
|
|
generator = req.sampling_params.generator or generator
|
|
height = req.sampling_params.height or height
|
|
width = req.sampling_params.width or width
|
|
num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
|
|
if req.sampling_params.guidance_scale_provided:
|
|
guidance_scale = req.sampling_params.guidance_scale
|
|
if guidance_scale <= 1.0:
|
|
logger.info("HunyuanImage3.0 runs without classifier-free guidance when guidance_scale <= 1.0.")
|
|
image_size = (height, width)
|
|
|
|
# ---- AR KV Reuse: extract injected KV from request ----
|
|
ar_kv_kwargs = self._extract_ar_kv_from_request(req)
|
|
|
|
model_inputs = self.prepare_model_inputs(
|
|
prompt=prompt,
|
|
cot_text=cot_text,
|
|
system_prompt=system_prompt,
|
|
mode="gen_image",
|
|
generator=generator,
|
|
image_size=image_size,
|
|
num_inference_steps=num_inference_steps,
|
|
guidance_scale=guidance_scale,
|
|
batch_cond_image_info=batch_cond_image_info,
|
|
)
|
|
|
|
model_inputs.update(ar_kv_kwargs)
|
|
|
|
outputs = self._generate(**model_inputs, **kwargs)
|
|
custom_output = {}
|
|
if any(t is not None for t in cot_text_list):
|
|
custom_output["ar_generated_text"] = cot_text_list[0] if len(cot_text_list) == 1 else cot_text_list
|
|
return DiffusionOutput(
|
|
output=outputs[0],
|
|
custom_output=custom_output,
|
|
stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
|
|
)
|