65 lines
2.4 KiB
Python
65 lines
2.4 KiB
Python
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()
|