Add RunPod Serverless support - Dockerfile, handler, and environment template

ComfyUI-master/.env.example

@@ -0,0 +1,143 @@
+# RunPod Configuration
+RUNPOD_API_KEY=your_runpod_api_key_here
+COMFYUI_ENDPOINT_ID=your_endpoint_id_here
+
+# Model Repository APIs
+CIVITAI_API_KEY=your_civitai_api_key_here
+HUGGINGFACE_USERNAME=your_huggingface_username_here
+HUGGINGFACE_PASSWORD=your_huggingface_password_here
+
+# New API Services (from latest ComfyUI update)
+# Minimax API
+MINIMAX_API_KEY=your_minimax_api_key_here
+MINIMAX_GROUP_ID=your_minimax_group_id_here
+
+# ByteDance API
+BYTEDANCE_API_KEY=your_bytedance_api_key_here
+
+# Ideogram API
+IDEOGRAM_API_KEY=your_ideogram_api_key_here
+
+# Kling API
+KLING_API_KEY=your_kling_api_key_here
+
+# Luma API
+LUMA_API_KEY=your_luma_api_key_here
+
+# OpenAI API
+OPENAI_API_KEY=your_openai_api_key_here
+
+# Pixverse API
+PIXVERSE_API_KEY=your_pixverse_api_key_here
+
+# Recraft API
+RECRAFT_API_KEY=your_recraft_api_key_here
+
+# Runway API
+RUNWAY_API_KEY=your_runway_api_key_here
+
+# Vidu API
+VIDU_API_KEY=your_vidu_api_key_here
+
+# RunPod GPU Performance Settings
+# Optimized for high-end GPUs (RTX 4090, A100, H100)
+COMFYUI_FAST_MODE=true
+COMFYUI_FP16_ACCUMULATION=true
+COMFYUI_FP8_OPTIMIZATIONS=true
+COMFYUI_PINNED_MEMORY=true
+COMFYUI_CUBLAS_OPS=true
+COMFYUI_AUTOTUNE=true
+
+# Cache Settings - Serverless Optimized
+COMFYUI_CACHE_TYPE=none  # No cache for serverless (fresh start each time)
+COMFYUI_CACHE_LRU_SIZE=0  # Disable LRU cache for serverless
+COMFYUI_CACHE_RAM_THRESHOLD=2.0  # Lower threshold for quick startup
+
+# Memory Management - Serverless Optimized
+COMFYUI_VRAM_MANAGEMENT=auto  # Auto management for variable workloads
+COMFYUI_ASYNC_OFFLOAD=false  # Disable for faster response times
+COMFYUI_NUM_STREAMS=2  # Fewer streams for faster startup
+
+# Model Loading - GPU Optimized
+COMFYUI_DISABLE_XFORMERS=false  # Keep xformers for GPU acceleration
+COMFYUI_USE_SPLIT_CROSS_ATTENTION=false  # Not needed with good GPU
+COMFYUI_ATTENTION_IMPLEMENTATION=xformers  # Force xformers for GPU
+COMFYUI_USE_SAGE_ATTENTION=true  # New sage attention for better performance
+
+# Quantization Settings - GPU Optimized
+COMFYUI_MIXED_PRECISION=true
+COMFYUI_SCALED_FP8=true
+COMFYUI_FORCE_FP16=true  # Force FP16 for GPU speed
+
+# GPU Specific Settings
+CUDA_VISIBLE_DEVICES=0  # Use first GPU
+PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:512,expandable_segments:True
+COMFYUI_FORCE_CUDA=true
+
+# Development and Debug
+COMFYUI_VERBOSE=false
+COMFYUI_DEBUG_MODE=false
+COMFYUI_LOG_LEVEL=INFO
+
+# RunPod Serverless Worker Settings
+RUNPOD_WORKER_ID=${RUNPOD_WORKER_ID}
+RUNPOD_JOB_ID=${RUNPOD_JOB_ID}
+RUNPOD_REQUEST_ID=${RUNPOD_REQUEST_ID}
+
+# Serverless Worker Performance
+COMFYUI_LISTEN=0.0.0.0
+COMFYUI_PORT=8000  # Standard serverless port
+COMFYUI_ENABLE_CORS_HEADER=*
+COMFYUI_SERVERLESS_MODE=true  # Enable serverless optimizations
+
+# GPU Memory Optimization for RunPod
+COMFYUI_NORMALVAE_TILED=true  # Use tiled VAE for memory efficiency
+COMFYUI_DISABLE_SMART_MEMORY=false  # Keep smart memory management
+COMFYUI_FORCE_UPCAST_ATTENTION=false  # Don't upcast on good GPUs
+
+# Model Storage Paths (Serverless optimized)
+COMFYUI_MODEL_PATH=/runpod-volume/models  # Network storage for serverless
+COMFYUI_OUTPUT_PATH=/tmp/outputs  # Temporary output for serverless
+COMFYUI_INPUT_PATH=/tmp/inputs  # Temporary input for serverless
+COMFYUI_TEMP_PATH=/tmp/comfyui
+
+# Serverless Processing Optimization
+COMFYUI_BATCH_SIZE=1  # Single request processing for serverless
+COMFYUI_MAX_BATCH_AREA=2097152  # 1024x1024 max area for faster processing
+COMFYUI_TIMEOUT=300  # 5 minute timeout for serverless jobs
+
+# Serverless Network Settings
+COMFYUI_EXTRA_MODEL_PATHS_CONFIG=/runpod-volume/extra_model_paths.yaml
+COMFYUI_AUTO_LAUNCH_BROWSER=false
+COMFYUI_DISABLE_METADATA=false  # Keep metadata for debugging
+COMFYUI_PREVIEW_METHOD=auto  # Auto preview method
+
+# RunPod Serverless Specific Settings
+RUNPOD_SERVERLESS=true
+RUNPOD_WEBHOOK_GET_JOB_URL=${RUNPOD_WEBHOOK_GET_JOB_URL}
+RUNPOD_WEBHOOK_POST_OUTPUT_URL=${RUNPOD_WEBHOOK_POST_OUTPUT_URL}
+
+# Worker Startup Optimization
+COMFYUI_PRELOAD_MODELS=true  # Preload common models for faster startup
+COMFYUI_WARMUP_ENABLED=true  # Enable warmup requests
+COMFYUI_STARTUP_TIMEOUT=120  # 2 minute startup timeout
+
+# Request Handling
+COMFYUI_MAX_CONCURRENT_REQUESTS=1  # Single request at a time for serverless
+COMFYUI_REQUEST_TIMEOUT=300  # 5 minute request timeout
+COMFYUI_CLEANUP_AFTER_REQUEST=true  # Clean up after each request
+
+# Model Management for Serverless
+COMFYUI_MODEL_CACHE_SIZE=2  # Cache only 2 models max
+COMFYUI_UNLOAD_MODELS_AFTER_USE=true  # Unload models to free VRAM
+COMFYUI_MODEL_DOWNLOAD_TIMEOUT=600  # 10 minute model download timeout
+
+# Output Handling
+COMFYUI_COMPRESS_OUTPUT=true  # Compress outputs for faster upload
+COMFYUI_OUTPUT_FORMAT=png  # Default output format
+COMFYUI_CLEANUP_TEMP_FILES=true  # Clean temp files after request
+
+# Logging for Serverless
+COMFYUI_LOG_REQUESTS=true  # Log all requests for debugging
+COMFYUI_LOG_PERFORMANCE=true  # Log performance metrics
+COMFYUI_STRUCTURED_LOGGING=true  # Use structured logging format

