From 703d2df687a13684321dbd528b815e5c2005f3e9 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 25 Jun 2026 16:06:38 +0000 Subject: [PATCH] Studio diffusion (Phase 6): img2img / inpaint / edit / LoRA / upscale on the native engine Builds on Phase 4's native stable-diffusion.cpp engine, extending it from text-to-image to the wider feature surface, since sd.cpp supports all of these through the binary already. Pure command-builder additions plus one engine method, so the txt2img path is unchanged. - sd_cpp_args.py: SdCppGenParams gains image-conditioning fields. init_img + strength make a run img2img, adding mask makes it inpaint, ref_images drives FLUX-Kontext / Qwen-Image-Edit style editing (repeated --ref-image), and lora_dir + the prompt syntax select LoRAs. New SdCppUpscaleParams + build_sd_cpp_upscale_command for the ESRGAN upscale run mode (input image + esrgan model, no prompt / text encoders). - sd_cpp_engine.py: the subprocess runner is factored into a shared _run() so generate() (now carrying the conditioning flags) and a new upscale() reuse the same streaming / error / output-check path. - scripts/sd_cpp_smoke.py: --task {txt2img,img2img,upscale} with --init-img / --strength / --upscale-model / --upscale-repeats. Tests: 10 new across the img2img / inpaint / edit / LoRA flag construction, the upscale builder and its validation, and the engine's img2img + upscale paths. Full diffusion suite 176 passing. Verified on a B200 box through SdCppEngine: img2img (Z-Image-Turbo Q4_K, the init image conditioned at strength 0.6, 4.8s) and ESRGAN upscale (512x512 -> 2048x2048 via RealESRGAN_x4plus_anime_6B, 2.7s), both producing coherent images. Video and the diffusers-path feature wiring are deferred. --- scripts/sd_cpp_smoke.py | 49 +++++++++-- studio/backend/core/inference/sd_cpp_args.py | 76 ++++++++++++++++- .../backend/core/inference/sd_cpp_engine.py | 68 ++++++++++++--- studio/backend/tests/test_sd_cpp_args.py | 85 +++++++++++++++++-- studio/backend/tests/test_sd_cpp_engine.py | 35 +++++++- 5 files changed, 283 insertions(+), 30 deletions(-) diff --git a/scripts/sd_cpp_smoke.py b/scripts/sd_cpp_smoke.py index 4398fa6cf3..ff5253caba 100644 --- a/scripts/sd_cpp_smoke.py +++ b/scripts/sd_cpp_smoke.py @@ -41,6 +41,7 @@ from core.inference.diffusion_memory import ( # noqa: E402 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 @@ -55,9 +56,15 @@ _MODE_TO_POLICY = { 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", required = True) + 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) @@ -90,6 +97,28 @@ def main(argv: list[str] | None = None) -> int: ) 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, @@ -98,20 +127,22 @@ def main(argv: list[str] | None = None) -> int: 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, + 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) - out = Path(args.out_image) t0 = time.time() result = engine.generate( files, diff --git a/studio/backend/core/inference/sd_cpp_args.py b/studio/backend/core/inference/sd_cpp_args.py index 876cd7e557..6a5f5e27c2 100644 --- a/studio/backend/core/inference/sd_cpp_args.py +++ b/studio/backend/core/inference/sd_cpp_args.py @@ -76,7 +76,14 @@ class SdCppModelFiles: @dataclass(frozen = True) class SdCppGenParams: - """Generation parameters, mapped 1:1 onto sd-cli's sampling flags.""" + """Generation parameters, mapped 1:1 onto sd-cli's sampling flags. + + The image-conditioning fields cover the img_gen variants: ``init_img`` + + ``strength`` make it img2img, adding ``mask`` makes it inpaint, and + ``ref_images`` drives FLUX Kontext / Qwen-Image-Edit style editing. ``lora_dir`` + points sd-cli at a LoRA directory; the LoRAs themselves are selected with + ```` tags inside ``prompt`` (sd.cpp's own syntax). + """ prompt: str negative_prompt: Optional[str] = None @@ -88,6 +95,24 @@ class SdCppGenParams: seed: Optional[int] = None sampling_method: Optional[str] = None batch_count: int = 1 + # image-to-image / inpaint / edit + init_img: Optional[str] = None + strength: Optional[float] = None + mask: Optional[str] = None + ref_images: tuple[str, ...] = () + # LoRA + lora_dir: Optional[str] = None + lora_apply_mode: Optional[str] = None + + +@dataclass(frozen = True) +class SdCppUpscaleParams: + """Inputs for sd-cli's ESRGAN upscale mode (a separate run mode).""" + + input_image: str + upscale_model: str + repeats: int = 1 + tile_size: Optional[int] = None def offload_flags( @@ -168,6 +193,20 @@ def build_sd_cpp_command( cmd += ["--prompt", params.prompt] if params.negative_prompt: cmd += ["--negative-prompt", params.negative_prompt] + # img2img / inpaint / edit conditioning (img_gen mode with an input image). + if params.init_img: + cmd += ["--init-img", params.init_img] + if params.strength is not None: + cmd += ["--strength", _fmt_float(params.strength)] + if params.mask: + cmd += ["--mask", params.mask] + for ref in params.ref_images: + cmd += ["--ref-image", ref] + # LoRA: the directory to scan; individual LoRAs are tags in prompt. + if params.lora_dir: + cmd += ["--lora-model-dir", params.lora_dir] + if params.lora_apply_mode: + cmd += ["--lora-apply-mode", params.lora_apply_mode] cmd += ["--width", str(int(params.width)), "--height", str(int(params.height))] if params.steps is not None: cmd += ["--steps", str(int(params.steps))] @@ -194,6 +233,41 @@ def build_sd_cpp_command( return cmd +def build_sd_cpp_upscale_command( + binary: str, + params: SdCppUpscaleParams, + *, + output_path: str, + verbose: bool = False, + extra_args: Optional[list[str]] = None, +) -> list[str]: + """Build the ``sd-cli --mode upscale`` argv (ESRGAN super-resolution). + + Upscale is a distinct run mode: it takes an input image and an ESRGAN model, + no prompt or text encoders. ``repeats`` runs the upscaler N times (each pass + is a fixed scale factor for the model). + """ + if not params.input_image: + raise ValueError("input_image is required for upscale") + if not params.upscale_model: + raise ValueError("upscale_model is required for upscale") + cmd: list[str] = [ + binary, "--mode", "upscale", + "--init-img", params.input_image, + "--upscale-model", params.upscale_model, + ] + if params.repeats and params.repeats != 1: + cmd += ["--upscale-repeats", str(int(params.repeats))] + if params.tile_size is not None: + cmd += ["--upscale-tile-size", str(int(params.tile_size))] + cmd += ["--output", output_path] + if verbose: + cmd += ["-v"] + if extra_args: + cmd += list(extra_args) + return cmd + + def _fmt_float(value: float) -> str: """Compact float -> str: drop a trailing ``.0`` so ``1.0`` -> ``1`` (sd-cli accepts both, but the tidy form keeps logged commands readable).""" diff --git a/studio/backend/core/inference/sd_cpp_engine.py b/studio/backend/core/inference/sd_cpp_engine.py index b3b921419b..1b85fd89fa 100644 --- a/studio/backend/core/inference/sd_cpp_engine.py +++ b/studio/backend/core/inference/sd_cpp_engine.py @@ -34,7 +34,9 @@ from typing import Callable, Optional from core.inference.sd_cpp_args import ( SdCppGenParams, SdCppModelFiles, + SdCppUpscaleParams, build_sd_cpp_command, + build_sd_cpp_upscale_command, ) logger = logging.getLogger(__name__) @@ -200,28 +202,68 @@ class SdCppEngine: nonzero, or no output file is produced. ``on_log`` (if given) receives each line of sd-cli's progress output as it arrives. """ + cmd = build_sd_cpp_command( + self._require_binary(), files, params, + output_path = str(self._prepare_out(output_path)), offload = offload, + threads = threads, verbose = verbose, extra_args = extra_args, + ) + return self._run(cmd, output_path, timeout = timeout, env = env, on_log = on_log) + + def upscale( + self, + params: "SdCppUpscaleParams", + *, + output_path: str, + verbose: bool = False, + extra_args: Optional[list[str]] = None, + timeout: Optional[float] = 1800.0, + env: Optional[dict[str, str]] = None, + on_log: Optional[Callable[[str], None]] = None, + ) -> Path: + """Upscale an image with an ESRGAN model; return the written path.""" + cmd = build_sd_cpp_upscale_command( + self._require_binary(), params, + output_path = str(self._prepare_out(output_path)), + verbose = verbose, extra_args = extra_args, + ) + return self._run(cmd, output_path, timeout = timeout, env = env, on_log = on_log) + + # ── internals ───────────────────────────────────────────────────────────── + + def _require_binary(self) -> str: if not self.is_available(): raise RuntimeError( "sd-cli (stable-diffusion.cpp) binary not found. Build it or set " "SD_CLI_PATH / UNSLOTH_SD_CPP_PATH." ) + return self.binary # type: ignore[return-value] + + @staticmethod + def _prepare_out(output_path: str) -> Path: 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, - ) + return out + + def _run( + self, + cmd: list[str], + output_path: str, + *, + timeout: Optional[float], + env: Optional[dict[str, str]], + on_log: Optional[Callable[[str], None]], + ) -> Path: + """Run an sd-cli argv, stream output, and return the produced image path. + + Raises ``RuntimeError`` on nonzero exit, timeout, or a missing output. + Shared by ``generate`` and ``upscale``. + """ + out = Path(output_path) base = dict(os.environ) if env: base.update(env) - run_env = runtime_env(self.binary, base) - logger.info("sd-cli generate: %s", " ".join(cmd)) + run_env = runtime_env(self._require_binary(), base) + logger.info("sd-cli run: %s", " ".join(cmd)) t0 = time.time() proc = subprocess.Popen( @@ -257,7 +299,7 @@ class SdCppEngine: f"sd-cli reported success but no image at {out}. Last output:\n" + "\n".join(tail[-12:]) ) - logger.info("sd-cli generate ok in %.1fs -> %s", time.time() - t0, out) + logger.info("sd-cli run ok in %.1fs -> %s", time.time() - t0, out) return out diff --git a/studio/backend/tests/test_sd_cpp_args.py b/studio/backend/tests/test_sd_cpp_args.py index 99e758e22e..f3287985df 100644 --- a/studio/backend/tests/test_sd_cpp_args.py +++ b/studio/backend/tests/test_sd_cpp_args.py @@ -20,7 +20,9 @@ from core.inference.diffusion_memory import ( from core.inference.sd_cpp_args import ( SdCppGenParams, SdCppModelFiles, + SdCppUpscaleParams, build_sd_cpp_command, + build_sd_cpp_upscale_command, offload_flags, text_encoder_flags_for_family, ) @@ -181,9 +183,80 @@ def test_build_requires_diffusion_model_and_prompt(): 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") + + +# ── img2img / inpaint / edit / LoRA (Phase 6) ─────────────────────────────── + + +def test_build_img2img_adds_init_and_strength(): + files = SdCppModelFiles(diffusion_model = "/m/z.gguf", vae = "/m/ae.sft", llm = "/m/q.gguf") + params = SdCppGenParams(prompt = "make it autumn", init_img = "/in/src.png", strength = 0.6) + cmd = build_sd_cpp_command("/bin/sd-cli", files, params, output_path = "/o.png") + assert _pair(cmd, "--init-img") == "/in/src.png" + assert _pair(cmd, "--strength") == "0.6" + assert _pair(cmd, "--mode") == "img_gen" # img2img is still img_gen mode + + +def test_build_inpaint_adds_mask(): + files = SdCppModelFiles(diffusion_model = "/m/z.gguf") + params = SdCppGenParams(prompt = "x", init_img = "/in/src.png", mask = "/in/mask.png", strength = 0.8) + cmd = build_sd_cpp_command("/bin/sd-cli", files, params, output_path = "/o.png") + assert _pair(cmd, "--mask") == "/in/mask.png" + assert _pair(cmd, "--init-img") == "/in/src.png" + + +def test_build_edit_repeats_ref_image(): + files = SdCppModelFiles(diffusion_model = "/m/flux.gguf") + params = SdCppGenParams(prompt = "add a hat", ref_images = ("/r/a.png", "/r/b.png")) + cmd = build_sd_cpp_command("/bin/sd-cli", files, params, output_path = "/o.png") + # each ref image gets its own --ref-image flag + idxs = [i for i, t in enumerate(cmd) if t == "--ref-image"] + assert len(idxs) == 2 + assert [cmd[i + 1] for i in idxs] == ["/r/a.png", "/r/b.png"] + + +def test_build_lora_dir_and_apply_mode(): + files = SdCppModelFiles(diffusion_model = "/m/z.gguf") + params = SdCppGenParams( + prompt = "a portrait ", lora_dir = "/loras", lora_apply_mode = "at_runtime", + ) + cmd = build_sd_cpp_command("/bin/sd-cli", files, params, output_path = "/o.png") + assert _pair(cmd, "--lora-model-dir") == "/loras" + assert _pair(cmd, "--lora-apply-mode") == "at_runtime" + # the tag rides in the prompt unchanged + assert _pair(cmd, "--prompt") == "a portrait " + + +def test_txt2img_omits_image_conditioning_flags(): + files = SdCppModelFiles(diffusion_model = "/m/z.gguf") + cmd = build_sd_cpp_command("/bin/sd-cli", files, SdCppGenParams(prompt = "x"), output_path = "/o.png") + for flag in ("--init-img", "--strength", "--mask", "--ref-image", "--lora-model-dir"): + assert flag not in cmd + + +# ── upscale mode ──────────────────────────────────────────────────────────── + + +def test_build_upscale_command(): + params = SdCppUpscaleParams(input_image = "/in/small.png", upscale_model = "/m/esrgan.pth", repeats = 2) + cmd = build_sd_cpp_upscale_command("/bin/sd-cli", params, output_path = "/out/big.png") + assert _pair(cmd, "--mode") == "upscale" + assert _pair(cmd, "--init-img") == "/in/small.png" + assert _pair(cmd, "--upscale-model") == "/m/esrgan.pth" + assert _pair(cmd, "--upscale-repeats") == "2" + assert _pair(cmd, "--output") == "/out/big.png" + # no prompt / text-encoder flags in upscale mode + assert "--prompt" not in cmd and "--llm" not in cmd + + +def test_build_upscale_requires_input_and_model(): + with pytest.raises(ValueError): + build_sd_cpp_upscale_command("/bin/sd-cli", + SdCppUpscaleParams(input_image = "", upscale_model = "/m/e.pth"), + output_path = "/o.png") + with pytest.raises(ValueError): + build_sd_cpp_upscale_command("/bin/sd-cli", + SdCppUpscaleParams(input_image = "/i.png", upscale_model = ""), + 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 25dd19cec6..f8080d183a 100644 --- a/studio/backend/tests/test_sd_cpp_engine.py +++ b/studio/backend/tests/test_sd_cpp_engine.py @@ -26,7 +26,7 @@ from core.inference.sd_cpp_engine import ( runtime_env, select_diffusion_engine, ) -from core.inference.sd_cpp_args import SdCppGenParams, SdCppModelFiles +from core.inference.sd_cpp_args import SdCppGenParams, SdCppModelFiles, SdCppUpscaleParams # ── binary discovery ──────────────────────────────────────────────────────── @@ -242,6 +242,39 @@ def test_generate_raises_when_binary_missing(): ) +def test_img2img_generate_passes_init_image(tmp_path, monkeypatch): + e = _engine(tmp_path) + out = tmp_path / "img.png" + src = tmp_path / "src.png" + src.write_bytes(b"\x89PNG\r\n") + _patch_popen(monkeypatch, lines = ["img2img"], returncode = 0, out_file = out) + e.generate(SdCppModelFiles(diffusion_model = "/m/z.gguf"), + SdCppGenParams(prompt = "x", init_img = str(src), strength = 0.5), + output_path = str(out)) + assert "--init-img" in _FakePopen.captured_cmd + assert str(src) == _FakePopen.captured_cmd[_FakePopen.captured_cmd.index("--init-img") + 1] + + +def test_upscale_runs_and_returns_path(tmp_path, monkeypatch): + e = _engine(tmp_path) + out = tmp_path / "big.png" + _patch_popen(monkeypatch, lines = ["upscaling", "done"], returncode = 0, out_file = out) + result = e.upscale( + SdCppUpscaleParams(input_image = "/in/small.png", upscale_model = "/m/esrgan.pth", repeats = 2), + output_path = str(out), + ) + assert result == out and out.is_file() + assert _FakePopen.captured_cmd[_FakePopen.captured_cmd.index("--mode") + 1] == "upscale" + assert "--upscale-model" in _FakePopen.captured_cmd + + +def test_upscale_raises_when_binary_missing(): + e = SdCppEngine(binary = None) + with pytest.raises(RuntimeError, match = "not found"): + e.upscale(SdCppUpscaleParams(input_image = "/i.png", upscale_model = "/m/e.pth"), + output_path = "/tmp/x.png") + + # ── engine routing ──────────────────────────────────────────────────────────