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 <lora:name:weight> 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.
This commit is contained in:
Daniel Han
parent
33977e0991
commit
703d2df687
5 changed files with 283 additions and 30 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
``<lora:name:weight>`` 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 <lora:name:w> 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)."""
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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:mystyle:0.8>", 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 <lora:...> tag rides in the prompt unchanged
|
||||
assert _pair(cmd, "--prompt") == "a portrait <lora:mystyle:0.8>"
|
||||
|
||||
|
||||
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")
|
||||
|
|
|
|||
|
|
@ -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 ──────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue