Add persistent HTTP daemon holding one model resident in VRAM

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

1 files changed, 131 insertions, 0 deletions

diff --git a/server.py b/server.py
new file mode 100644
index 0000000..f3be9e5
--- /dev/null
+++ b/server.py
@@ -0,0 +1,131 @@
+#!/usr/bin/env python3
+# Copyright (C) 2026 Danilo M. <danix@danix.xyz>
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License version 2 as
+# published by the Free Software Foundation.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU General Public License for more details.
+
+# Long-lived daemon that keeps ONE SDXL model resident in XPU VRAM so prompt
+# iteration skips the model-reload tax. Stdlib HTTP only. Serial by design:
+# one GPU, so concurrent generation would OOM. See README.md.
+
+import io
+import json
+import os
+import sys
+from http.server import BaseHTTPRequestHandler, HTTPServer
+
+import torch
+
+import generate
+
+PORT = int(os.environ.get("IMGGEN_PORT", "8765"))
+
+STATE = {"model": None, "pipe": None}
+
+
+def load(model_key):
+    if STATE["model"] == model_key and STATE["pipe"] is not None:
+        return
+    if STATE["pipe"] is not None:
+        del STATE["pipe"]
+        STATE["pipe"] = None
+        STATE["model"] = None
+        torch.xpu.empty_cache()
+    STATE["pipe"] = generate.build_pipe(model_key)
+    STATE["model"] = model_key
+
+
+class Handler(BaseHTTPRequestHandler):
+    def _json(self, code, obj):
+        body = json.dumps(obj).encode()
+        self.send_response(code)
+        self.send_header("Content-Type", "application/json")
+        self.send_header("Content-Length", str(len(body)))
+        self.end_headers()
+        self.wfile.write(body)
+
+    def _read_json(self):
+        length = int(self.headers.get("Content-Length", "0"))
+        raw = self.rfile.read(length) if length else b"{}"
+        return json.loads(raw)  # raises on bad JSON -> caught by caller
+
+    def do_GET(self):
+        if self.path == "/status":
+            self._json(200, {"model": STATE["model"], "ready": STATE["pipe"] is not None})
+        else:
+            self._json(404, {"error": "not found"})
+
+    def do_POST(self):
+        try:
+            data = self._read_json()
+        except (ValueError, json.JSONDecodeError):
+            self._json(400, {"error": "bad JSON"})
+            return
+
+        if self.path == "/model":
+            name = data.get("name")
+            try:
+                generate.resolve_model(name)
+            except SystemExit as e:
+                self._json(400, {"error": str(e)})
+                return
+            try:
+                load(name)
+            except Exception as e:  # download / load failure
+                self._json(500, {"error": f"load failed: {e}"})
+                return
+            self._json(200, {"model": STATE["model"], "ready": True})
+            return
+
+        if self.path == "/generate":
+            try:
+                image, elapsed = generate.generate(
+                    STATE["pipe"],
+                    data["prompt"],
+                    data.get("negative", ""),
+                    int(data.get("steps", 25)),
+                    float(data.get("guidance", 7.0)),
+                    data.get("seed"),
+                    int(data.get("width", 1024)),
+                    int(data.get("height", 1024)),
+                )
+            except torch.OutOfMemoryError:
+                torch.xpu.empty_cache()
+                self._json(500, {"error": "OOM, try fewer steps or smaller size"})
+                return
+            except Exception as e:
+                self._json(500, {"error": str(e)})
+                return
+            buf = io.BytesIO()
+            image.save(buf, format="PNG")
+            body = buf.getvalue()
+            self.send_response(200)
+            self.send_header("Content-Type", "image/png")
+            self.send_header("X-Elapsed", f"{elapsed:.2f}")
+            self.send_header("Content-Length", str(len(body)))
+            self.end_headers()
+            self.wfile.write(body)
+            return
+
+        self._json(404, {"error": "not found"})
+
+    def log_message(self, fmt, *args):
+        sys.stderr.write("%s - %s\n" % (self.address_string(), fmt % args))
+
+
+def main():
+    startup = os.environ.get("IMGGEN_MODEL", generate.DEFAULT_MODEL)
+    print(f"Loading {startup} on XPU ...", file=sys.stderr)
+    load(startup)
+    print(f"Ready on 127.0.0.1:{PORT} with model {startup}", file=sys.stderr)
+    HTTPServer(("0.0.0.0", PORT), Handler).serve_forever()
+
+
+if __name__ == "__main__":
+    main()