[Hunyuanimage 3.0] hunyuan accuracy test (#3655)

Signed-off-by: dengyunyang <584797741@qq.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

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."""

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

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中不同颜色区域之间的分界线，将被处理成纤细的、带有金属光泽的凸起边缘，这既能清晰地勾勒出图案，也符合珐琅工艺品的典型特征。整个冰箱贴会呈现出轻微的厚度和圆润的边缘，使其看起来像一个真实的、可触摸的物体。"
+           "最后，将这个制作完成的冰箱贴放置在图二所示的灰色织物背景上，并为其添加一个微妙的、柔和的阴影，以增强其立体感和与背景的融合度，最终呈现出一个精致、逼真的产品展示图。</think><recaption>这幅图像以产品摄影的精致风格，呈现了一枚根据`image_1`标志定制的圆形珐琅冰箱贴。最终图像使用`image_2`的分辨率。"
+           "冰箱贴居中放置在`image_2`的灰色织物背景上，其设计完美复刻了`image_1`中由深蓝、中蓝、浅蓝和白色构成的动态旋涡图案。整个冰箱贴被赋予了`image_2`中冰箱贴特有的高级质感：表面覆盖着一层光滑如镜的珐琅釉面，反射出柔和而清晰的高光；图案的每一个色块边缘都由纤细的抛光金属边框精确勾勒，增强了立体感。"
+           "柔和的顶光在冰箱贴的弧形边缘上形成平滑的过渡，并在其下方投下淡淡的、轮廓模糊的阴影，使其与织物背景无缝融合，营造出一种真实、静谧的视觉效果。<relation_1>最终图像完整保留了`image_1`中标志的全部设计元素。这包括其完美的圆形轮廓，以及内部由深蓝、中蓝、浅蓝和白色组成的精确旋涡状图案布局、形状和色彩关系。"
+           "</relation_1><relation_2>最终图像的分辨率、背景和材质均来自`image_2`。背景中灰色织物的纹理和质感被完整保留。冰箱贴的材质被完美重构，精确复刻了`image_2`中冰箱贴所展示的光滑珐琅质感、抛光金属边框的视觉效果，以及整体柔和、均匀的布光环境和由此产生的自然阴影。</relation_2></recaption><answer><boi>"
+           "<img_size_1024><img_ratio_36><timestep>[<img>]{3600}<eoi></answer>")
+# 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"