Configure Space to run ComfyUI with custom nodes and auto-GPU/CPU detection

Dockerfile

@@ -1,42 +1,46 @@
-FROM python:3.10-slim
-
-# Install system dependencies
-RUN apt-get update && apt-get install -y \
-    git \
-    wget \
-    libgl1 \
-    libglib2.0-0 \
-    && rm -rf /var/lib/apt/lists/*
-
-# Set working directory
-WORKDIR /app
-
-# Install Python packages
-RUN pip install --no-cache-dir \
-    gradio \
-    torch \
-    torchvision \
-    diffusers \
-    transformers \
-    accelerate \
-    safetensors \
-    pillow \
-    spaces
-
-# Copy application
-COPY app.py /app/
-COPY custom_nodes/ /app/custom_nodes/
-COPY workflows/ /app/workflows/
-
-# Create necessary directories
-RUN mkdir -p /app/models
-
-# Expose port
-EXPOSE 7860
-
-# Set environment for Zero GPU
-ENV PYTHONUNBUFFERED=1
-ENV HF_HOME=/app/models
-
-# Run the app
-CMD ["python", "app.py"]
+FROM python:3.10-slim
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    git \
+    wget \
+    curl \
+    libgl1 \
+    libglib2.0-0 \
+    build-essential \
+    && rm -rf /var/lib/apt/lists/*
+
+# Set working directory
+WORKDIR /app
+
+# Clone ComfyUI into /app
+RUN git clone https://github.com/comfyanonymous/ComfyUI.git .
+
+# Install ComfyUI requirements
+RUN pip install --no-cache-dir torch torchvision --extra-index-url https://download.pytorch.org/whl/cu121
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Install extra requirements for custom nodes
+RUN pip install --no-cache-dir \
+    gradio \
+    spaces \
+    diffusers \
+    transformers \
+    accelerate \
+    safetensors \
+    pillow \
+    opencv-python-headless \
+    scipy
+
+# Copy custom launcher script
+COPY app.py /app/app.py
+
+# Copy custom nodes and workflows
+COPY custom_nodes/ /app/custom_nodes/
+COPY workflows/ /app/workflows/
+
+# Expose port
+EXPOSE 7860
+
+# Run the launcher app
+CMD ["python", "app.py"]

app.py

@@ -1,250 +1,44 @@
-"""
-DGG ComfyUI API Wrapper for HuggingFace Spaces (Zero GPU)
-Provides Gradio interface and API endpoints for NWN character enhancement.
-Uses Zero GPU - GPU is only allocated during inference.
-"""
-import os
-import subprocess
-import threading
-import time
-import json
-import base64
-from pathlib import Path
-from io import BytesIO
-import random
-
-import gradio as gr
-from PIL import Image
-import numpy as np
-try:
-    import cv2
-except ImportError:
-    print("CV2 not found, installing headless...")
-    import subprocess
-    subprocess.check_call(["pip", "install", "opencv-python-headless"])
-    import cv2
-import spaces  # HuggingFace Zero GPU
-
-try:
-    import torch
-    from diffusers import StableDiffusionImg2ImgPipeline
-    DIFFUSERS_AVAILABLE = True
-except ImportError:
-    DIFFUSERS_AVAILABLE = False
-    print("Diffusers not available, will use fallback")
-
-# Global pipeline (loaded on first use)
-_pipeline = None
-_pipeline_lock = threading.Lock()
-
-def get_pipeline():
-    """Get or create the Stable Diffusion pipeline."""
-    global _pipeline
-    if _pipeline is not None:
-        return _pipeline
-
-    with _pipeline_lock:
-        if _pipeline is not None:
-            return _pipeline
-        
-        # Use v1-5 for general purpose (Terrain + Char)
-        model_id = "runwayml/stable-diffusion-v1-5"
-        
-        _pipeline = StableDiffusionImg2ImgPipeline.from_pretrained(
-            model_id,
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-            safety_checker=None,
-            requires_safety_checker=False
-        )
-        if torch.cuda.is_available():
-            _pipeline = _pipeline.to("cuda")
-    return _pipeline
-
-@spaces.GPU(duration=60)
-def enhance_image_gpu(
-    image: Image.Image,
-    prompt: str,
-    negative_prompt: str,
-    strength: float = 0.65,
-    guidance_scale: float = 7.5,
-    num_inference_steps: int = 25,
-    seed: int = -1
-) -> Image.Image:
-    if not DIFFUSERS_AVAILABLE:
-        return image
-    
-    pipe = get_pipeline()
-    if torch.cuda.is_available():
-        pipe = pipe.to("cuda")
-        
-    if image.mode != "RGB":
-        image = image.convert("RGB")
-        
-    # Resize Logic (Maintain aspect, Power of 8)
-    w, h = image.size
-    w = (w // 8) * 8
-    h = (h // 8) * 8
-    image = image.resize((w, h), Image.Resampling.LANCZOS)
-    
-    generator = None
-    if seed >= 0:
-        generator = torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu")
-        generator.manual_seed(seed)
-        
-    result = pipe(
-        prompt=prompt,
-        image=image,
-        strength=strength,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        negative_prompt=negative_prompt,
-        generator=generator
-    ).images[0]
-    
-    return result
-
-# --- CHARACTER LOGIC ---
-def enhance_nwn_character(input_image, character_type, denoise, steps, seed):
-    if input_image is None: return None
-    prompt = f"photorealistic {character_type}, highly detailed, 8k, cinematic lighting"
-    neg = "blurry, low quality, low poly, bad anatomy, watermark, text"
-    return enhance_image_gpu(input_image, prompt, neg, denoise, 7.5, steps, seed)
-
-CHARACTER_PRESETS = [
-    "female elf paladin in ornate silver armor",
-    "male human warrior in plate armor",
-    "female human mage in flowing robes"
-]
-
-# --- TERRAIN LOGIC ---
-def generate_noise_map(resolution=512, seed=-1):
-    if seed >= 0:
-        np.random.seed(seed)
-    # Simple fractal noise approximation
-    noise = np.random.rand(resolution, resolution).astype(np.float32)
-    # Blur to create "hills"
-    noise = cv2.GaussianBlur(noise, (101, 101), 0)
-    noise = (noise - noise.min()) / (noise.max() - noise.min())
-    return noise
-
-def erosion_sim(heightmap, iterations=10):
-    # Fast blur-based erosion
-    for _ in range(iterations):
-        blurred = cv2.GaussianBlur(heightmap, (3, 3), 0)
-        # Mix: Enhance valleys, sharpen peaks?
-        # Simple: H_new = H - (H - Blur) * strength
-        heightmap = heightmap - (heightmap - blurred) * 0.1
-    return heightmap
-
-def generate_terrain(seed, erosion_steps, ai_strength):
-    # 1. Base Noise
-    res = 512
-    h_map = generate_noise_map(res, seed)
-    
-    # 2. Convert to Image for AI
-    img_pil = Image.fromarray((h_map * 255).astype(np.uint8)).convert("RGB")
-    
-    # 3. AI Enhancement (Hallucinate details)
-    prompt = "high altitude aerial view of realistic mountain terrain heightmap, grayscale, erosion, geological details, 8k"
-    neg = "color, trees, water, buildings, roads, text, map overlay"
-    
-    enhanced = enhance_image_gpu(
-        img_pil, prompt, neg, strength=ai_strength, seed=seed
-    )
-    
-    # 4. Post-Process (16-bit conversion)
-    enhanced_np = np.array(enhanced.convert("L")).astype(np.float32) / 255.0
-    
-    # 5. Erosion on AI result
-    eroded = erosion_sim(enhanced_np, erosion_steps)
-    
-    # 6. Save as 16-bit
-    h_16 = (eroded * 65535).clip(0, 65535).astype(np.uint16)
-    
-    out_path = "output_terrain.png"
-    cv2.imwrite(out_path, h_16)
-    
-    # Return 8-bit preview and file path
-    preview = (eroded * 255).astype(np.uint8)
-    return Image.fromarray(preview), out_path
-
-def generate_courtyard_blueprint(seed, style, detail_level):
-    # 1. Create a base layout (Top-down blueprint style)
-    res = 512
-    # Create white canvas
-    canvas = np.ones((res, res, 3), dtype=np.uint8) * 255
-    
-    if seed >= 0:
-        np.random.seed(seed)
-    
-    # Draw a basic courtyard rectangle in center
-    cx, cy = res // 2, res // 2
-    w, h = 300, 200
-    cv2.rectangle(canvas, (cx - w//2, cy - h//2), (cx + w//2, cy + h//2), (0, 0, 0), 2)
-    # Draw a gate gap at south
-    cv2.rectangle(canvas, (cx - 30, cy + h//2 - 5), (cx + 30, cy + h//2 + 5), (255, 255, 255), -1)
-    # Draw central feature circle
-    cv2.circle(canvas, (cx, cy), 20, (50, 50, 50), -1)
-    
-    # Add some "blueprint" details via OpenCV
-    # Grid lines
-    for i in range(0, res, 50):
-        cv2.line(canvas, (i, 0), (i, res), (220, 220, 220), 1)
-        cv2.line(canvas, (0, i), (res, i), (220, 220, 220), 1)
-    
-    # Text labels
-    cv2.putText(canvas, f"STYLE: {style.upper()}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
-    cv2.putText(canvas, f"SEED: {seed}", (10, 50), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
-    cv2.putText(canvas, "TARGET: COURTYARD RECONSTRUCTION", (10, 70), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
-
-    return Image.fromarray(canvas)
-
-# --- APP UI ---
-with gr.Blocks(title="DGG Suite (Zero GPU)", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🛠️ DGG Content Suite")
-    
-    with gr.Tabs():
-        # TAB 1: CHARACTERS
-        with gr.Tab("Character Enhancer"):
-            with gr.Row():
-                with gr.Column():
-                    c_in = gr.Image(type="pil", label="Input")
-                    c_type = gr.Dropdown(CHARACTER_PRESETS, label="Type", value=CHARACTER_PRESETS[0], allow_custom_value=True)
-                    c_str = gr.Slider(0.3, 1.0, 0.65, label="Strength")
-                    c_seed = gr.Number(-1, label="Seed")
-                    c_btn = gr.Button("Enhance", variant="primary")
-                with gr.Column():
-                    c_out = gr.Image(label="Result")
-            c_btn.click(enhance_nwn_character, [c_in, c_type, c_str, gr.Number(25, visible=False), c_seed], c_out)
-
-        # TAB 2: TERRAIN
-        with gr.Tab("Terrain Builder"):
-            gr.Markdown("Generate 16-bit Heightmaps for UE5")
-            with gr.Row():
-                with gr.Column():
-                    t_seed = gr.Number(-1, label="Seed")
-                    t_iter = gr.Slider(0, 50, 10, label="Erosion Steps")
-                    t_ai = gr.Slider(0.0, 1.0, 0.5, label="AI Upscale Strength")
-                    t_btn = gr.Button("Generate Heightmap", variant="primary")
-                with gr.Column():
-                    t_prev = gr.Image(label="Preview (8-bit)")
-                    t_file = gr.File(label="Download 16-bit PNG")
-            
-            t_btn.click(generate_terrain, [t_seed, t_iter, t_ai], [t_prev, t_file])
-
-        # TAB 3: COURTYARD
-        with gr.Tab("Courtyard Architect"):
-            gr.Markdown("Generate blueprint reference for AI construction")
-            with gr.Row():
-                with gr.Column():
-                    cy_seed = gr.Number(-1, label="Seed")
-                    cy_style = gr.Dropdown(["Gothic", "Ancient", "Cybernetic", "Natural"], label="Style", value="Gothic")
-                    cy_detail = gr.Radio(["Low", "Medium", "High"], label="Detail Level", value="Medium")
-                    cy_btn = gr.Button("Generate Blueprint", variant="primary")
-                with gr.Column():
-                    cy_out = gr.Image(label="Blueprint Reference")
-            
-            cy_btn.click(generate_courtyard_blueprint, [cy_seed, cy_style, cy_detail], cy_out, api_name="generate_courtyard_blueprint")
-
-if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+import os
+import subprocess
+import urllib.request
+import torch
+
+def download_file(url, dest):
+    print(f"Downloading {url} to {dest}...")
+    os.makedirs(os.path.dirname(dest), exist_ok=True)
+    
+    # Simple console progress log
+    last_reported = -10
+    def progress(count, block_size, total_size):
+        nonlocal last_reported
+        if total_size <= 0:
+            return
+        percent = int(count * block_size * 100 / total_size)
+        if percent >= last_reported + 10:
+            print(f"Download progress: {percent}%")
+            last_reported = percent
+            
+    urllib.request.urlretrieve(url, dest, reporthook=progress)
+    print("Download complete!")
+
+# Check/download SD 1.5 model if missing
+checkpoint_path = "models/checkpoints/v1-5-pruned-emaonly.safetensors"
+if not os.path.exists(checkpoint_path):
+    sd15_url = "https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.safetensors"
+    try:
+        download_file(sd15_url, checkpoint_path)
+    except Exception as e:
+        print(f"Error downloading checkpoint: {e}")
+
+# Detect CPU vs GPU
+cuda_available = torch.cuda.is_available()
+cmd = ["python", "main.py", "--listen", "0.0.0.0", "--port", "7860"]
+if not cuda_available:
+    print("CUDA not available. Running ComfyUI in CPU mode...")
+    cmd.append("--cpu")
+else:
+    print("CUDA available. Running ComfyUI with GPU acceleration!")
+
+# Start ComfyUI
+print(f"Launching ComfyUI: {' '.join(cmd)}")
+subprocess.run(cmd)