Initial commit for ComfyUI Runpod worker

.dockerignore

@@ -0,0 +1,11 @@
+*.git
+*.gitignore
+
+env/
+.env
+
+__pycache__
+*.pyc
+*.pyo
+*.pyd
+.Python

.github/workflows/docker-publish.yml

@@ -0,0 +1,51 @@
+name: Docker Build and Push Container
+
+on:
+    push:
+        branches: [main]
+    pull_request:
+        branches: [main]
+
+jobs:
+    build:
+        runs-on: ubuntu-latest
+        steps:
+            - name: Free disk space
+              run: |
+                  df -h
+                  sudo rm -rf /usr/share/dotnet
+                  sudo rm -rf /usr/local/lib/android
+                  sudo rm -rf /opt/ghc
+                  sudo rm -rf /opt/hostedtoolcache
+                  df -h
+
+            - uses: actions/checkout@v3
+
+            - name: Generate version number
+              id: version
+              run: |
+                  echo "VERSION=v1.0.${{ github.run_number }}" >> $GITHUB_OUTPUT
+
+            - name: Set up Docker Buildx
+              uses: docker/setup-buildx-action@v2
+              with:
+                  driver: docker-container
+                  buildkitd-flags: --debug
+
+            - name: Login to Docker Hub
+              uses: docker/login-action@v2
+              with:
+                  username: ${{ secrets.DOCKERHUB_USERNAME }}
+                  password: ${{ secrets.DOCKERHUB_TOKEN }}
+
+            - name: Build container and push to Docker Hub
+              uses: docker/build-push-action@v4
+              with:
+                  context: .
+                  push: true
+                  tags: |
+                      ${{ secrets.DOCKERHUB_USERNAME }}/${{ secrets.DOCKERHUB_REPO_NAME }}:latest
+                      ${{ secrets.DOCKERHUB_USERNAME }}/${{ secrets.DOCKERHUB_REPO_NAME }}:${{ steps.version.outputs.VERSION }}
+                  cache-from: type=registry,ref=agwx2/${{ secrets.DOCKERHUB_REPO_NAME }}:buildcache
+                  cache-to: type=registry,ref=agwx2/${{ secrets.DOCKERHUB_REPO_NAME }}:buildcache,mode=max
+                  platforms: linux/amd64

.gitignore

@@ -0,0 +1,8 @@
+# git ignore
+
+# system
+.DS_Store
+
+# python
+venv
+__pycache__

Dockerfile

