handler.py

"""
RunPod Serverless Handler for ComfyUI with Wan Animate
Enhanced with model management, better error handling, and optimization
"""
import runpod
import json
import os
import sys
import uuid
import base64
import shutil
import time
from pathlib import Path
from io import BytesIO

# Add ComfyUI to path
sys.path.insert(0, '/comfyui')

# Global variables for optimization
COMFYUI_SERVER = None
EXECUTION_COUNT = 0
MAX_EXECUTIONS_PER_INSTANCE = 100  # Restart after this many executions

def setup_directories():
    """Ensure all required directories exist"""
    directories = [
        '/comfyui/input',
        '/comfyui/output',
        '/comfyui/temp'
    ]

    for directory in directories:
        Path(directory).mkdir(parents=True, exist_ok=True)

def clear_directories():
    """Clear temporary directories to free up space"""
    try:
        # Clear temp directory
        temp_dir = Path('/comfyui/temp')
        if temp_dir.exists():
            shutil.rmtree(temp_dir)
            temp_dir.mkdir(exist_ok=True)

        # Clear input directory (keep recent files)
        input_dir = Path('/comfyui/input')
        if input_dir.exists():
            for file in input_dir.glob('*'):
                if file.stat().st_mtime < time.time() - 3600:  # Older than 1 hour
                    file.unlink()

        print("🧹 Cleaned up temporary directories")
    except Exception as e:
        print(f"⚠️  Cleanup warning: {e}")

def validate_models():
    """Check if essential models are available"""
    essential_models = [
        '/comfyui/models/unet/wan2.2_animate_14B_bf16.safetensors',
        '/comfyui/models/vae/wan_2.1_vae.safetensors',
        '/comfyui/models/clip/umt5_xxl_fp8_e4m3fn_scaled.safetensors',
        '/comfyui/models/clip_vision/clip_vision_h.safetensors'
    ]

    missing_models = []
    for model_path in essential_models:
        if not Path(model_path).exists():
            missing_models.append(model_path)

    if missing_models:
        return False, f"Missing essential models: {', '.join(missing_models)}"
    return True, "All models available"

def save_uploaded_files(input_data):
    """Save uploaded images and videos to input directory"""
    images = input_data.get('images', {})
    videos = input_data.get('videos', {})

    saved_files = []

    # Save images
    for filename, base64_data in images.items():
        file_path = f'/comfyui/input/{filename}'
        try:
            with open(file_path, 'wb') as f:
                f.write(base64.b64decode(base64_data))
            saved_files.append(filename)
            print(f"✅ Saved image: {filename}")
        except Exception as e:
            print(f"❌ Failed to save image {filename}: {e}")

    # Save videos
    for filename, base64_data in videos.items():
        file_path = f'/comfyui/input/{filename}'
        try:
            with open(file_path, 'wb') as f:
                f.write(base64.b64decode(base64_data))
            saved_files.append(filename)
            print(f"✅ Saved video: {filename}")
        except Exception as e:
            print(f"❌ Failed to save video {filename}: {e}")

    return saved_files

def execute_workflow(workflow, prompt_id):
    """Execute ComfyUI workflow with error handling"""
    global COMFYUI_SERVER, EXECUTION_COUNT

    try:
        # Import ComfyUI modules
        import execution
        import server
        import folder_paths

        # Initialize server if not exists
        if COMFYUI_SERVER is None:
            print("🔧 Initializing ComfyUI server...")
            COMFYUI_SERVER = server.PromptServer.instance

        # Set up paths
        folder_paths.set_output_directory('/comfyui/output')
        folder_paths.set_input_directory('/comfyui/input')
        folder_paths.set_temp_directory('/comfyui/temp')

        # Validate workflow
        valid, error, unique_id = execution.validate_prompt(workflow)
        if not valid:
            return {
                "error": f"Workflow validation failed: {error}",
                "node_id": unique_id
            }

        # Execute workflow
        executor = execution.PromptExecutor(COMFYUI_SERVER)
        output_ui = {}

        print(f"🚀 Executing workflow: {prompt_id}")
        executor.execute(workflow, prompt_id, {})

        return {"status": "success"}

    except Exception as e:
        return {
            "error": f"Workflow execution failed: {str(e)}",
            "traceback": __import__('traceback').format_exc()
        }

