vllm-omni/vllm_omni/diffusion/models/hunyuan_image3/request_layout.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING, Any

import numpy as np
import torch
from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange

from vllm_omni.diffusion.diffusion_kv.kv_cache_utils import get_cache_namespace, hash_prefix_cache_value
from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
from vllm_omni.diffusion.request import OmniDiffusionRequest

from .hunyuan_image3_tokenizer import TokenizerEncodeOutput
from .hunyuan_image3_transformer import ImageInfo, JointImageInfo
from .system_prompt import get_system_prompt

if TYPE_CHECKING:
    from transformers.generation.configuration_utils import GenerationConfig

    from .hunyuan_image3_tokenizer import TokenizerWrapper
    from .hunyuan_image3_transformer import HunyuanImage3ImageProcessor


def _to_python_scalar(value: Any) -> Any:
    if isinstance(value, np.generic):
        return value.item()
    return value


def _to_tensor_if_needed(value: Any) -> Any:
    if isinstance(value, np.generic):
        return value.item()
    if isinstance(value, list):
        return torch.tensor(value)
    return value


def _image_info_to_payload(image_info: ImageInfo) -> dict[str, Any]:
    return {
        "image_type": image_info.image_type,
        "image_tensor": image_info.image_tensor,
        "image_width": _to_python_scalar(image_info.image_width),
        "image_height": _to_python_scalar(image_info.image_height),
        "token_width": _to_python_scalar(image_info.token_width),
        "token_height": _to_python_scalar(image_info.token_height),
        "image_token_length": _to_python_scalar(image_info.image_token_length),
        "base_size": _to_python_scalar(image_info.base_size),
        "ratio_index": _to_python_scalar(image_info.ratio_index),
        "add_timestep_token": image_info.add_timestep_token,
        "add_guidance_token": image_info.add_guidance_token,
        "add_timestep_r_token": image_info.add_timestep_r_token,
        "use_front_boi_token": image_info.use_front_boi_token,
        "add_image_shape_token": image_info.add_image_shape_token,
    }


def _image_info_from_payload(payload: dict[str, Any]) -> ImageInfo:
    return ImageInfo(
        image_type=payload.get("image_type"),
        image_tensor=_to_tensor_if_needed(payload.get("image_tensor")),
        image_width=payload.get("image_width"),
        image_height=payload.get("image_height"),
        token_width=payload.get("token_width"),
        token_height=payload.get("token_height"),
        image_token_length=payload.get("image_token_length"),
        base_size=payload.get("base_size"),
        ratio_index=payload.get("ratio_index"),
        add_timestep_token=payload.get("add_timestep_token", True),
        add_guidance_token=payload.get("add_guidance_token", False),
        add_timestep_r_token=payload.get("add_timestep_r_token", False),
        use_front_boi_token=payload.get("use_front_boi_token", True),
        add_image_shape_token=payload.get("add_image_shape_token", True),
    )


def joint_image_info_to_payload(joint_image_info: JointImageInfo) -> dict[str, Any]:
    return {
        "type": "joint_image_info",
        "vae_image_info": _image_info_to_payload(joint_image_info.vae_image_info),
        "vision_image_info": _image_info_to_payload(joint_image_info.vision_image_info),
        "vision_encoder_kwargs": joint_image_info.vision_encoder_kwargs,
    }


def _joint_image_info_from_payload(payload: Any) -> JointImageInfo:
    if isinstance(payload, JointImageInfo):
        return payload
    if not isinstance(payload, dict):
        raise TypeError(f"Expected dict or JointImageInfo for conditional image payload, got {type(payload)}.")

    vae_image_info = _image_info_from_payload(payload["vae_image_info"])
    vision_image_info = _image_info_from_payload(payload["vision_image_info"])
    vision_encoder_kwargs = payload.get("vision_encoder_kwargs") or {}
    if isinstance(vision_encoder_kwargs, dict):
        vision_encoder_kwargs = {key: _to_tensor_if_needed(value) for key, value in vision_encoder_kwargs.items()}
    return JointImageInfo(
        vae_image_info=vae_image_info,
        vision_image_info=vision_image_info,
        vision_encoder_kwargs=vision_encoder_kwargs,
    )


