215 lines
8.6 KiB
Python
215 lines
8.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from __future__ import annotations
|
|
|
|
import copy
|
|
import gc
|
|
import tempfile
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
import torch
|
|
import yaml
|
|
from PIL import Image
|
|
|
|
from tests.e2e.accuracy.helpers import CLIPScorer, assert_similarity, model_output_dir
|
|
from tests.helpers.runtime import OmniRunner
|
|
from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import build_prompt_tokens
|
|
|
|
pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
|
|
|
|
MODEL_NAME = "tencent/HunyuanImage-3.0-Instruct"
|
|
SEED = 42
|
|
NUM_INFERENCE_STEPS = 50
|
|
GUIDANCE_SCALE = 0
|
|
HEIGHT, WIDTH = 1024, 1024
|
|
PSNR_THRESHOLD = 0.0
|
|
SSIM_THRESHOLD = 0.0
|
|
CLIP_SCORE_THRESHOLD = 20.0
|
|
PROMPT = "A brown and white dog is running on the grass."
|
|
|
|
# fmt: off
|
|
_BASE_CONFIG = {
|
|
"stage_args": [
|
|
{
|
|
"stage_id": 0, "stage_type": "llm",
|
|
"runtime": {"process": True, "devices": "0,1", "max_batch_size": 1, "requires_multimodal_data": True},
|
|
"engine_args": {
|
|
"model_stage": "AR", "model_arch": "HunyuanImage3ForCausalMM",
|
|
"worker_cls": "vllm_omni.worker.gpu_ar_worker.GPUARWorker",
|
|
"scheduler_cls": "vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler",
|
|
"gpu_memory_utilization": 0.95, "enforce_eager": True, "trust_remote_code": True,
|
|
"engine_output_type": "latent", "enable_prefix_caching": False,
|
|
"max_num_batched_tokens": 32768, "tensor_parallel_size": 2, "pipeline_parallel_size": 1,
|
|
"hf_overrides": {"rope_parameters": {"mrope_section": [0, 32, 32], "rope_type": "default"}},
|
|
},
|
|
"is_comprehension": False, "final_output": True, "final_output_type": "text",
|
|
"default_sampling_params": {
|
|
"temperature": 0.0, "top_p": 1, "top_k": -1, "max_tokens": 8192,
|
|
"stop_token_ids": [128025], "detokenize": True, "skip_special_tokens": False,
|
|
},
|
|
"output_connectors": {"to_stage_1": "rdma_connector"},
|
|
},
|
|
{
|
|
"stage_id": 1, "stage_type": "diffusion",
|
|
"runtime": {"process": True, "devices": "2,3", "max_batch_size": 1, "requires_multimodal_data": True},
|
|
"engine_args": {
|
|
"model_stage": "dit", "model_arch": "HunyuanImage3ForCausalMM",
|
|
"enforce_eager": True, "trust_remote_code": True, "distributed_executor_backend": "mp",
|
|
"parallel_config": {"tensor_parallel_size": 2, "enable_expert_parallel": True},
|
|
},
|
|
"engine_input_source": [0],
|
|
"custom_process_input_func": "vllm_omni.model_executor.stage_input_processors.hunyuan_image3.ar2diffusion",
|
|
"final_output": True, "final_output_type": "image",
|
|
"default_sampling_params": {"num_inference_steps": NUM_INFERENCE_STEPS, "guidance_scale": GUIDANCE_SCALE},
|
|
"input_connectors": {"from_stage_0": "rdma_connector"},
|
|
},
|
|
],
|
|
"runtime": {
|
|
"enabled": True,
|
|
"connectors": {"rdma_connector": {
|
|
"name": "MooncakeTransferEngineConnector",
|
|
"extra": {"host": "auto", "zmq_port": 50051, "protocol": "rdma", "device_name": "", "memory_pool_size": 4294967296, "memory_pool_device": "cpu"},
|
|
}},
|
|
"edges": [{"from": 0, "to": 1}],
|
|
},
|
|
}
|
|
# fmt: on
|
|
|
|
|
|
def _make_config(enable_kv_reuse: bool, path: Path) -> None:
|
|
config = copy.deepcopy(_BASE_CONFIG)
|
|
config["stage_args"][0]["engine_args"]["omni_kv_config"] = {"need_send_cache": enable_kv_reuse}
|
|
config["stage_args"][1]["engine_args"]["omni_kv_config"] = {"need_recv_cache": enable_kv_reuse}
|
|
path.write_text(yaml.dump(config, default_flow_style=False, sort_keys=False))
|
|
|
|
|
|
def _run(stage_config_path: str, output_path: Path) -> tuple[Image.Image, str, float]:
|
|
from transformers import AutoTokenizer
|
|
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
|
|
from vllm_omni.platforms import current_omni_platform
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
|
|
token_ids = build_prompt_tokens(PROMPT, tokenizer, task="it2i_recaption", sys_type="en_unified")
|
|
|
|
with OmniRunner(MODEL_NAME, stage_configs_path=stage_config_path) as runner:
|
|
params_list = list(runner.omni.default_sampling_params_list)
|
|
for sp in params_list:
|
|
if isinstance(sp, OmniDiffusionSamplingParams):
|
|
sp.num_inference_steps = NUM_INFERENCE_STEPS
|
|
sp.guidance_scale = GUIDANCE_SCALE
|
|
sp.seed = SEED
|
|
sp.generator = torch.Generator(device=current_omni_platform.device_type or "cuda").manual_seed(SEED)
|
|
|
|
prompts: list[OmniPromptType] = [
|
|
{
|
|
"prompt_token_ids": token_ids,
|
|
"prompt": PROMPT,
|
|
"use_system_prompt": "en_unified",
|
|
"height": HEIGHT,
|
|
"width": WIDTH,
|
|
}
|
|
]
|
|
t0 = time.perf_counter()
|
|
outputs = list(runner.omni.generate(prompts=prompts, sampling_params_list=params_list))
|
|
elapsed = time.perf_counter() - t0
|
|
|
|
assert outputs, "Pipeline produced no outputs"
|
|
images = None
|
|
cot_text = ""
|
|
for out in outputs:
|
|
ro = getattr(out, "request_output", None)
|
|
if ro and getattr(ro, "outputs", None):
|
|
txt = "".join(getattr(o, "text", "") or "" for o in ro.outputs)
|
|
if txt:
|
|
cot_text = txt
|
|
|
|
imgs = getattr(out, "images", None)
|
|
if not imgs and ro and hasattr(ro, "images"):
|
|
imgs = ro.images
|
|
if imgs:
|
|
images = imgs
|
|
|
|
assert images, "Pipeline output had no images"
|
|
|
|
image = images[0].convert("RGB")
|
|
image.save(output_path)
|
|
gc.collect()
|
|
if torch.accelerator.is_available():
|
|
torch.accelerator.empty_cache()
|
|
return image, cot_text, elapsed
|
|
|
|
|
|
@pytest.mark.skipif(torch.accelerator.device_count() < 4, reason="Needs 4+ GPUs (2 AR + 2 DiT)")
|
|
def test_text_to_image_alignment(accuracy_artifact_root: Path) -> None:
|
|
"""KV reuse ON vs OFF: same pipeline, same seed → PSNR >= 40 dB."""
|
|
output_dir = model_output_dir(accuracy_artifact_root, MODEL_NAME + "-kv-reuse")
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
tmp = Path(tmpdir)
|
|
_make_config(False, tmp / "off.yaml")
|
|
image_no_reuse_1, cot_no_reuse_1, time_no_reuse_1 = _run(
|
|
str(tmp / "off.yaml"), output_dir / "without_kv_reuse_1.png"
|
|
)
|
|
image_no_reuse_2, cot_no_reuse_2, time_no_reuse_2 = _run(
|
|
str(tmp / "off.yaml"), output_dir / "without_kv_reuse_2.png"
|
|
)
|
|
_make_config(True, tmp / "on.yaml")
|
|
image_reuse, cot_reuse, time_reuse = _run(str(tmp / "on.yaml"), output_dir / "with_kv_reuse.png")
|
|
|
|
print("\n--- End-to-end time ---")
|
|
print(f" WITHOUT KV reuse 1: {time_no_reuse_1:.2f}s")
|
|
print(f" WITHOUT KV reuse 2: {time_no_reuse_2:.2f}s")
|
|
print(f" WITH KV reuse: {time_reuse:.2f}s")
|
|
|
|
(output_dir / "cot_with_kv_reuse.txt").write_text(cot_reuse, encoding="utf-8")
|
|
(output_dir / "cot_without_kv_reuse_1.txt").write_text(cot_no_reuse_1, encoding="utf-8")
|
|
(output_dir / "cot_without_kv_reuse_2.txt").write_text(cot_no_reuse_2, encoding="utf-8")
|
|
|
|
print(f"\n--- CoT WITH KV reuse (len={len(cot_reuse)}) ---\n{cot_reuse}")
|
|
print(f"\n--- CoT WITHOUT KV reuse 1 (len={len(cot_no_reuse_1)}) ---\n{cot_no_reuse_1}")
|
|
print(f"\n--- CoT WITHOUT KV reuse 2 (len={len(cot_no_reuse_2)}) ---\n{cot_no_reuse_2}")
|
|
|
|
clip_scorer = CLIPScorer()
|
|
clip_scorer.assert_score(
|
|
model_name=f"{MODEL_NAME} with-reuse",
|
|
image=image_reuse,
|
|
text=PROMPT,
|
|
threshold=CLIP_SCORE_THRESHOLD,
|
|
)
|
|
clip_scorer.assert_score(
|
|
model_name=f"{MODEL_NAME} without-reuse 1",
|
|
image=image_no_reuse_1,
|
|
text=PROMPT,
|
|
threshold=CLIP_SCORE_THRESHOLD,
|
|
)
|
|
clip_scorer.assert_score(
|
|
model_name=f"{MODEL_NAME} without-reuse 2",
|
|
image=image_no_reuse_2,
|
|
text=PROMPT,
|
|
threshold=CLIP_SCORE_THRESHOLD,
|
|
)
|
|
|
|
assert_similarity(
|
|
model_name=f"{MODEL_NAME} non-reuse run1 vs run2",
|
|
vllm_image=image_no_reuse_1,
|
|
diffusers_image=image_no_reuse_2,
|
|
ssim_threshold=SSIM_THRESHOLD,
|
|
psnr_threshold=PSNR_THRESHOLD,
|
|
width=WIDTH,
|
|
height=HEIGHT,
|
|
)
|
|
|
|
assert_similarity(
|
|
model_name=f"{MODEL_NAME} KV-reuse vs no-reuse",
|
|
vllm_image=image_reuse,
|
|
diffusers_image=image_no_reuse_1,
|
|
ssim_threshold=SSIM_THRESHOLD,
|
|
psnr_threshold=PSNR_THRESHOLD,
|
|
width=WIDTH,
|
|
height=HEIGHT,
|
|
)
|