Run sd-cli under a RAM watchdog and review candidates

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

1 files changed, 89 insertions, 0 deletions

diff --git a/fanfictioner b/fanfictioner
index 441d77a..f5b3e3c 100755
--- a/fanfictioner
+++ b/fanfictioner
@@ -59,6 +59,7 @@ MODEL_FLAGS = {"--diffusion-model", "--vae", "--llm", "--llm_vision"}
 SIZES = {"portrait": (832, 1216), "landscape": (1216, 832), "turnaround": (1216, 832)}
 # Edit inputs are shrunk first: full-size refs blow past 12 GB VRAM.
 SHRINK = {"portrait": "576x832", "landscape": "832x576", "ref": "560x384"}
+MIN_RAM_KB = 2 * 1024 * 1024  # kill sd-cli below 2 GB MemAvailable, before the OOM killer hits the desktop
 
 
 def load_conf(path=CONF):
@@ -230,6 +231,89 @@ def sd_args(profile, prompt, out, size, seed, refs=()):
     return args
 
 
+def mem_available_kb(meminfo):
+    return int(re.search(r"MemAvailable:\s+(\d+)", meminfo)[1])
+
+
+def run_sd(args, log):
+    """Run sd-cli with output in log; return an error string or None."""
+    with open(log, "w") as f:
+        f.write(shlex.join(args) + "\n")
+        f.flush()
+        proc = subprocess.Popen(args, stdout=f, stderr=subprocess.STDOUT)
+        try:
+            while proc.poll() is None:
+                time.sleep(2)
+                if proc.poll() is None and \
+                        mem_available_kb(Path("/proc/meminfo").read_text()) < MIN_RAM_KB:
+                    proc.kill()
+                    proc.wait()
+                    return "sd-cli killed by RAM watchdog (MemAvailable below 2 GB)"
+        except KeyboardInterrupt:
+            proc.kill()
+            proc.wait()
+            raise
+    if proc.returncode or "[ERROR" in Path(log).read_text(errors="replace"):
+        return f"sd-cli failed (exit {proc.returncode})"
+    return None
+
+
+def editor(path):
+    subprocess.run([*shlex.split(os.environ.get("EDITOR", "vi")), str(path)])
+
+
+def unload_llm():
+    pass
+
+
+def generate(d, profile, prompt, size, refs=()):
+    """One sd-cli call producing 3 candidates in d. Returns (files, seed, error)."""
+    unload_llm()
+    d.mkdir(parents=True, exist_ok=True)
+    seed = random.randint(0, 2**31 - 4)
+    log = d / f"{seed}.log"
+    print(f"sd-cli {profile}, seed {seed} ...", flush=True)
+    err = run_sd(sd_args(profile, prompt, d / f"{seed}_%d.png", size, seed, refs), log)
+    files = [d / f"{seed}_{i}.png" for i in (1, 2, 3)]
+    if not err and not all(f.exists() for f in files):
+        err = "sd-cli wrote no output"
+    if err:
+        print(err, *log.read_text(errors="replace").splitlines()[-15:], sep="\n")
+    return files, seed, err
+
+
+def show(files, seed, d):
+    m = d / f"{seed}_montage.png"
+    labelled = [x for i, f in enumerate(files) for x in ("-label", f"{i + 1}   seed {seed + i}", str(f))]
+    subprocess.run(["magick", "montage", "-pointsize", "32", *labelled,
+                    "-geometry", "+8+8", "-tile", "3x1", str(m)], check=True)
+    subprocess.run(["kitty", "+kitten", "icat", str(m)])
+
+
+def review(d, profile, prompt, size, refs=(), skip=False):
+    """Generate, show, ask. Returns the picked file, or None when the user skips."""
+    keys = "1-3 pick, r regen, p edit prompt" + (", s skip edit" if skip else "") + ", q quit"
+    while True:
+        files, seed, err = generate(d, profile, prompt, size, refs)
+        if not err:
+            show(files, seed, d)
+        while True:
+            k = input(f"[{keys}] > ").strip().lower()
+            if k in ("1", "2", "3") and not err:
+                return files[int(k) - 1]
+            if k == "r":
+                break
+            if k == "p":
+                (d / "prompt.txt").write_text(prompt)
+                editor(d / "prompt.txt")
+                prompt = (d / "prompt.txt").read_text().strip()
+                break
+            if k == "s" and skip:
+                return None
+            if k == "q":
+                sys.exit("quit, run again to resume")
+
+
 def selftest():
     plan = parse_plan(SAMPLE_PLAN)
     assert (plan["title"], plan["series"], plan["book"]) == ("Rooftop", "Test Series", "Book One"), plan
@@ -335,6 +419,11 @@ def selftest():
         keep(t / "refs" / "Mara.png", t / "deep" / "x.png")
         assert (t / "deep" / "x.png").read_bytes() == (t / "refs" / "Mara.png").read_bytes()
         assert not list((t / "deep").glob("*.tmp"))
+        assert mem_available_kb("MemTotal: 30000000 kB\nMemAvailable:   1500000 kB\n") == 1500000
+        assert run_sd(["true"], t / "ok.log") is None
+        assert (t / "ok.log").read_text().startswith("true")
+        assert run_sd(["sh", "-c", "echo '[ERROR] boom'"], t / "e.log") == "sd-cli failed (exit 0)"
+        assert run_sd(["false"], t / "f.log") == "sd-cli failed (exit 1)"
     print("selftest ok")