def normalize_hunyuan_single_stage_bot_task(bot_task: Any) -> str:
    if isinstance(bot_task, str) and bot_task.lower() == "none":
        bot_task = None
    tokenizer_bot_task = bot_task
    if tokenizer_bot_task == "think_recaption":
        tokenizer_bot_task = "think"
    elif tokenizer_bot_task == "vanilla":
        tokenizer_bot_task = "image"
    supported_bot_tasks = {"auto", "image", "think", "recaption", "img_ratio"}
    if tokenizer_bot_task is not None and tokenizer_bot_task not in supported_bot_tasks:
        raise ValueError(
            f"Unsupported HunyuanImage3 single-stage bot_task: {tokenizer_bot_task!r}. "
            f"Supported values are: {sorted(supported_bot_tasks)}."
        )
    return tokenizer_bot_task or "auto"


def normalize_hunyuan_cot_text(cot: str | None) -> str | None:
    """Restore an AR generation trigger tag omitted from generated text."""

    if not cot:
        return cot
    if "</think>" in cot and not cot.startswith("<think>"):
        return "<think>" + cot
    if "</recaption>" in cot and not cot.startswith("<recaption>"):
        return "<recaption>" + cot
    return cot


def extract_hunyuan_prompt_inputs(
    prompts: list[Any],
    extra_args: dict[str, Any],
    *,
    request_id: str,
    allow_cond_image: bool,
) -> tuple[list[str], list[str | None], str | None, list[list[JointImageInfo]] | None, str]:
    """Normalize request prompt fields shared by planning and execution."""

    is_dummy_warmup = OmniDiffusionRequest.is_dummy_run_request_id(request_id)
    bot_task = extra_args.get("bot_task")
    use_system_prompt = extra_args.get("use_system_prompt")
    system_prompt = extra_args.get("system_prompt")

    first_prompt = prompts[0] if prompts else None
    if isinstance(first_prompt, dict):
        if bot_task is None:
            bot_task = first_prompt.get("bot_task")
        if use_system_prompt is None:
            use_system_prompt = first_prompt.get("use_system_prompt")
        if system_prompt is None:
            system_prompt = first_prompt.get("system_prompt")
    tokenizer_bot_task = normalize_hunyuan_single_stage_bot_task(bot_task)
    if use_system_prompt is not None:
        system_prompt_bot_task = "image" if tokenizer_bot_task == "auto" else tokenizer_bot_task
        system_prompt = get_system_prompt(use_system_prompt, system_prompt_bot_task, system_prompt)
        system_prompt = system_prompt.strip() if system_prompt is not None else ""

    prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in prompts]
    cot_text_list = [
        (p.get("extra", {}).get("ar_generated_text") if isinstance(p, dict) else None) or None for p in prompts
    ]

    batch_cond_image_info: list[list[JointImageInfo]] | None = None
    if any(not isinstance(p, str) for p in prompts):
        batch_cond_image_info = []
        for prompt_item in prompts:
            if isinstance(prompt_item, str):
                batch_cond_image_info.append([])
                continue
            additional_info = prompt_item.get("additional_information") or {}
            cond_infos = additional_info.get("batch_cond_image_info", [])
            if isinstance(cond_infos, JointImageInfo | dict):
                cond_infos = [cond_infos]
            if cond_infos is None:
                cond_infos = []
            batch_cond_image_info.append([_joint_image_info_from_payload(item) for item in cond_infos])

        has_cond_image = [len(cond_infos) > 0 for cond_infos in batch_cond_image_info]
        if any(has_cond_image) and (not allow_cond_image) and not is_dummy_warmup:
            raise ValueError("HunyuanImage3 step execution does not support image editing requests yet.")
        if allow_cond_image and any(has_cond_image) and not all(has_cond_image):
            raise ValueError("When batching Hunyuan image editing requests, every prompt must include input image(s).")
        if not allow_cond_image or not any(has_cond_image):
            batch_cond_image_info = None

    return prompt, cot_text_list, system_prompt, batch_cond_image_info, tokenizer_bot_task


