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 c8bb050e7..507018ef7 100644 --- a/tests/entrypoints/openai_api/test_serving_chat_multistage_generation.py +++ b/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 diff --git a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py index 9cc2aafe5..988b91920 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py +++ b/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, ) diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py old mode 100644 new mode 100755 index 1be7aceba..7c6d7b9e5 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/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) diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py old mode 100644 new mode 100755 index c7b24e11f..9bac6c9ba --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/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: diff --git a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py old mode 100644 new mode 100755 index 7770bee95..7b975796f --- a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py +++ b/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)