3d pack on modal

3d-pack.py

import json
import subprocess
import uuid
from pathlib import Path
from typing import Dict
import asyncio
import aiohttp
import modal

# 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 = 1000
# Time to wait between poll attempts in milliseconds
COMFY_POLLING_INTERVAL_MS = 250
# Maximum number of poll attempts
COMFY_POLLING_MAX_RETRIES = 1000
# Host where ComfyUI is running
COMFY_HOST = "127.0.0.1:8188"

async def check_server(url, retries=50, delay=500):
    # for i in range(retries):
    while True:
        try:
            async with aiohttp.ClientSession() as session:
                async with session.get(url) as response:
                    # If the response status code is 200, the server is up and running
                    if response.status == 200:
                        print(f"comfy-modal - API is reachable")
                        return True
        except Exception as e:
            # If an exception occurs, the server may not be ready
            pass

        # Wait for the specified delay before retrying
        await asyncio.sleep(delay / 1000)

async def launch_comfy():
    server = await asyncio.subprocess.create_subprocess_shell("comfy launch", shell=True)
    
    ok = await check_server(
        f"http://{COMFY_HOST}",
        COMFY_API_AVAILABLE_MAX_RETRIES,
        COMFY_API_AVAILABLE_INTERVAL_MS,
    )
    
    server.terminate()

image = (  # build up a Modal Image to run ComfyUI, step by step
    # modal.Image.debian_slim(  # start from basic Linux with Python
    #     python_version="3.11"
    # )
    modal.Image.from_registry("nvidia/cuda:12.1.0-devel-ubuntu22.04", add_python="3.11")
    .apt_install("git")  # install git to clone ComfyUI
    .pip_install("comfy-cli==1.0.33")  # install comfy-cli
    .run_commands(  # use comfy-cli to install the ComfyUI repo and its dependencies
        "comfy --skip-prompt install --nvidia",
    )
    # .run_commands(  # download the inpainting model
    #     "comfy --skip-prompt model download --url https://huggingface.co/stabilityai/stable-diffusion-2-inpainting/resolve/main/512-inpainting-ema.safetensors --relative-path models/checkpoints"
    # )
    .apt_install(
        "git",
        "build-essential",
        "curl",
        "ffmpeg",
        "libegl1",
        "libegl1-mesa-dev",
        "libgl1",
        "libglib2.0-0",
        "libgl1-mesa-dev",
        "libgl1-mesa-glx",
        "libgles2",
        "libgles2-mesa-dev",
        "libglvnd-dev",
        "libglvnd0",
        "libglx0",
        "libsm6",
        "libxext6",
        "libxrender1",
        "ninja-build",
        "wget",
    )
    
    .env({
        "LD_LIBRARY_PATH": "/usr/lib64:$LD_LIBRARY_PATH",
        "NVIDIA_VISIBLE_DEVICES": "all",
        "NVIDIA_DRIVER_CAPABILITIES": "compute,utility,graphics",
        "PYOPENGL_PLATFORM": "egl",
    })
    
    .run_commands(  # download a custom node
        "comfy node install ComfyUI-3D-Pack", gpu="A10G"
    )
    
    .run_function(launch_comfy, gpu="A10G")
    # can layer additional models and custom node downloads as needed
)

app = modal.App(name="3d-pack", image=image)

@app.function(
    allow_concurrent_inputs=10,
    concurrency_limit=1,
    container_idle_timeout=30,
    timeout=3600,
    gpu="A10G",
)
@modal.web_server(8000, startup_timeout=60)
def ui():
    subprocess.Popen("comfy launch -- --listen 0.0.0.0 --port 8000", shell=True)