diff --git a/tests/e2e/accuracy/assets/hunyuan_image_ref.png b/tests/e2e/accuracy/assets/hunyuan_image_ref.png new file mode 100644 index 000000000..f3787ef0c Binary files /dev/null and b/tests/e2e/accuracy/assets/hunyuan_image_ref.png differ diff --git a/tests/e2e/accuracy/conftest.py b/tests/e2e/accuracy/conftest.py index a4804fc19..431a8cc9b 100644 --- a/tests/e2e/accuracy/conftest.py +++ b/tests/e2e/accuracy/conftest.py @@ -185,6 +185,12 @@ def accuracy_artifact_root() -> Path: return root +@pytest.fixture(scope="session") +def accuracy_assets_root() -> Path: + root = Path(__file__).resolve().parent / "assets" + return root + + @pytest.fixture(scope="session") def qwen_bear_image(accuracy_artifact_root: Path): """Download the Qwen bear image from the URL and save it to the accuracy artifact root.""" diff --git a/tests/e2e/accuracy/helpers.py b/tests/e2e/accuracy/helpers.py index 517fd9e13..a5044db76 100644 --- a/tests/e2e/accuracy/helpers.py +++ b/tests/e2e/accuracy/helpers.py @@ -1,9 +1,11 @@ import os +from io import BytesIO from pathlib import Path from typing import Any import numpy as np import pytest +import requests import torch from PIL import Image from torchmetrics.image import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure @@ -161,6 +163,16 @@ class CLIPScorer: similarity = (img_emb * txt_emb).sum(dim=-1) return float(similarity.item() * 100) + def image_image_score(self, image1: Image.Image, image2: Image.Image) -> float: + inputs = self._processor(text=[""], images=[image1, image2], return_tensors="pt", padding=True) + with torch.no_grad(): + outputs = self._model(**inputs) + img_emb1, img_emb2 = outputs.image_embeds.chunk(2) + img_emb1 = img_emb1 / img_emb1.norm(p=2, dim=-1, keepdim=True) + img_emb2 = img_emb2 / img_emb2.norm(p=2, dim=-1, keepdim=True) + similarity = (img_emb1 * img_emb2).sum(dim=-1) + return float(similarity.item() * 100) + def assert_score(self, *, model_name: str, image: Image.Image, text: str, threshold: float) -> None: value = self.score(image, text) print(f"{model_name} CLIP score:") @@ -180,3 +192,73 @@ def _pil_to_batched_tensor(image: Image.Image, *, compare_mode: str) -> torch.Te array = np.asarray(image.convert(compare_mode), dtype=np.float32) / 255.0 tensor = torch.from_numpy(array).permute(2, 0, 1).unsqueeze(0) return tensor + + +class SemanticSimilarityScorer: + """Semantic similarity scorer for text-text and image-text comparisons""" + + def __init__( + self, + text_model_name: str = "BAAI/bge-m3", + ): + """ + Initialize semantic similarity scorer + + Args: + text_model_name: Model for text-text similarity (BGE-M3) + """ + self._text_model = None + self._text_model_name = text_model_name + + def _load_text_model(self): + """Lazy load text similarity model (BGE-M3)""" + if self._text_model is None: + try: + from FlagEmbedding import BGEM3FlagModel + + print(f"Loading text similarity model: {self._text_model_name}") + self._text_model = BGEM3FlagModel(self._text_model_name, use_fp16=True) + except ImportError: + raise ImportError("FlagEmbedding not installed. Install with: pip install FlagEmbedding") + return self._text_model + + def text_similarity(self, text1: str, text2: str) -> float: + """ + Compute semantic similarity between two texts using BGE-M3 + """ + results = {} + # 1. Text prefix matching (maximum matching character length) + match_count = 0 + for char_vllm, char_ref in zip(text1, text2): + if char_vllm == char_ref: + match_count += 1 + else: + break + results["text_prefix_match_count"] = match_count + + # 2. Full CoT semantic similarity comparison + model = self._load_text_model() + emb1 = model.encode([text1], return_dense=True, return_sparse=False, return_colbert_vecs=False)["dense_vecs"][0] + emb2 = model.encode([text2], return_dense=True, return_sparse=False, return_colbert_vecs=False)["dense_vecs"][0] + similarity = np.dot(emb1, emb2) / (np.linalg.norm(emb1) * np.linalg.norm(emb2)) + results["cot_semantic_sim"] = float(similarity) + return results + + def release_models(self): + """Release loaded models to free memory""" + if self._text_model is not None: + del self._text_model + self._text_model = None + + +def download_images(urls: list[str]) -> list[Image.Image]: + """Download and cache multiple images""" + images = [] + for i, url in enumerate(urls): + print(f"Downloading image {i} from {url}") + response = requests.get(url, timeout=30) + response.raise_for_status() + image = Image.open(BytesIO(response.content)).convert("RGB") + images.append(image) + + return images diff --git a/tests/e2e/accuracy/test_hunyuan_image3.py b/tests/e2e/accuracy/test_hunyuan_image3.py index 0871793c5..bf230c07f 100644 --- a/tests/e2e/accuracy/test_hunyuan_image3.py +++ b/tests/e2e/accuracy/test_hunyuan_image3.py @@ -5,6 +5,8 @@ from __future__ import annotations import copy import gc +import importlib +import os import tempfile import time from pathlib import Path @@ -14,28 +16,59 @@ import torch import yaml from PIL import Image -from tests.e2e.accuracy.helpers import CLIPScorer, assert_similarity, model_output_dir +from tests.e2e.accuracy.helpers import ( + CLIPScorer, + SemanticSimilarityScorer, + compute_image_ssim_psnr, + download_images, + model_output_dir, +) from tests.helpers.runtime import OmniRunner -from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import build_prompt_tokens +from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import build_prompt_tokens, resolve_stop_token_ids + +os.environ["DIFFUSION_ATTENTION_BACKEND"] = "TORCH_SDPA" pytestmark = [pytest.mark.full_model, pytest.mark.diffusion] +# ============================================================================ +# Configurable Parameters +# ============================================================================ +AR_DEVICES = "0,1" +DIT_DEVICES = "2,3" 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." +GUIDANCE_SCALE = 2.5 +# ============================================================================ +# Constants +# ============================================================================ +MODEL_PATH = os.environ.get("HUNYUAN_MODEL_PATH", MODEL_NAME) +# Test input +PROMPT = "基于图一的logo,参考图二中冰箱贴的材质,制作一个新的冰箱贴" +TEST_IMAGE_URLS = [ + "https://raw.githubusercontent.com/Tencent-Hunyuan/HunyuanImage-3.0/main/assets/demo_instruct_imgs/input_1_0.png", + "https://raw.githubusercontent.com/Tencent-Hunyuan/HunyuanImage-3.0/main/assets/demo_instruct_imgs/input_1_1.png", +] +SEED = 42 +AR_TP_SIZE = len(AR_DEVICES.split(",")) +DIT_TP_SIZE = len(DIT_DEVICES.split(",")) + +# Precision thresholds +THRESHOLDS = { + # AR text comparison + "text_prefix_match": 10, # First 10 characters must match exactly + "cot_semantic_sim": 0.9, # Full CoT semantic similarity + # Image comparison + "clip_score": 85, # CLIP image semantic similarity + "ssim": 0.20, # Structural similarity + "psnr": 11.0, # Peak signal-to-noise ratio (dB) +} # 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}, + "runtime": {"process": True, "devices": AR_DEVICES, "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", @@ -50,11 +83,11 @@ _BASE_CONFIG = { "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"}, + "output_connectors": {"to_stage_1": "shared_memory_connector"}, }, { "stage_id": 1, "stage_type": "diffusion", - "runtime": {"process": True, "devices": "2,3", "max_batch_size": 1, "requires_multimodal_data": True}, + "runtime": {"process": True, "devices": DIT_DEVICES, "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", @@ -64,20 +97,32 @@ _BASE_CONFIG = { "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"}, + "input_connectors": {"from_stage_0": "shared_memory_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"}, + "connectors": {"shared_memory_connector": { + "name": "SharedMemoryConnector", + "extra": {"shm_threshold_bytes": 65536} }}, "edges": [{"from": 0, "to": 1}], }, } # fmt: on +# fmt: off +COT_REF = ("首先,我分析所有输入图像:图像1是一个圆形的logo,设计现代且抽象。它由不同色调的蓝色(深蓝、中蓝、浅蓝)和白色构成,这些色块以流畅的曲线相互交织,形成一个动态的、类似旋涡或波浪的图案。整个logo是扁平化的矢量图形,背景为纯黑色。图像2展示了四个并排摆放的卡通动物造型冰箱贴," + "它们被放置在灰色的织物背景上。这些冰箱贴的关键特征是其材质:它们具有光滑、高光的珐琅或烤漆质感,边缘有明显的金属包边,整体呈现出一种立体的、有厚度的实体感。用户的指令是“基于图一的logo,参考图二中冰箱贴的材质,制作一个新的冰箱贴”。这个指令要求将一个二维的平面设计(logo)" + "转化为一个具有特定物理属性(材质和立体感)的三维物体。核心任务是保留logo的视觉识别性,同时赋予其冰箱贴的实体质感。为了构建答案图像,我会将图一的圆形logo作为基础形状。然后,我会将图二中冰箱贴的材质特性应用到这个logo上。具体来说,logo中的每一个色块(深蓝、中蓝、浅蓝、白色)" + "都会被渲染成具有高光泽度的珐琅质感,表面会反射出柔和的环境光,形成自然的高光。logo中不同颜色区域之间的分界线,将被处理成纤细的、带有金属光泽的凸起边缘,这既能清晰地勾勒出图案,也符合珐琅工艺品的典型特征。整个冰箱贴会呈现出轻微的厚度和圆润的边缘,使其看起来像一个真实的、可触摸的物体。" + "最后,将这个制作完成的冰箱贴放置在图二所示的灰色织物背景上,并为其添加一个微妙的、柔和的阴影,以增强其立体感和与背景的融合度,最终呈现出一个精致、逼真的产品展示图。这幅图像以产品摄影的精致风格,呈现了一枚根据`image_1`标志定制的圆形珐琅冰箱贴。最终图像使用`image_2`的分辨率。" + "冰箱贴居中放置在`image_2`的灰色织物背景上,其设计完美复刻了`image_1`中由深蓝、中蓝、浅蓝和白色构成的动态旋涡图案。整个冰箱贴被赋予了`image_2`中冰箱贴特有的高级质感:表面覆盖着一层光滑如镜的珐琅釉面,反射出柔和而清晰的高光;图案的每一个色块边缘都由纤细的抛光金属边框精确勾勒,增强了立体感。" + "柔和的顶光在冰箱贴的弧形边缘上形成平滑的过渡,并在其下方投下淡淡的、轮廓模糊的阴影,使其与织物背景无缝融合,营造出一种真实、静谧的视觉效果。最终图像完整保留了`image_1`中标志的全部设计元素。这包括其完美的圆形轮廓,以及内部由深蓝、中蓝、浅蓝和白色组成的精确旋涡状图案布局、形状和色彩关系。" + "最终图像的分辨率、背景和材质均来自`image_2`。背景中灰色织物的纹理和质感被完整保留。冰箱贴的材质被完美重构,精确复刻了`image_2`中冰箱贴所展示的光滑珐琅质感、抛光金属边框的视觉效果,以及整体柔和、均匀的布光环境和由此产生的自然阴影。" + "[]{3600}") +# fmt: on + def _make_config(enable_kv_reuse: bool, path: Path) -> None: config = copy.deepcopy(_BASE_CONFIG) @@ -86,24 +131,25 @@ def _make_config(enable_kv_reuse: bool, path: Path) -> None: 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]: +def _run_offline(stage_configs_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 + build_kwargs: dict = {"task": "it2i", "bot_task": "think_recaption", "sys_type": "en_unified", "num_images": 2} + tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True) result = build_prompt_tokens( PROMPT, tokenizer, - task="it2i", - bot_task="recaption", - sys_type="en_unified", + **build_kwargs, ) token_ids = result.token_ids system_prompt_type = result.system_prompt_type - with OmniRunner(MODEL_NAME, stage_configs_path=stage_config_path) as runner: + ar_stop_token_ids = resolve_stop_token_ids(task="it2i", bot_task="think_recaption", tokenizer=tokenizer) + with OmniRunner(MODEL_NAME, stage_configs_path=stage_configs_path) as runner: params_list = list(runner.omni.default_sampling_params_list) for sp in params_list: if isinstance(sp, OmniDiffusionSamplingParams): @@ -111,14 +157,17 @@ def _run(stage_config_path: str, output_path: Path) -> tuple[Image.Image, str, f sp.guidance_scale = GUIDANCE_SCALE sp.seed = SEED sp.generator = torch.Generator(device=current_omni_platform.device_type or "cuda").manual_seed(SEED) + elif hasattr(sp, "stop_token_ids"): + sp.stop_token_ids = ar_stop_token_ids + images = download_images(TEST_IMAGE_URLS) prompts: list[OmniPromptType] = [ { "prompt_token_ids": token_ids, "prompt": PROMPT, "use_system_prompt": system_prompt_type, - "height": HEIGHT, - "width": WIDTH, + "modalities": ["image"], + "multi_modal_data": {"image": images}, } ] t0 = time.perf_counter() @@ -131,9 +180,13 @@ def _run(stage_config_path: str, output_path: Path) -> tuple[Image.Image, str, f 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 + cot_text = "".join(getattr(o, "text", "") or "" for o in ro.outputs) + if not cot_text: + ar_text = getattr(out, "custom_output", {}).get("ar_generated_text") + if isinstance(ar_text, list): + cot_text = "\n".join(text for text in ar_text if text) + else: + cot_text = ar_text or "" imgs = getattr(out, "images", None) if not imgs and ro and hasattr(ro, "images"): @@ -142,9 +195,11 @@ def _run(stage_config_path: str, output_path: Path) -> tuple[Image.Image, str, f images = imgs assert images, "Pipeline output had no images" + cot_text = cot_text.lstrip("\n") image = images[0].convert("RGB") - image.save(output_path) + image.save(output_path / "image_offline.png") + (output_path / "cot_offline.txt").write_text(cot_text, encoding="utf-8") gc.collect() if torch.accelerator.is_available(): torch.accelerator.empty_cache() @@ -152,71 +207,44 @@ def _run(stage_config_path: str, output_path: Path) -> tuple[Image.Image, str, f @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") +def test_image_to_image_alignment(accuracy_artifact_root: Path, accuracy_assets_root: Path) -> None: + if importlib.util.find_spec("FlagEmbedding") is None: + raise ImportError("Missing dependency: FlagEmbedding\nInstall with: pip install FlagEmbedding") + from tabulate import tabulate # lazy import + + """KV reuse ON vs OFF: same pipeline, same seed → PSNR >= 10 dB.""" + output_dir = model_output_dir(accuracy_artifact_root, MODEL_NAME + "-offline-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}") + omni_image, omni_cot, time_reuse = _run_offline(str(tmp / "on.yaml"), output_dir) + scorer = SemanticSimilarityScorer() 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, - ) + cot_results = scorer.text_similarity(omni_cot, COT_REF) + image_ref = Image.open(str(accuracy_assets_root / "hunyuan_image_ref.png")).convert("RGB") + image_clip_score = clip_scorer.image_image_score(omni_image, image_ref) + ssim_value, psnr_value = compute_image_ssim_psnr(prediction=omni_image, reference=image_ref, compare_mode="RGB") - 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, - ) + table = [ + ["COT similarity to reference", f"{cot_results['cot_semantic_sim']:.4f}", 0.9644], + ["COT prefix match", f"{cot_results['text_prefix_match_count']:.4f}", 29], + ["Image-Image similarity", f"{image_clip_score:.4f}", 94.5538], + ["SSIM", f"{ssim_value:.4f}", 0.242], + ["PSNR (dB)", f"{psnr_value:.2f}", 14.1], + ] - 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, + print(tabulate(table, headers=["Metric", "Value", "L20x Reference"], tablefmt="grid")) + + assert cot_results["cot_semantic_sim"] >= THRESHOLDS["cot_semantic_sim"], ( + f"COT semantic similarity {cot_results['cot_semantic_sim']:.4f} is below threshold {THRESHOLDS['cot_semantic_sim']}" ) + assert cot_results["text_prefix_match_count"] >= THRESHOLDS["text_prefix_match"], ( + f"COT prefix match count {cot_results['text_prefix_match_count']} is below threshold {THRESHOLDS['text_prefix_match']}" + ) + assert image_clip_score >= THRESHOLDS["clip_score"], ( + f"Image-Image similarity{image_clip_score:.4f} is below threshold {THRESHOLDS['clip_score']}" + ) + assert ssim_value >= THRESHOLDS["ssim"], f"SSIM {ssim_value:.4f} is below threshold {THRESHOLDS['ssim']}" + assert psnr_value >= THRESHOLDS["psnr"], f"PSNR {psnr_value:.2f} dB is below threshold {THRESHOLDS['psnr']} dB"