def resolve_hunyuan_guidance_scale(sampling: Any, default_scale: float = 5.0) -> float:
    if getattr(sampling, "guidance_scale_provided", False):
        return sampling.guidance_scale
    return default_scale


def hunyuan_num_image_tokens(image_info: ImageInfo) -> int:
    """Return the generated-image span overwritten on every denoise step."""

    return int(image_info.image_token_length) + hunyuan_num_special_tokens(image_info)


def hunyuan_num_special_tokens(image_info: ImageInfo) -> int:
    """Return the generated-image prefix tokens emitted before latent tokens."""

    return (
        int(image_info.add_timestep_token) + int(image_info.add_guidance_token) + int(image_info.add_timestep_r_token)
    )


def hunyuan_cfg_factor(image_info: ImageInfo, guidance_scale: float) -> int:
    """Return the execution branch count for standard or embedded CFG."""

    return 1 if image_info.add_guidance_token else 1 + int(guidance_scale > 1.0)


def build_hunyuan_batch_rope_image_info(
    output: TokenizerEncodeOutput,
    sections: list[list[dict[str, Any]]],
) -> list[list[tuple[slice, tuple[int, int]]]]:
    if output.all_image_slices is None:
        raise ValueError("Hunyuan tokenizer output is missing all_image_slices.")
    if len(output.all_image_slices) != len(sections):
        raise ValueError(
            "Hunyuan image-slice rows do not match template section rows: "
            f"image_slices={len(output.all_image_slices)}, sections={len(sections)}"
        )
    rope_image_info: list[list[tuple[slice, tuple[int, int]]]] = []
    for image_slices, sections_i in zip(output.all_image_slices, sections):
        image_shapes: list[tuple[int, int]] = []
        for section in sections_i:
            if "image" not in section["type"]:
                continue
            if isinstance(section["token_height"], list):
                if len(section["token_height"]) != len(section["token_width"]):
                    raise ValueError(
                        "token_height and token_width should have the same length, "
                        f"but got {len(section['token_height'])} and {len(section['token_width'])}"
                    )
                image_shapes.extend(
                    (int(token_height), int(token_width))
                    for token_height, token_width in zip(section["token_height"], section["token_width"])
                )
            else:
                image_shapes.append((int(section["token_height"]), int(section["token_width"])))
        if len(image_slices) != len(image_shapes):
            raise ValueError(f"Image slices({len(image_slices)}) do not match image shapes({len(image_shapes)}).")
        rope_image_info.append(list(zip(image_slices, image_shapes)))
    return rope_image_info


