main commit

app.py

@@ -1,80 +1,131 @@
-import gradio as gr
-import torch
-import torch.nn as nn
-import torchvision.utils as vutils
-from torchvision import transforms
-from PIL import Image
-
-# --- 1. Define the Brain Structure (Must match your training code!) ---
-class Generator(nn.Module):
-    def __init__(self, ngpu):
-        super(Generator, self).__init__()
-        self.ngpu = ngpu
-        nz = 100 # Noise size
-        self.main = nn.Sequential(
-            # Input is Z (Noise Vector)
-            nn.ConvTranspose2d(nz, 512, 4, 1, 0, bias=False),
-            nn.BatchNorm2d(512),
-            nn.ReLU(True),
-            # 4x4 -> 8x8
-            nn.ConvTranspose2d(512, 256, 4, 2, 1, bias=False),
-            nn.BatchNorm2d(256),
-            nn.ReLU(True),
-            # 8x8 -> 16x16
-            nn.ConvTranspose2d(256, 128, 4, 2, 1, bias=False),
-            nn.BatchNorm2d(128),
-            nn.ReLU(True),
-            # 16x16 -> 32x32
-            nn.ConvTranspose2d(128, 64, 4, 2, 1, bias=False),
-            nn.BatchNorm2d(64),
-            nn.ReLU(True),
-            # 32x32 -> 64x64
-            nn.ConvTranspose2d(64, 3, 4, 2, 1, bias=False),
-            nn.Tanh()
-        )
-
-    def forward(self, input):
-        return self.main(input)
-
-# --- 2. Load the Model ---
-device = torch.device("cpu") # Spaces run on CPU by default
-model = Generator(ngpu=0).to(device)
-
-# Load weights (map_location is vital for CPU spaces)
-model.load_state_dict(torch.load("generator_final.pth", map_location=device))
-model.eval()
-
-# --- 3. The Generation Function ---
-def generate_art(seed, noise_level):
-    torch.manual_seed(int(seed))
-    
-    # Create noise vector
-    noise = torch.randn(1, 100, 1, 1, device=device)
-    
-    # Generate image
-    with torch.no_grad():
-        fake = model(noise).detach().cpu()
-    
-    # Process image (Un-normalize from -1,1 to 0,1)
-    img_tensor = (fake[0] * 0.5) + 0.5
-    img_pil = transforms.ToPILImage()(img_tensor)
-    
-    # Upscale nicely for the web
-    img_pil = img_pil.resize((256, 256), resample=Image.NEAREST)
-    return img_pil
-
-# --- 4. The Interface ---
-with gr.Blocks() as demo:
-    gr.Markdown("# 👾 8-Bit Hero Generator (2025 Edition)")
-    gr.Markdown("I trained this AI from scratch on New Year's Eve 2025 using PyTorch & DCGAN.")
-    
-    with gr.Row():
-        with gr.Column():
-            seed_slider = gr.Slider(0, 9999, label="Character DNA (Seed)", value=42)
-            btn = gr.Button("Generate New Hero")
-        with gr.Column():
-            output_image = gr.Image(label="Result", type="pil")
-
-    btn.click(fn=generate_art, inputs=[seed_slider, seed_slider], outputs=output_image)
-
-demo.launch()
+import gradio as gr
+import torch
+import torch.nn as nn
+import torchvision.utils as vutils
+from torchvision import transforms
+from PIL import Image, ImageOps, ImageEnhance
+import numpy as np
+import random
+
+# --- 1. The Brain (Same as before) ---
+class Generator(nn.Module):
+    def __init__(self, ngpu):
+        super(Generator, self).__init__()
+        self.ngpu = ngpu
+        nz = 100
+        self.main = nn.Sequential(
+            nn.ConvTranspose2d(nz, 512, 4, 1, 0, bias=False),
+            nn.BatchNorm2d(512),
+            nn.ReLU(True),
+            nn.ConvTranspose2d(512, 256, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(256),
+            nn.ReLU(True),
+            nn.ConvTranspose2d(256, 128, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(128),
+            nn.ReLU(True),
+            nn.ConvTranspose2d(128, 64, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(64),
+            nn.ReLU(True),
+            nn.ConvTranspose2d(64, 3, 4, 2, 1, bias=False),
+            nn.Tanh()
+        )
+
+    def forward(self, input):
+        return self.main(input)
+
+# --- 2. Load Model ---
+device = torch.device("cpu")
+model = Generator(ngpu=0).to(device)
+model.load_state_dict(torch.load("generator_final.pth", map_location=device))
+model.eval()
+
+# --- 3. Advanced Generation Logic ---
+def generate_hero(seed_A, seed_B, morph_factor, clean_up, retro_mode, resolution):
+    # Set seeds
+    torch.manual_seed(int(seed_A))
+    noise_A = torch.randn(1, 100, 1, 1, device=device)
+    
+    torch.manual_seed(int(seed_B))
+    noise_B = torch.randn(1, 100, 1, 1, device=device)
+    
+    # MORPHING: Mix the two DNA seeds based on the slider
+    # If factor is 0, it's 100% A. If 1, it's 100% B.
+    final_noise = (1 - morph_factor) * noise_A + (morph_factor) * noise_B
+    
+    with torch.no_grad():
+        fake = model(final_noise).detach().cpu()
+    
+    # CLEANUP FILTER: Remove weak pixels (noise removal)
+    if clean_up:
+        # If a pixel is very dark, make it pure black (-1.0)
+        fake = torch.where(fake < -0.7, torch.tensor(-1.0), fake)
+
+    # Process to Image
+    img_tensor = (fake[0] * 0.5) + 0.5
+    img_pil = transforms.ToPILImage()(img_tensor)
+    
+    # RETRO GAMEBOY FILTER
+    if retro_mode:
+        # Convert to grayscale then tint green
+        img_pil = ImageOps.grayscale(img_pil)
+        img_pil = ImageOps.colorize(img_pil, black="#0f380f", white="#9bbc0f") # Gameboy Palette
+    
+    # RESOLUTION SCALING
+    size = int(resolution)
+    img_pil = img_pil.resize((size, size), resample=Image.NEAREST)
+    
+    return img_pil
+
+# --- 4. The "Pro" Interface ---
+css = """
+body {background-color: #121212; color: white;}
+.gradio-container {font-family: 'Courier New', monospace;}
+"""
+
+with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# ⚔️ 8-Bit Hero Forge (Ultimate Edition)")
+    gr.Markdown("Create, Morph, and Polish unique pixel characters. Trained on Dec 31, 2025.")
+    
+    with gr.Row():
+        # LEFT COLUMN: Controls
+        with gr.Column(scale=1):
+            gr.Markdown("### 🧬 DNA Mixer")
+            with gr.Row():
+                seed_A = gr.Number(label="Hero A ID", value=42, precision=0)
+                seed_B = gr.Number(label="Hero B ID", value=999, precision=0)
+            
+            morph = gr.Slider(0.0, 1.0, value=0.0, step=0.01, label="Morph Slider (Mix A & B)")
+            
+            gr.Markdown("### 🛠️ Polish Lab")
+            clean_btn = gr.Checkbox(label="Auto-Clean Noise", value=True)
+            retro_btn = gr.Checkbox(label="GameBoy Retro Mode", value=False)
+            res_drop = gr.Dropdown([64, 256, 512, 1024], value=512, label="Output Size (px)")
+            
+            gen_btn = gr.Button("✨ Forge Hero", variant="primary", size="lg")
+            
+            # Randomizer Button Logic
+            def randomize():
+                return random.randint(0, 9999), random.randint(0, 9999)
+            
+            rand_btn = gr.Button("🎲 Randomize IDs")
+            rand_btn.click(fn=randomize, inputs=None, outputs=[seed_A, seed_B])
+
+        # RIGHT COLUMN: Output
+        with gr.Column(scale=1):
+            output_img = gr.Image(label="Your Hero", type="pil", interactive=False)
+
+    # Connect the logic
+    gen_btn.click(
+        fn=generate_hero, 
+        inputs=[seed_A, seed_B, morph, clean_btn, retro_btn, res_drop], 
+        outputs=output_img
+    )
+    
+    # Live update when dragging morph slider
+    morph.change(
+        fn=generate_hero,
+        inputs=[seed_A, seed_B, morph, clean_btn, retro_btn, res_drop], 
+        outputs=output_img
+    )
+
+demo.launch()