# 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()}
# Siglip2VisionTransformer now returns a plain tensor (B, max_patches,
# hidden_size) after the vLLM-layers refactor, not an HF-style output.
image_embed = self.vision_model(image, **cur_kwargs)
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, generator=None):
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:
latents = vae_encode_result.latent_dist.sample(generator)
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,
generator=None,
):
# 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 batch_idx, cond_image_info_list in enumerate(batch_cond_image_info_list):
cond_vae_image_list, cond_t_list, cond_vit_image_list = [], [], []
cond_generator = generator[batch_idx] if isinstance(generator, list) else generator
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),
generator=cond_generator,
)
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 or 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 "" in cot and not cot.startswith(""):
cot = "" + cot
return cot
if "" in cot and not cot.startswith(""):
cot = "" + 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 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
# Skip encoding if AR KV reuse is enabled
has_ar_kv = kwargs.get("ar_kv_data")
if batch_cond_image_info is not None and len(batch_cond_image_info[0]) > 0 and not has_ar_kv:
cond_vae_images, cond_timestep, cond_vit_images = self._encode_cond_image(
batch_cond_image_info, cfg_factor[mode], generator=generator
)
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""] 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: , ![]()
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: , ![]()
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 {}
bot_task = extra_args.get("bot_task")
use_system_prompt = extra_args.get("use_system_prompt")
system_prompt = extra_args.get("system_prompt")
# Fall back to per-prompt prompt kwargs forwarded by ar2diffusion
if req.prompts:
first_prompt = req.prompts[0]
if isinstance(first_prompt, dict):
if bot_task is None:
bot_task = first_prompt.get("bot_task")
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 isinstance(bot_task, str) and bot_task.lower() == "none":
bot_task = None
if use_system_prompt is not None:
system_prompt = get_system_prompt(use_system_prompt, bot_task or "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)
tokenizer_bot_task = bot_task
if tokenizer_bot_task == "think_recaption":
tokenizer_bot_task = "think"
elif tokenizer_bot_task == "vanilla":
tokenizer_bot_task = "image"
supported_bot_tasks = {"auto", "image", "think", "recaption", "img_ratio"}
if tokenizer_bot_task is not None and tokenizer_bot_task not in supported_bot_tasks:
raise ValueError(
f"Unsupported HunyuanImage3 single-stage bot_task: {tokenizer_bot_task!r}. "
f"Supported values are: {sorted(supported_bot_tasks)}."
)
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,
bot_task=tokenizer_bot_task or "auto",
**ar_kv_kwargs,
)
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,
)