# SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 """End-to-end smoke for the native stable-diffusion.cpp engine. Drives the real ``SdCppEngine`` over a built ``sd-cli`` and a set of split GGUF assets (the same the diffusers path consumes), running one txt2img generation and reporting wall time. This is the GPU/native analogue of ``scripts/diffusion_bench.py``: it proves the engine wiring (finder -> command builder -> subprocess -> output PNG) works against real weights. Example (Z-Image-Turbo on one GPU): SD_CLI_PATH=.../sd-cli CUDA_VISIBLE_DEVICES=6 python scripts/sd_cpp_smoke.py \\ --family z-image \\ --diffusion-model .../z-image-turbo-Q4_K_M.gguf \\ --vae .../ae.safetensors \\ --llm .../Qwen3-4B-Instruct-2507-Q4_K_M.gguf \\ --memory-mode balanced --steps 8 --cfg-scale 1.0 --width 512 --height 512 """ from __future__ import annotations import argparse import sys import time from pathlib import Path _BACKEND_ROOT = Path(__file__).resolve().parent.parent / "studio" / "backend" if str(_BACKEND_ROOT) not in sys.path: sys.path.insert(0, str(_BACKEND_ROOT)) from core.inference.diffusion_memory import ( # noqa: E402 MEMORY_MODE_BALANCED, MEMORY_MODE_FAST, MEMORY_MODE_LOW_VRAM, OFFLOAD_GROUP, OFFLOAD_MODEL, OFFLOAD_NONE, ) from core.inference.sd_cpp_args import ( # noqa: E402 SdCppGenParams, SdCppModelFiles, SdCppUpscaleParams, offload_flags, ) from core.inference.sd_cpp_engine import SdCppEngine, find_sd_cpp_binary # noqa: E402 # memory-mode (user knob) -> sd.cpp offload policy, matching the diffusers planner. _MODE_TO_POLICY = { MEMORY_MODE_FAST: OFFLOAD_NONE, MEMORY_MODE_BALANCED: OFFLOAD_GROUP, MEMORY_MODE_LOW_VRAM: OFFLOAD_MODEL, } def main(argv: list[str] | None = None) -> int: p = argparse.ArgumentParser(description = "Native sd-cli engine smoke test.") p.add_argument("--task", default = "txt2img", choices = ["txt2img", "img2img", "upscale"]) p.add_argument("--binary", default = None, help = "sd-cli path (else env / finder)") p.add_argument("--family", default = "z-image") p.add_argument("--diffusion-model", default = None) # img2img + upscale inputs p.add_argument("--init-img", default = None) p.add_argument("--strength", type = float, default = 0.6) p.add_argument("--upscale-model", default = None) p.add_argument("--upscale-repeats", type = int, default = 1) p.add_argument("--vae", default = None) p.add_argument("--clip_l", default = None) p.add_argument("--t5xxl", default = None) p.add_argument("--llm", default = None) p.add_argument("--qwen2vl", default = None) p.add_argument( "--prompt", default = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed", ) p.add_argument("--negative-prompt", default = None) p.add_argument("--width", type = int, default = 512) p.add_argument("--height", type = int, default = 512) p.add_argument("--steps", type = int, default = 8) p.add_argument("--cfg-scale", type = float, default = 1.0) p.add_argument("--seed", type = int, default = 42) p.add_argument("--memory-mode", default = "balanced", choices = list(_MODE_TO_POLICY)) p.add_argument("--out-image", default = "outputs/sdcpp_verify/sdcpp_smoke.png") p.add_argument("--timeout", type = float, default = 1800.0) args = p.parse_args(argv) binary = args.binary or find_sd_cpp_binary() engine = SdCppEngine(binary = binary) print(f"binary: {engine.binary}", flush = True) print(f"available: {engine.is_available()}", flush = True) print(f"version: {engine.version()}", flush = True) if not engine.is_available(): print( "ERROR: sd-cli not found (set --binary / SD_CLI_PATH / UNSLOTH_SD_CPP_PATH).", flush = True, ) return 2 out = Path(args.out_image) if args.task == "upscale": if not args.init_img or not args.upscale_model: print("ERROR: upscale needs --init-img and --upscale-model.", flush = True) return 2 t0 = time.time() result = engine.upscale( SdCppUpscaleParams( input_image = args.init_img, upscale_model = args.upscale_model, repeats = args.upscale_repeats, ), output_path = str(out), verbose = True, timeout = args.timeout, on_log = lambda ln: print(f" [sd] {ln}", flush = True), ) dt = time.time() - t0 print( f"\nOK: upscaled {result} ({result.stat().st_size/1024:.0f} KB) in {dt:.1f}s", flush = True, ) print("SD-CPP-SMOKE-OK", flush = True) return 0 if not args.diffusion_model: print("ERROR: --diffusion-model is required for txt2img / img2img.", flush = True) return 2 files = SdCppModelFiles( diffusion_model = args.diffusion_model, vae = args.vae, clip_l = args.clip_l, t5xxl = args.t5xxl, llm = args.llm, qwen2vl = args.qwen2vl, ) is_img2img = args.task == "img2img" params = SdCppGenParams( prompt = args.prompt, negative_prompt = args.negative_prompt, width = args.width, height = args.height, steps = args.steps, cfg_scale = args.cfg_scale, seed = args.seed, init_img = args.init_img if is_img2img else None, strength = args.strength if is_img2img else None, ) if is_img2img and not args.init_img: print("ERROR: img2img needs --init-img.", flush = True) return 2 policy = _MODE_TO_POLICY[args.memory_mode] off = offload_flags(policy) print( f"task: {args.task}" + (f" (init={args.init_img}, strength={args.strength})" if is_img2img else ""), flush = True, ) print(f"memory: {args.memory_mode} -> policy={policy} -> flags={off}", flush = True) t0 = time.time() result = engine.generate( files, params, output_path = str(out), offload = off, verbose = True, timeout = args.timeout, on_log = lambda ln: print(f" [sd] {ln}", flush = True), ) dt = time.time() - t0 size_kb = result.stat().st_size / 1024 if result.is_file() else 0 print(f"\nOK: generated {result} ({size_kb:.0f} KB) in {dt:.1f}s", flush = True) print("SD-CPP-SMOKE-OK", flush = True) return 0 if __name__ == "__main__": raise SystemExit(main())