app.py

"""
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

# --- 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])

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)