ComfyUI-master/Dockerfile

@@ -0,0 +1,89 @@
+# RunPod Serverless ComfyUI Worker
+FROM runpod/pytorch:2.2.0-py3.11-cuda12.1.1-devel-ubuntu22.04
+
+# Set working directory
+WORKDIR /workspace
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    git \
+    wget \
+    curl \
+    unzip \
+    ffmpeg \
+    libsm6 \
+    libxext6 \
+    libxrender-dev \
+    libglib2.0-0 \
+    libgl1-mesa-glx \
+    && rm -rf /var/lib/apt/lists/*
+
+# Install Python dependencies
+RUN pip install --no-cache-dir \
+    runpod \
+    requests \
+    pillow \
+    numpy \
+    torch \
+    torchvision \
+    torchaudio \
+    xformers \
+    accelerate \
+    transformers \
+    diffusers \
+    opencv-python \
+    scipy \
+    scikit-image
+
+# Copy ComfyUI source code
+COPY . /workspace/ComfyUI
+
+# Set ComfyUI as working directory
+WORKDIR /workspace/ComfyUI
+
+# Install ComfyUI requirements
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Install additional dependencies for new features
+RUN pip install --no-cache-dir \
+    safetensors \
+    transformers[torch] \
+    accelerate \
+    bitsandbytes \
+    optimum
+
+# Copy environment example (users should provide their own .env)
+COPY .env.example /workspace/.env.example
+
+# Create necessary directories
+RUN mkdir -p /workspace/ComfyUI/models/checkpoints \
+    /workspace/ComfyUI/models/vae \
+    /workspace/ComfyUI/models/loras \
+    /workspace/ComfyUI/models/controlnet \
+    /workspace/ComfyUI/models/clip_vision \
+    /workspace/ComfyUI/models/upscale_models \
+    /workspace/ComfyUI/input \
+    /workspace/ComfyUI/output \
+    /tmp/inputs \
+    /tmp/outputs \
+    /tmp/comfyui
+
+# Set environment variables
+ENV PYTHONPATH="/workspace/ComfyUI:${PYTHONPATH}"
+ENV COMFYUI_SERVERLESS=true
+ENV NVIDIA_VISIBLE_DEVICES=all
+ENV NVIDIA_DRIVER_CAPABILITIES=compute,utility
+
+# Create startup script
+RUN echo '#!/bin/bash\n\
+cd /workspace/ComfyUI\n\
+python main.py --listen 0.0.0.0 --port 8000 --dont-print-server --disable-auto-launch &\n\
+sleep 10\n\
+cd /workspace/ComfyUI\n\
+python runpod_handler.py' > /workspace/start.sh && chmod +x /workspace/start.sh
+
+# Expose port
+EXPOSE 8000
+
+# Set the command
+CMD ["/workspace/start.sh"]

ComfyUI-master/runpod_handler.py

@@ -0,0 +1,245 @@
+#!/usr/bin/env python3
+"""
+RunPod Serverless Worker Handler for ComfyUI
+Optimized for the new ComfyUI features and performance improvements
+"""
+
+import os
+import json
+import time
+import logging
+import tempfile
+import requests
+from typing import Dict, Any, Optional
+import runpod
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+class ComfyUIServerlessHandler:
+    def __init__(self):
+        self.comfyui_url = "http://127.0.0.1:8000"
+        self.client_id = "runpod_serverless_worker"
+        self.setup_paths()
+        
+    def setup_paths(self):
+        """Setup required paths for serverless operation"""
+        os.makedirs("/tmp/inputs", exist_ok=True)
+        os.makedirs("/tmp/outputs", exist_ok=True)
+        os.makedirs("/tmp/comfyui", exist_ok=True)
+        
+    def wait_for_comfyui(self, timeout: int = 120) -> bool:
+        """Wait for ComfyUI to be ready"""
+        start_time = time.time()
+        while time.time() - start_time < timeout:
+            try:
+                response = requests.get(f"{self.comfyui_url}/system_stats")
+                if response.status_code == 200:
+                    logger.info("ComfyUI is ready")
+                    return True
+            except requests.exceptions.RequestException:
+                pass
+            time.sleep(2)
+        
+        logger.error(f"ComfyUI not ready after {timeout} seconds")
+        return False
+    
+    def download_input_files(self, input_data: Dict[str, Any]) -> Dict[str, str]:
+        """Download input files and return local paths"""
+        local_files = {}
+        
+        if "input_files" in input_data:
+            for file_key, file_url in input_data["input_files"].items():
+                try:
+                    response = requests.get(file_url, timeout=60)
+                    response.raise_for_status()
+                    
+                    # Create temporary file
+                    with tempfile.NamedTemporaryFile(
+                        delete=False, 
+                        dir="/tmp/inputs",
+                        suffix=os.path.splitext(file_url)[1]
+                    ) as tmp_file:
+                        tmp_file.write(response.content)
+                        local_files[file_key] = tmp_file.name
+                        
+                    logger.info(f"Downloaded {file_key} to {local_files[file_key]}")
+                    
+                except Exception as e:
+                    logger.error(f"Failed to download {file_key}: {str(e)}")
+                    raise
+        
+        return local_files
+    
+    def execute_workflow(self, workflow: Dict[str, Any]) -> Dict[str, Any]:
+        """Execute ComfyUI workflow"""
+        try:
+            # Queue the workflow
+            queue_response = requests.post(
+                f"{self.comfyui_url}/prompt",
+                json={
+                    "prompt": workflow,
+                    "client_id": self.client_id
+                },
+                timeout=30
+            )
+            queue_response.raise_for_status()
+            
+            prompt_id = queue_response.json()["prompt_id"]
+            logger.info(f"Queued workflow with prompt_id: {prompt_id}")
+            
+            # Wait for completion
+            return self.wait_for_completion(prompt_id)
+            
+        except Exception as e:
+            logger.error(f"Failed to execute workflow: {str(e)}")
+            raise
+    
+    def wait_for_completion(self, prompt_id: str, timeout: int = 300) -> Dict[str, Any]:
+        """Wait for workflow completion and return results"""
+        start_time = time.time()
+        
+        while time.time() - start_time < timeout:
+            try:
+                # Check queue status
+                queue_response = requests.get(f"{self.comfyui_url}/queue")
+                queue_data = queue_response.json()
+                
+                # Check if our job is still in queue
+                running = any(item[1]["prompt_id"] == prompt_id for item in queue_data.get("queue_running", []))
+                pending = any(item[1]["prompt_id"] == prompt_id for item in queue_data.get("queue_pending", []))
+                
+                if not running and not pending:
+                    # Job completed, get results
+                    history_response = requests.get(f"{self.comfyui_url}/history/{prompt_id}")
+                    if history_response.status_code == 200:
+                        history_data = history_response.json()
+                        if prompt_id in history_data:
+                            return self.process_results(history_data[prompt_id])
+                
+                time.sleep(2)
+                
+            except Exception as e:
+                logger.error(f"Error checking completion: {str(e)}")
+                time.sleep(5)
+        
+        raise TimeoutError(f"Workflow {prompt_id} timed out after {timeout} seconds")
+    
+    def process_results(self, history_data: Dict[str, Any]) -> Dict[str, Any]:
+        """Process and upload results"""
+        results = {
+            "status": "completed",
+            "outputs": [],
+            "metadata": {}
+        }
+        
+        if "outputs" in history_data:
+            for node_id, node_output in history_data["outputs"].items():
+                if "images" in node_output:
+                    for image_info in node_output["images"]:
+                        # Download image from ComfyUI
+                        image_url = f"{self.comfyui_url}/view"
+                        params = {
+                            "filename": image_info["filename"],
+                            "subfolder": image_info.get("subfolder", ""),
+                            "type": image_info.get("type", "output")
+                        }
+                        
+                        try:
+                            image_response = requests.get(image_url, params=params)
+                            image_response.raise_for_status()
+                            
+                            # Save to temp file for upload
+                            output_path = f"/tmp/outputs/{image_info['filename']}"
+                            with open(output_path, "wb") as f:
+                                f.write(image_response.content)
+                            
+                            results["outputs"].append({
+                                "type": "image",
+                                "filename": image_info["filename"],
+                                "path": output_path,
+                                "node_id": node_id
+                            })
+                            
+                        except Exception as e:
+                            logger.error(f"Failed to process image {image_info['filename']}: {str(e)}")
+        
+        return results
+    
+    def cleanup(self):
+        """Clean up temporary files"""
+        try:
+            import shutil
+            shutil.rmtree("/tmp/inputs", ignore_errors=True)
+            shutil.rmtree("/tmp/outputs", ignore_errors=True)
+            os.makedirs("/tmp/inputs", exist_ok=True)
+            os.makedirs("/tmp/outputs", exist_ok=True)
+            logger.info("Cleaned up temporary files")
+        except Exception as e:
+            logger.error(f"Cleanup failed: {str(e)}")
+
+def handler(job: Dict[str, Any]) -> Dict[str, Any]:
+    """Main serverless handler function"""
+    handler_instance = ComfyUIServerlessHandler()
+    
+    try:
+        # Wait for ComfyUI to be ready
+        if not handler_instance.wait_for_comfyui():
+            return {"error": "ComfyUI failed to start"}
+        
+        # Get job input
+        job_input = job.get("input", {})
+        
+        # Download input files if any
+        local_files = handler_instance.download_input_files(job_input)
+        
+        # Update workflow with local file paths
+        workflow = job_input.get("workflow", {})
+        if local_files and "file_mappings" in job_input:
+            for node_id, mappings in job_input["file_mappings"].items():
+                if node_id in workflow:
+                    for input_key, file_key in mappings.items():
+                        if file_key in local_files:
+                            workflow[node_id]["inputs"][input_key] = local_files[file_key]
+        
+        # Execute workflow
+        results = handler_instance.execute_workflow(workflow)
+        
+        # Upload output files to RunPod storage or return base64
+        output_urls = []
+        for output in results.get("outputs", []):
+            if output["type"] == "image":
+                # For serverless, we typically return base64 or upload to storage
+                with open(output["path"], "rb") as f:
+                    import base64
+                    image_data = base64.b64encode(f.read()).decode()
+                    output_urls.append({
+                        "filename": output["filename"],
+                        "data": image_data,
+                        "node_id": output["node_id"]
+                    })
+        
+        return {
+            "status": "success",
+            "outputs": output_urls,
+            "execution_time": time.time() - job.get("start_time", time.time())
+        }
+        
+    except Exception as e:
+        logger.error(f"Handler error: {str(e)}")
+        return {
+            "error": str(e),
+            "status": "failed"
+        }
+    
+    finally:
+        # Always cleanup
+        handler_instance.cleanup()
+
+if __name__ == "__main__":
+    # Start the serverless worker
+    runpod.serverless.start({"handler": handler})