Create app.py

app.py

@@ -0,0 +1,78 @@
+import gradio as gr
+import cv2
+import numpy as np
+import torch
+import os
+from pathlib import Path
+
+# Placeholder for the actual model loading and processing logic
+# Replace this with the actual code from generate.py
+def load_model(model_path, use_cpu=False):
+    # Example: Load your ESRGAN model here
+    # This is a placeholder; adapt it based on the actual model loading in generate.py
+    model = torch.load(model_path)
+    if not use_cpu and torch.cuda.is_available():
+        model = model.cuda()
+    model.eval()
+    return model
+
+def process_image(input_image, tile_size=512, seamless=False, use_cpu=False):
+    # Convert Gradio input (PIL image) to OpenCV format
+    img = np.array(input_image)
+    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+
+    # Load models (adjust paths to match your uploaded model files)
+    normal_model = load_model("models/normal_model.pth", use_cpu)
+    disp_model = load_model("models/displacement_model.pth", use_cpu)
+    rough_model = load_model("models/roughness_model.pth", use_cpu)
+
+    # Placeholder processing logic (replace with actual generate.py logic)
+    # For example, apply the model to the image
+    with torch.no_grad():
+        # Convert image to tensor
+        img_tensor = torch.from_numpy(img.transpose(2, 0, 1)).float() / 255.0
+        if not use_cpu and torch.cuda.is_available():
+            img_tensor = img_tensor.cuda()
+
+        # Generate maps (simplified example)
+        normal_map = normal_model(img_tensor.unsqueeze(0)).cpu().numpy().squeeze().transpose(1, 2, 0)
+        disp_map = disp_model(img_tensor.unsqueeze(0)).cpu().numpy().squeeze().transpose(1, 2, 0)
+        rough_map = rough_model(img_tensor.unsqueeze(0)).cpu().numpy().squeeze().transpose(1, 2, 0)
+
+    # Convert back to uint8 for display
+    normal_map = (normal_map * 255).astype(np.uint8)
+    disp_map = (disp_map * 255).astype(np.uint8)
+    rough_map = (rough_map * 255).astype(np.uint8)
+
+    # Convert to RGB for Gradio output
+    normal_map = cv2.cvtColor(normal_map, cv2.COLOR_BGR2RGB)
+    disp_map = cv2.cvtColor(disp_map, cv2.COLOR_BGR2RGB)
+    rough_map = cv2.cvtColor(rough_map, cv2.COLOR_BGR2RGB)
+
+    return normal_map, disp_map, rough_map
+
+# Gradio interface
+def generate_maps(input_image, tile_size, seamless, use_cpu):
+    normal_map, disp_map, rough_map = process_image(input_image, tile_size, seamless, use_cpu)
+    return input_image, normal_map, disp_map, rough_map
+
+interface = gr.Interface(
+    fn=generate_maps,
+    inputs=[
+        gr.Image(type="pil", label="Diffuse Texture"),
+        gr.Slider(minimum=256, maximum=1024, step=64, value=512, label="Tile Size"),
+        gr.Checkbox(label="Seamless", value=False),
+        gr.Checkbox(label="Use CPU", value=False),
+    ],
+    outputs=[
+        gr.Image(type="numpy", label="Input Diffuse Texture"),
+        gr.Image(type="numpy", label="Normal Map"),
+        gr.Image(type="numpy", label="Displacement Map"),
+        gr.Image(type="numpy", label="Roughness Map"),
+    ],
+    title="Material Map Generator",
+    description="Upload a diffuse texture to generate AI-generated Normal, Displacement, and Roughness maps."
+)
+
+if __name__ == "__main__":
+    interface.launch()