add config

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"]

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

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)

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:
+comfyui:
-#     base_path: path/to/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
+    # 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
+    #is_default: true
-#     checkpoints: models/checkpoints/
+    checkpoints: models/checkpoints/
-#     clip: models/clip/
+    clip: models/clip/
-#     clip_vision: models/clip_vision/
+    clip_vision: models/clip_vision/
-#     configs: models/configs/
+    configs: models/configs/
-#     controlnet: models/controlnet/
+    controlnet: models/controlnet/
-#     diffusion_models: |
+    diffusion_models: |
-#                  models/diffusion_models
+                 models/diffusion_models
-#                  models/unet
+                 models/unet
-#     embeddings: models/embeddings/
+    embeddings: models/embeddings/
-#     loras: models/loras/
+    loras: models/loras/
-#     upscale_models: models/upscale_models/
+    upscale_models: models/upscale_models/
-#     vae: models/vae/
+    vae: models/vae/
 
 #other_ui:
 #    base_path: path/to/ui

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"

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})

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

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"
+      }
+    }
+  }
+}