diff --git a/scripts/sd_cpp_smoke.py b/scripts/sd_cpp_smoke.py index 8693481d2b..4398fa6cf3 100644 --- a/scripts/sd_cpp_smoke.py +++ b/scripts/sd_cpp_smoke.py @@ -63,7 +63,10 @@ def main(argv: list[str] | None = None) -> int: 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( + "--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) @@ -81,17 +84,28 @@ def main(argv: list[str] | None = None) -> int: 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) + print( + "ERROR: sd-cli not found (set --binary / SD_CLI_PATH / UNSLOTH_SD_CPP_PATH).", + 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, + diffusion_model = args.diffusion_model, + vae = args.vae, + clip_l = args.clip_l, + t5xxl = args.t5xxl, + llm = args.llm, + qwen2vl = args.qwen2vl, ) 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, + 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, ) policy = _MODE_TO_POLICY[args.memory_mode] off = offload_flags(policy) @@ -100,8 +114,13 @@ def main(argv: list[str] | None = None) -> int: out = Path(args.out_image) 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), + 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 diff --git a/studio/backend/core/inference/sd_cpp_engine.py b/studio/backend/core/inference/sd_cpp_engine.py index cf645a91c5..b3b921419b 100644 --- a/studio/backend/core/inference/sd_cpp_engine.py +++ b/studio/backend/core/inference/sd_cpp_engine.py @@ -208,9 +208,14 @@ class SdCppEngine: out = Path(output_path) out.parent.mkdir(parents = True, exist_ok = True) cmd = build_sd_cpp_command( - self.binary, files, params, - output_path = str(out), offload = offload, threads = threads, - verbose = verbose, extra_args = extra_args, + self.binary, + files, + params, + output_path = str(out), + offload = offload, + threads = threads, + verbose = verbose, + extra_args = extra_args, ) base = dict(os.environ) if env: @@ -246,9 +251,7 @@ class SdCppEngine: proc.kill() if ret != 0: - raise RuntimeError( - f"sd-cli exited {ret}. Last output:\n" + "\n".join(tail[-12:]) - ) + raise RuntimeError(f"sd-cli exited {ret}. Last output:\n" + "\n".join(tail[-12:])) if not out.is_file(): raise RuntimeError( f"sd-cli reported success but no image at {out}. Last output:\n" diff --git a/studio/backend/tests/test_sd_cpp_args.py b/studio/backend/tests/test_sd_cpp_args.py index 9d86a4cb8b..99e758e22e 100644 --- a/studio/backend/tests/test_sd_cpp_args.py +++ b/studio/backend/tests/test_sd_cpp_args.py @@ -60,7 +60,13 @@ def test_offload_group_streams_with_flash_attention(): def test_offload_model_pushes_everything_to_cpu_and_tiles(): flags = offload_flags(OFFLOAD_MODEL) - for expected in ("--offload-to-cpu", "--clip-on-cpu", "--vae-on-cpu", "--vae-tiling", "--diffusion-fa"): + for expected in ( + "--offload-to-cpu", + "--clip-on-cpu", + "--vae-on-cpu", + "--vae-tiling", + "--diffusion-fa", + ): assert expected in flags # sequential maps the same as model assert offload_flags(OFFLOAD_SEQUENTIAL) == flags @@ -81,7 +87,9 @@ def test_offload_forced_flags_dedup(): def test_build_zimage_command_minimal(): files = SdCppModelFiles( - diffusion_model = "/m/z.gguf", vae = "/m/ae.sft", llm = "/m/qwen3.gguf", + diffusion_model = "/m/z.gguf", + vae = "/m/ae.sft", + llm = "/m/qwen3.gguf", ) params = SdCppGenParams(prompt = "a cat", width = 512, height = 768, steps = 8, cfg_scale = 1.0, seed = 42) cmd = build_sd_cpp_command("/bin/sd-cli", files, params, output_path = "/out/x.png") @@ -105,8 +113,10 @@ def test_build_zimage_command_minimal(): def test_build_flux1_dual_text_encoders(): files = SdCppModelFiles( - diffusion_model = "/m/flux.gguf", vae = "/m/ae.sft", - clip_l = "/m/clip_l.sft", t5xxl = "/m/t5.gguf", + diffusion_model = "/m/flux.gguf", + vae = "/m/ae.sft", + clip_l = "/m/clip_l.sft", + t5xxl = "/m/t5.gguf", ) params = SdCppGenParams(prompt = "x", guidance = 3.5) cmd = build_sd_cpp_command("/bin/sd-cli", files, params, output_path = "/o.png") @@ -121,8 +131,14 @@ def test_build_appends_offload_and_extra_args_last(): params = SdCppGenParams(prompt = "x") off = offload_flags(OFFLOAD_GROUP) cmd = build_sd_cpp_command( - "/bin/sd-cli", files, params, output_path = "/o.png", - offload = off, threads = 8, verbose = True, extra_args = ["--rng", "cuda"], + "/bin/sd-cli", + files, + params, + output_path = "/o.png", + offload = off, + threads = 8, + verbose = True, + extra_args = ["--rng", "cuda"], ) assert "--offload-to-cpu" in cmd assert _pair(cmd, "--threads") == "8" @@ -143,14 +159,31 @@ def test_build_omits_unset_optional_params(): files = SdCppModelFiles(diffusion_model = "/m/z.gguf") params = SdCppGenParams(prompt = "x") # no steps/cfg/seed/sampler cmd = build_sd_cpp_command("/bin/sd-cli", files, params, output_path = "/o.png") - for flag in ("--steps", "--cfg-scale", "--guidance", "--seed", "--sampling-method", "--batch-count", "--threads", "-v"): + for flag in ( + "--steps", + "--cfg-scale", + "--guidance", + "--seed", + "--sampling-method", + "--batch-count", + "--threads", + "-v", + ): assert flag not in cmd def test_build_requires_diffusion_model_and_prompt(): with pytest.raises(ValueError): - build_sd_cpp_command("/bin/sd-cli", SdCppModelFiles(diffusion_model = ""), - SdCppGenParams(prompt = "x"), output_path = "/o.png") + build_sd_cpp_command( + "/bin/sd-cli", + SdCppModelFiles(diffusion_model = ""), + SdCppGenParams(prompt = "x"), + output_path = "/o.png", + ) with pytest.raises(ValueError): - build_sd_cpp_command("/bin/sd-cli", SdCppModelFiles(diffusion_model = "/m/z.gguf"), - SdCppGenParams(prompt = " "), output_path = "/o.png") + build_sd_cpp_command( + "/bin/sd-cli", + SdCppModelFiles(diffusion_model = "/m/z.gguf"), + SdCppGenParams(prompt = " "), + output_path = "/o.png", + ) diff --git a/studio/backend/tests/test_sd_cpp_engine.py b/studio/backend/tests/test_sd_cpp_engine.py index 8a20c73963..25dd19cec6 100644 --- a/studio/backend/tests/test_sd_cpp_engine.py +++ b/studio/backend/tests/test_sd_cpp_engine.py @@ -61,7 +61,9 @@ def test_find_custom_install_dir_build_layout(tmp_path, monkeypatch): def test_find_falls_back_to_path(tmp_path, monkeypatch): _clear_env(monkeypatch) monkeypatch.setattr(eng.Path, "home", staticmethod(lambda: tmp_path / "nohome")) - monkeypatch.setattr(eng.shutil, "which", lambda stem: "/usr/bin/sd-cli" if stem == "sd-cli" else None) + monkeypatch.setattr( + eng.shutil, "which", lambda stem: "/usr/bin/sd-cli" if stem == "sd-cli" else None + ) assert find_sd_cpp_binary() == "/usr/bin/sd-cli" @@ -124,7 +126,16 @@ class _FakePopen: captured_cmd: list[str] = [] captured_env: dict = {} - def __init__(self, cmd, *, lines, returncode, out_file, write, env = None): + def __init__( + self, + cmd, + *, + lines, + returncode, + out_file, + write, + env = None, + ): type(self).captured_cmd = list(cmd) type(self).captured_env = dict(env or {}) self._lines = list(lines) @@ -148,10 +159,23 @@ class _FakePopen: pass -def _patch_popen(monkeypatch, *, lines, returncode, out_file, write = True): +def _patch_popen( + monkeypatch, + *, + lines, + returncode, + out_file, + write = True, +): def _factory(cmd, **kw): - return _FakePopen(cmd, lines = lines, returncode = returncode, out_file = out_file, - write = write, env = kw.get("env")) + return _FakePopen( + cmd, + lines = lines, + returncode = returncode, + out_file = out_file, + write = write, + env = kw.get("env"), + ) monkeypatch.setattr(eng.subprocess, "Popen", _factory) @@ -165,7 +189,9 @@ def _engine(tmp_path): def test_generate_success_returns_path_and_collects_logs(tmp_path, monkeypatch): e = _engine(tmp_path) out = tmp_path / "img.png" - _patch_popen(monkeypatch, lines = ["loading model", "step 1/8", "done"], returncode = 0, out_file = out) + _patch_popen( + monkeypatch, lines = ["loading model", "step 1/8", "done"], returncode = 0, out_file = out + ) seen: list[str] = [] files = SdCppModelFiles(diffusion_model = "/m/z.gguf", vae = "/m/ae.sft", llm = "/m/q.gguf") params = SdCppGenParams(prompt = "a cat", steps = 8, seed = 1) @@ -187,8 +213,11 @@ def test_generate_raises_on_nonzero_exit(tmp_path, monkeypatch): out = tmp_path / "img.png" _patch_popen(monkeypatch, lines = ["boom: bad gguf"], returncode = 1, out_file = out, write = False) with pytest.raises(RuntimeError, match = "exited 1"): - e.generate(SdCppModelFiles(diffusion_model = "/m/z.gguf"), - SdCppGenParams(prompt = "x"), output_path = str(out)) + e.generate( + SdCppModelFiles(diffusion_model = "/m/z.gguf"), + SdCppGenParams(prompt = "x"), + output_path = str(out), + ) def test_generate_raises_when_no_output_despite_success(tmp_path, monkeypatch): @@ -196,15 +225,21 @@ def test_generate_raises_when_no_output_despite_success(tmp_path, monkeypatch): out = tmp_path / "img.png" _patch_popen(monkeypatch, lines = ["ok"], returncode = 0, out_file = out, write = False) with pytest.raises(RuntimeError, match = "no image"): - e.generate(SdCppModelFiles(diffusion_model = "/m/z.gguf"), - SdCppGenParams(prompt = "x"), output_path = str(out)) + e.generate( + SdCppModelFiles(diffusion_model = "/m/z.gguf"), + SdCppGenParams(prompt = "x"), + output_path = str(out), + ) def test_generate_raises_when_binary_missing(): e = SdCppEngine(binary = None) with pytest.raises(RuntimeError, match = "not found"): - e.generate(SdCppModelFiles(diffusion_model = "/m/z.gguf"), - SdCppGenParams(prompt = "x"), output_path = "/tmp/x.png") + e.generate( + SdCppModelFiles(diffusion_model = "/m/z.gguf"), + SdCppGenParams(prompt = "x"), + output_path = "/tmp/x.png", + ) # ── engine routing ────────────────────────────────────────────────────────── @@ -225,6 +260,11 @@ def test_routing_cpu_falls_back_to_diffusers_without_binary(): def test_routing_prefer_native_overrides_gpu(): - assert select_diffusion_engine("cuda", native_available = True, prefer_native = True) == ENGINE_SD_CPP + assert ( + select_diffusion_engine("cuda", native_available = True, prefer_native = True) == ENGINE_SD_CPP + ) # but only if a binary is actually available - assert select_diffusion_engine("cuda", native_available = False, prefer_native = True) == ENGINE_DIFFUSERS + assert ( + select_diffusion_engine("cuda", native_available = False, prefer_native = True) + == ENGINE_DIFFUSERS + ) diff --git a/studio/backend/tests/test_sd_cpp_install.py b/studio/backend/tests/test_sd_cpp_install.py index ae5fc4e5d3..171ea3e810 100644 --- a/studio/backend/tests/test_sd_cpp_install.py +++ b/studio/backend/tests/test_sd_cpp_install.py @@ -37,7 +37,11 @@ _ASSETS = [ ] -def _resolve(system, machine, accelerator = "auto"): +def _resolve( + system, + machine, + accelerator = "auto", +): return resolve_release_asset(_ASSETS, system = system, machine = machine, accelerator = accelerator) @@ -47,7 +51,9 @@ def _resolve(system, machine, accelerator = "auto"): def test_macos_arm64_picks_darwin_arm64(): assert _resolve("Darwin", "arm64") == "sd-master-8caa3f9-bin-Darwin-macOS-15.7.7-arm64.zip" # aarch64 spelling resolves the same - assert _resolve("Darwin", "aarch64").startswith("sd-master") and "arm64" in _resolve("Darwin", "aarch64") + assert _resolve("Darwin", "aarch64").startswith("sd-master") and "arm64" in _resolve( + "Darwin", "aarch64" + ) def test_macos_intel_has_no_prebuilt(): @@ -64,7 +70,10 @@ def test_linux_x86_64_auto_picks_plain_cpu_build(): def test_linux_vulkan_and_rocm_select_accelerator_builds(): - assert _resolve("Linux", "x86_64", "vulkan") == "sd-master-8caa3f9-bin-Linux-Ubuntu-24.04-x86_64-vulkan.zip" + assert ( + _resolve("Linux", "x86_64", "vulkan") + == "sd-master-8caa3f9-bin-Linux-Ubuntu-24.04-x86_64-vulkan.zip" + ) assert "rocm" in _resolve("Linux", "x86_64", "rocm") diff --git a/studio/install_sd_cpp_prebuilt.py b/studio/install_sd_cpp_prebuilt.py index 2156995356..3850412369 100644 --- a/studio/install_sd_cpp_prebuilt.py +++ b/studio/install_sd_cpp_prebuilt.py @@ -40,7 +40,13 @@ RELEASES_API = f"https://api.github.com/repos/{REPO}/releases/latest" # accelerator -> the token that must appear in a Linux/Windows asset name. _LINUX_ACCEL_TOKEN = {"rocm": "rocm", "vulkan": "vulkan"} -_WINDOWS_ACCEL_TOKEN = {"cuda": "cuda12", "vulkan": "vulkan", "rocm": "rocm", "cpu": "avx2", "auto": "avx2"} +_WINDOWS_ACCEL_TOKEN = { + "cuda": "cuda12", + "vulkan": "vulkan", + "rocm": "rocm", + "cpu": "avx2", + "auto": "avx2", +} # Tokens that mark an accelerator-specific Linux build; "auto"/"cpu" want none of them. _LINUX_ACCEL_MARKERS = ("rocm", "vulkan", "cuda", "sycl", "musa") @@ -73,11 +79,16 @@ def resolve_release_asset( system = system.lower() accel = accelerator.lower() arch = _arch_tokens(machine) - zips = [a for a in asset_names if a.lower().endswith(".zip") and not a.lower().startswith("cudart")] + zips = [ + a for a in asset_names if a.lower().endswith(".zip") and not a.lower().startswith("cudart") + ] if system == "darwin": - pool = [a for a in zips if ("darwin" in a.lower() or "macos" in a.lower()) - and any(t in a.lower() for t in arch)] + pool = [ + a + for a in zips + if ("darwin" in a.lower() or "macos" in a.lower()) and any(t in a.lower() for t in arch) + ] return pool[0] if pool else None if system == "windows": @@ -144,7 +155,10 @@ def install( release = _fetch_latest_release(token = token) names = [a["name"] for a in release.get("assets", [])] chosen = resolve_release_asset( - names, system = platform.system(), machine = platform.machine(), accelerator = accelerator, + names, + system = platform.system(), + machine = platform.machine(), + accelerator = accelerator, ) if not chosen: raise RuntimeError( @@ -172,16 +186,23 @@ def install( def main(argv: Optional[list[str]] = None) -> int: p = argparse.ArgumentParser(description = "Install a prebuilt sd-cli (stable-diffusion.cpp).") - p.add_argument("--accelerator", default = "auto", choices = ["auto", "cpu", "vulkan", "rocm", "cuda"]) + p.add_argument( + "--accelerator", default = "auto", choices = ["auto", "cpu", "vulkan", "rocm", "cuda"] + ) p.add_argument("--install-dir", default = None) - p.add_argument("--print-asset", action = "store_true", help = "resolve + print the asset, don't download") + p.add_argument( + "--print-asset", action = "store_true", help = "resolve + print the asset, don't download" + ) args = p.parse_args(argv) if args.print_asset: release = _fetch_latest_release() names = [a["name"] for a in release.get("assets", [])] chosen = resolve_release_asset( - names, system = platform.system(), machine = platform.machine(), accelerator = args.accelerator, + names, + system = platform.system(), + machine = platform.machine(), + accelerator = args.accelerator, ) print(chosen or "(no matching prebuilt; build from source)") return 0 if chosen else 2