From 4d87ff4b228b92127dd5b310f78b6cee68134fe0 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Mon, 29 Jun 2026 10:33:09 +0000
Subject: [PATCH] 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.
---
 .../backend/core/inference/sd_cpp_engine.py   | 15 ++++++++---
 studio/backend/tests/test_sd_cpp_engine.py    | 25 +++++++++++++++++++
 2 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/studio/backend/core/inference/sd_cpp_engine.py b/studio/backend/core/inference/sd_cpp_engine.py
index 7809661827..80078ef6fc 100644
--- a/studio/backend/core/inference/sd_cpp_engine.py
+++ b/studio/backend/core/inference/sd_cpp_engine.py
@@ -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 ──────────────────────────────────────────────────────────
diff --git a/studio/backend/tests/test_sd_cpp_engine.py b/studio/backend/tests/test_sd_cpp_engine.py
index ac63596a97..ca6e875ea6 100644
--- a/studio/backend/tests/test_sd_cpp_engine.py
+++ b/studio/backend/tests/test_sd_cpp_engine.py
@@ -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"