anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity
Files
094abd70c1d98edd965d5e2fd392993318092bd3
vllm-omni/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py
T

2181 lines
95 KiB
Python
Raw Blame History

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import copy
import logging
from collections.abc import Iterable
from typing import TYPE_CHECKING, Any, ClassVar
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.diffusion_kv.config import DiffusionKVCacheMode
from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery
from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import (
DiffusionPipelineProfilerMixin,
)
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
from vllm_omni.inputs.data import OmniTextPrompt
from vllm_omni.model_executor.models.hunyuan_image3.siglip2 import Siglip2VisionTransformer
from . import request_layout as request_layout_utils
from .hunyuan_image3_tokenizer import TokenizerEncodeOutput, 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,
retrieve_timesteps,
)
from .request_layout import HunyuanPreparedLayout
if TYPE_CHECKING:
from vllm_omni.diffusion.worker.input_batch import InputBatch
from vllm_omni.diffusion.worker.utils import StepRequestState
logger = logging.getLogger(__name__)
BatchRaggedImages = torch.Tensor | list[torch.Tensor | list[torch.Tensor]]
BatchRaggedTensor = torch.Tensor | list[torch.Tensor]
_STEP_MODEL_KWARGS = "hunyuan_model_kwargs"
_STEP_INPUT_IDS = "hunyuan_input_ids"
_STEP_GENERATOR = "hunyuan_generator"
_STEP_GUIDANCE_SCALE = "hunyuan_guidance_scale"
_STEP_CFG_FACTOR = "hunyuan_cfg_factor"
_STEP_OUTPUT_SIZE = "hunyuan_output_size"
_STEP_COT_TEXT_LIST = "hunyuan_cot_text_list"
_STEP_AR_KV = "hunyuan_ar_kv"
_STEP_PROMPT_KV = "hunyuan_prompt_kv"
_HUNYUAN_DEFAULT_OUTPUT_TYPE = "pil"
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 _shift_full_attn_spans(
spans: list[tuple[int, int]],
prefix_len: int | None,
max_prefix_len: int | None,
) -> list[tuple[int, int]]:
if prefix_len is None or max_prefix_len is None:
return list(spans)
shift = max_prefix_len - prefix_len
if shift == 0:
return list(spans)
return [(start + shift, end + shift) if start >= prefix_len else (start, end) for start, end in spans]
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)
prepare_diffusion_kv_layout = od_config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
tokenizer_wrapper = TokenizerWrapper(od_config.model) if prepare_diffusion_kv_layout else None
generation_config = GenerationConfig.from_pretrained(od_config.model) if prepare_diffusion_kv_layout else None
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
# Use index-based primary resolution lookup for conditioning images
# (matching official HunyuanImage-3.0 behavior). This avoids the
# area-based sub-resolution matching via get_target_size().match(),
# which would shrink small inputs (e.g. 512x512 → 512x512) and
# diverge from the output generation size determined by AR ratio_idx.
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)
reso = image_processor.reso_group[ratio_idx]
target_width = int(reso.width)
target_height = int(reso.height)
vae_input = _resize_and_crop_center(pil_image, target_width, target_height)
vae_tensor = image_processor.vae_processor(vae_input)
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 request_layout_utils.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):
prompt = request.prompt
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.prompt = prompt
if prepare_diffusion_kv_layout:
assert tokenizer_wrapper is not None and generation_config is not None
prepared_layout = request_layout_utils.prepare_hunyuan_layout(
request,
tokenizer_wrapper=tokenizer_wrapper,
image_processor=image_processor,
generation_config=generation_config,
image_base_size=hf_config.image_base_size,
)
request.prepared_layout = prepared_layout
request.diffusion_kv_requests = request_layout_utils.build_hunyuan_diffusion_kv_requests(
request,
prepared_layout,
)
return request
return pre_process_func
def get_hunyuan_image3_post_process_func(od_config: OmniDiffusionConfig):
"""GPU tensor → PIL, runs outside pipeline_forward to overlap with next request."""
from diffusers.image_processor import VaeImageProcessor
image_processor = VaeImageProcessor()
def post_process_func(images: torch.Tensor | dict):
# Handle dict envelope format (upstream): {"payload": {"image": ...}, "metadata": ...}
if isinstance(images, dict) and "payload" in images:
tensor = images["payload"].get("image")
if tensor is None:
return images
if tensor.dim() == 3:
tensor = tensor.unsqueeze(0)
do_denormalize = [True] * tensor.shape[0]
images["payload"]["image"] = image_processor.postprocess(
tensor,
output_type=_HUNYUAN_DEFAULT_OUTPUT_TYPE,
do_denormalize=do_denormalize,
)
return images
# Legacy raw tensor format
if images.dim() == 3:
images = images.unsqueeze(0)
do_denormalize = [True] * images.shape[0]
return image_processor.postprocess(
images,
output_type=_HUNYUAN_DEFAULT_OUTPUT_TYPE,
do_denormalize=do_denormalize,
)
return post_process_func
class HunyuanImage3Pipeline(
HunyuanImage3PreTrainedModel,
GenerationMixin,
SupportImageInput,
SupportsComponentDiscovery,
DiffusionPipelineProfilerMixin,
):
supports_step_execution: ClassVar[bool] = True
supports_request_batch = False
support_image_input = True
_dit_modules: ClassVar[list[str]] = ["model"]
_encoder_modules: ClassVar[list[str]] = ["vision_model"]
_vae_modules: ClassVar[list[str]] = ["vae"]
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
def _validate_step_request(self, state: "StepRequestState") -> None:
prompt = state.prompt
sampling = state.sampling
if prompt is None:
raise ValueError("HunyuanImage3 step execution currently requires exactly one prompt per request.")
if sampling.timesteps is not None or sampling.sigmas is not None:
raise ValueError("HunyuanImage3 step execution does not support custom timesteps or sigmas yet.")
if sampling.num_outputs_per_prompt not in (None, 0, 1):
raise ValueError("HunyuanImage3 step execution supports only one output per prompt.")
parallel_config = self.od_config.parallel_config
if parallel_config.sequence_parallel_size > 1:
raise ValueError("HunyuanImage3 step execution does not support sequence parallel yet.")
if parallel_config.cfg_parallel_size > 1:
raise ValueError("HunyuanImage3 step execution does not support CFG parallel yet.")
cache_backend = getattr(self.od_config, "cache_backend", "none")
if cache_backend not in (None, "", "none"):
raise ValueError("HunyuanImage3 step execution does not support diffusion cache backends yet.")
if getattr(self.od_config, "diffusion_kv_cache_skip_step_indices", None):
raise ValueError("HunyuanImage3 step execution does not support diffusion KV cache step skips yet.")
def _snapshot_injected_ar_kv(self) -> list[list[tuple[torch.Tensor, torch.Tensor]] | None] | None:
snapshot: list[list[tuple[torch.Tensor, torch.Tensor]] | None] = []
found = False
for layer in self.model.layers:
mgr = layer.self_attn.image_attn
injected = mgr._injected_ar_kv
if injected is None:
snapshot.append(None)
continue
found = True
snapshot.append([(k.detach().clone(), v.detach().clone()) for k, v in injected])
mgr._injected_ar_kv = None
return snapshot if found else None
def _restore_injected_ar_kv(
self,
states: list["StepRequestState"],
row_state_indexes: list[int],
row_branches: list[int],
) -> None:
snapshots = [state.extra.get(_STEP_AR_KV) for state in states]
if any(snapshot is None for snapshot in snapshots):
if any(snapshot is not None for snapshot in snapshots):
raise ValueError("Cannot mix Hunyuan AR-KV reuse and non-reuse requests in one step batch yet.")
for layer in self.model.layers:
layer.self_attn.image_attn._injected_ar_kv = None
return
for layer_idx, layer in enumerate(self.model.layers):
injected_rows = []
for state_idx, branch in zip(row_state_indexes, row_branches):
layer_snapshot = snapshots[state_idx][layer_idx]
if layer_snapshot is None:
raise ValueError("Hunyuan AR-KV snapshot is incomplete for grouped step execution.")
injected_rows.append(layer_snapshot[branch])
layer.self_attn.image_attn._injected_ar_kv = injected_rows
@staticmethod
def _pad_tensor_rows(rows: list[torch.Tensor], pad_value: bool | int | float = 0) -> torch.Tensor:
if not rows:
raise ValueError("Cannot merge an empty Hunyuan step batch.")
if all(tuple(row.shape[1:]) == tuple(rows[0].shape[1:]) for row in rows):
return torch.cat(rows, dim=0)
max_shape = [max(int(row.shape[dim]) for row in rows) for dim in range(1, rows[0].ndim)]
padded = []
for row in rows:
out = row.new_full((row.shape[0], *max_shape), pad_value)
slices = (slice(None),) + tuple(slice(0, int(size)) for size in row.shape[1:])
out[slices] = row
padded.append(out)
return torch.cat(padded, dim=0)
@staticmethod
def _merge_attention_masks(
rows: list[torch.Tensor],
prefix_lens: list[int] | None,
) -> torch.Tensor:
if prefix_lens is None:
return HunyuanImage3Pipeline._pad_tensor_rows(rows, False)
max_q = max(int(row.shape[-2]) for row in rows)
current_lens = [int(row.shape[-1]) - prefix_len for row, prefix_len in zip(rows, prefix_lens)]
max_prefix = max(prefix_lens)
max_current = max(current_lens)
padded = []
for row, prefix_len, current_len in zip(rows, prefix_lens, current_lens):
out = row.new_zeros((row.shape[0], row.shape[1], max_q, max_prefix + max_current))
q_len = int(row.shape[-2])
out[:, :, :q_len, :prefix_len] = row[:, :, :q_len, :prefix_len]
out[:, :, :q_len, max_prefix : max_prefix + current_len] = row[
:,
:,
:q_len,
prefix_len : prefix_len + current_len,
]
padded.append(out)
return torch.cat(padded, dim=0)
@staticmethod
def _shift_full_attn_spans(
spans: list[tuple[int, int]],
prefix_len: int | None,
max_prefix_len: int | None,
) -> list[tuple[int, int]]:
return _shift_full_attn_spans(spans, prefix_len, max_prefix_len)
def _row_from_value(self, value: Any, branch: int, pad_value: int = 0) -> Any:
if isinstance(value, torch.Tensor):
return value[branch : branch + 1]
if isinstance(value, TokenizerEncodeOutput):
return type(value)(**{key: self._row_from_value(value[key], branch, pad_value) for key in value.keys()})
if isinstance(value, tuple):
return tuple(self._row_from_value(item, branch, pad_value) for item in value)
if isinstance(value, list):
return [value[branch]]
if isinstance(value, dict):
return {key: self._row_from_value(item, branch, pad_value) for key, item in value.items()}
return value
def _merge_rows(
self,
values: list[Any],
branches: list[int],
key: str,
*,
first_step: bool,
prefix_lens: list[int] | None,
) -> Any:
first = values[0]
if first is None:
return None
if isinstance(first, torch.Tensor):
rows = [self._row_from_value(value, branch) for value, branch in zip(values, branches)]
if key == "attention_mask":
return self._merge_attention_masks(rows, None if first_step else prefix_lens)
pad_value = self._tkwrapper.pad_token_id if key == "input_ids" else 0
return self._pad_tensor_rows(rows, pad_value)
if isinstance(first, TokenizerEncodeOutput):
return type(first)(
**{
output_key: self._merge_rows(
[value[output_key] for value in values],
branches,
output_key,
first_step=first_step,
prefix_lens=prefix_lens,
)
for output_key in first.keys()
}
)
if isinstance(first, tuple):
return tuple(
self._merge_rows(
[value[i] for value in values],
branches,
key,
first_step=first_step,
prefix_lens=prefix_lens,
)
for i in range(len(first))
)
if isinstance(first, dict):
return {
sub_key: self._merge_rows(
[value[sub_key] for value in values],
branches,
sub_key,
first_step=first_step,
prefix_lens=prefix_lens,
)
for sub_key in first
}
if isinstance(first, list):
return [value[branch] for value, branch in zip(values, branches)]
if all(value == first for value in values):
return first
raise ValueError(f"Cannot batch heterogeneous Hunyuan model kwarg {key!r}: {values}")
def _prompt_kv_prefix_lens(
self,
states: list["StepRequestState"],
row_state_indexes: list[int],
row_branches: list[int],
) -> list[int] | None:
if not states or states[0].step_index == 0:
return None
prefix_lens: list[int] = []
for state_idx, branch in zip(row_state_indexes, row_branches):
prompt_kv = states[state_idx].extra.get(_STEP_PROMPT_KV)
if prompt_kv is None:
raise ValueError(
f"Missing Hunyuan prompt KV cache for request {states[state_idx].request_id} "
"during later denoise step."
)
prefix_lens.append(int(prompt_kv[0]["lens"][branch].item()))
return prefix_lens
def _merge_step_model_inputs(
self,
states: list["StepRequestState"],
row_state_indexes: list[int],
row_branches: list[int],
first_step: bool,
) -> tuple[torch.Tensor | None, dict[str, Any]]:
prefix_lens = self._prompt_kv_prefix_lens(states, row_state_indexes, row_branches)
max_prefix_len = max(prefix_lens) if prefix_lens is not None else None
per_row_kwargs = [states[state_idx].extra[_STEP_MODEL_KWARGS] for state_idx in row_state_indexes]
merged: dict[str, Any] = {}
for key in per_row_kwargs[0]:
if key in {
"batch_gen_image_info",
"eos_token_id",
"generator",
"guidance_rescale",
"guidance_scale",
"max_new_tokens",
"num_inference_steps",
}:
continue
values = [kwargs.get(key) for kwargs in per_row_kwargs]
if key == "full_attn_spans" and values[0] is not None:
merged[key] = [
self._shift_full_attn_spans(value[branch], prefix_len, max_prefix_len)
for value, branch, prefix_len in zip(
values,
row_branches,
prefix_lens or [None] * len(row_branches),
)
]
continue
merged[key] = self._merge_rows(
values,
row_branches,
key,
first_step=first_step,
prefix_lens=prefix_lens,
)
if "attention_mask" in merged:
b, _, q_len, seq_len = merged["attention_mask"].shape
merged["query_lens"] = [q_len] * b
merged["seq_lens"] = [seq_len] * b
input_values = [states[state_idx].extra.get(_STEP_INPUT_IDS) for state_idx in row_state_indexes]
if input_values[0] is None:
input_ids = None
else:
rows = [value[branch : branch + 1] for value, branch in zip(input_values, row_branches)]
input_ids = self._pad_tensor_rows(rows, self._tkwrapper.pad_token_id)
return input_ids, merged
def _restore_prompt_kv_cache(
self,
states: list["StepRequestState"],
row_state_indexes: list[int],
row_branches: list[int],
) -> None:
if states[0].step_index == 0:
for layer in self.model.layers:
mgr = layer.self_attn.image_attn
mgr.image_kv_cache_map = None
mgr.image_kv_cache_lens = None
return
for layer_idx, layer in enumerate(self.model.layers):
rows_k, rows_v, lens = [], [], []
for state_idx, branch in zip(row_state_indexes, row_branches):
cache = states[state_idx].extra[_STEP_PROMPT_KV][layer_idx]
rows_k.append(cache["key"][branch : branch + 1])
rows_v.append(cache["value"][branch : branch + 1])
lens.append(cache["lens"][branch : branch + 1])
mgr = layer.self_attn.image_attn
mgr.image_kv_cache_map = (
self._pad_tensor_rows(rows_k),
self._pad_tensor_rows(rows_v),
)
mgr.image_kv_cache_lens = torch.cat(lens, dim=0).to(device=mgr.image_kv_cache_map[0].device)
def _capture_prompt_kv_cache(
self,
states: list["StepRequestState"],
row_state_indexes: list[int],
row_branches: list[int],
) -> None:
by_state: dict[int, list[dict[str, torch.Tensor]]] = {i: [] for i in range(len(states))}
for layer in self.model.layers:
mgr = layer.self_attn.image_attn
if mgr.image_kv_cache_map is None or mgr.image_kv_cache_lens is None:
raise ValueError("Hunyuan first step did not produce prompt KV cache.")
key, value = mgr.image_kv_cache_map
lens = mgr.image_kv_cache_lens
for state_idx in by_state:
branch_rows = [
row_idx for row_idx, (idx, _) in enumerate(zip(row_state_indexes, row_branches)) if idx == state_idx
]
state_lens = lens[branch_rows]
max_len = int(state_lens.max().item())
by_state[state_idx].append(
{
"key": key[branch_rows, :max_len].detach().clone(),
"value": value[branch_rows, :max_len].detach().clone(),
"lens": state_lens.detach().clone(),
}
)
for state_idx, cache in by_state.items():
states[state_idx].extra[_STEP_PROMPT_KV] = cache
@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.to(x.device).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
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)}."
)
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,
prepared_layout: HunyuanPreparedLayout | None = 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":
if prepared_layout is not None:
# TODO: Support batch_size > 1 once prepared layouts can
# preserve the public-request-to-CFG-branch mapping.
if batch_size != 1:
raise ValueError("Hunyuan prepared layout currently supports one public request at a time.")
batch_gen_image_info = [prepared_layout.generated_image_info]
else:
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 prepared_layout is None:
# Pull template options from generation_config only when the
# Worker must tokenize locally.
drop_think = kwargs.get("drop_think", self.generation_config.drop_think)
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 = out["output"]
rope_image_info = request_layout_utils.build_hunyuan_batch_rope_image_info(output, out["sections"])
else:
if mode != "gen_image" or batch_message_list is not None:
raise ValueError("Hunyuan prepared layout only supports prompt-based gen_image execution.")
output = prepared_layout.tokenizer_output
rope_image_info = prepared_layout.rope_image_info
expected_rows = batch_size * cfg_factor[mode]
if prepared_layout.num_branches != expected_rows:
raise ValueError(
"Hunyuan prepared branches do not match execution CFG layout: "
f"expected_rows={expected_rows}, prepared_branches={prepared_layout.num_branches}"
)
# 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
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=output.gen_timestep_scatter_index,
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
device = inputs_tensor.device
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, device=device).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]
kwargs["num_image_tokens"] = request_layout_utils.hunyuan_num_image_tokens(image_info)
# 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 _extract_ar_kv_from_sampling(self, sampling: Any) -> dict[str, Any]:
kv = getattr(sampling, "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
}
return {"ar_kv_data": ar_kv_data} if ar_kv_data else {}
def prepare_encode(
self,
state: "StepRequestState",
**kwargs: Any,
) -> "StepRequestState":
del kwargs
self._validate_step_request(state)
pipe = self.pipeline
sampling = state.sampling
(
prompt,
cot_text_list,
system_prompt,
batch_cond_image_info,
tokenizer_bot_task,
) = request_layout_utils.extract_hunyuan_prompt_inputs(
[state.prompt] if state.prompt is not None else [],
getattr(sampling, "extra_args", {}) or {},
request_id=state.request_id,
allow_cond_image=False,
)
cot_text = (
[request_layout_utils.normalize_hunyuan_cot_text(text) for text in cot_text_list]
if any(text is not None for text in cot_text_list)
else None
)
height = sampling.height or 1024
width = sampling.width or 1024
image_size = (height, width)
num_inference_steps = sampling.num_inference_steps or 50
guidance_scale = request_layout_utils.resolve_hunyuan_guidance_scale(sampling)
if guidance_scale <= 1.0:
logger.info("HunyuanImage3.0 step execution runs without classifier-free guidance.")
pipe._guidance_scale = guidance_scale
pipe._guidance_rescale = getattr(sampling, "guidance_rescale", 0.0)
model_kwargs = self.prepare_model_inputs(
prompt=prompt,
cot_text=cot_text,
system_prompt=system_prompt,
mode="gen_image",
generator=sampling.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,
prepared_layout=request_layout_utils.get_hunyuan_prepared_layout(state),
)
model_kwargs.update(self._extract_ar_kv_from_sampling(sampling))
model_kwargs["use_cache"] = False
input_ids = model_kwargs.pop("input_ids")
batch_gen_image_info: list[ImageInfo] = model_kwargs["batch_gen_image_info"]
image_info = batch_gen_image_info[0]
target_height = int(image_info.image_height)
target_width = int(image_info.image_width)
model_kwargs["num_image_tokens"] = request_layout_utils.hunyuan_num_image_tokens(image_info)
timesteps, _ = retrieve_timesteps(
self.scheduler,
num_inference_steps,
self.device,
None,
None,
)
pipe._num_timesteps = len(timesteps)
req_scheduler = copy.deepcopy(self.scheduler)
if hasattr(req_scheduler, "set_begin_index"):
req_scheduler.set_begin_index(0)
latents = pipe.prepare_latents(
batch_size=1,
latent_channel=self.config.vae["latent_channels"],
image_size=[target_height, target_width],
dtype=torch.bfloat16,
device=self.device,
generator=model_kwargs["generator"],
)
attention_mask = self._prepare_attention_mask_for_generation(
input_ids,
self.generation_config,
model_kwargs=model_kwargs,
)
bsz, _, query_len, seq_len = attention_mask.shape
model_kwargs["query_lens"] = [query_len] * bsz
model_kwargs["seq_lens"] = [seq_len] * bsz
model_kwargs["attention_mask"] = attention_mask.to(latents.device)
pipe._guidance_scale = guidance_scale
pipe._guidance_rescale = 0.0
input_ids, ar_kv_reuse_len = pipe._maybe_handle_ar_kv_reuse(
input_ids,
model_kwargs,
batch_size=1,
cfg_parallel_ready=False,
cfg_rank=None,
device=self.device,
)
model_kwargs["ar_kv_reuse_len"] = ar_kv_reuse_len
state.latents = latents
state.timesteps = timesteps
state.step_index = 0
state.scheduler = req_scheduler
state.do_true_cfg = guidance_scale > 1.0
state.extra = {
_STEP_MODEL_KWARGS: model_kwargs,
_STEP_INPUT_IDS: input_ids,
_STEP_GENERATOR: model_kwargs["generator"],
_STEP_GUIDANCE_SCALE: guidance_scale,
_STEP_CFG_FACTOR: 1 + int(guidance_scale > 1.0),
_STEP_OUTPUT_SIZE: (target_height, target_width),
_STEP_COT_TEXT_LIST: cot_text_list,
_STEP_AR_KV: self._snapshot_injected_ar_kv(),
}
return state
def _step_group_key(self, state: "StepRequestState") -> tuple[Any, ...]:
if _STEP_CFG_FACTOR not in state.extra:
raise ValueError(f"Missing Hunyuan CFG factor for request {state.request_id}.")
if _STEP_MODEL_KWARGS not in state.extra:
raise ValueError(f"Missing Hunyuan step model kwargs for request {state.request_id}.")
if state.latents is None:
raise ValueError(f"Missing Hunyuan latents for request {state.request_id}.")
model_kwargs = state.extra[_STEP_MODEL_KWARGS]
return (
state.step_index == 0,
state.extra[_STEP_CFG_FACTOR],
tuple(state.latents.shape[1:]),
model_kwargs.get("num_image_tokens"),
model_kwargs.get("ar_kv_reuse_len", 0),
state.extra.get(_STEP_AR_KV) is not None,
)
def _ensure_grouped_attention_backend_supported(self, num_states: int) -> None:
if num_states <= 1:
return
spec, source = self.od_config.diffusion_attention_config.resolve_with_source(role="self")
backend = spec.backend.upper() if spec is not None else "AUTO"
if backend == "TORCH_SDPA":
return
source_msg = f" from {source}" if source is not None else ""
raise ValueError(
"HunyuanImage3 grouped DiT batching only supports TORCH_SDPA attention. "
f"Current diffusion attention backend is {backend}{source_msg}. "
"Set DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA or configure "
"diffusion_attention_config.default.backend=TORCH_SDPA when max_num_seqs > 1."
)
def _split_step_groups(
self,
states: list["StepRequestState"],
) -> list[list["StepRequestState"]]:
grouped: dict[tuple[Any, ...], list[StepRequestState]] = {}
for state in states:
grouped.setdefault(self._step_group_key(state), []).append(state)
return list(grouped.values())
@staticmethod
def _step_row_order(
states: list["StepRequestState"],
cfg_factor: int,
) -> tuple[list[int], list[int]]:
row_state_indexes = [state_idx for _ in range(cfg_factor) for state_idx in range(len(states))]
row_branches = [branch for branch in range(cfg_factor) for _ in states]
return row_state_indexes, row_branches
@staticmethod
def _validate_step_group_states(states: list["StepRequestState"]) -> tuple[bool, int]:
if not states:
raise ValueError("HunyuanImage3 denoise_step received an empty group.")
first_step = states[0].step_index == 0
if _STEP_CFG_FACTOR not in states[0].extra:
raise ValueError(f"Missing Hunyuan CFG factor for request {states[0].request_id}.")
cfg_factor = int(states[0].extra[_STEP_CFG_FACTOR])
if cfg_factor not in (1, 2):
raise ValueError(
f"Unsupported Hunyuan CFG factor {cfg_factor} for request {states[0].request_id}; "
"expected no-CFG or cond/uncond CFG rows."
)
for state in states:
if (state.step_index == 0) != first_step:
raise ValueError("Hunyuan step group mixed first-step and later-step requests.")
if _STEP_CFG_FACTOR not in state.extra:
raise ValueError(f"Missing Hunyuan CFG factor for request {state.request_id}.")
if int(state.extra[_STEP_CFG_FACTOR]) != cfg_factor:
raise ValueError("Hunyuan step group mixed requests with different CFG factors.")
if _STEP_MODEL_KWARGS not in state.extra:
raise ValueError(f"Missing Hunyuan step model kwargs for request {state.request_id}.")
if state.latents is None:
raise ValueError(f"Missing Hunyuan latents for request {state.request_id}.")
if first_step:
if _STEP_AR_KV not in state.extra:
raise ValueError(f"Missing Hunyuan AR KV snapshot for request {state.request_id}.")
if _STEP_INPUT_IDS not in state.extra:
raise ValueError(f"Missing Hunyuan step input ids for request {state.request_id}.")
elif _STEP_PROMPT_KV not in state.extra:
raise ValueError(f"Missing Hunyuan prompt KV cache for request {state.request_id}.")
return first_step, cfg_factor
def _split_merged_kwargs_to_states(
self,
states: list["StepRequestState"],
merged_kwargs: dict[str, Any],
row_state_indexes: list[int],
row_branches: list[int],
) -> None:
cfg_factor = states[0].extra[_STEP_CFG_FACTOR]
for state_idx, state in enumerate(states):
state_rows = [
row_idx for row_idx, (idx, _) in enumerate(zip(row_state_indexes, row_branches)) if idx == state_idx
]
next_kwargs: dict[str, Any] = {}
for key, value in merged_kwargs.items():
if key in {"query_lens", "seq_lens"}:
continue
if key == "full_attn_spans":
next_kwargs[key] = state.extra[_STEP_MODEL_KWARGS].get(key)
elif key == "attention_mask" and isinstance(value, torch.Tensor):
cache = state.extra.get(_STEP_PROMPT_KV)
if cache is None:
next_kwargs[key] = value[state_rows]
else:
compact_masks = []
for row_idx in state_rows:
branch = row_branches[row_idx]
row_mask = value[row_idx : row_idx + 1]
query_len = int(row_mask.shape[-2])
prefix_len = int(cache[0]["lens"][branch].item())
max_prefix_len = int(row_mask.shape[-1]) - query_len
prefix_mask = row_mask[:, :, :, :prefix_len]
current_mask = row_mask[:, :, :, max_prefix_len : max_prefix_len + query_len]
compact_masks.append(torch.cat([prefix_mask, current_mask], dim=-1))
next_kwargs[key] = torch.cat(compact_masks, dim=0)
elif isinstance(value, torch.Tensor):
next_kwargs[key] = value[state_rows]
elif isinstance(value, tuple):
next_kwargs[key] = tuple(
item[state_rows] if isinstance(item, torch.Tensor) else item for item in value
)
elif isinstance(value, list) and len(value) == len(row_state_indexes):
next_kwargs[key] = [value[row_idx] for row_idx in state_rows]
else:
next_kwargs[key] = value
if len(state_rows) != cfg_factor:
raise ValueError("Hunyuan step split lost CFG rows while updating request state.")
if "attention_mask" in next_kwargs:
bsz, _, query_len, seq_len = next_kwargs["attention_mask"].shape
next_kwargs["query_lens"] = [query_len] * bsz
next_kwargs["seq_lens"] = [seq_len] * bsz
state.extra[_STEP_MODEL_KWARGS] = next_kwargs
state.extra[_STEP_INPUT_IDS] = None
def _denoise_step_group(self, states: list["StepRequestState"]) -> torch.Tensor:
first_step, cfg_factor = self._validate_step_group_states(states)
row_state_indexes, row_branches = self._step_row_order(states, cfg_factor)
latents = torch.cat([state.latents for state in states], dim=0)
latent_model_input = torch.cat([latents] * cfg_factor, dim=0)
timestep = torch.cat(
[
state.current_timestep.reshape(1).to(device=latents.device)
for _ in range(cfg_factor)
for state in states
],
dim=0,
)
input_ids, model_kwargs = self._merge_step_model_inputs(
states,
row_state_indexes,
row_branches,
first_step,
)
if first_step:
self._restore_injected_ar_kv(states, row_state_indexes, row_branches)
else:
self._restore_prompt_kv_cache(states, row_state_indexes, row_branches)
model_inputs = self.prepare_inputs_for_generation(
input_ids,
images=latent_model_input,
timestep=timestep,
**model_kwargs,
)
step_indexes = {state.step_index for state in states}
context_step_index = next(iter(step_indexes)) if len(step_indexes) == 1 else None
set_forward_context_denoise_step_idx(context_step_index)
try:
with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16, enabled=True):
model_output = self.forward_call(**model_inputs, first_step=first_step)
pred = model_output["diffusion_prediction"]
finally:
set_forward_context_denoise_step_idx(None)
pred = pred.to(dtype=torch.float32)
if first_step:
self._capture_prompt_kv_cache(states, row_state_indexes, row_branches)
if cfg_factor > 1:
pred_cond, pred_uncond = pred.chunk(2)
pred = torch.cat(
[
self.pipeline.cfg_operator(
pred_cond[state_idx : state_idx + 1],
pred_uncond[state_idx : state_idx + 1],
state.extra[_STEP_GUIDANCE_SCALE],
step=state.step_index,
)
for state_idx, state in enumerate(states)
],
dim=0,
)
updated_kwargs = self._update_model_kwargs_for_generation(model_output, model_kwargs)
self._split_merged_kwargs_to_states(states, updated_kwargs, row_state_indexes, row_branches)
return pred
def denoise_step(
self,
input_batch: "InputBatch",
**kwargs: Any,
) -> torch.Tensor | None:
del kwargs
if getattr(self, "interrupt", False):
return None
states = list(input_batch.states)
if not states:
raise ValueError("HunyuanImage3 denoise_step received an empty batch.")
self._ensure_grouped_attention_backend_supported(len(states))
outputs: dict[str, torch.Tensor] = {}
for group in self._split_step_groups(states):
pred = self._denoise_step_group(group)
for state, state_pred in zip(group, pred):
outputs[state.request_id] = state_pred.unsqueeze(0)
return torch.cat([outputs[state.request_id] for state in states], dim=0)
def step_scheduler(
self,
state: "StepRequestState",
noise_pred: torch.Tensor,
**kwargs: Any,
) -> None:
del kwargs
if getattr(self, "interrupt", False):
return
generator = state.extra.get(_STEP_GENERATOR)
step_kwargs = self.pipeline.prepare_extra_func_kwargs(state.scheduler.step, {"generator": generator})
latent_dtype = state.latents.dtype
state.latents = state.scheduler.step(
noise_pred,
state.current_timestep,
state.latents,
**step_kwargs,
return_dict=False,
)[0].to(dtype=latent_dtype)
state.step_index += 1
def post_decode(
self,
state: "StepRequestState",
**kwargs: Any,
) -> DiffusionOutput:
output_type = kwargs.get("output_type", "pil")
generator = state.extra.get(_STEP_GENERATOR)
latents = state.latents
if output_type == "latent":
return DiffusionOutput(
output=latents,
stage_durations=getattr(self, "stage_durations", None),
)
if hasattr(self.vae.config, "scaling_factor") and self.vae.config.scaling_factor:
latents = latents / self.vae.config.scaling_factor
if hasattr(self.vae.config, "shift_factor") and self.vae.config.shift_factor:
latents = latents + self.vae.config.shift_factor
if hasattr(self.vae, "ffactor_temporal"):
latents = latents.unsqueeze(2)
with torch.autocast(device_type=self.device.type, dtype=torch.float16, enabled=True):
image = self.vae.decode(latents, return_dict=False, generator=generator)[0]
if hasattr(self.vae, "ffactor_temporal"):
assert image.shape[2] == 1, "image should have shape [B, C, T, H, W] and T should be 1"
image = image.squeeze(2)
# Postprocess deferred to engine post_process_func for overlap with next request.
cot_text_list = state.extra.get(_STEP_COT_TEXT_LIST) or []
metadata = {}
if any(text is not None for text in cot_text_list):
metadata["text"] = {"ar_generated_text": cot_text_list[0]}
return DiffusionOutput(
output={
"payload": {"image": image[0]},
"metadata": metadata,
},
stage_durations=getattr(self, "stage_durations", None),
)
def forward(
self,
req: DiffusionRequestBatch,
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 {}
(
prompt_from_req,
cot_text_list,
system_prompt,
batch_cond_image_info,
tokenizer_bot_task,
) = request_layout_utils.extract_hunyuan_prompt_inputs(
req.prompts,
extra_args,
request_id=req.request_id,
allow_cond_image=True,
)
prompt = prompt_from_req or prompt
cot_text = (
[request_layout_utils.normalize_hunyuan_cot_text(t) for t in cot_text_list]
if any(t is not None for t in cot_text_list)
else 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
guidance_scale = request_layout_utils.resolve_hunyuan_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,
bot_task=tokenizer_bot_task,
prepared_layout=request_layout_utils.get_hunyuan_prepared_layout(req.requests[0]),
**ar_kv_kwargs,
)
model_inputs.update(ar_kv_kwargs)
outputs = self._generate(**model_inputs, **kwargs)
image = outputs[0]
metadata = {}
if any(t is not None for t in cot_text_list):
metadata["text"] = {"ar_generated_text": cot_text_list[0] if len(cot_text_list) == 1 else cot_text_list}
stage_durations = self.stage_durations if hasattr(self, "stage_durations") else None
return DiffusionOutput(
output={
"payload": {"image": image},
"metadata": metadata,
},
stage_durations=stage_durations,
)
Reference in new issue View Git Blame Copy Permalink