[Bugfix]Fix HunyuanImage3 conditional image and prompt kwargs alignment (#4145)

Signed-off-by: reandeng <reandeng@tencent.com>
Signed-off-by: Yaegaki1Erika <70182590+Yaegaki1Erika@users.noreply.github.com>
Co-authored-by: reandeng <reandeng@tencent.com>

tests/entrypoints/openai_api/test_serving_chat_multistage_generation.py

@@ -61,7 +61,7 @@ def test_build_multistage_generation_inputs_applies_stage_specific_overrides(ser
     assert engine_prompt["prompt"] == "draw a robot"
     assert engine_prompt["modalities"] == ["img2img"]
     assert engine_prompt["negative_prompt"] == "blurry"
-    assert engine_prompt["mm_processor_kwargs"] == {"target_h": 768, "target_w": 1024}
+    assert engine_prompt["mm_processor_kwargs"] == {"target_h": 768, "target_w": 1024, "vae_generator_seed": 0}
     assert engine_prompt["multi_modal_data"]["img2img"].size == (24, 24)
 
     assert len(sampling_params_list) == 3

vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py

@@ -626,7 +626,7 @@ class HunyuanImage3Pipeline(
             rope_image_info.append(list(zip(image_slices, image_shapes)))
         return rope_image_info
 
-    def vae_encode(self, image, cfg_factor=1):
+    def vae_encode(self, image, cfg_factor=1, generator=None):
         config = self.vae.config
 
         if image.ndim == 3:
@@ -641,11 +641,7 @@ class HunyuanImage3Pipeline(
             if isinstance(vae_encode_result, torch.Tensor):
                 latents = vae_encode_result
             else:
-                # Match HunyuanImage-3's cond encode path: sample the
-                # posterior, but use a fixed generator so repeated online
-                # requests are deterministic.
-                _cond_vae_gen = torch.Generator(device=image.device).manual_seed(0)
-                latents = vae_encode_result.latent_dist.sample(_cond_vae_gen)
+                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:
@@ -668,14 +664,17 @@ class HunyuanImage3Pipeline(
         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 cond_image_info_list in batch_cond_image_info_list:
+        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))
@@ -869,7 +868,7 @@ class HunyuanImage3Pipeline(
         has_ar_kv = kwargs.get("ar_kv_data")
         if batch_cond_image_info is not None and len(batch_cond_image_info[0]) > 0 and not has_ar_kv:
             cond_vae_images, cond_timestep, cond_vit_images = self._encode_cond_image(
-                batch_cond_image_info, cfg_factor[mode]
+                batch_cond_image_info, cfg_factor[mode], generator=generator
             )
             vit_kwargs = {"spatial_shapes": [], "attention_mask": []}
             for cond_image_info in batch_cond_image_info:
@@ -1393,18 +1392,23 @@ class HunyuanImage3Pipeline(
         **kwargs,
     ) -> DiffusionOutput:
         extra_args = getattr(getattr(req, "sampling_params", None), "extra_args", {}) or {}
+        bot_task = extra_args.get("bot_task")
         use_system_prompt = extra_args.get("use_system_prompt")
         system_prompt = extra_args.get("system_prompt")
-        # Fall back to per-prompt use_system_prompt forwarded by ar2diffusion
+        # Fall back to per-prompt prompt kwargs forwarded by ar2diffusion
         if req.prompts:
             first_prompt = req.prompts[0]
             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")
+        if isinstance(bot_task, str) and bot_task.lower() == "none":
+            bot_task = None
         if use_system_prompt is not None:
-            system_prompt = get_system_prompt(use_system_prompt, "image", system_prompt)
+            system_prompt = get_system_prompt(use_system_prompt, bot_task or "image", 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 req.prompts] or prompt
 
@@ -1451,6 +1455,18 @@ class HunyuanImage3Pipeline(
         # ---- AR KV Reuse: extract injected KV from request ----
         ar_kv_kwargs = self._extract_ar_kv_from_request(req)
 
+        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)}."
+            )
+
         model_inputs = self.prepare_model_inputs(
             prompt=prompt,
             cot_text=cot_text,
@@ -1461,6 +1477,7 @@ class HunyuanImage3Pipeline(
             num_inference_steps=num_inference_steps,
             guidance_scale=guidance_scale,
             batch_cond_image_info=batch_cond_image_info,
+            bot_task=tokenizer_bot_task or "auto",
             **ar_kv_kwargs,
         )
 

vllm_omni/entrypoints/openai/api_server.py

@@ -2015,6 +2015,16 @@ async def edit_images(
         _update_if_not_none(gen_params, "layers", layers)
         _update_if_not_none(gen_params, "resolution", resolution)
 
+        extra_args = dict(getattr(gen_params, "extra_args", {}) or {})
+        if bot_task is not None:
+            extra_args["bot_task"] = bot_task
+        if sys_type is not None:
+            extra_args["use_system_prompt"] = sys_type
+        if system_prompt is not None:
+            extra_args["system_prompt"] = system_prompt
+        if extra_args:
+            gen_params.extra_args = extra_args
+
         # 4. Generate images
         request_id = f"img_edit-{random_uuid()}"
         raw_request.state.request_metadata = RequestResponseMetadata(request_id=request_id)

vllm_omni/entrypoints/openai/serving_chat.py

@@ -2599,6 +2599,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
             mm_processor_kwargs["target_h"] = height
         if width is not None:
             mm_processor_kwargs["target_w"] = width
+        if seed is not None and engine_prompt_data is not None:
+            mm_processor_kwargs["vae_generator_seed"] = int(seed)
         if mm_processor_kwargs:
             engine_prompt["mm_processor_kwargs"] = mm_processor_kwargs
         if engine_prompt_data is not None:

vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py

@@ -1043,6 +1043,9 @@ class HunyuanImage3MultiModalProcessor(BaseMultiModalProcessor[HunyuanImage3Proc
         images = mm_data.get("images", [])
         logger.debug(f"process image count: {len(images)}")
         batch_feature = image_processor(prompt, images, **tok_kwargs)
+        vae_generator_seed = mm_kwargs.get("vae_generator_seed")
+        if vae_generator_seed is not None and images:
+            batch_feature["vae_generator_seed"] = torch.full((len(images),), int(vae_generator_seed), dtype=torch.long)
         return batch_feature
 
     def _hf_processor_applies_updates(
@@ -1080,6 +1083,8 @@ class HunyuanImage3MultiModalProcessor(BaseMultiModalProcessor[HunyuanImage3Proc
             config["base_size"] = MultiModalFieldConfig.batched("image")
         if "ratio_index" in hf_inputs:
             config["ratio_index"] = MultiModalFieldConfig.batched("image")
+        if "vae_generator_seed" in hf_inputs:
+            config["vae_generator_seed"] = MultiModalFieldConfig.batched("image")
         return config
 
     def _get_prompt_updates(
@@ -1732,6 +1737,7 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
         # we reconstruct per-image shapes from vae_token_grid_hw below.
         kwargs.pop("vae_pixel_size", None)
         vae_token_grid_hw = kwargs.pop("vae_token_grid_hw", None)
+        vae_generator_seed = kwargs.pop("vae_generator_seed", None)
 
         if vit_pixel_values is None or vae_pixel_values is None:
             return None
@@ -1771,6 +1777,7 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
                 "vit_spatial_shapes": vit_spatial_shapes,
                 "vae_pixel_values": vae_image_list,
                 "vae_token_grid_hw": vae_token_grid_hw,
+                "vae_generator_seed": vae_generator_seed,
             },
         )
 
@@ -1778,6 +1785,7 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
         self,
         images: torch.Tensor,
         cfg_factor: int = 1,
+        generator: torch.Generator | None = None,
     ) -> tuple[tuple[torch.Tensor, torch.Tensor], ...]:
         """
         Encode images through VAE encoder.
@@ -1797,11 +1805,7 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
             images = images.to(dtype=self.vae.dtype)
 
         vae_encode_result = self.vae.encode(images)
-        # Match HunyuanImage-3's cond encode path: sample the posterior, but
-        # use a fixed generator so online requests do not consume the global
-        # RNG and drift across a long-running server.
-        _cond_vae_gen = torch.Generator(device=images.device).manual_seed(0)
-        latents = vae_encode_result.latent_dist.sample(_cond_vae_gen)
+        latents = vae_encode_result.latent_dist.sample(generator)
 
         # Apply shift and scaling factors if present
         if hasattr(config, "shift_factor") and config.shift_factor:
@@ -1933,6 +1937,12 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
         vit_pixel_attention_mask = pixel_values["vit_pixel_attention_mask"]
         vit_spatial_shapes = pixel_values["vit_spatial_shapes"]
         vae_pixel_values = pixel_values["vae_pixel_values"]
+        vae_generator_seed = pixel_values.get("vae_generator_seed")
+        vae_generator = None
+        if vae_generator_seed is not None and vae_generator_seed.numel() > 0:
+            vae_generator = torch.Generator(device=vae_pixel_values[0].device).manual_seed(
+                int(vae_generator_seed.reshape(-1)[0].item())
+            )
 
         # Perform ViT encoding
         vit_embeddings = self._vit_encode(vit_pixel_values, vit_pixel_attention_mask, vit_spatial_shapes)
@@ -1941,7 +1951,7 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
         # `reso_group` bucket so shapes are ragged across the image-batch dim.
         vae_token_embeddings = []
         for vae_image_i in vae_pixel_values:
-            t_i, latents_i = self._vae_encode(vae_image_i.unsqueeze(0), vae_cfg_factor)
+            t_i, latents_i = self._vae_encode(vae_image_i.unsqueeze(0), vae_cfg_factor, generator=vae_generator)
             t_emb = self.time_embed(t_i[0])
             vae_tokens, _, _ = self.patch_embed(latents_i, t_emb)
             vae_token_embeddings.append(vae_tokens)