vllm-omni/probe/issue-2/offline_fp8v.py

#!/usr/bin/env python
"""Issue 2 probe: offline Qwen-Image-2.1 edit driver with prefix-KV extras.

The served path cannot reach `extras.prefix_kv_cache_dtype` on commit
8484d85d (the structured-config bridge drops deploy-YAML extras), so this
driver exercises the knob through the offline `Omni(...)` engine kwargs —
the other route the recipe names. It mirrors image_edit.py's request
construction: same prompt builder, same sampling fields, one request at a
time.

Usage: offline_fp8v.py --out RUN_DIR [--kv-dtype fp8_v] [--fixture FID]
"""

from __future__ import annotations

import argparse
import hashlib
import json
import pathlib
import time

import torch
from PIL import Image

from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.model_extras import (
    build_image_to_image_prompt,
    get_model_class_name,
)
from vllm_omni.platforms import current_omni_platform

W = pathlib.Path.home() / "Projects/vllm-omni-val/issue2"
SUITE = (
    pathlib.Path.home()
    / ".local/share/multi-model-helper-server/data/evaluation/media/image/references-suite-2"
)
SNAP = (
    pathlib.Path.home()
    / ".local/state/multi-model-helper-server/state/image-service/weights/hub/models--Qwen--Qwen-Image-2.1/snapshots/790c92633540aa0cb11d9abf19eb46d861714758"
)
STEPS = 20
CFG = 1.0


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--out", required=True, type=pathlib.Path)
    parser.add_argument("--kv-dtype", default="fp8_v")
    parser.add_argument(
        "--quant",
        default="none",
        choices=["none", "fp8", "fp8-imgmlp", "fp8-dit", "fp8-encoder"],
        help="weight quantization: none, fp8 (all), fp8 minus img_mlp, transformer-only, or text-encoder-only",
    )
    parser.add_argument("--fixture", default="edit-compose-four")
    parser.add_argument("--seeds", type=lambda s: tuple(int(x) for x in s.split(",")), default=(2301,))
    parser.add_argument("--no-extras", action="store_true", help="run without the extras knob (control)")
    args = parser.parse_args()

    fixture = json.loads((W / "fixtures.json").read_text())[args.fixture]
    refs = [Image.open(SUITE / f"{r}.png").convert("RGB") for r in fixture["references"]]

    omni_kwargs = dict(
        model=str(SNAP),
        enforce_eager=True,
        enable_cpu_offload=True,
        init_timeout=1200,
    )
    if args.quant == "fp8":
        omni_kwargs["quantization"] = "fp8"
    elif args.quant == "fp8-imgmlp":
        omni_kwargs["quantization_config"] = {
            "method": "fp8",
            "ignored_layers": ["img_mlp"],
        }
    elif args.quant == "fp8-dit":
        omni_kwargs["quantization_config"] = {"transformer": {"method": "fp8"}}
    elif args.quant == "fp8-encoder":
        omni_kwargs["quantization_config"] = {"text_encoder": {"method": "fp8"}}
    if not args.no_extras:
        omni_kwargs["extras"] = {"prefix_kv_cache_dtype": args.kv_dtype}
    t0 = time.time()
    omni = Omni(**omni_kwargs)
    print(
        f"load_s={time.time() - t0:.1f} quant={args.quant} extras={omni_kwargs.get('extras')}",
        flush=True,
    )
    model_class_name = get_model_class_name(omni)

    args.out.mkdir(parents=True, exist_ok=True)
    records = []
    for seed in args.seeds:
        generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
        prompt_dict = build_image_to_image_prompt(
            model_class_name=model_class_name,
            prompt=fixture["prompt"],
            negative_prompt=None,
            input_image=refs[0] if len(refs) == 1 else refs,
            height=fixture["height"],
            width=fixture["width"],
        )
        params = OmniDiffusionSamplingParams(
            generator=generator,
            seed=seed,
            true_cfg_scale=CFG,
            num_inference_steps=STEPS,
            height=fixture["height"],
            width=fixture["width"],
        )
        torch.cuda.reset_peak_memory_stats()
        t1 = time.time()
        rec = {"fixture_id": args.fixture, "seed": seed, "n_references": len(refs), "kv_dtype": args.kv_dtype, "quant": args.quant}
        try:
            outputs = omni.generate(prompt_dict, sampling_params_list=[params])
            wall = time.time() - t1
            images = None
            for output in outputs:
                images = getattr(output, "images", None)
                if images:
                    break
            if not images:
                raise RuntimeError("no images in output")
            out_path = args.out / "measured" / f"{args.fixture}.s{seed}.png"
            out_path.parent.mkdir(parents=True, exist_ok=True)
            images[0].save(out_path)
            png = out_path.read_bytes()
            rec.update(
                wall_s=round(wall, 2),
                peak_gpu_mb=round(torch.cuda.max_memory_allocated() / 2**20, 1),
                bytes=len(png),
                sha256=hashlib.sha256(png).hexdigest(),
            )
        except Exception as exc:  # noqa: BLE001 - recorded as the probe result
            rec.update(wall_s=round(time.time() - t1, 2), error=f"{type(exc).__name__}: {exc}"[:2000])
        records.append(rec)
        print(json.dumps(rec), flush=True)

    (args.out / f"offline-{args.fixture}.json").write_text(json.dumps(records, indent=1) + "\n")
    ok = all("sha256" in r for r in records)
    print(("PASS" if ok else "FAIL") + f" {len(records)} requests", flush=True)
    return 0 if ok else 1


if __name__ == "__main__":
    raise SystemExit(main())