Update App: Add Terrain Builder Tab

app.py

@@ -11,43 +11,37 @@ 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
+import cv2
 import spaces  # HuggingFace Zero GPU
 
-# For Zero GPU spaces, we use a simpler approach:
-# Direct diffusers pipeline instead of ComfyUI backend
-# This avoids ComfyUI's CUDA initialization issues
-
 try:
     import torch
-    from diffusers import StableDiffusionImg2ImgPipeline, AutoPipelineForImage2Image
+    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
         
-        print("Loading Stable Diffusion pipeline...")
-        
-        # Use a lightweight model for Zero GPU
+        # Use v1-5 for general purpose (Terrain + Char)
         model_id = "runwayml/stable-diffusion-v1-5"
         
         _pipeline = StableDiffusionImg2ImgPipeline.from_pretrained(
@@ -56,17 +50,11 @@ def get_pipeline():
             safety_checker=None,
             requires_safety_checker=False
         )
-        
         if torch.cuda.is_available():
             _pipeline = _pipeline.to("cuda")
-            print("Pipeline loaded on CUDA")
-        else:
-            print("Pipeline loaded on CPU")
-        
-        return _pipeline
-
+    return _pipeline
 
-@spaces.GPU(duration=60)  # Request GPU for up to 60 seconds
+@spaces.GPU(duration=60)
 def enhance_image_gpu(
     image: Image.Image,
     prompt: str,
@@ -76,43 +64,27 @@ def enhance_image_gpu(
     num_inference_steps: int = 25,
     seed: int = -1
 ) -> Image.Image:
-    """
-    Enhance an image using Stable Diffusion img2img.
-    This function runs on GPU via Zero GPU.
-    """
     if not DIFFUSERS_AVAILABLE:
-        return image  # Return original if diffusers not available
+        return image
     
     pipe = get_pipeline()
-    
-    # Move to GPU if available
     if torch.cuda.is_available():
         pipe = pipe.to("cuda")
-    
-    # Prepare image
+        
     if image.mode != "RGB":
         image = image.convert("RGB")
-    
-    # Resize to reasonable size for inference
-    max_size = 512
-    ratio = min(max_size / image.width, max_size / image.height)
-    if ratio < 1:
-        new_size = (int(image.width * ratio), int(image.height * ratio))
-        image = image.resize(new_size, Image.LANCZOS)
-    
-    # Ensure dimensions are multiples of 8
+        
+    # 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.LANCZOS)
+    image = image.resize((w, h), Image.Resampling.LANCZOS)
     
-    # Set seed
     generator = None
     if seed >= 0:
         generator = torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu")
         generator.manual_seed(seed)
-    
-    # Run inference
+        
     result = pipe(
         prompt=prompt,
         image=image,
@@ -125,120 +97,103 @@ def enhance_image_gpu(
     
     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)
 
-def enhance_nwn_character(
-    input_image: Image.Image,
-    character_type: str,
-    denoise: float,
-    steps: int,
-    seed: int
-) -> Image.Image:
-    """
-    Enhance an NWN character screenshot to photorealistic quality.
-    """
-    if input_image is None:
-        return None
-    
-    prompt = f"photorealistic {character_type}, highly detailed facial features, 8k, cinematic lighting, professional character art, sharp focus, intricate armor details, fantasy portrait"
-    
-    negative_prompt = "blurry, low quality, low poly, pixelated, cartoonish, anime style, deformed, bad anatomy, extra limbs, watermark, signature, text, simple, flat lighting"
-    
-    result = enhance_image_gpu(
-        image=input_image,
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        strength=denoise,
-        guidance_scale=7.5,
-        num_inference_steps=steps,
-        seed=seed
-    )
-    
-    return result
-
-
-# Character presets for NWN
 CHARACTER_PRESETS = [
     "female elf paladin in ornate silver armor",
     "male human warrior in plate armor",
-    "female human mage in flowing robes",
-    "male dwarf fighter in heavy armor",
-    "female half-elf ranger in leather armor",
-    "male elf wizard in arcane robes",
-    "female human cleric in white robes",
-    "male half-orc barbarian in fur armor",
+    "female human mage in flowing robes"
 ]
 
-
-# Create Gradio app
-with gr.Blocks(title="DGG NWN Character Enhancer", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("""
-    # 🎮 DGG NWN Character Enhancer
+# --- 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
+    )
     
-    Transform low-poly Neverwinter Nights characters into photorealistic images for MetaHuman conversion.
+    # 4. Post-Process (16-bit conversion)
+    enhanced_np = np.array(enhanced.convert("L")).astype(np.float32) / 255.0
     
-    **Powered by Stable Diffusion on Zero GPU**
-    """)
+    # 5. Erosion on AI result
+    eroded = erosion_sim(enhanced_np, erosion_steps)
     
-    with gr.Row():
-        with gr.Column(scale=1):
-            input_image = gr.Image(type="pil", label="NWN Screenshot", height=400)
-            
-            character_type = gr.Dropdown(
-                choices=CHARACTER_PRESETS,
-                value=CHARACTER_PRESETS[0],
-                label="Character Type",
-                allow_custom_value=True
-            )
-            
-            with gr.Row():
-                denoise = gr.Slider(
-                    minimum=0.3, maximum=0.9, value=0.65, step=0.05,
-                    label="Enhancement Strength",
-                    info="Lower = more faithful to original"
-                )
-                steps = gr.Slider(
-                    minimum=15, maximum=40, value=25, step=5,
-                    label="Quality Steps"
-                )
-            
-            seed = gr.Number(value=-1, label="Seed (-1 for random)", precision=0)
-            
-            enhance_btn = gr.Button("✨ Enhance Character", variant="primary", size="lg")
-        
-        with gr.Column(scale=1):
-            output_image = gr.Image(type="pil", label="Enhanced Result", height=400)
+    # 6. Save as 16-bit
+    h_16 = (eroded * 65535).clip(0, 65535).astype(np.uint16)
     
-    enhance_btn.click(
-        fn=enhance_nwn_character,
-        inputs=[input_image, character_type, denoise, steps, seed],
-        outputs=output_image,
-        api_name="enhance"  # This enables API access at /enhance
-    )
-    
-    gr.Markdown("""
-    ---
-    ## API Usage
-    ```python
-    from gradio_client import Client, handle_file
+    out_path = "output_terrain.png"
+    cv2.imwrite(out_path, h_16)
     
-    client = Client("BenjaminPittsley/DGG-ComfyUI")
-    result = client.predict(
-        input_image=handle_file("nwn_screenshot.png"),
-        character_type="female elf paladin in ornate silver armor",
-        denoise=0.65,
-        steps=25,
-        seed=-1,
-        api_name="/enhance"
-    )
-    ```
-    
-    ## Workflow
-    1. Take screenshot of NWN character
-    2. Upload and enhance here
-    3. Use enhanced image with Meshy.ai for 3D generation
-    4. Import 3D mesh to Unreal → Mesh to MetaHuman
-    """)
+    # 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)