def prepare_hunyuan_layout(
    request: OmniDiffusionRequest,
    *,
    tokenizer_wrapper: TokenizerWrapper,
    image_processor: HunyuanImage3ImageProcessor,
    generation_config: GenerationConfig,
    image_base_size: int,
    cfg_distilled: bool = False,
    use_meanflow: bool = False,
) -> HunyuanPreparedLayout:
    """Build the CPU token/image layout reused by Scheduler and Worker."""

    sampling = request.sampling_params
    extra_args = getattr(sampling, "extra_args", {}) or {}
    prompt, cot_text_list, system_prompt, batch_cond_image_info, tokenizer_bot_task = extract_hunyuan_prompt_inputs(
        [request.prompt],
        extra_args,
        request_id=request.request_id,
        allow_cond_image=True,
    )
    cot_text = (
        [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
    guidance_scale = resolve_hunyuan_guidance_scale(sampling)
    image_info_kwargs: dict[str, bool] = {}
    if cfg_distilled:
        image_info_kwargs["add_guidance_token"] = True
    if use_meanflow:
        image_info_kwargs["add_timestep_r_token"] = True
    generated_image_info = image_processor.build_image_info((height, width), **image_info_kwargs)
    result = tokenizer_wrapper.apply_chat_template(
        batch_prompt=prompt,
        mode="gen_image",
        batch_gen_image_info=[generated_image_info],
        batch_cond_image_info=batch_cond_image_info,
        batch_system_prompt=[system_prompt],
        batch_cot_text=cot_text,
        max_length=None,
        bot_task=tokenizer_bot_task,
        image_base_size=image_base_size,
        sequence_template=getattr(generation_config, "sequence_template", "pretrain"),
        cfg_factor=hunyuan_cfg_factor(generated_image_info, guidance_scale),
        drop_think=getattr(generation_config, "drop_think", False),
    )
    tokenizer_output = result["output"]
    return HunyuanPreparedLayout(
        tokenizer_output=tokenizer_output,
        rope_image_info=build_hunyuan_batch_rope_image_info(tokenizer_output, result["sections"]),
        generated_image_info=generated_image_info,
    )


@dataclass(frozen=True)
class HunyuanPreparedLayout:
    """CPU execution layout prepared once before Scheduler admission."""

    tokenizer_output: TokenizerEncodeOutput
    rope_image_info: list[list[tuple[slice, tuple[int, int]]]]
    generated_image_info: ImageInfo

    def __post_init__(self) -> None:
        tokens = self.tokenizer_output.tokens
        if tokens is None or tokens.ndim != 2:
            shape = None if tokens is None else tuple(tokens.shape)
            raise ValueError(f"Hunyuan prepared tokenizer tokens must be 2-D, got {shape}")
        if len(self.rope_image_info) != int(tokens.shape[0]):
            raise ValueError(
                "Hunyuan prepared RoPE rows do not match tokenizer rows: "
                f"rope_rows={len(self.rope_image_info)}, tokens={int(tokens.shape[0])}"
            )
        prefix_positions = self.tokenizer_output.gen_timestep_scatter_index
        if prefix_positions is None or prefix_positions.ndim != 2 or prefix_positions.shape[1] == 0:
            raise ValueError("Hunyuan prepared layout requires a non-empty 2-D gen_timestep_scatter_index")
        if prefix_positions.shape[0] != int(tokens.shape[0]):
            raise ValueError(
                "Hunyuan prepared prefix-position rows do not match tokenizer rows: "
                f"prefix_positions={prefix_positions.shape[0]}, tokens={int(tokens.shape[0])}"
            )
        real_pos = self.tokenizer_output.real_pos
        if real_pos is None or real_pos.ndim != 2 or real_pos.shape[1] == 0:
            raise ValueError("Hunyuan prepared layout requires a non-empty 2-D real_pos")
        if real_pos.shape[0] != int(tokens.shape[0]):
            raise ValueError(
                "Hunyuan prepared valid-length rows do not match tokenizer rows: "
                f"real_pos={real_pos.shape[0]}, tokens={int(tokens.shape[0])}"
            )
        seq_lens = [int(row[-1].item()) for row in real_pos]
        if any(seq_len <= 0 or seq_len > int(tokens.shape[1]) for seq_len in seq_lens):
            raise ValueError(
                "Hunyuan prepared valid sequence lengths must fit the padded tokenizer width: "
                f"seq_lens={seq_lens}, padded_width={int(tokens.shape[1])}"
            )

    @property
    def num_branches(self) -> int:
        return int(self.tokenizer_output.tokens.shape[0])


def build_hunyuan_diffusion_kv_requests(
    request: OmniDiffusionRequest,
    prepared_layout: HunyuanPreparedLayout,
) -> tuple[DiffusionKVRequest, ...]:
    """Build allocation-only KV requests, even when prefix caching is disabled."""

    tokenizer_output = prepared_layout.tokenizer_output
    cfg_factor = hunyuan_cfg_factor(
        prepared_layout.generated_image_info,
        resolve_hunyuan_guidance_scale(request.sampling_params),
    )
    if prepared_layout.num_branches != cfg_factor:
        raise ValueError(
            "Hunyuan tokenizer sequence count does not match CFG execution: "
            f"rows={prepared_layout.num_branches}, cfg_factor={cfg_factor}"
        )
    prefix_positions = tokenizer_output.gen_timestep_scatter_index
    real_pos = tokenizer_output.real_pos
    assert prefix_positions is not None and real_pos is not None

    def boundary_rows(field_name: str) -> list[int | None]:
        positions = getattr(tokenizer_output, field_name)
        if positions is None:
            return [None] * cfg_factor
        if not isinstance(positions, list) or not positions:
            raise ValueError(f"Hunyuan {field_name} must be a non-empty boundary list")
        # Batched tokenizer output contains a separate boundary for each CFG
        # row. Never broadcast a conditional boundary into a negative row.
        if cfg_factor == 1 and len(positions) == 1 and (positions[0] is None or type(positions[0]) is int):
            return positions
        if len(positions) == cfg_factor and all(
            isinstance(row, list) and len(row) == 1 and (row[0] is None or type(row[0]) is int) for row in positions
        ):
            return [row[0] for row in positions]
        raise ValueError(f"Hunyuan {field_name} must contain one boundary per CFG row")

    think_boundaries = boundary_rows("think_recaption_end_pos")
    uncond_boundaries = boundary_rows("uncond_cfg_start_pos")

    target_len = hunyuan_num_image_tokens(prepared_layout.generated_image_info)
    reusable_lens = [0] * cfg_factor
    if think_boundaries[0] is not None:
        reusable_lens[0] = think_boundaries[0] or 0
    if any(boundary is not None for boundary in uncond_boundaries):
        for sequence_id in range(1, cfg_factor):
            reusable_lens[sequence_id] = min(reusable_lens[0], uncond_boundaries[sequence_id] or 0)
    return tuple(
        DiffusionKVRequest(
            f"{request.request_id}/diffusion-kv/{sequence_id}",
            sequence_id=sequence_id,
            # The generated-image timestep position terminates the reusable
            # prompt/reference-image prefix for this execution row.
            prefix_len=int(prefix_row[-1].item()),
            target_len=target_len,
            seq_len=int(valid_row[-1].item()),
            # Native AR -> DiT transfer needs token IDs to describe the
            # transferred prefix.  Local paged prefix caching does not: its
            # identity is attached later by ``prepare_hunyuan_prefix_cache``.
            # Keep this conversion out of the disabled/local-only path.
            prompt_token_ids=(
                tokenizer_output.tokens[sequence_id, : reusable_lens[sequence_id]].tolist()
                if request.kv_transfer_params is not None
                else None
            ),
            # Prompt and reference-image tokens are already embedded in this
            # row's primary self-attention sequence. Hunyuan therefore has no
            # independently projected cross/joint-attention KV context.
            kv_contexts=(),
        )
        for sequence_id, (prefix_row, valid_row) in enumerate(zip(prefix_positions, real_pos))
    )


def prepare_hunyuan_prefix_cache(request: OmniDiffusionRequest) -> None:
    """Attach native MM identities / positions when Engine enables caching.

    This is model input adaptation, not a hashing framework. All reference
    inputs are hashed once, then reused by VAE/ViT subspans and CFG rows. The
    native block hasher consumes these ranges directly, without token extras.
    """

    layout = get_hunyuan_prepared_layout(request)
    if layout is None or not request.diffusion_kv_requests:
        raise ValueError("Hunyuan prefix caching requires prepared layout and KV requests")
    sampling = request.sampling_params
    _, _, _, images, _ = extract_hunyuan_prompt_inputs(
        [request.prompt], sampling.extra_args or {}, request_id=request.request_id, allow_cond_image=True
    )
    reference_digest = None
    if images:
        # Match Worker: an explicitly supplied generator overrides the seed.
        # Conditional VAE samples latents, so image bytes alone are not enough.
        if sampling.generator is not None:
            generators = sampling.generator if isinstance(sampling.generator, list) else [sampling.generator]
            random_state: object = (
                "generator-state",
                tuple((str(generator.device), generator.get_state()) for generator in generators),
            )
        elif sampling.seed is not None:
            random_state = ("seed", int(sampling.seed), str(sampling.generator_device or "worker-default"))
        else:
            random_state = ("request-local-random-state", request.request_id)
        # Preserve the conservative whole-reference-set policy. Splitting
        # independent image identities also needs the VAE RNG sequence modeled.
        reference_digest = hash_prefix_cache_value(
            (
                "hunyuan-reference-v1",
                [[joint_image_info_to_payload(image) for image in row] for row in images],
                random_state,
            )
        )

    namespace = get_cache_namespace("hunyuan-image3-primary-v4", sampling)
    output = layout.tokenizer_output
    joint_rows = output.joint_image_slices or [[] for _ in range(layout.num_branches)]
    prepared = []
    for row in request.diffusion_kv_requests:
        if row.block_hashes or row.num_computed_tokens:
            raise ValueError("Hunyuan cache inputs must be prepared before KV execution")
        joint_spans = [(int(span.start or 0), int(span.stop or 0)) for span in joint_rows[row.sequence_id]]
        features = []
        for image_slice, (height, width) in layout.rope_image_info[row.sequence_id]:
            start = int(image_slice.start or 0)
            end = int(image_slice.stop or output.tokens.shape[1])
            if start >= row.prefix_len or end <= 0:
                continue
            # Joint spans enclose VAE + separator + ViT; RoPE lists subspans.
            is_reference = any(joint_start <= start and end <= joint_end for joint_start, joint_end in joint_spans)
            if is_reference and reference_digest is None:
                raise ValueError("Hunyuan reference-image prefix is missing its canonical inputs")
            identifier = hash_prefix_cache_value(
                (
                    "hunyuan-image-span-v2",
                    start,
                    end,
                    int(height),
                    int(width),
                    is_reference,
                    reference_digest if is_reference else None,
                )
            ).hex()
            features.append(
                MultiModalFeatureSpec(
                    data=None,
                    modality="image",
                    identifier=identifier,
                    mm_position=PlaceholderRange(offset=start, length=end - start),
                )
            )
        # Raw model inputs remain in request.prompt/prepared_layout; data=None
        # here does not imply an encoder hit. These are Scheduler-only KV keys.
        tokens = tuple(output.tokens[row.sequence_id, : row.prefix_len].tolist())
        prepared.append((row, tokens, sorted(features, key=lambda feature: feature.mm_position.offset)))

    for row, tokens, features in prepared:
        row.cache_token_ids = tokens
        row.mm_features = features
        row.cache_namespace = namespace


def get_hunyuan_prepared_layout(source: Any) -> HunyuanPreparedLayout | None:
    prepared_layout = getattr(source, "prepared_layout", None)
    if prepared_layout is None:
        return None
    if not isinstance(prepared_layout, HunyuanPreparedLayout):
        raise TypeError(
            f"HunyuanImage3 expected prepared_layout to be HunyuanPreparedLayout, got {type(prepared_layout).__name__}"
        )
    return prepared_layout


def native_kv_covers_cond_images(output: TokenizerEncodeOutput, computed_tokens: tuple[int, ...]) -> bool:
    """Skip image encoding only when every CFG row already contains its image KV."""
    image_slices = output.joint_image_slices
    if not image_slices or len(image_slices) != len(computed_tokens):
        return False
    return all(
        slices and all(image_slice.stop <= computed for image_slice in slices)
        for slices, computed in zip(image_slices, computed_tokens, strict=True)
    )