Add worker.pys

2025-11-25 13:33:12 -08:00
parent 0339b471c5
commit 2ce0450809
5 changed files with 256 additions and 3 deletions
@@ -0,0 +1,84 @@
+import random
+import sys
+
+from vastai import Worker, WorkerConfig, HandlerConfig, LogActionConfig, BenchmarkConfig
+
+# ComyUI model configuration
+MODEL_SERVER_URL           = 'http://127.0.0.1'
+MODEL_SERVER_PORT          = 18288
+MODEL_LOG_FILE             = '/var/log/portal/comfyui.log'
+MODEL_HEALTHCHECK_ENDPOINT = "/health"
+
+# ComyUI-specific log messages
+MODEL_LOAD_LOG_MSG = [
+    "To see the GUI go to: "
+]
+
+MODEL_ERROR_LOG_MSGS = [
+    "MetadataIncompleteBuffer",
+    "Value not in list: ",
+    "[ERROR] Provisioning Script failed"
+]
+
+MODEL_INFO_LOG_MSGS = [
+    '"message":"Downloading'
+]
+
+benchmark_prompts = [
+    "Cartoon hoodie hero; orc, anime cat, bunny; black goo; buff; vector on white.",
+    "Cozy farming-game scene with fine details.",
+    "2D vector child with soccer ball; airbrush chrome; swagger; antique copper.",
+    "Realistic futuristic downtown of low buildings at sunset.",
+    "Perfect wave front view; sunny seascape; ultra-detailed water; artful feel.",
+    "Clear cup with ice, fruit, mint; creamy swirls; fluid-sim CGI; warm glow.",
+    "Male biker with backpack on motorcycle; oilpunk; award-worthy magazine cover.",
+    "Collage for textile; surreal cartoon cat in cap/jeans before poster; crisp.",
+    "Medieval village inside glass sphere; volumetric light; macro focus.",
+    "Iron Man with glowing axe; mecha sci-fi; jungle scene; dynamic light.",
+    "Pope Francis DJ in leather jacket, mixing on giant console; dramatic.",
+]
+
+
+def parse_request(json_msg):
+    return {"input" : json_msg}
+
+benchmark_dataset = [
+    {
+        "input": {
+            "request_id": f"test-{random.randint(1000, 99999)}",
+            "modifier": "Text2Image",
+            "modifications": {
+                "prompt": prompt,
+                "width": 512,
+                "height": 512,
+                "steps": 20,
+                "seed": random.randint(0, sys.maxsize)
+            }
+        }
+    } for prompt in benchmark_prompts
+]
+
+worker_config = WorkerConfig(
+    model_server_url=MODEL_SERVER_URL,
+    model_server_port=MODEL_SERVER_PORT,
+    model_log_file=MODEL_LOG_FILE,
+    model_healthcheck_url=MODEL_HEALTHCHECK_ENDPOINT,
+    handlers=[
+        HandlerConfig(
+            route="/generate/sync",
+            allow_parallel_requests=False,
+            max_queue_time=10.0,
+            request_parser=parse_request,
+            benchmark_config=BenchmarkConfig(
+                dataset=benchmark_dataset,
+            )
+        )
+    ],
+    log_action_config=LogActionConfig(
+        on_load=MODEL_LOAD_LOG_MSG,
+        on_error=MODEL_ERROR_LOG_MSGS,
+        on_info=MODEL_INFO_LOG_MSGS
+    )
+)
+
+Worker(worker_config).run()