From 40aed9b5f8d85f3f5589cf10cc528901f57b8976 Mon Sep 17 00:00:00 2001 From: Colter Downing Date: Thu, 4 Dec 2025 10:52:57 -0800 Subject: [PATCH] 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)