Fix hunyuan resolve stop token ids (#3896)

Signed-off-by: zzh <943967662@qq.com>

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

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["</think>"]
+    end_of_recaption_id = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["</recaption>"]
+
+    # 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["<img_ratio_0>"]
+    end = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_32>"]
+    other_start = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_33>"]
+    other_end = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_36>"]
+    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["<answer>"]
+    assert resolve_stop_token_ids(task="i2t", bot_task=None, tokenizer=tok, image_size="1024x1024") == [answer_id]
+
 
 @pytest.mark.parametrize(
     "task,bot_task",

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,
+            "<img>": 2,
+            "<recaption>": 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"
+    )

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["</think>"]
+            end_of_recaption_id = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["</recaption>"]
+            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: <img_ratio_0> .. <img_ratio_32>.
         start = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_0>"]
         end = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_32>"]

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