vllm-omni/tests/e2e/accuracy/test_joy_image_edit.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

from __future__ import annotations

import gc
import os
from pathlib import Path
from typing import Any

import pytest
import requests
import torch
from PIL import Image
from vllm.config.load import LoadConfig

from benchmarks.accuracy.common import decode_base64_image, pil_to_png_bytes
from tests.e2e.accuracy.helpers import assert_similarity, model_output_dir
from tests.helpers.clean import cleanup_test_environment
from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniServer
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.inputs.data import OmniDiffusionSamplingParams

pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]


MODEL_ID = "jdopensource/JoyAI-Image-Edit-Diffusers"
MODEL_ENV_VAR = "JOY_IMAGE_EDIT_MODEL"
INPUT_IMAGE_PATH = Path(__file__).resolve().parents[3] / "docs" / "assets" / "WeChat.jpg"
PROMPT = "Change the background to a clean studio while preserving the subject."
NEGATIVE_PROMPT = ""
WIDTH = 1024
HEIGHT = 1024
NUM_INFERENCE_STEPS = 20
GUIDANCE_SCALE = 4.0
SEED = 42
SSIM_THRESHOLD = 0.90
PSNR_THRESHOLD = 25.0
SERVER_ARGS = [
    "--num-gpus",
    "1",
    "--enforce-eager",
    "--init-timeout",
    "1200",
]


def _model_name() -> str:
    return os.environ.get(MODEL_ENV_VAR, MODEL_ID)


def _local_files_only(model: str) -> bool:
    return Path(model).exists()


def _joy_image_edit_pipeline_cls() -> type[Any]:
    diffusers = pytest.importorskip("diffusers")
    pipeline_cls = getattr(diffusers, "JoyImageEditPipeline", None)
    if pipeline_cls is None:
        pytest.skip("diffusers.JoyImageEditPipeline is required for JoyAI-Image-Edit parity testing.")
    return pipeline_cls


def _load_input_image() -> Image.Image:
    if not INPUT_IMAGE_PATH.exists():
        raise AssertionError(f"JoyAI input image does not exist: {INPUT_IMAGE_PATH}")
    image = Image.open(INPUT_IMAGE_PATH).convert("RGB")
    image.load()
    return image


def _latent_metrics(
    vllm_latents: torch.Tensor,
    diffusers_latents: torch.Tensor,
) -> dict[str, float]:
    vllm_flat = vllm_latents.detach().float().cpu().flatten()
    diffusers_flat = diffusers_latents.detach().float().cpu().flatten()
    if vllm_flat.shape != diffusers_flat.shape:
        raise AssertionError(
            "Latent outputs have different shapes: "
            f"vllm_omni={tuple(vllm_latents.shape)}, diffusers={tuple(diffusers_latents.shape)}"
        )
    delta = vllm_flat - diffusers_flat
    return {
        "max_abs": delta.abs().max().item(),
        "mean_abs": delta.abs().mean().item(),
        "rmse": torch.sqrt(torch.mean(delta.square())).item(),
        "cosine_similarity": torch.nn.functional.cosine_similarity(
            vllm_flat.unsqueeze(0),
            diffusers_flat.unsqueeze(0),
        ).item(),
    }


def _print_latent_metrics(metrics: dict[str, float]) -> None:
    print(f"{MODEL_ID} latent metrics:")
    for name, value in metrics.items():
        print(f"  {name}: {value:.8f}")


def _ensure_diffusion_distributed_initialized() -> None:
    from vllm_omni.diffusion.distributed.parallel_state import (
        init_distributed_environment,
        initialize_model_parallel,
        model_parallel_is_initialized,
    )

    if not torch.distributed.is_initialized():
        os.environ.setdefault("MASTER_ADDR", "localhost")
        os.environ.setdefault("MASTER_PORT", "29513")
        os.environ.setdefault("RANK", "0")
        os.environ.setdefault("LOCAL_RANK", "0")
        os.environ.setdefault("WORLD_SIZE", "1")
        init_distributed_environment(world_size=1, rank=0, local_rank=0)
    if not model_parallel_is_initialized():
        initialize_model_parallel()


def _run_vllm_omni_joy_image_edit(
    *,
    model: str,
    input_image: Image.Image,
    output_path: Path,
) -> Image.Image:
    with OmniServer(model=model, serve_args=SERVER_ARGS) as server:
        response = requests.post(
            f"http://{server.host}:{server.port}/v1/images/edits",
            data={
                "model": server.model,
                "prompt": PROMPT,
                "size": f"{WIDTH}x{HEIGHT}",
                "n": 1,
                "response_format": "b64_json",
                "negative_prompt": NEGATIVE_PROMPT,
                "num_inference_steps": NUM_INFERENCE_STEPS,
                "guidance_scale": GUIDANCE_SCALE,
                "seed": SEED,
            },
            files=[("image", ("input.png", pil_to_png_bytes(input_image), "image/png"))],
            timeout=1200,
        )
        response.raise_for_status()
        payload = response.json()
        assert len(payload["data"]) == 1
        image = decode_base64_image(payload["data"][0]["b64_json"])
        image.load()
        image.save(output_path)
        return image


