2479 lines
109 KiB
Python
2479 lines
109 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni 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 (
|
|
get_forward_context,
|
|
is_forward_context_available,
|
|
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 _config_flag(config: Any, name: str) -> bool:
|
|
"""Read an optional Hunyuan config flag with legacy-safe defaults."""
|
|
|
|
return bool(getattr(config, name, False))
|
|
|
|
|
|
class HunyuanImage3FlowMatchEulerDiscreteScheduler(FlowMatchEulerDiscreteScheduler):
|
|
def get_timestep_r(self, timestep: float | torch.FloatTensor):
|
|
if self.step_index is None:
|
|
self._init_step_index(timestep)
|
|
if self.step_index + 1 < len(self.timesteps):
|
|
return self.timesteps[self.step_index + 1]
|
|
return self.timesteps.new_zeros(())
|
|
|
|
|
|
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 get_hunyuan_image_3_prefix_cache_func(od_config: OmniDiffusionConfig):
|
|
return request_layout_utils.prepare_hunyuan_prefix_cache
|
|
|
|
|
|
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,
|
|
cfg_distilled=_config_flag(hf_config, "cfg_distilled"),
|
|
use_meanflow=_config_flag(hf_config, "use_meanflow"),
|
|
)
|
|
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
|
|
# The warmup request's blank image routes this AR+diffusion model down its
|
|
# image-edit path, where the token block cannot match the latent patch grid.
|
|
dummy_run_num_frames: ClassVar[int] = 0
|
|
_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)
|
|
self.hf_config.cfg_distilled = _config_flag(self.hf_config, "cfg_distilled")
|
|
self.hf_config.use_meanflow = _config_flag(self.hf_config, "use_meanflow")
|
|
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)
|
|
# Distillation-specific embedding modules
|
|
if _config_flag(self.hf_config, "cfg_distilled"):
|
|
self.guidance_emb = TimestepEmbedder(hidden_size=self.hf_config.hidden_size)
|
|
if _config_flag(self.hf_config, "use_meanflow"):
|
|
self.timestep_r_emb = 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",
|
|
"guidance_emb",
|
|
"timestep_r_emb",
|
|
"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)
|
|
|
|
# Note: guidance_emb and timestep_r_emb are no longer skipped
|
|
# to support HunyuanImage-3.0-Distil and MeanFlow distilled models
|
|
loader = AutoWeightsLoader(self)
|
|
return loader.load_weights(
|
|
weights, mapper=WeightsMapper(orig_to_new_prefix={name: None for name in (skip_prefixes or ())})
|
|
)
|
|
|
|
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 = HunyuanImage3FlowMatchEulerDiscreteScheduler(
|
|
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:
|
|
if self._uses_scheduler_paged_kv():
|
|
raise ValueError(
|
|
"HunyuanImage3 paged_scheduler currently supports request-level execution only; "
|
|
"disable step execution or use dense_legacy."
|
|
)
|
|
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 _uses_scheduler_paged_kv(self) -> bool:
|
|
return (
|
|
getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
|
|
is DiffusionKVCacheMode.PAGED_SCHEDULER
|
|
)
|
|
|
|
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),
|
|
strict=True,
|
|
)
|
|
]
|
|
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_guidance_tokens(
|
|
self,
|
|
x: torch.Tensor,
|
|
guidance: BatchRaggedTensor | None = None,
|
|
guidance_scatter_index: torch.Tensor | None = None,
|
|
):
|
|
"""Instantiate guidance embedding tokens for CFG distilled models.
|
|
|
|
Args:
|
|
x: Input sequence tensor (batch_size, seq_len, n_embd)
|
|
guidance: Guidance scale values (as tensor, typically guidance_scale * 1000)
|
|
guidance_scatter_index: Index positions to scatter guidance embeddings
|
|
"""
|
|
if guidance is None or guidance_scatter_index is None:
|
|
return x
|
|
if not hasattr(self, "guidance_emb"):
|
|
return x
|
|
|
|
batch_size, seq_len, n_embd = x.shape
|
|
guidance_src = self.guidance_emb(guidance.reshape(-1)) # (bsz * n, n_embd)
|
|
x.scatter_(
|
|
dim=1,
|
|
index=guidance_scatter_index.unsqueeze(-1).repeat(1, 1, n_embd),
|
|
src=guidance_src.reshape(batch_size, -1, n_embd),
|
|
)
|
|
return x
|
|
|
|
def instantiate_timestep_r_tokens(
|
|
self,
|
|
x: torch.Tensor,
|
|
timesteps_r: BatchRaggedTensor | None = None,
|
|
timesteps_r_scatter_index: torch.Tensor | None = None,
|
|
):
|
|
"""Instantiate timestep_r embedding tokens for MeanFlow distilled models.
|
|
|
|
Args:
|
|
x: Input sequence tensor (batch_size, seq_len, n_embd)
|
|
timesteps_r: Timestep r values for MeanFlow
|
|
timesteps_r_scatter_index: Index positions to scatter embeddings
|
|
"""
|
|
if timesteps_r is None or timesteps_r_scatter_index is None:
|
|
return x
|
|
if not hasattr(self, "timestep_r_emb"):
|
|
return x
|
|
|
|
batch_size, seq_len, n_embd = x.shape
|
|
|
|
if isinstance(timesteps_r, list):
|
|
for i, timestep_r in enumerate(timesteps_r):
|
|
timestep_r_src = self.timestep_r_emb(timestep_r) # (n, n_embd)
|
|
x[i : i + 1].scatter_(
|
|
dim=1,
|
|
index=timesteps_r_scatter_index[i].unsqueeze(0).unsqueeze(-1).repeat(1, 1, n_embd),
|
|
src=timestep_r_src.reshape(1, -1, n_embd),
|
|
)
|
|
else:
|
|
timesteps_r_src = self.timestep_r_emb(timesteps_r.reshape(-1)) # (bsz * n, n_embd)
|
|
x.scatter_(
|
|
dim=1,
|
|
index=timesteps_r_scatter_index.unsqueeze(-1).repeat(1, 1, n_embd),
|
|
src=timesteps_r_src.reshape(batch_size, -1, n_embd),
|
|
)
|
|
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, num_special_tokens: int = None):
|
|
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[:, num_special_tokens:, :]
|
|
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,
|
|
add_guidance_token=_config_flag(getattr(self, "hf_config", None), "cfg_distilled"),
|
|
add_timestep_r_token=_config_flag(getattr(self, "hf_config", None), "use_meanflow"),
|
|
)
|
|
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
|
|
# For CFG distilled models, cfg_factor is always 1 (CFG is embedded in the model)
|
|
# For non-distilled models, cfg_factor depends on guidance_scale
|
|
cfg_factor = {"gen_text": 1, "gen_image": 1 + int(guidance_scale > 1.0)}
|
|
if _config_flag(getattr(self, "hf_config", None), "cfg_distilled"):
|
|
cfg_factor["gen_image"] = 1
|
|
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") or request_layout_utils.native_kv_covers_cond_images(
|
|
output, kwargs.get("kv_computed_tokens", ())
|
|
)
|
|
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 CFG distilled models
|
|
guidance_scatter_index=to_device(output.guidance_scatter_index, device)
|
|
if hasattr(output, "guidance_scatter_index") and output.guidance_scatter_index is not None
|
|
else None,
|
|
# for MeanFlow distilled models
|
|
timesteps_r_scatter_index=to_device(output.timesteps_r_scatter_index, device)
|
|
if hasattr(output, "timesteps_r_scatter_index") and output.timesteps_r_scatter_index is not None
|
|
else None,
|
|
# 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 reference_slice in tokenizer_output.joint_image_slices[i]:
|
|
logger.debug(
|
|
"Hunyuan reference span: sequence_id=%d start=%d end=%d",
|
|
i,
|
|
int(reference_slice.start or 0),
|
|
int(reference_slice.stop or seq_len),
|
|
)
|
|
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
|
|
|
|
@staticmethod
|
|
def _slice_cached_prefix_inputs(
|
|
inputs_embeds: torch.Tensor,
|
|
attention_mask: torch.Tensor | None,
|
|
position_ids: torch.Tensor | None,
|
|
custom_pos_emb: tuple[torch.Tensor, ...] | None,
|
|
image_mask: torch.Tensor | None,
|
|
gen_timestep_scatter_index: torch.Tensor | None,
|
|
query_lens: list[int] | None,
|
|
seq_lens: list[int] | None,
|
|
cached_prefix_len: int,
|
|
) -> tuple[
|
|
torch.Tensor,
|
|
torch.Tensor,
|
|
torch.Tensor,
|
|
tuple[torch.Tensor, ...],
|
|
torch.Tensor,
|
|
torch.Tensor,
|
|
list[int],
|
|
int,
|
|
]:
|
|
"""Remove a cached first-step prefix while preserving full KV lengths."""
|
|
|
|
seq_len = int(inputs_embeds.shape[1])
|
|
if (
|
|
query_lens is None
|
|
or seq_lens is None
|
|
or len(query_lens) != inputs_embeds.shape[0]
|
|
or len(seq_lens) != len(query_lens)
|
|
):
|
|
raise ValueError("Hunyuan paged prefix execution requires one query/sequence length per row")
|
|
if any(query_len != seq_len for query_len in query_lens) or any(
|
|
logical_len != seq_len for logical_len in seq_lens
|
|
):
|
|
raise ValueError(
|
|
"Hunyuan paged prefix slicing requires an untrimmed first-step tensor: "
|
|
f"query_lens={query_lens}, tensor_len={seq_len}"
|
|
)
|
|
if not 0 < cached_prefix_len < seq_len:
|
|
raise ValueError(
|
|
"Hunyuan cached prefix must leave a non-empty first-step suffix: "
|
|
f"cached={cached_prefix_len}, seq_len={seq_len}"
|
|
)
|
|
if (
|
|
attention_mask is None
|
|
or position_ids is None
|
|
or custom_pos_emb is None
|
|
or image_mask is None
|
|
or gen_timestep_scatter_index is None
|
|
):
|
|
raise ValueError("Hunyuan paged prefix slicing is missing first-step layout tensors")
|
|
|
|
return (
|
|
inputs_embeds[:, cached_prefix_len:],
|
|
attention_mask[:, :, cached_prefix_len:, :],
|
|
position_ids[:, cached_prefix_len:],
|
|
tuple(pos[:, cached_prefix_len:] for pos in custom_pos_emb),
|
|
image_mask[:, cached_prefix_len:],
|
|
gen_timestep_scatter_index - cached_prefix_len,
|
|
[query_len - cached_prefix_len for query_len in query_lens],
|
|
seq_len - cached_prefix_len,
|
|
)
|
|
|
|
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"),
|
|
# For CFG distilled models
|
|
"guidance": kwargs.get("guidance"),
|
|
"guidance_scatter_index": kwargs.get("guidance_scatter_index"),
|
|
# For MeanFlow distilled models
|
|
"timesteps_r": kwargs.get("timesteps_r"),
|
|
"timesteps_r_scatter_index": kwargs.get("timesteps_r_scatter_index"),
|
|
"query_lens": kwargs.get("query_lens"),
|
|
"seq_lens": kwargs.get("seq_lens"),
|
|
"num_image_tokens": kwargs.get("num_image_tokens"),
|
|
"num_special_tokens": kwargs.get("num_special_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"],
|
|
"num_special_tokens": model_kwargs["num_special_tokens"],
|
|
"guidance": model_kwargs.get("guidance"),
|
|
"guidance_scatter_index": model_kwargs.get("guidance_scatter_index"),
|
|
"timesteps_r": model_kwargs.get("timesteps_r"),
|
|
"timesteps_r_scatter_index": model_kwargs.get("timesteps_r_scatter_index"),
|
|
}
|
|
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)
|
|
image_position_ids = index.masked_select(image_mask.bool()).reshape(bsz, -1)
|
|
special_position_ids = [
|
|
index[torch.arange(bsz), model_kwargs["gen_timestep_scatter_index"][:, -1]].unsqueeze(-1)
|
|
]
|
|
if model_kwargs.get("guidance_scatter_index") is not None:
|
|
special_position_ids.append(
|
|
index[torch.arange(bsz), model_kwargs["guidance_scatter_index"][:, -1]].unsqueeze(-1)
|
|
)
|
|
if model_kwargs.get("timesteps_r_scatter_index") is not None:
|
|
special_position_ids.append(
|
|
index[torch.arange(bsz), model_kwargs["timesteps_r_scatter_index"][:, -1]].unsqueeze(-1)
|
|
)
|
|
special_position_ids = torch.cat(special_position_ids, dim=1)
|
|
updated_model_kwargs["position_ids"] = torch.cat([special_position_ids, image_position_ids], dim=1)
|
|
|
|
# attention mask
|
|
mask_list = []
|
|
current_starts = special_position_ids[:, 0]
|
|
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)
|
|
kwargs["num_special_tokens"] = request_layout_utils.hunyuan_num_special_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,
|
|
num_special_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,
|
|
# for CFG distilled models
|
|
guidance: torch.Tensor | None = None,
|
|
guidance_scatter_index: torch.Tensor | None = None,
|
|
# for MeanFlow distilled models
|
|
timesteps_r: torch.Tensor | None = None,
|
|
timesteps_r_scatter_index: torch.Tensor | 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)
|
|
# Instantiate guidance and timestep_r tokens for distilled models
|
|
inputs_embeds = self.instantiate_guidance_tokens(inputs_embeds, guidance, guidance_scatter_index)
|
|
inputs_embeds = self.instantiate_timestep_r_tokens(
|
|
inputs_embeds, timesteps_r, timesteps_r_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)
|
|
extra_embeds = []
|
|
if hasattr(self, "guidance_emb") and guidance is not None:
|
|
guidance_emb = self.guidance_emb(guidance.reshape(-1)).reshape(bsz, -1, n_embd)
|
|
extra_embeds.append(guidance_emb)
|
|
if hasattr(self, "timestep_r_emb") and timesteps_r is not None:
|
|
timesteps_r_emb = self.timestep_r_emb(timesteps_r.reshape(-1)).reshape(bsz, -1, n_embd)
|
|
extra_embeds.append(timesteps_r_emb)
|
|
inputs_embeds = torch.cat([timestep_emb, *extra_embeds, 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
|
|
|
|
paged_kv_cached_prefix_len = 0
|
|
if first_step and self._uses_scheduler_paged_kv():
|
|
if not is_forward_context_available():
|
|
raise RuntimeError("Hunyuan paged prefix execution requires an Omni ForwardContext")
|
|
paged_kv_cached_prefix_len = int(get_forward_context().paged_kv_cached_prefix_len)
|
|
if paged_kv_cached_prefix_len:
|
|
if ar_kv_reuse_len:
|
|
raise ValueError("Hunyuan cannot combine imported AR KV and Scheduler prefix-cache hits")
|
|
(
|
|
inputs_embeds,
|
|
attention_mask,
|
|
position_ids,
|
|
custom_pos_emb,
|
|
image_mask,
|
|
gen_timestep_scatter_index,
|
|
query_lens,
|
|
seq_len,
|
|
) = self._slice_cached_prefix_inputs(
|
|
inputs_embeds,
|
|
attention_mask,
|
|
position_ids,
|
|
custom_pos_emb,
|
|
image_mask,
|
|
gen_timestep_scatter_index,
|
|
query_lens,
|
|
seq_lens,
|
|
paged_kv_cached_prefix_len,
|
|
)
|
|
logger.debug(
|
|
"Hunyuan prefix slice: cached_prefix_len=%d query_len=%d",
|
|
paged_kv_cached_prefix_len,
|
|
seq_len,
|
|
)
|
|
|
|
# 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,
|
|
paged_kv_cached_prefix_len=paged_kv_cached_prefix_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, num_special_tokens
|
|
)
|
|
|
|
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)
|
|
|
|
ar_kv_kwargs = self._extract_ar_kv_from_sampling(sampling)
|
|
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),
|
|
kv_computed_tokens=state.extra.get("kv_computed_tokens", ()),
|
|
**ar_kv_kwargs,
|
|
)
|
|
model_kwargs.update(ar_kv_kwargs)
|
|
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)
|
|
model_kwargs["num_special_tokens"] = request_layout_utils.hunyuan_num_special_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
|
|
cfg_factor = request_layout_utils.hunyuan_cfg_factor(image_info, guidance_scale)
|
|
state.do_true_cfg = cfg_factor > 1
|
|
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: cfg_factor,
|
|
_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
|
|
|
|
def _validate_step_group_states(self, 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 _config_flag(self.hf_config, "cfg_distilled"):
|
|
model_kwargs["guidance"] = torch.tensor(
|
|
[1000.0 * state.extra[_STEP_GUIDANCE_SCALE] for state in states] * cfg_factor,
|
|
device=latents.device,
|
|
dtype=torch.bfloat16,
|
|
)
|
|
if _config_flag(self.hf_config, "use_meanflow"):
|
|
model_kwargs["timesteps_r"] = torch.cat(
|
|
[
|
|
state.scheduler.get_timestep_r(state.current_timestep).reshape(1).to(device=latents.device)
|
|
for _ in range(cfg_factor)
|
|
for state in states
|
|
],
|
|
dim=0,
|
|
)
|
|
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.")
|
|
if self._uses_scheduler_paged_kv():
|
|
raise ValueError(
|
|
"HunyuanImage3 paged_scheduler currently supports request-level execution only; "
|
|
"disable step execution or use dense_legacy."
|
|
)
|
|
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]),
|
|
kv_computed_tokens=getattr(req.requests[0], "kv_computed_tokens", ()),
|
|
**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,
|
|
)
|