File size: 1,067 Bytes
cd02163 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 | """RunPod Serverless worker: one Jev/SystemOne request (or a list) per job, billed per second.
Job input: {"input": <request>} or {"input": {"requests": [<request>, ...]}}.
The model loads once per worker from DECIDER_MODEL (repo id or folder).
"""
import os
import sys
import runpod
import torch
from huggingface_hub import snapshot_download
MODEL = os.environ.get("DECIDER_MODEL", "Jeetu456/decider-2b")
path = MODEL if os.path.isdir(MODEL) else snapshot_download(MODEL)
sys.path.insert(0, path)
from decider_model import SchemaError, load_release_model, systemone # noqa: E402
model = load_release_model(path, device="cuda" if torch.cuda.is_available() else "cpu")
def handler(job):
request = job["input"]
try:
if isinstance(request, dict) and "requests" in request:
return [systemone(model, r) for r in request["requests"]]
return systemone(model, request)
except (SchemaError, ValueError) as error:
return {"error": str(error)}
if __name__ == "__main__":
runpod.serverless.start({"handler": handler})
|