From f04138e13bee6835e9237fd129f4f7a40c1550ff Mon Sep 17 00:00:00 2001 From: Colter Downing Date: Wed, 3 Dec 2025 20:16:25 -0800 Subject: [PATCH 1/5] update to be able to get images --- workers/comfyui-json/client.py | 351 ++++++++++++++++++++++++++++++--- workers/comfyui-json/server.py | 32 +++ 2 files changed, 359 insertions(+), 24 deletions(-) diff --git a/workers/comfyui-json/client.py b/workers/comfyui-json/client.py index 93e184c..b80a9ba 100644 --- a/workers/comfyui-json/client.py +++ b/workers/comfyui-json/client.py @@ -1,35 +1,338 @@ -from .data_types import count_workload +import os +import sys +import json import uuid import random +import base64 import asyncio -import random +import logging +import argparse from vastai import Serverless -async def main(): - async with Serverless() as client: - endpoint = await client.get_endpoint(name="my-comfy-endpoint") # Change this to your endpoint name +# ---------------------- Config ---------------------- +DEFAULT_PROMPT = "a beautiful sunset over mountains, digital art, highly detailed" +ENDPOINT_NAME = "Comfy-Prod2" +DEFAULT_WIDTH = 512 +DEFAULT_HEIGHT = 512 +DEFAULT_STEPS = 20 +COST = 100 # Fixed cost for ComfyUI requests - payload = { - "input": { - "request_id": str(uuid.uuid4()), - "modifier": "Text2Image", - "modifications": { - "prompt": "a beautiful landscape with mountains and lakes", - "width": 1024, - "height": 1024, - "steps": 20, - "seed": random.randint(0, 2**32 - 1) - }, - "workflow_json": {} # Empty since using modifier approach - } +logging.basicConfig(level=logging.INFO, format="%(levelname)s - %(message)s") +log = logging.getLogger(__name__) + + +# ---------------------- API Functions ---------------------- +async def call_generate( + client: Serverless, + *, + endpoint_name: str, + prompt: str, + width: int, + height: int, + steps: int, + seed: int, +) -> dict: + """Generate image using Text2Image modifier""" + endpoint = await client.get_endpoint(name=endpoint_name) + payload = { + "input": { + "request_id": str(uuid.uuid4()), + "modifier": "Text2Image", + "modifications": { + "prompt": prompt, + "width": width, + "height": height, + "steps": steps, + "seed": seed, + }, } - - response = await endpoint.request("/generate/sync", payload, cost=count_workload()) + } + return await endpoint.request("/generate/sync", payload, cost=COST) + + +async def call_generate_workflow( + client: Serverless, + *, + endpoint_name: str, + workflow_json: dict, +) -> dict: + """Generate using custom workflow JSON""" + endpoint = await client.get_endpoint(name=endpoint_name) + payload = { + "input": { + "request_id": str(uuid.uuid4()), + "workflow_json": workflow_json, + } + } + return await endpoint.request("/generate/sync", payload, cost=COST) + + +# ---------------------- Demo Class ---------------------- +class APIDemo: + def __init__(self, client: Serverless, endpoint_name: str): + self.client = client + self.endpoint_name = endpoint_name + + def extract_images(self, response: dict) -> list: + """Extract image info from ComfyUI response""" + images = [] + + # Check for output array (S3/webhook configured) + if "output" in response: + for item in response["output"]: + if "url" in item: + images.append({"type": "url", "path": item["url"]}) + elif "local_path" in item: + images.append({"type": "local", "path": item["local_path"]}) + elif "base64" in item: + images.append({"type": "base64", "data": item["base64"]}) + + # Check for comfyui_response format (default) + if "comfyui_response" in response: + for prompt_id, data in response["comfyui_response"].items(): + if isinstance(data, dict) and "outputs" in data: + for node_id, node_output in data["outputs"].items(): + if "images" in node_output: + for img in node_output["images"]: + images.append({ + "type": "remote", + "filename": img.get("filename"), + "subfolder": img.get("subfolder", ""), + }) + + return images + + async def save_images(self, images: list, worker_url: str, prefix: str = "comfy") -> list: + """Save images locally by fetching from remote server""" + os.makedirs("generated_images", exist_ok=True) + saved = [] + seen = set() + + for i, img in enumerate(images): + if img["type"] == "base64": + data = img["data"] + if data.startswith("data:"): + data = data.split(",", 1)[-1] + path = f"generated_images/{prefix}_{i}.png" + with open(path, "wb") as f: + f.write(base64.b64decode(data)) + print(f" šŸ’¾ Saved: {path}") + saved.append(path) + + elif img["type"] == "url": + url = img["path"] + if url in seen: + continue + seen.add(url) + try: + import urllib.request + path = f"generated_images/{prefix}_{len(saved)}.png" + urllib.request.urlretrieve(url, path) + print(f" šŸ’¾ Downloaded: {path}") + saved.append(path) + except Exception as e: + print(f" šŸ”— URL: {url}") + saved.append(url) + + elif img["type"] == "local": + remote_path = img["path"] + if remote_path in seen: + continue + seen.add(remote_path) + filename = os.path.basename(remote_path) + # Try to fetch via /view endpoint + local_path = await self._fetch_image(worker_url, filename, "", f"{prefix}_{len(saved)}.png") + if local_path: + saved.append(local_path) + else: + print(f" šŸ“‚ Remote: {remote_path}") + saved.append(remote_path) + + elif img["type"] == "remote": + filename = img["filename"] + if filename in seen: + continue + seen.add(filename) + subfolder = img.get("subfolder", "") + # Try to fetch via /view endpoint + local_path = await self._fetch_image(worker_url, filename, subfolder, f"{prefix}_{len(saved)}.png") + if local_path: + saved.append(local_path) + else: + print(f" šŸ–¼ļø Remote: {filename}") + saved.append(filename) + + return saved + + async def _fetch_image(self, worker_url: str, filename: str, subfolder: str, local_name: str) -> str | None: + """Fetch image directly from worker's /view endpoint""" + if not worker_url: + print(f" āš ļø No worker URL available") + return None + + try: + import aiohttp + + params = {"filename": filename, "type": "output"} + if subfolder: + params["subfolder"] = subfolder + + url = f"{worker_url}/view" + print(f" šŸ”— Fetching from: {url}") + + async with aiohttp.ClientSession() as session: + async with session.get(url, params=params, ssl=False) as resp: + if resp.status == 200: + raw_bytes = await resp.read() + path = f"generated_images/{local_name}" + with open(path, "wb") as f: + f.write(raw_bytes) + print(f" šŸ’¾ Saved: {path}") + return path + else: + text = await resp.text() + print(f" āŒ HTTP {resp.status}: {text[:100]}") + return None + except Exception as e: + print(f" āŒ Fetch error: {e}") + return None + + async def demo_prompt( + self, + prompt: str, + width: int, + height: int, + steps: int, + seed: int | None, + ): + """Demo: Generate image from text prompt""" + print("=" * 60) + print("COMFYUI TEXT-TO-IMAGE DEMO") + print("=" * 60) + + if seed is None: + seed = random.randint(0, 2**32 - 1) + + print(f"Prompt: {prompt[:100]}..." if len(prompt) > 100 else f"Prompt: {prompt}") + print(f"Size: {width}x{height}, Steps: {steps}, Seed: {seed}") + print("\nšŸŽØ Generating image...") + + response = await call_generate( + self.client, + endpoint_name=self.endpoint_name, + prompt=prompt, + width=width, + height=height, + steps=steps, + seed=seed, + ) + + print("\nāœ… Generation complete!") + + # Get worker URL for fetching images + worker_url = response.get("url", "") + print(f"Worker URL: {worker_url}") + + # Extract and handle images + if "response" in response: + images = self.extract_images(response["response"]) + if images: + print(f"\nšŸ“ {len(images)} image(s) generated:") + await self.save_images(images, worker_url, prefix=f"comfy_{seed}") + else: + print("\nNo images found in response") + print(json.dumps(response, indent=2, default=str)[:2000]) + else: + print("\nUnexpected response format") + print(json.dumps(response, indent=2, default=str)[:2000]) + + async def demo_workflow(self, workflow_file: str): + """Demo: Generate using custom workflow file""" + print("=" * 60) + print("COMFYUI CUSTOM WORKFLOW DEMO") + print("=" * 60) + + if not os.path.exists(workflow_file): + log.error(f"Workflow file not found: {workflow_file}") + return + + with open(workflow_file, "r") as f: + workflow_json = json.load(f) + + print(f"Workflow: {workflow_file}") + print("\nšŸŽØ Generating...") + + response = await call_generate_workflow( + self.client, + endpoint_name=self.endpoint_name, + workflow_json=workflow_json, + ) + + print("\nāœ… Generation complete!") + + worker_url = response.get("url", "") + + if "response" in response: + images = self.extract_images(response["response"]) + if images: + print(f"\nšŸ“ {len(images)} image(s) generated:") + await self.save_images(images, worker_url, prefix="workflow") + else: + print("\nNo images found in response") + print(json.dumps(response, indent=2, default=str)[:2000]) + else: + print("\nUnexpected response format") + print(json.dumps(response, indent=2, default=str)[:2000]) + + +# ---------------------- CLI ---------------------- +def build_arg_parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser(description="Vast ComfyUI-JSON Demo (Serverless SDK)") + p.add_argument("--endpoint", default=ENDPOINT_NAME, help=f"Vast endpoint name (default: {ENDPOINT_NAME})") + p.add_argument("--prompt", type=str, default=DEFAULT_PROMPT, metavar="TEXT", + help=f"Prompt text (default: '{DEFAULT_PROMPT[:30]}...')") + p.add_argument("--workflow", type=str, metavar="FILE", help="Use custom workflow JSON file instead") + p.add_argument("--width", type=int, default=DEFAULT_WIDTH, help=f"Image width (default: {DEFAULT_WIDTH})") + p.add_argument("--height", type=int, default=DEFAULT_HEIGHT, help=f"Image height (default: {DEFAULT_HEIGHT})") + p.add_argument("--steps", type=int, default=DEFAULT_STEPS, help=f"Steps (default: {DEFAULT_STEPS})") + p.add_argument("--seed", type=int, default=None, help="Seed (default: random)") + return p + + +async def main_async(): + args = build_arg_parser().parse_args() + + print("=" * 60) + print(f"Using endpoint: {args.endpoint}") + + try: + async with Serverless() as client: + demo = APIDemo(client, args.endpoint) + + if args.workflow: + await demo.demo_workflow(workflow_file=args.workflow) + else: + await demo.demo_prompt( + prompt=args.prompt, + width=args.width, + height=args.height, + steps=args.steps, + seed=args.seed, + ) + + except AttributeError as e: + if "API key" in str(e): + log.error("API key missing. Set VAST_API_KEY environment variable.") + else: + log.error(f"Error: {e}") + sys.exit(1) + except Exception as e: + log.error(f"Error: {e}") + import traceback + traceback.print_exc() + sys.exit(1) - # Get the file from the path on the local machine using SCP or SFTP - # or configure S3 to upload to cloud storage. - print(response["response"]["output"][0]["local_path"]) if __name__ == "__main__": - asyncio.run(main()) \ No newline at end of file + asyncio.run(main_async()) diff --git a/workers/comfyui-json/server.py b/workers/comfyui-json/server.py index ed4e578..7998e71 100644 --- a/workers/comfyui-json/server.py +++ b/workers/comfyui-json/server.py @@ -4,6 +4,7 @@ import dataclasses import base64 from typing import Optional, Union, Type +import aiohttp from aiohttp import web, ClientResponse from lib.backend import Backend, LogAction @@ -108,8 +109,39 @@ async def handle_ping(_): return web.Response(body="pong") +async def handle_view(request: web.Request) -> web.Response: + """Proxy /view requests to ComfyUI to fetch generated images""" + # Forward query params to ComfyUI + query_string = request.query_string + url = f"{MODEL_SERVER_URL}/view?{query_string}" + + log.debug(f"Proxying /view request to: {url}") + + try: + async with aiohttp.ClientSession() as session: + async with session.get(url) as resp: + if resp.status == 200: + content = await resp.read() + return web.Response( + body=content, + status=200, + content_type=resp.content_type or "image/png" + ) + else: + text = await resp.text() + return web.Response( + text=text, + status=resp.status, + content_type="text/plain" + ) + except Exception as e: + log.error(f"Error proxying /view: {e}") + return web.Response(text=str(e), status=500) + + routes = [ web.post("/generate/sync", backend.create_handler(ComfyWorkflowHandler())), + web.get("/view", handle_view), web.get("/ping", handle_ping), ] From e839cfc6e8fa3a32eef152381359df27bf15a953 Mon Sep 17 00:00:00 2001 From: Colter Downing Date: Wed, 3 Dec 2025 20:22:45 -0800 Subject: [PATCH 2/5] include view in API wrapper --- workers/comfyui-json/server.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/workers/comfyui-json/server.py b/workers/comfyui-json/server.py index 7998e71..daf35e5 100644 --- a/workers/comfyui-json/server.py +++ b/workers/comfyui-json/server.py @@ -14,6 +14,7 @@ from .data_types import ComfyWorkflowData MODEL_SERVER_URL = os.getenv("MODEL_SERVER_URL", "http://127.0.0.1:18288") +COMFYUI_URL = os.getenv("COMFYUI_URL", "http://127.0.0.1:18188") # Raw ComfyUI server # This is the last log line that gets emitted once comfyui+extensions have been fully loaded MODEL_SERVER_START_LOG_MSG = "To see the GUI go to: " @@ -110,10 +111,10 @@ async def handle_ping(_): async def handle_view(request: web.Request) -> web.Response: - """Proxy /view requests to ComfyUI to fetch generated images""" - # Forward query params to ComfyUI + """Proxy /view requests to raw ComfyUI server to fetch generated images""" + # Forward query params to raw ComfyUI (not the API wrapper) query_string = request.query_string - url = f"{MODEL_SERVER_URL}/view?{query_string}" + url = f"{COMFYUI_URL}/view?{query_string}" log.debug(f"Proxying /view request to: {url}") From d4d36bf86e03f40f727179975dfb8d53518e9ed2 Mon Sep 17 00:00:00 2001 From: Colter Downing Date: Wed, 3 Dec 2025 20:45:55 -0800 Subject: [PATCH 3/5] done with comfy updates --- workers/comfyui-json/README.md | 71 ++++++++++++--- workers/comfyui-json/client.py | 156 +++++++-------------------------- 2 files changed, 94 insertions(+), 133 deletions(-) diff --git a/workers/comfyui-json/README.md b/workers/comfyui-json/README.md index 7aa1ba3..5306a23 100644 --- a/workers/comfyui-json/README.md +++ b/workers/comfyui-json/README.md @@ -1,8 +1,16 @@ # ComfyUI PyWorker -This is the base PyWorker for ComfyUI. It provides a unified interface for running any ComfyUI workflow through a proxy-based architecture. +This is the base PyWorker for ComfyUI. It provides a unified interface for running any ComfyUI workflow through a proxy-based architecture. See the [Serverless documentation](https://docs.vast.ai/serverless) for guides and how-to's. -The cost for each request has a static value of `1`. ComfyUI does not handle concurrent workloads and there is no current provision to load multiple instances of ComfyUI per worker node. +The cost for each request has a static value of `1`. ComfyUI does not handle concurrent workloads and there is no current provision to load multiple instances of ComfyUI per worker node. + +## Instance Setup + +1. Pick a template + +- [ComfyUI (Serverless)](https://cloud.vast.ai/?ref_id=62897&creator_id=62897&name=ComfyUI%20(Serverless)) + +2. Follow the [getting started guide](https://docs.vast.ai/documentation/serverless/quickstart) for help with configuring your serverless setup. For testing, we recommend that you use the default options presented by the web interface. ## Requirements @@ -10,6 +18,57 @@ This worker requires both [ComfyUI](https://github.com/comfyanonymous/ComfyUI) a A docker image is provided but you may use any if the above requirements are met. +## Client + +The client demonstrates how to use the Vast Serverless SDK to generate images and save them locally. + +### Setup + +1. Clone the PyWorker repository to your local machine and install the necessary requirements for running the test client. + +```bash +git clone https://github.com/vast-ai/pyworker +cd pyworker +pip install uv +uv venv -p 3.12 +source .venv/bin/activate +uv pip install -r requirements.txt +``` + +2. Set your API key: + +```bash +export VAST_API_KEY= +``` + +### Usage + +```bash +# Default prompt +python -m workers.comfyui-json.client + +# Custom prompt +python -m workers.comfyui-json.client --prompt "a cat sitting on a rainbow" + +# With options +python -m workers.comfyui-json.client --prompt "sunset" --width 1024 --height 1024 --steps 30 +``` + +### CLI Flags + +| Flag | Default | Description | +|------|---------|-------------| +| `--endpoint` | `my-comfyui-endpoint` | Vast endpoint name | +| `--prompt` | (default) | Text prompt for image generation | +| `--width` | 512 | Image width in pixels | +| `--height` | 512 | Image height in pixels | +| `--steps` | 20 | Number of denoising steps | +| `--seed` | (random) | Random seed for reproducibility | + +### Output + +Images are saved to `./generated_images/comfy_{seed}.png`. + ## Benchmarking ### Custom Benchmark Workflows @@ -212,11 +271,3 @@ WEBHOOK_TIMEOUT=30 # Webhook timeout in seconds } } ``` - -## Client Libraries - -See the test client examples for implementation details on how to integrate with the ComfyUI worker. - ---- - -See Vast's serverless documentation for more details on how to use ComfyUI with autoscaler. \ No newline at end of file diff --git a/workers/comfyui-json/client.py b/workers/comfyui-json/client.py index b80a9ba..a243183 100644 --- a/workers/comfyui-json/client.py +++ b/workers/comfyui-json/client.py @@ -3,16 +3,16 @@ import sys import json import uuid import random -import base64 import asyncio import logging import argparse +import aiohttp from vastai import Serverless # ---------------------- Config ---------------------- DEFAULT_PROMPT = "a beautiful sunset over mountains, digital art, highly detailed" -ENDPOINT_NAME = "Comfy-Prod2" +ENDPOINT_NAME = "my-comfyui-endpoint" DEFAULT_WIDTH = 512 DEFAULT_HEIGHT = 512 DEFAULT_STEPS = 20 @@ -74,128 +74,40 @@ class APIDemo: self.client = client self.endpoint_name = endpoint_name - def extract_images(self, response: dict) -> list: - """Extract image info from ComfyUI response""" - images = [] - - # Check for output array (S3/webhook configured) - if "output" in response: - for item in response["output"]: - if "url" in item: - images.append({"type": "url", "path": item["url"]}) - elif "local_path" in item: - images.append({"type": "local", "path": item["local_path"]}) - elif "base64" in item: - images.append({"type": "base64", "data": item["base64"]}) - - # Check for comfyui_response format (default) + def extract_filename(self, response: dict) -> str | None: + """Extract the generated image filename from ComfyUI response""" if "comfyui_response" in response: - for prompt_id, data in response["comfyui_response"].items(): + for data in response["comfyui_response"].values(): if isinstance(data, dict) and "outputs" in data: - for node_id, node_output in data["outputs"].items(): - if "images" in node_output: - for img in node_output["images"]: - images.append({ - "type": "remote", - "filename": img.get("filename"), - "subfolder": img.get("subfolder", ""), - }) - - return images + for node_output in data["outputs"].values(): + if "images" in node_output and node_output["images"]: + return node_output["images"][0].get("filename") + return None - async def save_images(self, images: list, worker_url: str, prefix: str = "comfy") -> list: - """Save images locally by fetching from remote server""" + async def save_image(self, worker_url: str, filename: str, local_name: str) -> str | None: + """Fetch and save image locally from the worker""" os.makedirs("generated_images", exist_ok=True) - saved = [] - seen = set() + return await self._fetch_image(worker_url, filename, local_name) - for i, img in enumerate(images): - if img["type"] == "base64": - data = img["data"] - if data.startswith("data:"): - data = data.split(",", 1)[-1] - path = f"generated_images/{prefix}_{i}.png" - with open(path, "wb") as f: - f.write(base64.b64decode(data)) - print(f" šŸ’¾ Saved: {path}") - saved.append(path) - - elif img["type"] == "url": - url = img["path"] - if url in seen: - continue - seen.add(url) - try: - import urllib.request - path = f"generated_images/{prefix}_{len(saved)}.png" - urllib.request.urlretrieve(url, path) - print(f" šŸ’¾ Downloaded: {path}") - saved.append(path) - except Exception as e: - print(f" šŸ”— URL: {url}") - saved.append(url) - - elif img["type"] == "local": - remote_path = img["path"] - if remote_path in seen: - continue - seen.add(remote_path) - filename = os.path.basename(remote_path) - # Try to fetch via /view endpoint - local_path = await self._fetch_image(worker_url, filename, "", f"{prefix}_{len(saved)}.png") - if local_path: - saved.append(local_path) - else: - print(f" šŸ“‚ Remote: {remote_path}") - saved.append(remote_path) - - elif img["type"] == "remote": - filename = img["filename"] - if filename in seen: - continue - seen.add(filename) - subfolder = img.get("subfolder", "") - # Try to fetch via /view endpoint - local_path = await self._fetch_image(worker_url, filename, subfolder, f"{prefix}_{len(saved)}.png") - if local_path: - saved.append(local_path) - else: - print(f" šŸ–¼ļø Remote: {filename}") - saved.append(filename) - - return saved - - async def _fetch_image(self, worker_url: str, filename: str, subfolder: str, local_name: str) -> str | None: - """Fetch image directly from worker's /view endpoint""" + async def _fetch_image(self, worker_url: str, filename: str, local_name: str) -> str | None: + """Fetch image from worker's /view endpoint and save locally""" if not worker_url: - print(f" āš ļø No worker URL available") return None try: - import aiohttp - - params = {"filename": filename, "type": "output"} - if subfolder: - params["subfolder"] = subfolder - url = f"{worker_url}/view" - print(f" šŸ”— Fetching from: {url}") + params = {"filename": filename, "type": "output"} async with aiohttp.ClientSession() as session: async with session.get(url, params=params, ssl=False) as resp: if resp.status == 200: - raw_bytes = await resp.read() path = f"generated_images/{local_name}" with open(path, "wb") as f: - f.write(raw_bytes) + f.write(await resp.read()) print(f" šŸ’¾ Saved: {path}") return path - else: - text = await resp.text() - print(f" āŒ HTTP {resp.status}: {text[:100]}") - return None - except Exception as e: - print(f" āŒ Fetch error: {e}") + return None + except Exception: return None async def demo_prompt( @@ -234,18 +146,17 @@ class APIDemo: worker_url = response.get("url", "") print(f"Worker URL: {worker_url}") - # Extract and handle images + # Fetch and save image if "response" in response: - images = self.extract_images(response["response"]) - if images: - print(f"\nšŸ“ {len(images)} image(s) generated:") - await self.save_images(images, worker_url, prefix=f"comfy_{seed}") + filename = self.extract_filename(response["response"]) + if filename: + path = await self.save_image(worker_url, filename, f"comfy_{seed}.png") + if not path: + print(f"āŒ Failed to fetch image") else: - print("\nNo images found in response") - print(json.dumps(response, indent=2, default=str)[:2000]) + print("āŒ No image in response") else: - print("\nUnexpected response format") - print(json.dumps(response, indent=2, default=str)[:2000]) + print("āŒ Unexpected response format") async def demo_workflow(self, workflow_file: str): """Demo: Generate using custom workflow file""" @@ -274,16 +185,15 @@ class APIDemo: worker_url = response.get("url", "") if "response" in response: - images = self.extract_images(response["response"]) - if images: - print(f"\nšŸ“ {len(images)} image(s) generated:") - await self.save_images(images, worker_url, prefix="workflow") + filename = self.extract_filename(response["response"]) + if filename: + path = await self.save_image(worker_url, filename, "workflow.png") + if not path: + print(f"āŒ Failed to fetch image") else: - print("\nNo images found in response") - print(json.dumps(response, indent=2, default=str)[:2000]) + print("āŒ No image in response") else: - print("\nUnexpected response format") - print(json.dumps(response, indent=2, default=str)[:2000]) + print("āŒ Unexpected response format") # ---------------------- CLI ---------------------- From 40aed9b5f8d85f3f5589cf10cc528901f57b8976 Mon Sep 17 00:00:00 2001 From: Colter Downing Date: Thu, 4 Dec 2025 10:52:57 -0800 Subject: [PATCH 4/5] adding s3 as an option --- workers/comfyui-json/README.md | 33 ++++++++++++++- workers/comfyui-json/client.py | 74 +++++++++++++++++++++++++++++++--- 2 files changed, 101 insertions(+), 6 deletions(-) diff --git a/workers/comfyui-json/README.md b/workers/comfyui-json/README.md index 5306a23..9517dbb 100644 --- a/workers/comfyui-json/README.md +++ b/workers/comfyui-json/README.md @@ -20,7 +20,7 @@ A docker image is provided but you may use any if the above requirements are met ## Client -The client demonstrates how to use the Vast Serverless SDK to generate images and save them locally. +The client demonstrates how to use the Vast Serverless SDK to generate images, save them locally, and optionally upload to S3-compatible storage. ### Setup @@ -52,6 +52,12 @@ python -m workers.comfyui-json.client --prompt "a cat sitting on a rainbow" # With options python -m workers.comfyui-json.client --prompt "sunset" --width 1024 --height 1024 --steps 30 + +# Using a custom workflow file +python -m workers.comfyui-json.client --workflow my_workflow.json + +# With S3 upload +python -m workers.comfyui-json.client --s3 ``` ### CLI Flags @@ -60,15 +66,40 @@ python -m workers.comfyui-json.client --prompt "sunset" --width 1024 --height 10 |------|---------|-------------| | `--endpoint` | `my-comfyui-endpoint` | Vast endpoint name | | `--prompt` | (default) | Text prompt for image generation | +| `--workflow` | (none) | Path to custom workflow JSON file | | `--width` | 512 | Image width in pixels | | `--height` | 512 | Image height in pixels | | `--steps` | 20 | Number of denoising steps | | `--seed` | (random) | Random seed for reproducibility | +| `--s3` | (disabled) | Upload generated images to S3 | ### Output Images are saved to `./generated_images/comfy_{seed}.png`. +### S3 Upload (Optional) + +You can optionally upload generated images to an S3-compatible storage service (AWS S3, Cloudflare R2, Backblaze B2, etc.) by using the `--s3` flag. + +**1. Set environment variables:** + +```bash +export S3_ENDPOINT_URL="https://your-account.r2.cloudflarestorage.com" +export S3_BUCKET_NAME="my-bucket" +export S3_ACCESS_KEY_ID="your-access-key-id" +export S3_SECRET_ACCESS_KEY="your-secret-access-key" +``` + +**2. Run with S3 upload enabled:** + +```bash +python -m workers.comfyui-json.client --prompt "a beautiful landscape" --s3 +``` + +Images will be saved locally AND uploaded to `s3://{bucket}/comfyui/{filename}`. + +**Note:** Requires `boto3` (`pip install boto3`). + ## Benchmarking ### Custom Benchmark Workflows diff --git a/workers/comfyui-json/client.py b/workers/comfyui-json/client.py index a243183..10a1d91 100644 --- a/workers/comfyui-json/client.py +++ b/workers/comfyui-json/client.py @@ -12,16 +12,45 @@ from vastai import Serverless # ---------------------- Config ---------------------- DEFAULT_PROMPT = "a beautiful sunset over mountains, digital art, highly detailed" -ENDPOINT_NAME = "my-comfyui-endpoint" +ENDPOINT_NAME = "Comfy-Prod" DEFAULT_WIDTH = 512 DEFAULT_HEIGHT = 512 DEFAULT_STEPS = 20 COST = 100 # Fixed cost for ComfyUI requests +# Optional S3 Configuration (from environment variables) +S3_ENDPOINT_URL = os.getenv("S3_ENDPOINT_URL") +S3_BUCKET_NAME = os.getenv("S3_BUCKET_NAME") +S3_ACCESS_KEY_ID = os.getenv("S3_ACCESS_KEY_ID") +S3_SECRET_ACCESS_KEY = os.getenv("S3_SECRET_ACCESS_KEY") + logging.basicConfig(level=logging.INFO, format="%(levelname)s - %(message)s") log = logging.getLogger(__name__) +def get_s3_client(): + """Create and return an S3 client configured for the S3-compatible endpoint""" + try: + import boto3 + from botocore.config import Config + except ImportError: + log.error("boto3 is required for S3 uploads. Install with: pip install boto3") + return None + + if not all([S3_ENDPOINT_URL, S3_BUCKET_NAME, S3_ACCESS_KEY_ID, S3_SECRET_ACCESS_KEY]): + log.error("S3 environment variables not fully configured. Required:") + log.error(" S3_ENDPOINT_URL, S3_BUCKET_NAME, S3_ACCESS_KEY_ID, S3_SECRET_ACCESS_KEY") + return None + + return boto3.client( + "s3", + endpoint_url=S3_ENDPOINT_URL, + aws_access_key_id=S3_ACCESS_KEY_ID, + aws_secret_access_key=S3_SECRET_ACCESS_KEY, + config=Config(signature_version="s3v4"), + ) + + # ---------------------- API Functions ---------------------- async def call_generate( client: Serverless, @@ -70,9 +99,14 @@ async def call_generate_workflow( # ---------------------- Demo Class ---------------------- class APIDemo: - def __init__(self, client: Serverless, endpoint_name: str): + def __init__(self, client: Serverless, endpoint_name: str, upload_s3: bool = False): self.client = client self.endpoint_name = endpoint_name + self.upload_s3 = upload_s3 + self.s3_client = get_s3_client() if upload_s3 else None + + if upload_s3 and not self.s3_client: + log.warning("S3 upload requested but client creation failed. Images will only be saved locally.") def extract_filename(self, response: dict) -> str | None: """Extract the generated image filename from ComfyUI response""" @@ -85,10 +119,29 @@ class APIDemo: return None async def save_image(self, worker_url: str, filename: str, local_name: str) -> str | None: - """Fetch and save image locally from the worker""" + """Fetch and save image locally from the worker, optionally upload to S3""" os.makedirs("generated_images", exist_ok=True) return await self._fetch_image(worker_url, filename, local_name) + def _upload_to_s3(self, local_path: str, s3_key: str) -> str | None: + """Upload a local file to S3 and return the S3 URL""" + if not self.s3_client: + return None + + try: + self.s3_client.upload_file( + local_path, + S3_BUCKET_NAME, + s3_key, + ExtraArgs={"ContentType": "image/png"} + ) + s3_url = f"{S3_ENDPOINT_URL}/{S3_BUCKET_NAME}/{s3_key}" + print(f" ā˜ļø Uploaded to S3: {s3_key}") + return s3_url + except Exception as e: + log.error(f"Failed to upload to S3: {e}") + return None + async def _fetch_image(self, worker_url: str, filename: str, local_name: str) -> str | None: """Fetch image from worker's /view endpoint and save locally""" if not worker_url: @@ -102,9 +155,16 @@ class APIDemo: async with session.get(url, params=params, ssl=False) as resp: if resp.status == 200: path = f"generated_images/{local_name}" + image_data = await resp.read() with open(path, "wb") as f: - f.write(await resp.read()) + f.write(image_data) print(f" šŸ’¾ Saved: {path}") + + # Upload to S3 if enabled + if self.upload_s3 and self.s3_client: + s3_key = f"comfyui/{local_name}" + self._upload_to_s3(path, s3_key) + return path return None except Exception: @@ -207,6 +267,8 @@ def build_arg_parser() -> argparse.ArgumentParser: p.add_argument("--height", type=int, default=DEFAULT_HEIGHT, help=f"Image height (default: {DEFAULT_HEIGHT})") p.add_argument("--steps", type=int, default=DEFAULT_STEPS, help=f"Steps (default: {DEFAULT_STEPS})") p.add_argument("--seed", type=int, default=None, help="Seed (default: random)") + p.add_argument("--s3", action="store_true", + help="Upload generated images to S3 (requires S3_ENDPOINT_URL, S3_BUCKET_NAME, S3_ACCESS_KEY_ID, S3_SECRET_ACCESS_KEY env vars)") return p @@ -215,10 +277,12 @@ async def main_async(): print("=" * 60) print(f"Using endpoint: {args.endpoint}") + if args.s3: + print(f"S3 upload: enabled (bucket: {S3_BUCKET_NAME})") try: async with Serverless() as client: - demo = APIDemo(client, args.endpoint) + demo = APIDemo(client, args.endpoint, upload_s3=args.s3) if args.workflow: await demo.demo_workflow(workflow_file=args.workflow) From 222ac2a0ddfe77c96abd5036bb78af6534274d85 Mon Sep 17 00:00:00 2001 From: Colter Downing Date: Thu, 4 Dec 2025 10:54:55 -0800 Subject: [PATCH 5/5] default endpoint name --- workers/comfyui-json/client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/workers/comfyui-json/client.py b/workers/comfyui-json/client.py index 10a1d91..d79b30d 100644 --- a/workers/comfyui-json/client.py +++ b/workers/comfyui-json/client.py @@ -12,7 +12,7 @@ from vastai import Serverless # ---------------------- Config ---------------------- DEFAULT_PROMPT = "a beautiful sunset over mountains, digital art, highly detailed" -ENDPOINT_NAME = "Comfy-Prod" +ENDPOINT_NAME = "my-comfyui-endpoint" DEFAULT_WIDTH = 512 DEFAULT_HEIGHT = 512 DEFAULT_STEPS = 20