init commit

scripts/serve.py

@@ -0,0 +1,64 @@
+import zmq
+import argparse
+import pickle
+
+from nitrogen.inference_session import InferenceSession
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Model inference server")
+    parser.add_argument("ckpt", type=str, help="Path to checkpoint file")
+    parser.add_argument("--port", type=int, default=5555, help="Port to serve on")
+    parser.add_argument("--old-layout", action="store_true", help="Use old layout")
+    parser.add_argument("--cfg", type=float, default=1.0, help="CFG scale")
+    parser.add_argument("--ctx", type=int, default=1, help="Context length")
+    args = parser.parse_args()
+
+    session = InferenceSession.from_ckpt(args.ckpt, old_layout=args.old_layout, cfg_scale=args.cfg, context_length=args.ctx)
+
+    # Setup ZeroMQ
+    context = zmq.Context()
+    socket = context.socket(zmq.REP)
+    socket.bind(f"tcp://*:{args.port}")
+
+    # Create poller
+    poller = zmq.Poller()
+    poller.register(socket, zmq.POLLIN)
+
+    print(f"\n{'='*60}")
+    print(f"Server running on port {args.port}")
+    print(f"Waiting for requests...")
+    print(f"{'='*60}\n")
+
+    try:
+        while True:
+            # Poll with 100ms timeout to allow interrupt handling
+            events = dict(poller.poll(timeout=100))
+            if socket in events and events[socket] == zmq.POLLIN:
+                # Receive request only when data is available
+                request = socket.recv()
+                request = pickle.loads(request)
+                if request["type"] == "reset":
+                    session.reset()
+                    response = {"status": "ok"}
+                    print("Session reset")
+                elif request["type"] == "info":
+                    info = session.info()
+                    response = {"status": "ok", "info": info}
+                    print("Sent session info")
+                elif request["type"] == "predict":
+                    raw_image = request["image"]
+                    result = session.predict(raw_image)
+                    response = {
+                        "status": "ok",
+                        "pred": result
+                    }
+                else:
+                    response = {"status": "error", "message": f"Unknown request type: {request['type']}"}
+                # Send response
+                socket.send(pickle.dumps(response))
+    except KeyboardInterrupt:
+        print("\nShutting down server...")
+        exit(0)
+    finally:
+        socket.close()
+        context.term()