Prepare isolated Qwen Image 2.1 evaluation

2026-09-21 08:52:36 +02:00
parent 302b08e051
commit 1a660c20ba
8 changed files with 351 additions and 1 deletions
@@ -0,0 +1,35 @@
+#!/usr/bin/env bash
+# Download pinned official weights, build the isolated worker and leave it stopped.
+set -Eeuo pipefail
+
+ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+ENV_FILE=${STACK_ENV:-/etc/mike-ai/stack.env}
+MODEL_DIR=${QWEN_IMAGE_21_MODEL_DIR:-/data/models/Qwen-Image-2.1-ComfyUI}
+OUTPUT_DIR=${QWEN_IMAGE_21_OUTPUT_DIR:-/data/qwen-image-2.1-test-output}
+BASE=https://huggingface.co/Comfy-Org/Qwen-Image-2.1/resolve/ace0edeb3791a594ddfa36ed5f41a178a394e921
+
+download() {
+  local relative=$1 expected=$2 target="$MODEL_DIR/$1"
+  install -d -m 0755 "$(dirname "$target")"
+  if [[ -f $target ]] && echo "$expected  $target" | sha256sum -c --status; then
+    echo "OK: $relative"
+    return
+  fi
+  curl -fL --retry 5 --retry-all-errors -C - -o "$target.part" "$BASE/$relative"
+  echo "$expected  $target.part" | sha256sum -c --status
+  mv "$target.part" "$target"
+}
+
+download diffusion_models/qwen_image_2.1_int8_convrot.safetensors cb74113cb03faecd79611b01fd7fd642f0aa60d6f0b95086abee214d75eaa57d
+download text_encoders/qwen3vl_8b_int8_convrot.safetensors 8bfd0f6e12abf2d2d697ecc888e5e90b0d6741d6708f05799f53afa560452e8f
+download vae/qwen_image_2.1_vae_bf16.safetensors bb21f7473051e1ac368515dd3f2e15cd44d7a11748ee8823e1ddca3e4876b7c9
+install -d -m 0755 "$OUTPUT_DIR"
+
+compose=(docker compose --env-file "$ENV_FILE" -f "$ROOT_DIR/compose.yaml" --profile qwen-image-test)
+"${compose[@]}" build qwen-image-21-test
+"${compose[@]}" up -d --build --no-deps profile-controller
+"${compose[@]}" create qwen-image-21-test
+"${compose[@]}" stop qwen-image-21-test
+state=$(docker inspect -f '{{.State.Status}}' mike-ai-qwen-image-2.1-test)
+[[ $state == exited || $state == created ]]
+echo "Qwen-Image-2.1 test is prepared and stopped. No GPU model was loaded."
@@ -0,0 +1,136 @@
+#!/usr/bin/env python3
+"""One-shot Qwen-Image-2.1 evaluation with guaranteed profile restoration."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import pathlib
+import random
+import time
+import urllib.parse
+import urllib.request
+
+
+CONTROLLER = os.environ.get("CONTROLLER_URL", "http://profile-controller:8090")
+TOKEN = os.environ["CONTROLLER_TOKEN"]
+COMFY = os.environ.get("COMFY_URL", "http://qwen-image-21-test:8188")
+OUTPUT = pathlib.Path(os.environ.get("OUTPUT_DIR", "/output"))
+
+
+def request_json(url: str, *, payload: dict | None = None,
+                 authenticated: bool = False, timeout: float = 30) -> dict:
+    body = None if payload is None else json.dumps(payload).encode()
+    headers = {"Content-Type": "application/json"}
+    if authenticated:
+        headers["Authorization"] = f"Bearer {TOKEN}"
+    method = "POST" if payload is not None else "GET"
+    req = urllib.request.Request(url, data=body, headers=headers, method=method)
+    with urllib.request.urlopen(req, timeout=timeout) as response:
+        return json.load(response)
+
+
+def controller(path: str, *, post: bool = False) -> dict:
+    return request_json(CONTROLLER + path, payload={} if post else None,
+                        authenticated=True, timeout=900)
+
+
+def wait_comfy(timeout: int = 900) -> None:
+    deadline = time.monotonic() + timeout
+    while time.monotonic() < deadline:
+        try:
+            request_json(COMFY + "/system_stats", timeout=3)
+            return
+        except Exception:
+            time.sleep(2)
+    raise TimeoutError("Qwen-Image worker did not become ready")
+
+
+def workflow(prompt: str, seed: int, steps: int, width: int, height: int) -> dict:
+    return {
+        "1": {"class_type": "UNETLoader", "inputs": {
+            "unet_name": "qwen_image_2.1_int8_convrot.safetensors",
+            "weight_dtype": "default"}},
+        "2": {"class_type": "CLIPLoader", "inputs": {
+            "clip_name": "qwen3vl_8b_int8_convrot.safetensors",
+            "type": "qwen_image", "device": "default"}},
+        "3": {"class_type": "VAELoader", "inputs": {
+            "vae_name": "qwen_image_2.1_vae_bf16.safetensors"}},
+        "4": {"class_type": "TextEncodeQwenImage21", "inputs": {
+            "clip": ["2", 0], "prompt": prompt, "negative_prompt": "",
+            "resolution": max(width, height)}},
+        "5": {"class_type": "EmptyLatentImage", "inputs": {
+            "width": width, "height": height, "batch_size": 1}},
+        "6": {"class_type": "KSampler", "inputs": {
+            "model": ["1", 0], "positive": ["4", 0], "negative": ["4", 1],
+            "latent_image": ["5", 0], "seed": seed, "steps": steps,
+            "cfg": 1.0, "sampler_name": "euler", "scheduler": "simple",
+            "denoise": 1.0}},
+        "7": {"class_type": "VAEDecode", "inputs": {
+            "samples": ["6", 0], "vae": ["3", 0]}},
+        "8": {"class_type": "SaveImage", "inputs": {
+            "filename_prefix": "qwen-image-2.1-test", "images": ["7", 0]}},
+    }
+
+
+def wait_result(prompt_id: str, timeout: int = 3600) -> dict:
+    deadline = time.monotonic() + timeout
+    while time.monotonic() < deadline:
+        history = request_json(f"{COMFY}/history/{prompt_id}", timeout=10)
+        if prompt_id in history:
+            result = history[prompt_id]
+            status = result.get("status", {})
+            if status.get("status_str") == "error" or not status.get("completed", True):
+                raise RuntimeError("generation failed: " + json.dumps(status))
+            return result
+        time.sleep(2)
+    raise TimeoutError("Qwen-Image generation timed out")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("prompt")
+    parser.add_argument("--seed", type=int, default=None)
+    parser.add_argument("--steps", type=int, default=25)
+    parser.add_argument("--width", type=int, default=1024)
+    parser.add_argument("--height", type=int, default=1024)
+    args = parser.parse_args()
+    if args.width % 32 or args.height % 32:
+        parser.error("width and height must be multiples of 32")
+    seed = args.seed if args.seed is not None else random.randrange(2**53)
+    previous = controller("/profiles/status").get("active_profile")
+    started = time.monotonic()
+    try:
+        controller("/workers/qwen-image-test/start", post=True)
+        wait_comfy()
+        queued = request_json(COMFY + "/prompt", payload={
+            "prompt": workflow(args.prompt, seed, args.steps, args.width, args.height),
+            "client_id": "athena-qwen-image-21-test"}, timeout=30)
+        prompt_id = queued["prompt_id"]
+        result = wait_result(prompt_id)
+        images = []
+        for node in result.get("outputs", {}).values():
+            images.extend(node.get("images", []))
+        if not images:
+            raise RuntimeError("generation completed without an image")
+        image = images[0]
+        query = urllib.parse.urlencode({
+            "filename": image["filename"], "subfolder": image.get("subfolder", ""),
+            "type": image.get("type", "output")})
+        target = OUTPUT / image["filename"]
+        target.parent.mkdir(parents=True, exist_ok=True)
+        with urllib.request.urlopen(COMFY + "/view?" + query, timeout=120) as src:
+            target.write_bytes(src.read())
+        print(json.dumps({"status": "ok", "file": str(target), "seed": seed,
+                          "seconds": round(time.monotonic() - started, 1)}, indent=2))
+    finally:
+        try:
+            controller("/workers/qwen-image-test/stop", post=True)
+        finally:
+            if previous:
+                controller(f"/profiles/{previous}/activate", post=True)
+
+
+if __name__ == "__main__":
+    main()