Studio diffusion (Phase 4) review round 2: collect sd-cli batch outputs

Codex review: when batch_count > 1, stable-diffusion.cpp's save_results() writes
the numbered files <stem>_<idx><suffix> (base_0.png, base_1.png, ...) instead of
the literal --output path. SdCppEngine.generate checked only the literal path, so
a batch generation would exit 0 and then raise 'no image' (or return a stale
file). generate now returns the literal path when present and otherwise falls
back to the numbered siblings; single-image behavior is unchanged.

Test: a fake sd-cli that writes img_0.png/img_1.png (not img.png) is collected
without error.
This commit is contained in:
Daniel Han 2026-06-29 10:33:09 +00:00
commit 4d87ff4b22
2 changed files with 37 additions and 3 deletions

View file

@ -280,13 +280,22 @@ class SdCppEngine:
if ret != 0:
raise RuntimeError(f"sd-cli exited {ret}. Last output:\n" + "\n".join(tail[-12:]))
if not out.is_file():
if out.is_file():
produced: Optional[Path] = out
else:
# For batch_count > 1, stable-diffusion.cpp's save_results() writes
# "<stem>_<idx><suffix>" (base_0.png, base_1.png, ...) rather than the
# literal --output path, so the single-path check above misses them.
# Fall back to the numbered siblings and return the first.
batch = sorted(out.parent.glob(f"{out.stem}_*{out.suffix}"))
produced = batch[0] if batch else None
if produced is None:
raise RuntimeError(
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)
return out
logger.info("sd-cli generate ok in %.1fs -> %s", time.time() - t0, produced)
return produced
# ── engine routing ──────────────────────────────────────────────────────────

View file

@ -211,6 +211,31 @@ def test_generate_success_returns_path_and_collects_logs(tmp_path, monkeypatch):
assert str(Path(e.binary).resolve().parent) in _FakePopen.captured_env.get(var, "")
def test_generate_collects_batch_output_paths(tmp_path, monkeypatch):
# batch_count > 1: stable-diffusion.cpp writes "<stem>_<idx><suffix>"
# (img_0.png, img_1.png, ...) rather than the literal --output path, so the
# single-path check must fall back to the numbered siblings.
e = _engine(tmp_path)
out = tmp_path / "img.png"
def _factory(cmd, **kw):
# Emulate batch save_results(): write the numbered files, NOT the literal path.
(tmp_path / "img_0.png").write_bytes(b"\x89PNG\r\n")
(tmp_path / "img_1.png").write_bytes(b"\x89PNG\r\n")
return _FakePopen(
cmd, lines = ["done"], returncode = 0, out_file = out, write = False, env = kw.get("env")
)
monkeypatch.setattr(eng.subprocess, "Popen", _factory)
result = e.generate(
SdCppModelFiles(diffusion_model = "/m/z.gguf"),
SdCppGenParams(prompt = "x", batch_count = 2),
output_path = str(out),
)
assert result == tmp_path / "img_0.png" and result.is_file()
assert not out.exists() # the literal --output path was never written
def test_generate_raises_on_nonzero_exit(tmp_path, monkeypatch):
e = _engine(tmp_path)
out = tmp_path / "img.png"