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