From d4d36bf86e03f40f727179975dfb8d53518e9ed2 Mon Sep 17 00:00:00 2001 From: Colter Downing Date: Wed, 3 Dec 2025 20:45:55 -0800 Subject: [PATCH] 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 ----------------------