class _JoyImageEditRequest:
    def __init__(self, input_image: Image.Image) -> None:
        self.prompts = [
            {
                "prompt": PROMPT,
                "negative_prompt": NEGATIVE_PROMPT,
                "multi_modal_data": {"image": input_image},
            }
        ]
        self.sampling_params = OmniDiffusionSamplingParams(
            height=HEIGHT,
            width=WIDTH,
            num_inference_steps=NUM_INFERENCE_STEPS,
            guidance_scale=GUIDANCE_SCALE,
            guidance_scale_provided=True,
            seed=SEED,
        )
        self.request_id = "joy-image-edit-latent-parity"

    def is_dummy_run(self) -> bool:
        return False


def _load_vllm_omni_joy_image_edit_pipeline(model: str):
    from vllm_omni.diffusion.models.joy_image.pipeline_joy_image_edit import (
        JoyImageEditPipeline,
    )

    _ensure_diffusion_distributed_initialized()
    od_config = OmniDiffusionConfig(
        model=model,
        model_class_name="JoyImageEditPipeline",
        dtype=torch.bfloat16,
        enforce_eager=True,
    )
    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
    return loader.load_model(
        load_device="cuda",
        load_format="custom_pipeline",
        custom_pipeline_name=JoyImageEditPipeline,
    )


def _extract_diffusers_latents(result: Any) -> torch.Tensor:
    if isinstance(result, torch.Tensor):
        return result
    images = getattr(result, "images", None)
    if isinstance(images, torch.Tensor):
        return images
    raise TypeError(f"Unexpected Diffusers latent output type: {type(result)!r}")


def _run_latent_parity_comparison(
    *,
    model: str,
    input_image: Image.Image,
) -> dict[str, float]:
    pipeline_cls = _joy_image_edit_pipeline_cls()
    cleanup_test_environment()
    diffusers_pipe = None
    try:
        diffusers_pipe = pipeline_cls.from_pretrained(
            model,
            torch_dtype=torch.bfloat16,
            trust_remote_code=True,
            local_files_only=_local_files_only(model),
        ).to("cuda")
        if hasattr(diffusers_pipe, "set_progress_bar_config"):
            diffusers_pipe.set_progress_bar_config(disable=False)

        processed_image = diffusers_pipe.vae_image_processor.resize_center_crop(
            input_image,
            (HEIGHT, WIDTH),
        )
        prompt_embeds, prompt_embeds_mask = diffusers_pipe.encode_prompt_multiple_images(
            prompt=PROMPT,
            images=processed_image,
            device=torch.device("cuda"),
            num_images_per_prompt=1,
            max_sequence_length=4096,
        )
        negative_prompt_embeds, negative_prompt_embeds_mask = diffusers_pipe.encode_prompt_multiple_images(
            prompt=NEGATIVE_PROMPT,
            images=processed_image,
            device=torch.device("cuda"),
            num_images_per_prompt=1,
            max_sequence_length=4096,
        )
        num_channels_latents = int(diffusers_pipe.transformer.config.in_channels)
        generator = torch.Generator(device="cuda").manual_seed(SEED)
        fixed_latents = torch.randn(
            1,
            1,
            num_channels_latents,
            1,
            HEIGHT // diffusers_pipe.vae_scale_factor_spatial,
            WIDTH // diffusers_pipe.vae_scale_factor_spatial,
            generator=generator,
            device="cuda",
            dtype=prompt_embeds.dtype,
        )
        diffusers_result = diffusers_pipe(
            prompt=None,
            image=input_image,
            height=HEIGHT,
            width=WIDTH,
            num_inference_steps=NUM_INFERENCE_STEPS,
            guidance_scale=GUIDANCE_SCALE,
            latents=fixed_latents,
            prompt_embeds=prompt_embeds,
            prompt_embeds_mask=prompt_embeds_mask,
            negative_prompt_embeds=negative_prompt_embeds,
            negative_prompt_embeds_mask=negative_prompt_embeds_mask,
            output_type="latent",
        )
        diffusers_latents = _extract_diffusers_latents(diffusers_result).detach().cpu()
        prompt_embeds = prompt_embeds.detach().cpu()
        prompt_embeds_mask = prompt_embeds_mask.detach().cpu()
        negative_prompt_embeds = negative_prompt_embeds.detach().cpu()
        negative_prompt_embeds_mask = negative_prompt_embeds_mask.detach().cpu()
        fixed_latents = fixed_latents.detach().cpu()
    finally:
        if diffusers_pipe is not None and hasattr(diffusers_pipe, "maybe_free_model_hooks"):
            diffusers_pipe.maybe_free_model_hooks()
        del diffusers_pipe
        gc.collect()
        if torch.cuda.is_available():
            torch.accelerator.empty_cache()
        cleanup_test_environment()

    cleanup_test_environment()
    vllm_pipe = None
    try:
        from vllm_omni.diffusion.models.joy_image.pipeline_joy_image_edit import (
            get_joy_image_edit_pre_process_func,
        )

        vllm_pipe = _load_vllm_omni_joy_image_edit_pipeline(model)
        vllm_pipe.eval()
        request = _JoyImageEditRequest(input_image)
        request = get_joy_image_edit_pre_process_func(vllm_pipe.od_config)(request)
        with torch.inference_mode():
            vllm_result = vllm_pipe(
                request,
                prompt_embeds=prompt_embeds.to(device=vllm_pipe.device, dtype=torch.bfloat16),
                prompt_embeds_mask=prompt_embeds_mask.to(device=vllm_pipe.device),
                negative_prompt_embeds=negative_prompt_embeds.to(device=vllm_pipe.device, dtype=torch.bfloat16),
                negative_prompt_embeds_mask=negative_prompt_embeds_mask.to(device=vllm_pipe.device),
                latents=fixed_latents.to(device=vllm_pipe.device, dtype=torch.bfloat16),
                output_type="latent",
            )
        vllm_latents = vllm_result.output
    finally:
        if vllm_pipe is not None and hasattr(vllm_pipe, "maybe_free_model_hooks"):
            vllm_pipe.maybe_free_model_hooks()
        del vllm_pipe
        gc.collect()
        if torch.cuda.is_available():
            torch.accelerator.empty_cache()
        cleanup_test_environment()

    return _latent_metrics(vllm_latents, diffusers_latents)


