From f97df737b8bb7428c81e06d3a1aac354f2ed572b Mon Sep 17 00:00:00 2001 From: minkhant1996-dev Date: Thu, 31 Oct 2024 17:39:58 +0700 Subject: [PATCH] add config --- Dockerfile | 47 +++ docker-compose.yml | 16 + folder_paths.py | 3 +- ...hs.yaml.example => min-comfyui-config.yaml | 32 +- run_docker.sh | 25 ++ src_min/rp_handler.py | 350 ++++++++++++++++++ src_min/start.sh | 20 + src_min/test_input.json | 63 ++++ 8 files changed, 538 insertions(+), 18 deletions(-) create mode 100644 Dockerfile create mode 100644 docker-compose.yml rename extra_model_paths.yaml.example => min-comfyui-config.yaml (59%) create mode 100644 run_docker.sh create mode 100644 src_min/rp_handler.py create mode 100644 src_min/start.sh create mode 100644 src_min/test_input.json diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 000000000..e78e7d454 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,47 @@ +# Stage 1: Base image with common dependencies +FROM nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04 as base + +# Prevents prompts from packages asking for user input during installation +ENV DEBIAN_FRONTEND=noninteractive +# Prefer binary wheels over source distributions for faster pip installations +ENV PIP_PREFER_BINARY=1 +# Ensures output from python is printed immediately to the terminal without buffering +ENV PYTHONUNBUFFERED=1 + +# Install Python, git and other necessary tools +RUN apt-get update && apt-get install -y \ + python3.10 \ + python3-pip \ + git \ + wget + +# Clean up to reduce image size +RUN apt-get autoremove -y && apt-get clean -y && rm -rf /var/lib/apt/lists/* + +# Clone ComfyUI repository +RUN git clone https://github.com/comfyanonymous/ComfyUI.git /comfyui + +# Change working directory to ComfyUI +WORKDIR /comfyui + +# Install ComfyUI dependencies +RUN pip3 install --upgrade --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121 \ + && pip3 install --upgrade -r requirements.txt + +# Install runpod +RUN pip3 install runpod requests + +# Support for the network volume +ADD src_min/extra_model_paths.yaml ./ + +# Go back to the root +WORKDIR / + +# Add the start and the handler +ADD src_min/start.sh src_min/rp_handler.py src_min/test_input.json ./ + +RUN chmod +x /start.sh + +# WORKDIR /comfyui + +CMD ["python3", "comfyui/main.py", "--listen", "0.0.0.0", "--port", "8188"] diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 000000000..74613dffa --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,16 @@ +version: "3.8" + +services: + comfyui: + build: . + container_name: comfyui-worker-min + environment: + - NVIDIA_VISIBLE_DEVICES=all + - SERVE_API_LOCALLY=true + ports: + - "8000:8000" + - "8188:8188" + runtime: nvidia + volumes: + - /home/minkhant/Documents/BrookAI/AI_MODELS/output:/comfyui/output + - /home/minkhant/Documents/BrookAI/AI_MODELS:/runpod-volume diff --git a/folder_paths.py b/folder_paths.py index 01ae821de..7801cbf70 100644 --- a/folder_paths.py +++ b/folder_paths.py @@ -12,10 +12,9 @@ supported_pt_extensions: set[str] = {'.ckpt', '.pt', '.bin', '.pth', '.safetenso folder_names_and_paths: dict[str, tuple[list[str], set[str]]] = {} base_path = os.path.dirname(os.path.realpath(__file__)) -models_dir = os.path.join(base_path, "models") +models_dir = "/home/minkhant/Documents/BrookAI/AI_MODELS" #os.path.join(base_path, "models") folder_names_and_paths["checkpoints"] = ([os.path.join(models_dir, "checkpoints")], supported_pt_extensions) folder_names_and_paths["configs"] = ([os.path.join(models_dir, "configs")], [".yaml"]) - folder_names_and_paths["loras"] = ([os.path.join(models_dir, "loras")], supported_pt_extensions) folder_names_and_paths["vae"] = ([os.path.join(models_dir, "vae")], supported_pt_extensions) folder_names_and_paths["clip"] = ([os.path.join(models_dir, "clip")], supported_pt_extensions) diff --git a/extra_model_paths.yaml.example b/min-comfyui-config.yaml similarity index 59% rename from extra_model_paths.yaml.example rename to min-comfyui-config.yaml index b55913a5a..ff65dc9ec 100644 --- a/extra_model_paths.yaml.example +++ b/min-comfyui-config.yaml @@ -23,22 +23,22 @@ a111: #config for comfyui #your base path should be either an existing comfy install or a central folder where you store all of your models, loras, etc. -#comfyui: -# base_path: path/to/comfyui/ -# # You can use is_default to mark that these folders should be listed first, and used as the default dirs for eg downloads -# #is_default: true -# checkpoints: models/checkpoints/ -# clip: models/clip/ -# clip_vision: models/clip_vision/ -# configs: models/configs/ -# controlnet: models/controlnet/ -# diffusion_models: | -# models/diffusion_models -# models/unet -# embeddings: models/embeddings/ -# loras: models/loras/ -# upscale_models: models/upscale_models/ -# vae: models/vae/ +comfyui: + base_path: "/home/minkhant/Documents/BrookAI/AI_MODELS" + # You can use is_default to mark that these folders should be listed first, and used as the default dirs for eg downloads + #is_default: true + checkpoints: models/checkpoints/ + clip: models/clip/ + clip_vision: models/clip_vision/ + configs: models/configs/ + controlnet: models/controlnet/ + diffusion_models: | + models/diffusion_models + models/unet + embeddings: models/embeddings/ + loras: models/loras/ + upscale_models: models/upscale_models/ + vae: models/vae/ #other_ui: # base_path: path/to/ui diff --git a/run_docker.sh b/run_docker.sh new file mode 100644 index 000000000..ea26f88d0 --- /dev/null +++ b/run_docker.sh @@ -0,0 +1,25 @@ +#!/bin/bash + +# Start time +start_time=$(date +%s) + +# Run Docker Compose +echo "Starting Docker Compose..." +docker-compose up -d + +# Check if Docker Compose ran successfully +if [ $? -eq 0 ]; then + echo "Docker Compose ran successfully." +else + echo "Docker Compose failed to start." + exit 1 +fi + +# End time +end_time=$(date +%s) + +# Calculate time taken +time_taken=$((end_time - start_time)) + +# Print time taken in seconds +echo "Time taken: ${time_taken} seconds" diff --git a/src_min/rp_handler.py b/src_min/rp_handler.py new file mode 100644 index 000000000..2150284a6 --- /dev/null +++ b/src_min/rp_handler.py @@ -0,0 +1,350 @@ +import runpod +from runpod.serverless.utils import rp_upload +import json +import urllib.request +import urllib.parse +import time +import os +import requests +import base64 +from io import BytesIO + +# Time to wait between API check attempts in milliseconds +COMFY_API_AVAILABLE_INTERVAL_MS = 50 +# Maximum number of API check attempts +COMFY_API_AVAILABLE_MAX_RETRIES = 500 +# Time to wait between poll attempts in milliseconds +COMFY_POLLING_INTERVAL_MS = os.environ.get("COMFY_POLLING_INTERVAL_MS", 250) +# Maximum number of poll attempts +COMFY_POLLING_MAX_RETRIES = os.environ.get("COMFY_POLLING_MAX_RETRIES", 500) +# Host where ComfyUI is running +COMFY_HOST = "127.0.0.1:8188" +# Enforce a clean state after each job is done +# see https://docs.runpod.io/docs/handler-additional-controls#refresh-worker +REFRESH_WORKER = os.environ.get("REFRESH_WORKER", "false").lower() == "true" + + +def validate_input(job_input): + """ + Validates the input for the handler function. + + Args: + job_input (dict): The input data to validate. + + Returns: + tuple: A tuple containing the validated data and an error message, if any. + The structure is (validated_data, error_message). + """ + # Validate if job_input is provided + if job_input is None: + return None, "Please provide input" + + # Check if input is a string and try to parse it as JSON + if isinstance(job_input, str): + try: + job_input = json.loads(job_input) + except json.JSONDecodeError: + return None, "Invalid JSON format in input" + + # Validate 'workflow' in input + workflow = job_input.get("workflow") + if workflow is None: + return None, "Missing 'workflow' parameter" + + # Validate 'images' in input, if provided + images = job_input.get("images") + if images is not None: + if not isinstance(images, list) or not all( + "name" in image and "image" in image for image in images + ): + return ( + None, + "'images' must be a list of objects with 'name' and 'image' keys", + ) + + # Return validated data and no error + return {"workflow": workflow, "images": images}, None + + +def check_server(url, retries=500, delay=50): + """ + Check if a server is reachable via HTTP GET request + + Args: + - url (str): The URL to check + - retries (int, optional): The number of times to attempt connecting to the server. Default is 50 + - delay (int, optional): The time in milliseconds to wait between retries. Default is 500 + + Returns: + bool: True if the server is reachable within the given number of retries, otherwise False + """ + + for i in range(retries): + try: + response = requests.get(url) + + # If the response status code is 200, the server is up and running + if response.status_code == 200: + print(f"runpod-worker-comfy - API is reachable") + return True + except requests.RequestException as e: + # If an exception occurs, the server may not be ready + pass + + # Wait for the specified delay before retrying + time.sleep(delay / 1000) + + print( + f"runpod-worker-comfy - Failed to connect to server at {url} after {retries} attempts." + ) + return False + + +def upload_images(images): + """ + Upload a list of base64 encoded images to the ComfyUI server using the /upload/image endpoint. + + Args: + images (list): A list of dictionaries, each containing the 'name' of the image and the 'image' as a base64 encoded string. + server_address (str): The address of the ComfyUI server. + + Returns: + list: A list of responses from the server for each image upload. + """ + if not images: + return {"status": "success", "message": "No images to upload", "details": []} + + responses = [] + upload_errors = [] + + print(f"runpod-worker-comfy - image(s) upload") + + for image in images: + name = image["name"] + image_data = image["image"] + blob = base64.b64decode(image_data) + + # Prepare the form data + files = { + "image": (name, BytesIO(blob), "image/png"), + "overwrite": (None, "true"), + } + + # POST request to upload the image + response = requests.post(f"http://{COMFY_HOST}/upload/image", files=files) + if response.status_code != 200: + upload_errors.append(f"Error uploading {name}: {response.text}") + else: + responses.append(f"Successfully uploaded {name}") + + if upload_errors: + print(f"runpod-worker-comfy - image(s) upload with errors") + return { + "status": "error", + "message": "Some images failed to upload", + "details": upload_errors, + } + + print(f"runpod-worker-comfy - image(s) upload complete") + return { + "status": "success", + "message": "All images uploaded successfully", + "details": responses, + } + + +def queue_workflow(workflow): + """ + Queue a workflow to be processed by ComfyUI + + Args: + workflow (dict): A dictionary containing the workflow to be processed + + Returns: + dict: The JSON response from ComfyUI after processing the workflow + """ + + # The top level element "prompt" is required by ComfyUI + data = json.dumps({"prompt": workflow}).encode("utf-8") + + req = urllib.request.Request(f"http://{COMFY_HOST}/prompt", data=data) + return json.loads(urllib.request.urlopen(req).read()) + + +def get_history(prompt_id): + """ + Retrieve the history of a given prompt using its ID + + Args: + prompt_id (str): The ID of the prompt whose history is to be retrieved + + Returns: + dict: The history of the prompt, containing all the processing steps and results + """ + with urllib.request.urlopen(f"http://{COMFY_HOST}/history/{prompt_id}") as response: + return json.loads(response.read()) + + +def base64_encode(img_path): + """ + Returns base64 encoded image. + + Args: + img_path (str): The path to the image + + Returns: + str: The base64 encoded image + """ + with open(img_path, "rb") as image_file: + encoded_string = base64.b64encode(image_file.read()).decode("utf-8") + return f"{encoded_string}" + + +def process_output_images(outputs, job_id): + """ + This function takes the "outputs" from image generation and the job ID, + then determines the correct way to return the image, either as a direct URL + to an AWS S3 bucket or as a base64 encoded string, depending on the + environment configuration. + + Args: + outputs (dict): A dictionary containing the outputs from image generation, + typically includes node IDs and their respective output data. + job_id (str): The unique identifier for the job. + + Returns: + dict: A dictionary with the status ('success' or 'error') and the message, + which is either the URL to the image in the AWS S3 bucket or a base64 + encoded string of the image. In case of error, the message details the issue. + + The function works as follows: + - It first determines the output path for the images from an environment variable, + defaulting to "/comfyui/output" if not set. + - It then iterates through the outputs to find the filenames of the generated images. + - After confirming the existence of the image in the output folder, it checks if the + AWS S3 bucket is configured via the BUCKET_ENDPOINT_URL environment variable. + - If AWS S3 is configured, it uploads the image to the bucket and returns the URL. + - If AWS S3 is not configured, it encodes the image in base64 and returns the string. + - If the image file does not exist in the output folder, it returns an error status + with a message indicating the missing image file. + """ + + # The path where ComfyUI stores the generated images + COMFY_OUTPUT_PATH = os.environ.get("COMFY_OUTPUT_PATH", "/comfyui/output") + + output_images = {} + + for node_id, node_output in outputs.items(): + if "images" in node_output: + for image in node_output["images"]: + output_images = os.path.join(image["subfolder"], image["filename"]) + + print(f"runpod-worker-comfy - image generation is done") + + # expected image output folder + local_image_path = f"{COMFY_OUTPUT_PATH}/{output_images}" + + print(f"runpod-worker-comfy - {local_image_path}") + + # The image is in the output folder + if os.path.exists(local_image_path): + if os.environ.get("BUCKET_ENDPOINT_URL", False): + # URL to image in AWS S3 + image = rp_upload.upload_image(job_id, local_image_path) + print( + "runpod-worker-comfy - the image was generated and uploaded to AWS S3" + ) + else: + # base64 image + image = base64_encode(local_image_path) + print( + "runpod-worker-comfy - the image was generated and converted to base64" + ) + + return { + "status": "success", + "message": image, + } + else: + print("runpod-worker-comfy - the image does not exist in the output folder") + return { + "status": "error", + "message": f"the image does not exist in the specified output folder: {local_image_path}", + } + + +def handler(job): + """ + The main function that handles a job of generating an image. + + This function validates the input, sends a prompt to ComfyUI for processing, + polls ComfyUI for result, and retrieves generated images. + + Args: + job (dict): A dictionary containing job details and input parameters. + + Returns: + dict: A dictionary containing either an error message or a success status with generated images. + """ + job_input = job["input"] + + # Make sure that the input is valid + validated_data, error_message = validate_input(job_input) + if error_message: + return {"error": error_message} + + # Extract validated data + workflow = validated_data["workflow"] + images = validated_data.get("images") + + # Make sure that the ComfyUI API is available + check_server( + f"http://{COMFY_HOST}", + COMFY_API_AVAILABLE_MAX_RETRIES, + COMFY_API_AVAILABLE_INTERVAL_MS, + ) + + # Upload images if they exist + upload_result = upload_images(images) + + if upload_result["status"] == "error": + return upload_result + + # Queue the workflow + try: + queued_workflow = queue_workflow(workflow) + prompt_id = queued_workflow["prompt_id"] + print(f"runpod-worker-comfy - queued workflow with ID {prompt_id}") + except Exception as e: + return {"error": f"Error queuing workflow: {str(e)}"} + + # Poll for completion + print(f"runpod-worker-comfy - wait until image generation is complete") + retries = 0 + try: + while retries < COMFY_POLLING_MAX_RETRIES: + history = get_history(prompt_id) + + # Exit the loop if we have found the history + if prompt_id in history and history[prompt_id].get("outputs"): + break + else: + # Wait before trying again + time.sleep(COMFY_POLLING_INTERVAL_MS / 1000) + retries += 1 + else: + return {"error": "Max retries reached while waiting for image generation"} + except Exception as e: + return {"error": f"Error waiting for image generation: {str(e)}"} + + # Get the generated image and return it as URL in an AWS bucket or as base64 + images_result = process_output_images(history[prompt_id].get("outputs"), job["id"]) + + result = {**images_result, "refresh_worker": REFRESH_WORKER} + + return result + + +# Start the handler only if this script is run directly +if __name__ == "__main__": + runpod.serverless.start({"handler": handler}) diff --git a/src_min/start.sh b/src_min/start.sh new file mode 100644 index 000000000..6cb5cb84a --- /dev/null +++ b/src_min/start.sh @@ -0,0 +1,20 @@ +#!/usr/bin/env bash + +# Use libtcmalloc for better memory management +TCMALLOC="$(ldconfig -p | grep -Po "libtcmalloc.so.\d" | head -n 1)" +export LD_PRELOAD="${TCMALLOC}" + +# Serve the API and don't shutdown the container +if [ "$SERVE_API_LOCALLY" == "true" ]; then + echo "runpod-worker-comfy: Starting ComfyUI" + python3 /comfyui/main.py --disable-auto-launch --disable-metadata --listen & + + echo "runpod-worker-comfy: Starting RunPod Handler" + python3 -u /rp_handler.py --rp_serve_api --rp_api_host=0.0.0.0 +else + echo "runpod-worker-comfy: Starting ComfyUI" + python3 /comfyui/main.py --disable-auto-launch --disable-metadata & + + echo "runpod-worker-comfy: Starting RunPod Handler" + python3 -u /rp_handler.py +fi \ No newline at end of file diff --git a/src_min/test_input.json b/src_min/test_input.json new file mode 100644 index 000000000..7fdb0f200 --- /dev/null +++ b/src_min/test_input.json @@ -0,0 +1,63 @@ +{ + "input": { + "workflow": { + "3": { + "inputs": { + "seed": 234234, + "steps": 20, + "cfg": 8, + "sampler_name": "euler", + "scheduler": "normal", + "denoise": 1, + "model": ["4", 0], + "positive": ["6", 0], + "negative": ["7", 0], + "latent_image": ["5", 0] + }, + "class_type": "KSampler" + }, + "4": { + "inputs": { + "ckpt_name": "sd_xl_base_1.0.safetensors" + }, + "class_type": "CheckpointLoaderSimple" + }, + "5": { + "inputs": { + "width": 512, + "height": 512, + "batch_size": 1 + }, + "class_type": "EmptyLatentImage" + }, + "6": { + "inputs": { + "text": "beautiful scenery nature glass bottle landscape, purple galaxy bottle,", + "clip": ["4", 1] + }, + "class_type": "CLIPTextEncode" + }, + "7": { + "inputs": { + "text": "text, watermark", + "clip": ["4", 1] + }, + "class_type": "CLIPTextEncode" + }, + "8": { + "inputs": { + "samples": ["3", 0], + "vae": ["4", 2] + }, + "class_type": "VAEDecode" + }, + "9": { + "inputs": { + "filename_prefix": "ComfyUI/test", + "images": ["8", 0] + }, + "class_type": "SaveImage" + } + } + } +}