def collect_output_files():
    """Collect and encode output files"""
    output_dir = Path('/comfyui/output')
    output_files = []

    if not output_dir.exists():
        return output_files

    # Get all video and image files
    file_extensions = ['*.mp4', '*.avi', '*.mov', '*.gif', '*.png', '*.jpg', '*.jpeg']
    files_found = []

    for ext in file_extensions:
        files_found.extend(output_dir.glob(ext))

    # Sort by modification time (newest first)
    files_found.sort(key=lambda x: x.stat().st_mtime, reverse=True)

    for file_path in files_found[:10]:  # Limit to 10 most recent files
        try:
            # Read and encode file (with size limit for serverless)
            file_size = file_path.stat().st_size
            max_size = 100 * 1024 * 1024  # 100MB limit

            if file_size > max_size:
                print(f"⚠️  Skipping {file_path.name} - file too large ({file_size / (1024*1024):.1f}MB)")
                continue

            with open(file_path, 'rb') as f:
                file_data = base64.b64encode(f.read()).decode('utf-8')

            output_files.append({
                "filename": file_path.name,
                "data": file_data,
                "size": file_size,
                "type": "video" if file_path.suffix.lower() in ['.mp4', '.avi', '.mov', '.gif'] else "image"
            })

            print(f"✅ Processed output: {file_path.name} ({file_size / (1024*1024):.1f}MB)")

        except Exception as e:
            print(f"❌ Failed to process {file_path.name}: {e}")

    return output_files

def handler(event):
    """
    RunPod handler function

    Expected input format:
    {
        "input": {
            "workflow": {...},  # ComfyUI workflow JSON
            "images": {         # Optional: base64 encoded images
                "filename.png": "base64_data..."
            },
            "videos": {         # Optional: base64 encoded videos
                "filename.mp4": "base64_data..."
            },
            "action": "execute" | "check_models" | "cleanup"  # Optional action
        }
    }
    """
    global EXECUTION_COUNT

    start_time = time.time()
    action = event.get('input', {}).get('action', 'execute')

    try:
        # Setup directories
        setup_directories()

        # Handle different actions
        if action == "check_models":
            models_valid, message = validate_models()
            return {
                "action": "check_models",
                "models_valid": models_valid,
                "message": message
            }

        elif action == "cleanup":
            clear_directories()
            return {
                "action": "cleanup",
                "status": "completed"
            }

        # Main execution path
        input_data = event.get('input', {})
        workflow = input_data.get('workflow', {})

        if not workflow:
            return {
                "error": "No workflow provided"
            }

        # Validate models before execution
        models_valid, message = validate_models()
        if not models_valid:
            return {
                "error": f"Model validation failed: {message}"
            }

        # Generate prompt ID
        prompt_id = str(uuid.uuid4())

        # Save uploaded files
        saved_files = save_uploaded_files(input_data)
        print(f"📁 Saved {len(saved_files)} files")

        # Execute workflow
        result = execute_workflow(workflow, prompt_id)

        if result.get("error"):
            return result

        # Collect output files
        output_files = collect_output_files()

        # Increment execution counter
        EXECUTION_COUNT += 1
        print(f"📈 Execution count: {EXECUTION_COUNT}/{MAX_EXECUTIONS_PER_INSTANCE}")

        # Schedule cleanup for next request
        if EXECUTION_COUNT >= MAX_EXECUTIONS_PER_INSTANCE:
            clear_directories()
            print("🔄 Instance restart recommended due to execution limit")

        execution_time = time.time() - start_time
        print(f"⏱️  Total execution time: {execution_time:.2f} seconds")

        return {
            "status": "success",
            "action": "execute",
            "prompt_id": prompt_id,
            "output_files": output_files,
            "execution_time": execution_time,
            "execution_count": EXECUTION_COUNT,
            "saved_files": saved_files
        }

    except Exception as e:
        execution_time = time.time() - start_time
        return {
            "error": f"Handler failed: {str(e)}",
            "execution_time": execution_time,
            "traceback": __import__('traceback').format_exc()
        }

def handler_wrapper(event):
    """Wrapper to handle RunPod serverless lifecycle"""
    try:
        result = handler(event)

        # Add memory usage info
        try:
            import psutil
            process = psutil.Process()
            memory_info = process.memory_info()
            result["memory_usage_mb"] = memory_info.rss / (1024 * 1024)
        except:
            pass

        return result

    except Exception as e:
        return {
            "error": f"Fatal error: {str(e)}",
            "traceback": __import__('traceback').format_exc()
        }

if __name__ == "__main__":
    # Start RunPod serverless
    print("🎯 Starting RunPod serverless handler...")
    runpod.serverless.start({
        "handler": handler_wrapper,
        "return_aggregate_stream": True
    })