def _run_diffusers_joy_image_edit(
    *,
    model: str,
    input_image: Image.Image,
    output_path: Path,
) -> Image.Image:
    pipeline_cls = _joy_image_edit_pipeline_cls()
    cleanup_test_environment()
    pipe = None
    try:
        pipe = pipeline_cls.from_pretrained(
            model,
            torch_dtype=torch.bfloat16,
            trust_remote_code=True,
            local_files_only=_local_files_only(model),
        ).to("cuda")
        if hasattr(pipe, "set_progress_bar_config"):
            pipe.set_progress_bar_config(disable=False)
        generator = torch.Generator(device="cuda").manual_seed(SEED)
        result = pipe(
            prompt=PROMPT,
            image=input_image,
            negative_prompt=NEGATIVE_PROMPT,
            height=HEIGHT,
            width=WIDTH,
            num_inference_steps=NUM_INFERENCE_STEPS,
            guidance_scale=GUIDANCE_SCALE,
            generator=generator,
        )
        output_image = result.images[0].convert("RGB")
        output_image.save(output_path)
        return output_image
    finally:
        if pipe is not None and hasattr(pipe, "maybe_free_model_hooks"):
            pipe.maybe_free_model_hooks()
        del pipe
        gc.collect()
        if torch.cuda.is_available():
            torch.accelerator.empty_cache()
        cleanup_test_environment()


@pytest.mark.benchmark
@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=1)
def test_joy_image_edit_matches_diffusers(accuracy_artifact_root: Path) -> None:
    _joy_image_edit_pipeline_cls()
    model = _model_name()
    output_dir = model_output_dir(accuracy_artifact_root, MODEL_ID)
    input_image = _load_input_image()
    input_path = output_dir / "input.png"
    vllm_output_path = output_dir / "vllm_omni.png"
    diffusers_output_path = output_dir / "diffusers.png"
    input_image.save(input_path)

    vllm_image = _run_vllm_omni_joy_image_edit(
        model=model,
        input_image=input_image,
        output_path=vllm_output_path,
    )
    diffusers_image = _run_diffusers_joy_image_edit(
        model=model,
        input_image=input_image,
        output_path=diffusers_output_path,
    )

    print(f"{MODEL_ID} generated images:")
    print(f"  input: {input_path}")
    print(f"  vllm_omni: {vllm_output_path}")
    print(f"  diffusers: {diffusers_output_path}")

    assert_similarity(
        model_name=MODEL_ID,
        vllm_image=vllm_image,
        diffusers_image=diffusers_image,
        width=WIDTH,
        height=HEIGHT,
        ssim_threshold=SSIM_THRESHOLD,
        psnr_threshold=PSNR_THRESHOLD,
    )
    _print_latent_metrics(
        _run_latent_parity_comparison(
            model=model,
            input_image=input_image,
        )
    )