@@ -0,0 +1,55 @@
+# dockerfile
+
+FROM nvidia/cuda:11.8.0-runtime-ubuntu22.04
+
+WORKDIR /
+
+# Set environment variables
+ENV DEBIAN_FRONTEND=noninteractive \
+    PIP_PREFER_BINARY=1 \
+    PYTHONUNBUFFERED=1 \
+    CUDA_HOME=/usr/local/cuda \
+    PATH=/usr/local/cuda/bin:/opt/venv/bin:$PATH \
+    ENVIRONMENT=PRODUCTION \
+    COMFYUI_ROOT=/comfyui \
+    APP_ROOT=/app
+
+# Install Python, system dependencies, and CUDA
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends \
+    bash \
+    git \
+    python3.10 \
+    python3.10-venv \
+    python3-pip \
+    libgl1 \
+    libglib2.0-0 \
+    wget && \
+    rm -rf /var/lib/apt/lists/*
+
+# Create virtual environment for Python and install requirements
+RUN python3.10 -m venv /opt/venv && \
+    /opt/venv/bin/pip install --upgrade pip
+
+# Copy the Python dependencies (requirements.txt) and install
+RUN /opt/venv/bin/pip install --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
+
+# Clone ComfyUI and install its dependencies
+RUN git clone https://github.com/comfyanonymous/ComfyUI.git ${COMFYUI_ROOT} && \
+    cd ${COMFYUI_ROOT} && \
+    /opt/venv/bin/pip install --no-cache-dir -r requirements.txt
+
+# Copy the Python dependencies (requirements.txt) and install
+COPY requirements.txt /requirements.txt
+RUN /opt/venv/bin/pip install --no-cache-dir -r /requirements.txt
+
+# Copy application code
+COPY src ${APP_ROOT}
+RUN chmod +x ${APP_ROOT}/main.sh
+
+# download models
+RUN chmod +x ${APP_ROOT}/download.sh
+RUN bash ${APP_ROOT}/download.sh
+
+# Set the final command
+CMD ["${APP_ROOT}/main.sh"]

README.md

@@ -0,0 +1,55 @@
+<div align="center">
+
+<h1>ComfyUI | Worker</h1>
+
+This worker is a RunPod worker that uses the Stable Diffusion model for AI tasks. The worker is built upon the ComfyUI Stable Diffusion Web UI, which is a user interface for Stable Diffusion AI models.
+
+</div>
+
+## Model
+
+Models are auto downloaded in the Dockerfile.
+
+## Building the Worker
+
+The worker is built using a Dockerfile. The Dockerfile specifies the base image, environment variables, system package dependencies, Python dependencies, and the steps to install and setup the ComfyUi Stable Diffusion Web UI. It also downloads models and sets up the API server.
+
+The Python dependencies are specified in requirements.txt. The primary dependency is runpod.
+
+## Running the Worker
+
+The worker can be run using the start.sh script. This script starts the system and runs the serverless handler script.
+
+### build and run
+
+To test locally, activate venv for ComfyUI/Torch, install reqs, then run:
+
+```
+ENV=DEVELOPMENT src/main.sh
+```
+
+### github actions
+
+Please add these secret vars in your Github account's settings for actions:
+
+```
+DOCKERHUB_USERNAME="your_dockerhub_username"
+DOCKERHUB_REPO_NAME="your_dockerhub_repo_name
+DOCKERHUB_TOKEN="your_dockerhub_personal_access_token"
+
+# if you want to upload results to S3
+AWS_ACCESS_KEY="your_aws_access_key"
+AWS_SECRET_KEY="your_aws_secret_key"
+AWS_BUCKET_NAME="your_aws_s3_bucket_name"
+AWS_REGION="your_aws_region"
+```
+
+## API
+
+<!-- The worker provides an API for inference. The API is set up using supervisor, and the configuration is specified in webui_api.conf. The API runs on port 3000. -->
+
+TODO
+
+## Serverless Handler
+
+The serverless handler (handler.py) is a Python script that handles inference requests. It defines a function handler(event) that takes an inference request, runs the inference using the Stable Diffusion model, and returns the output.

download.sh

@@ -0,0 +1,9 @@
+#!/usr/bin/env bash
+
+
+# SD Classic 1.5
+wget --show-progress --progress=dot:giga -O /comfyui/models/checkpoints/v1-5-pruned-emaonly.safetensors https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.safetensors
+
+
+# SDXL Lightning 4-step
+wget --show-progress --progress=dot:giga -O /comfyui/models/checkpoints/sdxl_lightning_4step.safetensors https://huggingface.co/ByteDance/SDXL-Lightning/resolve/main/sdxl_lightning_4step.safetensors

main.sh

@@ -0,0 +1,13 @@
+#!/usr/bin/env bash
+
+ENV="${ENV:-DEVELOPMENT}"
+
+if [ "$ENV" = "PRODUCTION" ]; then
+    echo "Worker Initiated"
+    echo "Starting Runpod Handler"
+    python -u "/$APP_ROOT/handler.py"
+else
+    echo "Local Test Worker Initiated"
+    echo "Starting Test Runpod Handler"
+    python -u "src/handler.py" --rp_serve_api --rp_api_host="0.0.0.0"
+fi

requirements.txt

@@ -0,0 +1,7 @@
+# requirements
+
+runpod
+requests
+urllib3
+websockets
+boto3