From ffd9771c54d23ccf7383c86c5c7a776f29ba588f Mon Sep 17 00:00:00 2001 From: zzh <67720593+BLANKETusers@users.noreply.github.com> Date: Sun, 31 May 2026 07:47:36 +0800 Subject: [PATCH] Fix hunyuan resolve stop token ids (#3896) Signed-off-by: zzh <943967662@qq.com> --- .../hunyuan_image3/end2end.py | 20 +++- .../hunyuan_image3/test_prompt_utils.py | 52 +++++++++ ...test_serving_chat_multistage_generation.py | 101 ++++++++++++++++++ .../models/hunyuan_image3/prompt_utils.py | 9 ++ vllm_omni/entrypoints/openai/serving_chat.py | 20 ++++ 5 files changed, 199 insertions(+), 3 deletions(-) diff --git a/examples/offline_inference/hunyuan_image3/end2end.py b/examples/offline_inference/hunyuan_image3/end2end.py index 537379b89..e8cd8a36f 100644 --- a/examples/offline_inference/hunyuan_image3/end2end.py +++ b/examples/offline_inference/hunyuan_image3/end2end.py @@ -61,8 +61,8 @@ def parse_args(): parser.add_argument("--steps", type=int, default=50, help="Number of inference steps.") parser.add_argument("--guidance-scale", type=float, default=5.0, help="Classifier-free guidance scale.") parser.add_argument("--seed", type=int, default=42, help="Random seed.") - parser.add_argument("--height", type=int, default=1024, help="Output image height.") - parser.add_argument("--width", type=int, default=1024, help="Output image width.") + parser.add_argument("--height", type=int, default=None, help="Output image height.") + parser.add_argument("--width", type=int, default=None, help="Output image width.") parser.add_argument("--vae-use-tiling", action="store_true", help="Enable VAE tiling.") parser.add_argument( "--bot-task", @@ -229,7 +229,21 @@ def main(): from vllm_omni.inputs.data import OmniDiffusionSamplingParams - ar_stop_token_ids = resolve_stop_token_ids(task=task, bot_task=bot_task, tokenizer=tokenizer) + if (args.height is None) != (args.width is None): + raise ValueError("--height and --width must both be specified or both omitted.") + user_specified_size = args.height is not None and args.width is not None + if args.modality in ("img2text", "text2text"): + ar_image_size = "auto" + elif user_specified_size: + ar_image_size = f"{args.width}x{args.height}" + else: + ar_image_size = None + ar_stop_token_ids = resolve_stop_token_ids( + task=task, bot_task=bot_task, tokenizer=tokenizer, image_size=ar_image_size + ) + print( + f"[AR Config] task={task}, bot_task={bot_task}, image_size={ar_image_size}, stop_token_ids={ar_stop_token_ids}" + ) for sp in params_list: if isinstance(sp, OmniDiffusionSamplingParams): sp.num_inference_steps = args.steps diff --git a/tests/diffusion/models/hunyuan_image3/test_prompt_utils.py b/tests/diffusion/models/hunyuan_image3/test_prompt_utils.py index 7c3256eee..7cb7ea83a 100644 --- a/tests/diffusion/models/hunyuan_image3/test_prompt_utils.py +++ b/tests/diffusion/models/hunyuan_image3/test_prompt_utils.py @@ -124,6 +124,58 @@ def test_resolve_stop_token_ids_image_tasks_stop_on_ratio_range(): assert resolve_stop_token_ids(task="i2t", bot_task=None, tokenizer=tok) == [answer_id] assert resolve_stop_token_ids(task="t2t", bot_task=None, tokenizer=tok) == [answer_id] + # image_size="auto" is equivalent to None (need_ratio=True). + for bot in ("think", "recaption", "think_recaption", "vanilla"): + assert resolve_stop_token_ids(task="t2i", bot_task=bot, tokenizer=tok, image_size="auto") == expected + assert resolve_stop_token_ids(task="it2i", bot_task=bot, tokenizer=tok, image_size="auto") == expected + + +def test_resolve_stop_token_ids_explicit_image_size_stops_on_terminator(): + """When image_size is an explicit "WxH" string, need_ratio=False: AR + already knows the size and stops at the cot terminator instead of the + ratio range. + + Mirrors upstream logic where ``need_ratio`` is false when the user + supplies a target size, so ``final_stop_tokens`` is set to the + recaption/think terminator (not the ratio tokens). + """ + tok = FakeTokenizer() + + end_of_think_id = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""] + end_of_recaption_id = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""] + + # think_recaption / recaption -> [end_of_recaption] + for bot in ("think_recaption", "recaption"): + for task in ("t2i", "it2i"): + result = resolve_stop_token_ids(task=task, bot_task=bot, tokenizer=tok, image_size="1024x1024") + assert result == [end_of_recaption_id], ( + f"task={task}, bot_task={bot}, image_size='1024x1024': expected [{end_of_recaption_id}], got {result}" + ) + + # think -> [end_of_think, end_of_recaption] + for task in ("t2i", "it2i"): + result = resolve_stop_token_ids(task=task, bot_task="think", tokenizer=tok, image_size="1024x1024") + assert result == [end_of_think_id, end_of_recaption_id], ( + f"task={task}, bot_task='think', image_size='1024x1024': " + f"expected [{end_of_think_id}, {end_of_recaption_id}], got {result}" + ) + + # vanilla -> no early return, falls through to ratio range + start = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""] + end = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""] + other_start = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""] + other_end = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""] + expected_ratio = list(range(start, end + 1)) + list(range(other_start, other_end + 1)) + for task in ("t2i", "it2i"): + result = resolve_stop_token_ids(task=task, bot_task="vanilla", tokenizer=tok, image_size="1024x1024") + assert result == expected_ratio, ( + f"task={task}, bot_task='vanilla', image_size='1024x1024': expected ratio range, got {result}" + ) + + # Text tasks are unaffected by image_size. + answer_id = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""] + assert resolve_stop_token_ids(task="i2t", bot_task=None, tokenizer=tok, image_size="1024x1024") == [answer_id] + @pytest.mark.parametrize( "task,bot_task", diff --git a/tests/entrypoints/openai_api/test_serving_chat_multistage_generation.py b/tests/entrypoints/openai_api/test_serving_chat_multistage_generation.py index 4b63588ba..c8bb050e7 100644 --- a/tests/entrypoints/openai_api/test_serving_chat_multistage_generation.py +++ b/tests/entrypoints/openai_api/test_serving_chat_multistage_generation.py @@ -363,3 +363,104 @@ def test_build_multistage_generation_inputs_custom_system_prompt(serving_chat): f"custom system_prompt content must reach the rendered prompt; " f"marker {marker!r} not found in prompt of length {len(out['prompt'])}" ) + + +def test_build_multistage_generation_inputs_sets_ar_stop_token_ids_with_explicit_size(serving_chat): + """When height+width are provided with bot_task, the AR (llm) stage + must receive stop_token_ids from resolve_stop_token_ids so AR stops + at the correct terminator instead of generating until max_tokens. + + Without this, AR would generate past the cot boundary and produce + garbage that the DiT bridge cannot parse. + """ + from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat + + class FakeTokenizer: + SPECIAL = { + "<|startoftext|>": 1, + "": 2, + "": 4, + } + + def convert_tokens_to_ids(self, tok: str) -> int: + return self.SPECIAL.get(tok, 0) + + def encode(self, text: str, add_special_tokens: bool = False) -> list[int]: + return list(range(100, 100 + len(text))) + + engine = SimpleNamespace( + stage_configs=[ + SimpleNamespace(stage_type="llm", is_comprehension=True), + SimpleNamespace(stage_type="diffusion", is_comprehension=False), + ], + default_sampling_params_list=[ + SamplingParams(temperature=0.0), + OmniDiffusionSamplingParams(), + ], + ) + images = [Image.new("RGB", (32, 32), color="red")] + + # With height+width and bot_task=think, ar_image_size="1024x768" -> + # need_ratio=False -> stop_token_ids=[end_of_think, end_of_recaption] + _, sampling_params_list = OmniOpenAIServingChat._build_multistage_generation_inputs( + serving_chat, + engine=engine, + prompt="draw a cat", + extra_body={"bot_task": "think"}, + reference_images=images, + gen_params=OmniDiffusionSamplingParams(height=768, width=1024), + tokenizer=FakeTokenizer(), + ) + + # AR stage (index 0) must have stop_token_ids set. + ar_params = sampling_params_list[0] + assert ar_params.stop_token_ids is not None, ( + "AR stage must have stop_token_ids set when height+width and bot_task are provided" + ) + assert len(ar_params.stop_token_ids) > 0, "stop_token_ids must be non-empty" + + # Diffusion stage (index 1) must NOT have stop_token_ids set. + diff_params = sampling_params_list[1] + assert getattr(diff_params, "stop_token_ids", None) is None, "Diffusion stage must not have stop_token_ids set" + + +def test_build_multistage_generation_inputs_no_stop_token_ids_without_size(serving_chat): + """Without height+width, ar_image_size=None -> need_ratio=True -> + stop_token_ids is set to the ratio range (AR predicts ratio). + This still assigns stop_token_ids, but the set is the ratio range, + not the terminator. + + Without bot_task at all, ar_stop_token_ids stays None and + stop_token_ids is not set on any stage. + """ + from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat + + engine = SimpleNamespace( + stage_configs=[ + SimpleNamespace(stage_type="llm", is_comprehension=True), + SimpleNamespace(stage_type="diffusion", is_comprehension=False), + ], + default_sampling_params_list=[ + SamplingParams(temperature=0.0), + OmniDiffusionSamplingParams(), + ], + ) + images = [Image.new("RGB", (32, 32), color="red")] + + # No height/width, no bot_task -> ar_stop_token_ids=None -> + # stop_token_ids not set on any stage. + _, sampling_params_list = OmniOpenAIServingChat._build_multistage_generation_inputs( + serving_chat, + engine=engine, + prompt="draw a cat", + extra_body={}, + reference_images=images, + gen_params=OmniDiffusionSamplingParams(), + ) + + # SamplingParams defaults stop_token_ids=[], not None. + # The key contract: it was NOT set by resolve_stop_token_ids, + # so it stays as the SamplingParams default (empty list). + assert sampling_params_list[0].stop_token_ids == [], ( + "Without bot_task, AR stage stop_token_ids must be the default empty list" + ) diff --git a/vllm_omni/diffusion/models/hunyuan_image3/prompt_utils.py b/vllm_omni/diffusion/models/hunyuan_image3/prompt_utils.py index f054fef26..b64fe2e97 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/prompt_utils.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/prompt_utils.py @@ -146,6 +146,7 @@ def resolve_stop_token_ids( task: str = "it2i", bot_task: str | None | _DefaultBotTask = _DEFAULT_BOT_TASK, tokenizer: Any | None = None, + image_size: str | None = None, ) -> list[int]: """AR stop-token ids for a given (task, bot_task) generation request. @@ -171,6 +172,14 @@ def resolve_stop_token_ids( if bot_task not in _BOT_TASK_PRESETS: raise ValueError(f"Unknown bot_task {bot_task!r}. Choose from: {available_bot_tasks()}") if task in ("it2i", "t2i"): + need_ratio = image_size is None or image_size == "auto" + if not need_ratio: + end_of_think_id = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""] + end_of_recaption_id = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""] + if bot_task in ("think_recaption", "recaption"): + return [end_of_recaption_id] + if bot_task == "think": + return [end_of_think_id, end_of_recaption_id] # Main ratio range: .. . start = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""] end = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""] diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 06ccfcd70..f98732984 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -2335,11 +2335,13 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): prompt_token_ids: list[int] | None = None system_prompt_type: str | None = None build_kwargs: dict[str, Any] = {} + ar_stop_token_ids: list[int] | None = None if bot_task is not None or use_system_prompt is not None or custom_system_prompt is not None: from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import ( build_prompt, build_prompt_tokens, + resolve_stop_token_ids, ) build_kwargs: dict[str, Any] = { @@ -2368,6 +2370,19 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): engine_prompt_data = {"image": reference_images[0]} modalities = ["image"] + ar_task = "it2i" if reference_images else "t2i" + # ar_image_size: None -> need_ratio=True (AR predicts ratio); + # explicit size -> need_ratio=False (AR stops at terminator). + ar_image_size: str | None = None + if height is not None and width is not None: + ar_image_size = f"{width}x{height}" + ar_stop_token_ids = resolve_stop_token_ids( + task=ar_task, + bot_task=bot_task, + tokenizer=tokenizer, + image_size=ar_image_size, + ) + engine_prompt: OmniTextPrompt = {"prompt": prompt} if prompt_token_ids is not None: engine_prompt["prompt_token_ids"] = prompt_token_ids @@ -2413,6 +2428,11 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): stage_type = get_stage_type(stage_cfg) default_stage_params = sampling_params_list[idx] + # AR stop tokens: use stage_type=="llm" instead of comprehension_idx + # (None for DictConfig where is_comprehension is nested in engine_args). + if stage_type == "llm" and ar_stop_token_ids is not None: + default_stage_params.stop_token_ids = ar_stop_token_ids + if ( comprehension_idx is not None and idx == comprehension_idx