[image][bugfix] Fix HunyuanImage3 bridge request batching (#4849)
Signed-off-by: Celeste-jq <591998922@qq.com>
This commit is contained in:
Jiaqian Liu authored and GitHub committed
1 parent
431a7c71ac
commit
f00caf35c4
3 files changed
+151
-121
No files matched your search
@@ -524,11 +524,11 @@ class TestAr2DiffusionIntegration:
|
||||
self,
|
||||
token_ids: list[int],
|
||||
prompt_kwargs: dict[str, Any] | None = None,
|
||||
) -> tuple[Any, Any]:
|
||||
) -> dict[str, Any] | None:
|
||||
"""Invoke ``ar2diffusion`` with a single fake AR output.
|
||||
|
||||
Returns ``(diffusion_inputs, ar_predicted_sizes)`` so callers can
|
||||
assert on the resolved height/width.
|
||||
Returns the diffusion request payload so callers can assert on the
|
||||
resolved height/width.
|
||||
"""
|
||||
from vllm_omni.model_executor.stage_input_processors.hunyuan_image3 import (
|
||||
ar2diffusion,
|
||||
@@ -601,8 +601,8 @@ class TestAr2DiffusionIntegration:
|
||||
token_ids=[150000, 150001, token_id],
|
||||
prompt_kwargs={"height": 1024, "width": 1024, "image_base_size": 1024},
|
||||
)
|
||||
assert len(result) == 1
|
||||
diffusion_input = result[0]
|
||||
assert result is not None
|
||||
diffusion_input = result
|
||||
assert diffusion_input["height"] == expected_h, (
|
||||
f"ratio_idx={ratio_idx}: expected height {expected_h}, got {diffusion_input['height']}"
|
||||
)
|
||||
@@ -617,8 +617,8 @@ class TestAr2DiffusionIntegration:
|
||||
token_ids=[150000, 150001], # no ratio tokens
|
||||
prompt_kwargs={"height": 768, "width": 1024, "image_base_size": 1024},
|
||||
)
|
||||
assert len(result) == 1
|
||||
diffusion_input = result[0]
|
||||
assert result is not None
|
||||
diffusion_input = result
|
||||
assert diffusion_input["height"] == 768
|
||||
assert diffusion_input["width"] == 1024
|
||||
|
||||
@@ -633,9 +633,9 @@ class TestAr2DiffusionIntegration:
|
||||
token_ids=[150000],
|
||||
prompt_kwargs={"height": 512, "width": 512, "image_base_size": 1024},
|
||||
)
|
||||
assert len(result) == 1
|
||||
assert result[0]["height"] == 512
|
||||
assert result[0]["width"] == 512
|
||||
assert result is not None
|
||||
assert result["height"] == 512
|
||||
assert result["width"] == 512
|
||||
|
||||
# -- cache behaviour -----------------------------------------------------
|
||||
|
||||
|
||||
@@ -45,6 +45,33 @@ def test_truncate_at_cot_end_strips_tail_after_recaption_marker():
|
||||
assert text == "body text</recaption>"
|
||||
|
||||
|
||||
def test_ar2diffusion_returns_one_request_payload_for_request_level_batching():
|
||||
result = ar2diffusion(
|
||||
[_source_output([100], text="thought")],
|
||||
prompt={"prompt": "edit"},
|
||||
)
|
||||
|
||||
assert isinstance(result, dict)
|
||||
assert result["prompt"] == "edit"
|
||||
|
||||
|
||||
def test_ar2diffusion_uses_parent_output_when_companions_are_present():
|
||||
result = ar2diffusion(
|
||||
[
|
||||
_source_output([100], text="parent thought"),
|
||||
_source_output([200], text="companion thought"),
|
||||
],
|
||||
prompt={"prompt": "edit"},
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert result["extra"]["ar_generated_text"] == "parent thought"
|
||||
|
||||
|
||||
def test_ar2diffusion_returns_none_without_parent_output():
|
||||
assert ar2diffusion([], prompt={"prompt": "edit"}) is None
|
||||
|
||||
|
||||
def test_ar2diffusion_applies_ratio_and_truncates_tail_without_tokenizer(monkeypatch: pytest.MonkeyPatch):
|
||||
real_import = builtins.__import__
|
||||
|
||||
@@ -67,10 +94,10 @@ def test_ar2diffusion_applies_ratio_and_truncates_tail_without_tokenizer(monkeyp
|
||||
prompt=[{"prompt": "edit", "height": 64, "width": 64}],
|
||||
)
|
||||
|
||||
assert len(result) == 1
|
||||
assert (result[0]["height"], result[0]["width"]) == (512, 2048)
|
||||
assert result[0]["extra"]["ar_generated_text"] == "decoded without special tokens"
|
||||
assert "ar_token_ids" not in result[0]["extra"]
|
||||
assert result is not None
|
||||
assert (result["height"], result["width"]) == (512, 2048)
|
||||
assert result["extra"]["ar_generated_text"] == "decoded without special tokens"
|
||||
assert "ar_token_ids" not in result["extra"]
|
||||
|
||||
|
||||
def test_ar2diffusion_forwards_custom_system_prompt_body():
|
||||
@@ -88,5 +115,6 @@ def test_ar2diffusion_forwards_custom_system_prompt_body():
|
||||
],
|
||||
)
|
||||
|
||||
assert result[0]["use_system_prompt"] == "custom"
|
||||
assert result[0]["system_prompt"] == marker
|
||||
assert result is not None
|
||||
assert result["use_system_prompt"] == "custom"
|
||||
assert result["system_prompt"] == marker
|
||||
@@ -95,128 +95,130 @@ def ar2diffusion(
|
||||
source_outputs: list[Any],
|
||||
prompt: OmniTokensPrompt | TextPrompt | list | None = None,
|
||||
requires_multimodal_data: bool = False,
|
||||
) -> list[dict[str, Any]]:
|
||||
) -> dict[str, Any] | None:
|
||||
"""Process AR stage outputs to create Diffusion stage inputs.
|
||||
|
||||
HunyuanImage3 produces one downstream diffusion request per parent AR
|
||||
request. ``source_outputs`` may include companion outputs, but the parent
|
||||
output is expected first: the orchestrator bundles diffusion inputs as
|
||||
``[parent_output, *cfg_companion_outputs]`` before invoking this bridge.
|
||||
|
||||
Args:
|
||||
prompt: Original user prompt (may contain multimodal data).
|
||||
requires_multimodal_data: Whether to forward multimodal data.
|
||||
|
||||
Returns:
|
||||
List of dicts, each consumable by the HunyuanImage3 diffusion pipeline.
|
||||
One prompt dict consumable by the HunyuanImage3 diffusion pipeline, or
|
||||
``None`` when no parent output is available.
|
||||
"""
|
||||
ar_outputs = source_outputs
|
||||
diffusion_inputs = []
|
||||
if not source_outputs:
|
||||
return None
|
||||
|
||||
# Normalize prompt to list
|
||||
if not isinstance(prompt, list):
|
||||
prompt = [prompt] if prompt is not None else [{}]
|
||||
# The orchestrator constructs this list as [parent, *cfg_companions].
|
||||
ar_output = source_outputs[0]
|
||||
output = ar_output.outputs[0]
|
||||
generated_token_ids = output.cumulative_token_ids
|
||||
# Prefer cumulative_text, fallback to text if aggregation dropped it
|
||||
generated_text = getattr(output, "cumulative_text", None) or getattr(output, "text", "") or ""
|
||||
|
||||
for i, ar_output in enumerate(ar_outputs):
|
||||
output = ar_output.outputs[0]
|
||||
generated_token_ids = output.cumulative_token_ids
|
||||
# Prefer cumulative_text, fallback to text if aggregation dropped it
|
||||
generated_text = getattr(output, "cumulative_text", None) or getattr(output, "text", "") or ""
|
||||
if isinstance(prompt, list):
|
||||
original_prompt = prompt[0] if prompt else {}
|
||||
elif prompt is not None:
|
||||
original_prompt = prompt
|
||||
else:
|
||||
original_prompt = {}
|
||||
if isinstance(original_prompt, dict):
|
||||
pass
|
||||
elif hasattr(original_prompt, "_asdict"):
|
||||
original_prompt = original_prompt._asdict()
|
||||
elif hasattr(original_prompt, "__dict__"):
|
||||
original_prompt = vars(original_prompt)
|
||||
else:
|
||||
original_prompt = {}
|
||||
|
||||
# Get original prompt info
|
||||
original_prompt = prompt[i] if i < len(prompt) else {}
|
||||
if isinstance(original_prompt, dict):
|
||||
pass
|
||||
elif hasattr(original_prompt, "_asdict"):
|
||||
original_prompt = original_prompt._asdict()
|
||||
elif hasattr(original_prompt, "__dict__"):
|
||||
original_prompt = vars(original_prompt)
|
||||
else:
|
||||
original_prompt = {}
|
||||
height = original_prompt.get("height", 1024)
|
||||
width = original_prompt.get("width", 1024)
|
||||
text_prompt = original_prompt.get("prompt", "")
|
||||
use_system_prompt = original_prompt.get("use_system_prompt")
|
||||
custom_system_prompt = original_prompt.get("system_prompt")
|
||||
|
||||
height = original_prompt.get("height", 1024)
|
||||
width = original_prompt.get("width", 1024)
|
||||
text_prompt = original_prompt.get("prompt", "")
|
||||
use_system_prompt = original_prompt.get("use_system_prompt")
|
||||
custom_system_prompt = original_prompt.get("system_prompt")
|
||||
# Prefer the AR's predicted output aspect (`<img_size_*><img_ratio_*>`
|
||||
# tail emitted by `HunyuanImage3ForCausalMM.sample` under the
|
||||
# ratio-restriction logits processor) over the carried-through
|
||||
# height/width, which the serving layer fills with the first
|
||||
# reference image's bucket and so collapses non-square targets to
|
||||
# square in the multi-image / mismatched-aspect case. Mirrors the
|
||||
# official upstream where `reso_group[ratio_index]` is the
|
||||
# canonical source of the diffusion target shape.
|
||||
ratio_idx = _extract_ratio_index(generated_token_ids)
|
||||
ar_predicted = False
|
||||
if ratio_idx is not None:
|
||||
base_size = int(original_prompt.get("image_base_size", 1024))
|
||||
reso_group: ResolutionGroup = get_cached_resolution_group(base_size=base_size)
|
||||
try:
|
||||
reso: Resolution = reso_group[ratio_idx]
|
||||
height = reso.height
|
||||
width = reso.width
|
||||
ar_predicted = True
|
||||
except IndexError:
|
||||
logger.warning(
|
||||
"[ar2diffusion] Request 0: ratio_index=%d out of range [0,%d), keeping prompt size %dx%d",
|
||||
ratio_idx,
|
||||
len(reso_group),
|
||||
height,
|
||||
width,
|
||||
)
|
||||
|
||||
# Prefer the AR's predicted output aspect (`<img_size_*><img_ratio_*>`
|
||||
# tail emitted by `HunyuanImage3ForCausalMM.sample` under the
|
||||
# ratio-restriction logits processor) over the carried-through
|
||||
# height/width, which the serving layer fills with the first
|
||||
# reference image's bucket and so collapses non-square targets to
|
||||
# square in the multi-image / mismatched-aspect case. Mirrors the
|
||||
# official upstream where `reso_group[ratio_index]` is the
|
||||
# canonical source of the diffusion target shape.
|
||||
ratio_idx = _extract_ratio_index(generated_token_ids)
|
||||
ar_predicted = False
|
||||
if ratio_idx is not None:
|
||||
base_size = int(original_prompt.get("image_base_size", 1024))
|
||||
reso_group: ResolutionGroup = get_cached_resolution_group(base_size=base_size)
|
||||
try:
|
||||
reso: Resolution = reso_group[ratio_idx]
|
||||
height = reso.height
|
||||
width = reso.width
|
||||
ar_predicted = True
|
||||
except IndexError:
|
||||
logger.warning(
|
||||
"[ar2diffusion] Request %d: ratio_index=%d out of range [0,%d), keeping prompt size %dx%d",
|
||||
i,
|
||||
ratio_idx,
|
||||
len(reso_group),
|
||||
height,
|
||||
width,
|
||||
)
|
||||
cot_text_for_dit = _truncate_at_cot_end(generated_text)
|
||||
|
||||
cot_text_for_dit = _truncate_at_cot_end(generated_text)
|
||||
logger.info(
|
||||
"[ar2diffusion] Request 0: AR generated %d tokens, text length=%d, cot_text length=%d, target size=%dx%d (%s)",
|
||||
len(generated_token_ids),
|
||||
len(generated_text),
|
||||
len(cot_text_for_dit),
|
||||
height,
|
||||
width,
|
||||
f"AR ratio_idx={ratio_idx}" if ar_predicted else "from prompt (no AR ratio token)",
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"[ar2diffusion] Request %d: AR generated %d tokens, text length=%d, "
|
||||
"cot_text length=%d, target size=%dx%d (%s)",
|
||||
i,
|
||||
len(generated_token_ids),
|
||||
len(generated_text),
|
||||
len(cot_text_for_dit),
|
||||
height,
|
||||
width,
|
||||
f"AR ratio_idx={ratio_idx}" if ar_predicted else "from prompt (no AR ratio token)",
|
||||
)
|
||||
diffusion_input: dict[str, Any] = {
|
||||
"prompt": text_prompt,
|
||||
"height": height,
|
||||
"width": width,
|
||||
"extra": {
|
||||
"ar_generated_text": cot_text_for_dit,
|
||||
},
|
||||
}
|
||||
|
||||
diffusion_input: dict[str, Any] = {
|
||||
"prompt": text_prompt,
|
||||
"height": height,
|
||||
"width": width,
|
||||
"extra": {
|
||||
"ar_generated_text": cot_text_for_dit,
|
||||
},
|
||||
}
|
||||
# Forward use_system_prompt so the DiT can build the same system prefix.
|
||||
# Also forward the custom system prompt body when sys_type=custom so
|
||||
# DiT's `get_system_prompt(use, "image", body)` doesn't fall back to
|
||||
# an empty prefix and silently diverge from AR.
|
||||
if use_system_prompt is not None:
|
||||
diffusion_input["use_system_prompt"] = use_system_prompt
|
||||
if custom_system_prompt is not None:
|
||||
diffusion_input["system_prompt"] = custom_system_prompt
|
||||
|
||||
# Forward use_system_prompt so the DiT can build the same system prefix.
|
||||
# Also forward the custom system prompt body when sys_type=custom so
|
||||
# DiT's `get_system_prompt(use, "image", body)` doesn't fall back to
|
||||
# an empty prefix and silently diverge from AR.
|
||||
if use_system_prompt is not None:
|
||||
diffusion_input["use_system_prompt"] = use_system_prompt
|
||||
if custom_system_prompt is not None:
|
||||
diffusion_input["system_prompt"] = custom_system_prompt
|
||||
# Forward multimodal data (original image for IT2I conditioning).
|
||||
# The diffusion pre_process_func reads multi_modal_data["image"], which
|
||||
# matches vLLM's standard prompt schema, so we only need to pass it once.
|
||||
mm_data = original_prompt.get("multi_modal_data")
|
||||
if mm_data:
|
||||
prompt_images = mm_data.get("image")
|
||||
if prompt_images is None:
|
||||
prompt_images = mm_data.get("images")
|
||||
if prompt_images is not None:
|
||||
diffusion_input["multi_modal_data"] = {"image": prompt_images}
|
||||
|
||||
# Forward multimodal data (original image for IT2I conditioning).
|
||||
# The diffusion pre_process_func reads multi_modal_data["image"], which
|
||||
# matches vLLM's standard prompt schema, so we only need to pass it once.
|
||||
mm_data = original_prompt.get("multi_modal_data")
|
||||
if mm_data:
|
||||
prompt_images = mm_data.get("image")
|
||||
if prompt_images is None:
|
||||
prompt_images = mm_data.get("images")
|
||||
if prompt_images is not None:
|
||||
diffusion_input["multi_modal_data"] = {"image": prompt_images}
|
||||
# Forward multimodal output from AR (if any)
|
||||
if hasattr(ar_output, "multimodal_output") and ar_output.multimodal_output:
|
||||
mm_output = ar_output.multimodal_output
|
||||
if isinstance(mm_output, Mapping):
|
||||
diffusion_input["extra"]["ar_multimodal_output"] = mm_output
|
||||
|
||||
# Forward multimodal output from AR (if any)
|
||||
if hasattr(ar_output, "multimodal_output") and ar_output.multimodal_output:
|
||||
mm_output = ar_output.multimodal_output
|
||||
if isinstance(mm_output, Mapping):
|
||||
diffusion_input["extra"]["ar_multimodal_output"] = mm_output
|
||||
# Forward sampling params
|
||||
for key in ["seed", "num_inference_steps", "guidance_scale", "negative_prompt"]:
|
||||
if key in original_prompt:
|
||||
diffusion_input[key] = original_prompt[key]
|
||||
|
||||
# Forward sampling params
|
||||
for key in ["seed", "num_inference_steps", "guidance_scale", "negative_prompt"]:
|
||||
if key in original_prompt:
|
||||
diffusion_input[key] = original_prompt[key]
|
||||
|
||||
diffusion_inputs.append(diffusion_input)
|
||||
|
||||
return diffusion_inputs
|
||||
return diffusion_input
|
||||
Reference in new issue
Block a user