"""Calibrate the eleven-view Turbo workflow on a local CUDA GPU. No Spaces/Gradio import, quota timer, or captions are used. The angle LoRA is off by default and can be enabled for a comparison using an existing --image. An omitted --image uses the bundled mannequin for a speed smoke test only; --create-sample first generates a synthetic character reference with the same loaded model. A memory limit caps PyTorch's allocator; it does not emulate MIG, half a GPU, or ZeroGPU latency. """ import argparse from datetime import datetime, timezone import importlib.metadata import json import math from pathlib import Path import sys import time from PIL import Image from core import prepare_reference, conditioning_reference, rgb_on_white from pose_guides import get_guide from views import BY_KEY, DEFAULT_KEYS, SIZES, make_prompt, PROMPT_VERSION, REFERENCE_RESOLUTION # Keep in sync with backend.py without importing its ZeroGPU startup code. MODEL_ID = "Qwen/Qwen-Image-2.1-Turbo" MODEL_REVISION = "d65dbc9a7e8f6b5479e33dee6030eaab2a906509" ANGLES_REPO = "akhaliq/Qwen-Image-2.1-Multiple-Angles-LoRA" ANGLES_WEIGHT = ("checkpoints_v2/angles_v2_full_qwen21_multiple_angles_v2_" "qwen21_multiple_angles_v2_000001500.safetensors") ANGLES_STRENGTH = 0.8 GIB = 1024**3 SAMPLE_SEED = 20261012 SAMPLE_PROMPT = ( "A single clean anime-style character illustration of a clearly adult woman in her late twenties, " "full body, straight front view at eye level, standing in a relaxed neutral A-pose with her hands " "lowered and legs narrowly spaced. Short teal bob haircut ending at her jaw, straight bangs, " "amber eyes and a calm neutral expression. She wears a bright orange hip-length zip-front jacket " "with a large cobalt-blue rectangular patch on the left chest and a yellow stripe around the " "right forearm, a cream crew-neck shirt, straight dark navy trousers, and white ankle boots " "with black soles. Opaque everyday clothing, natural adult proportions, clearly defined clothing " "seams and consistent flat colors with restrained cel shading. Show the complete top of her " "head and both boots with generous margins. Plain white background, soft even lighting. " "One character only, no text, no letters, no logo, no reference panels, no collage." ) def positive_int(value): number = int(value) if number < 1: raise argparse.ArgumentTypeError("must be at least 1") return number def positive_float(value): number = float(value) if not math.isfinite(number) or number <= 0: raise argparse.ArgumentTypeError("must be a finite number greater than 0") return number def parse_args(argv=None): parser = argparse.ArgumentParser(description=__doc__) source = parser.add_mutually_exclusive_group() source.add_argument("--image", type=Path, help="Character reference. Omit both source options for mannequin-only speed smoke test.") source.add_argument("--create-sample", action="store_true", help="Generate a synthetic adult anime character reference with the same loaded Turbo model.") parser.add_argument("--output-dir", type=Path, required=True, help="Save PNGs and metrics.json here; existing metrics are never overwritten.") parser.add_argument("--resolution", choices=tuple(SIZES), default="fast") parser.add_argument("--pose-guides", choices=("on", "off"), default="on") parser.add_argument("--reference-crop", choices=("on", "off"), default="off", help="Use close-view crops only for an upright full-body reference on white.") parser.add_argument("--angles-lora", choices=("on", "off"), default="off", help="Compare the optional Qwen2.1 angle adapter at strength0.8. Requires --image.") parser.add_argument("--repeat", type=positive_int, default=1, help="Number of complete eleven-view batches (default: 1).") parser.add_argument("--warmup", type=int, choices=(0, 1), default=0, help="Optionally generate one extra front view before the measured batches. " "Its time, memory and image are also recorded (default: 0).") parser.add_argument("--memory-limit-gib", type=positive_float, help="PyTorch allocator cap in GiB, e.g. 48. This is not a MIG/compute cap.") parser.add_argument("--seed", type=int, default=0, help="Base seed; each view adds its fixed index, identical across repeats.") args = parser.parse_args(argv) if args.angles_lora == "on" and not args.image: parser.error("--angles-lora on requires --image; reuse an existing reference " "instead of --create-sample for a comparable source image") return args def write_metrics(path, metrics): """Replace the JSON atomically, retaining completed results after a late error.""" temporary = path.with_suffix(".json.tmp") temporary.write_text(json.dumps(metrics, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") temporary.replace(path) def cuda_memory(torch): return { "allocated_gib": round(torch.cuda.memory_allocated() / GIB, 4), "reserved_gib": round(torch.cuda.memory_reserved() / GIB, 4), "peak_allocated_gib": round(torch.cuda.max_memory_allocated() / GIB, 4), "peak_reserved_gib": round(torch.cuda.max_memory_reserved() / GIB, 4), } def package_versions(): versions = {} for name in ("torch", "diffusers", "transformers", "accelerate", "peft", "pillow"): try: versions[name] = importlib.metadata.version(name) except importlib.metadata.PackageNotFoundError: versions[name] = None return versions def benchmark(args): output_dir = args.output_dir.expanduser().resolve() output_dir.mkdir(parents=True, exist_ok=True) metrics_path = output_dir / "metrics.json" if metrics_path.exists(): raise FileExistsError(f"{metrics_path} already exists; choose a new --output-dir") pose_enabled = args.pose_guides == "on" angles_enabled = args.angles_lora == "on" metrics = { "status": "running", "started_at_utc": datetime.now(timezone.utc).isoformat(), "model": MODEL_ID, "revision": MODEL_REVISION, "input": str(args.image.expanduser().resolve()) if args.image else None, "input_kind": ("user_reference" if args.image else "synthetic_character" if args.create_sample else "mannequin_speed_smoke_only"), "quality_validation": "not performed by this benchmark", "resolution": args.resolution, "pose_guides": pose_enabled, "angles_lora": angles_enabled, "prompt_version": PROMPT_VERSION, "conditioning": ("original reference with white-background close-view crops v1" if args.reference_crop == "on" else "original reference unchanged"), "crop_close_views": args.reference_crop == "on", "angles_adapter": ({"repo": ANGLES_REPO, "weight": ANGLES_WEIGHT, "strength": ANGLES_STRENGTH, "adapter_name": "angles", "strategy": "fuse on CPU, unload adapter, then move pipeline to CUDA"} if angles_enabled else None), "base_seed": args.seed, "requested_views": DEFAULT_KEYS, "repeat": args.repeat, "warmup": args.warmup, "memory_limit_gib": args.memory_limit_gib, "memory_limit_scope": "PyTorch allocator only; not MIG or a compute limit. " "Non-PyTorch CUDA allocations can be outside this cap.", "timing_scope": "Each inference uses CUDA synchronization before and after the pipeline call; " "PNG writing is excluded. Batch wall time includes saving PNGs and metrics. " "Cold means the first pipeline call in this process. A synthetic reference " "or warmup prewarms the pipeline before the measured edit batches. " "No ZeroGPU startup, queue, transport, or quota is simulated.", "pipeline": {"true_cfg_scale": 1.0, "use_kv_cache": True, "output_resolution": REFERENCE_RESOLUTION[args.resolution], "num_images_per_prompt": 1, "sampling": "checkpoint-configured eight sigmas"}, "packages": package_versions(), "generations": [], "batches": [], "completed_images": 0, } write_metrics(metrics_path, metrics) process_started = time.perf_counter() torch = None stage = "input_preparation" active_generation = None active_batch = None pipeline_call_count = 0 edit_call_count = 0 exit_code = 0 try: max_edge = 768 if args.resolution == "fast" else 1216 if args.image: with Image.open(args.image.expanduser()) as source: reference = prepare_reference(source, max_edge) elif not args.create_sample: reference = prepare_reference(get_guide("full_front", 512, 768), max_edge) print("No --image: using the mannequin for a SPEED SMOKE TEST ONLY; " "these images do not validate character identity or pose quality.", flush=True) else: reference = None if reference is not None: metrics["reference_size"] = list(reference.size) guides = None if pose_enabled: guides = {key: get_guide(key, *SIZES[args.resolution][BY_KEY[key].shape]) for key in DEFAULT_KEYS} stage = "cuda_initialization" import torch as torch_module torch = torch_module if not torch.cuda.is_available(): raise RuntimeError("CUDA is unavailable; this benchmark requires a local CUDA GPU") device = torch.cuda.current_device() properties = torch.cuda.get_device_properties(device) metrics["gpu"] = {"name": torch.cuda.get_device_name(device), "device_index": device, "total_memory_gib": round(properties.total_memory / GIB, 4), "compute_capability": f"{properties.major}.{properties.minor}", "torch_cuda_version": torch.version.cuda} if args.memory_limit_gib is not None: fraction = args.memory_limit_gib * GIB / properties.total_memory if fraction > 1: raise ValueError("--memory-limit-gib exceeds the selected GPU's total memory") # Set before model loading or CUDA model placement. torch.cuda.set_per_process_memory_fraction(fraction, device) metrics["allocator_fraction"] = fraction print(f"PyTorch allocator cap: {args.memory_limit_gib:g} GiB. " "The GPU still has its full compute capacity; this is not MIG or half-GPU emulation.", flush=True) torch.cuda.synchronize() torch.cuda.reset_peak_memory_stats() write_metrics(metrics_path, metrics) stage = "model_load" from diffusers import QwenImage21Pipeline load_started = time.perf_counter() pipe = QwenImage21Pipeline.from_pretrained( MODEL_ID, revision=MODEL_REVISION, dtype=torch.bfloat16, ) pipe.set_progress_bar_config(disable=True) if angles_enabled: from huggingface_hub import hf_hub_download from angles_adapter import load_angles_adapter checkpoint = hf_hub_download(ANGLES_REPO, filename=ANGLES_WEIGHT) metrics["angles_adapter"]["conversion"] = load_angles_adapter( pipe, checkpoint, adapter_name="angles", strength=ANGLES_STRENGTH, ) pipe.fuse_lora(adapter_names=["angles"]) pipe.unload_lora_weights() pipe.vae.enable_tiling() pipe.to("cuda") torch.cuda.synchronize() metrics["model_load_seconds"] = round(time.perf_counter() - load_started, 4) metrics["model_load_scope"] = ("from_pretrained through CUDA placement; includes optional " "CPU adapter fusion and downloads if uncached") metrics["model_load_memory"] = cuda_memory(torch) sample_sigmas = pipe.config.get("sample_sigmas", []) if len(sample_sigmas) != 8: raise RuntimeError("The pinned Turbo checkpoint's eight-sigma sampling schedule was not loaded") metrics["sample_sigmas"] = list(sample_sigmas) metrics["peak_allocated_gib"] = metrics["model_load_memory"]["peak_allocated_gib"] metrics["peak_reserved_gib"] = metrics["model_load_memory"]["peak_reserved_gib"] print(f"Loaded {metrics['gpu']['name']} in {metrics['model_load_seconds']:.3f}s; " f"model allocation {metrics['model_load_memory']['allocated_gib']:.2f} GiB.", flush=True) write_metrics(metrics_path, metrics) if args.create_sample: stage = "sample_generation" active_generation = { "phase": "sample_reference", "key": "reference", "status": "running", "seed": SAMPLE_SEED, "prompt": SAMPLE_PROMPT, "width": 832, "height": 1216, "output_resolution": 1024, "references": 0, "cache_state": "cold", "is_first_pipeline_call": True, "is_first_edit": False, "pipeline_call_index": 1, } metrics["sample_generation"] = active_generation metrics["active_generation"] = active_generation write_metrics(metrics_path, metrics) torch.cuda.synchronize() torch.cuda.reset_peak_memory_stats() sample_started = time.perf_counter() pipeline_call_count += 1 sample_image = pipe( prompt=SAMPLE_PROMPT, width=832, height=1216, output_resolution=1024, true_cfg_scale=1.0, use_kv_cache=True, num_images_per_prompt=1, generator=torch.Generator("cpu").manual_seed(SAMPLE_SEED), ).images[0] torch.cuda.synchronize() sample_seconds = round(time.perf_counter() - sample_started, 4) active_generation["seconds"] = sample_seconds active_generation.update(cuda_memory(torch)) stage = "sample_save" rgb_on_white(sample_image).save(output_dir / "reference.png", format="PNG") active_generation["image"] = "reference.png" active_generation["status"] = "complete" metrics["sample_generation_seconds"] = sample_seconds metrics["cold_generation"] = dict(active_generation) metrics["first_pipeline_generation"] = dict(active_generation) for field in ("peak_allocated_gib", "peak_reserved_gib"): metrics[field] = max(metrics[field], active_generation[field]) reference = prepare_reference(sample_image, max_edge) metrics["reference_size"] = list(reference.size) metrics["input"] = "reference.png" metrics.pop("active_generation", None) active_generation = None write_metrics(metrics_path, metrics) print(f"Synthetic reference: {sample_seconds:.3f}s; saved reference.png. " "This time is separate from edit batches; image quality has not been validated.", flush=True) def generate(key, *, phase, repeat_index): nonlocal stage, active_generation, pipeline_call_count, edit_call_count view = BY_KEY[key] width, height = SIZES[args.resolution][view.shape] inputs = [conditioning_reference(reference, key, width, height) if args.reference_crop == "on" else reference] if guides is not None: inputs.append(guides[key]) seed = (args.seed + DEFAULT_KEYS.index(key)) % (2**31) directory = output_dir / ("warmup" if phase == "warmup" else f"run_{repeat_index:02d}") directory.mkdir(exist_ok=True) target = directory / f"{DEFAULT_KEYS.index(key) + 1:02d}_{key}.png" active_generation = {"key": key, "phase": phase, "repeat_index": repeat_index, "cache_state": "cold" if pipeline_call_count == 0 else "warm", "is_first_pipeline_call": pipeline_call_count == 0, "is_first_edit": edit_call_count == 0, "pipeline_call_index": pipeline_call_count + 1, "seed": seed, "width": width, "height": height, "references": len(inputs), "status": "running"} metrics["active_generation"] = active_generation write_metrics(metrics_path, metrics) stage = "generation" torch.cuda.synchronize() torch.cuda.reset_peak_memory_stats() inference_started = time.perf_counter() pipeline_call_count += 1 edit_call_count += 1 # Match backend.py. No budget callback: calibration must attempt all eleven views. edit_prompt = make_prompt(view, "", pose_enabled, angles_enabled) active_generation["prompt"] = edit_prompt image = pipe( image=inputs, prompt=edit_prompt, width=width, height=height, output_resolution=REFERENCE_RESOLUTION[args.resolution], true_cfg_scale=1.0, use_kv_cache=True, num_images_per_prompt=1, generator=torch.Generator("cpu").manual_seed(seed), ).images[0] torch.cuda.synchronize() active_generation["seconds"] = round(time.perf_counter() - inference_started, 4) active_generation.update(cuda_memory(torch)) stage = "image_save" rgb_on_white(image).save(target, format="PNG") active_generation["image"] = str(target.relative_to(output_dir)) active_generation["status"] = "complete" metrics["generations"].append(active_generation) if active_generation["cache_state"] == "cold": metrics["cold_generation"] = dict(active_generation) if active_generation["is_first_pipeline_call"]: metrics["first_pipeline_generation"] = dict(active_generation) if active_generation["is_first_edit"]: metrics["first_edit_generation"] = dict(active_generation) for field in ("peak_allocated_gib", "peak_reserved_gib"): metrics[field] = max(metrics[field], active_generation[field]) metrics["completed_images"] += 1 if phase == "batch": active_batch["completed_views"].append(key) active_batch["inference_seconds"] = round( active_batch["inference_seconds"] + active_generation["seconds"], 4) metrics.pop("active_generation", None) completed = active_generation active_generation = None write_metrics(metrics_path, metrics) print(f"{phase} {repeat_index or '-'} / {key}: {completed['seconds']:.3f}s, " f"peak allocated {completed['peak_allocated_gib']:.2f} GiB, " f"peak reserved {completed['peak_reserved_gib']:.2f} GiB " f"({completed['cache_state']})", flush=True) if args.warmup: generate("full_front", phase="warmup", repeat_index=0) for repeat_index in range(1, args.repeat + 1): active_batch = {"repeat_index": repeat_index, "status": "running", "completed_views": [], "inference_seconds": 0.0, "includes_cold_generation": pipeline_call_count == 0} metrics["batches"].append(active_batch) batch_started = time.perf_counter() for key in DEFAULT_KEYS: generate(key, phase="batch", repeat_index=repeat_index) active_batch["wall_seconds"] = round(time.perf_counter() - batch_started, 4) active_batch["status"] = "complete" print(f"Batch {repeat_index}: 11/11 views, inference " f"{active_batch['inference_seconds']:.3f}s, wall {active_batch['wall_seconds']:.3f}s.", flush=True) write_metrics(metrics_path, metrics) active_batch = None metrics["status"] = "complete" except (Exception, KeyboardInterrupt) as exc: exit_code = 130 if isinstance(exc, KeyboardInterrupt) else 1 metrics["status"] = "interrupted" if exit_code == 130 else "failed" metrics["error"] = {"stage": stage, "type": type(exc).__name__, "message": str(exc)} if active_generation is not None: active_generation["status"] = metrics["status"] if active_batch is not None: active_batch["status"] = metrics["status"] active_batch["wall_seconds"] = round(time.perf_counter() - batch_started, 4) if torch is not None: try: metrics["memory_at_failure"] = cuda_memory(torch) except Exception: pass print(f"Benchmark {metrics['status']} during {stage}: {type(exc).__name__}: {exc}\n" "Completed PNGs and metrics are retained.", file=sys.stderr, flush=True) finally: metrics["process_seconds"] = round(time.perf_counter() - process_started, 4) metrics["finished_at_utc"] = datetime.now(timezone.utc).isoformat() write_metrics(metrics_path, metrics) print(f"Metrics: {metrics_path}", flush=True) return exit_code def main(argv=None): args = parse_args(argv) try: return benchmark(args) except (OSError, ValueError) as exc: print(f"Cannot start benchmark: {exc}", file=sys.stderr) return 1 if __name__ == "__main__": raise SystemExit(main())