Update app.py

app.py

@@ -1,7 +1,94 @@
 import gradio as gr
+from diffusers import ControlNetModel, StableDiffusionXLControlNetPipeline, AutoencoderKL
+from diffusers import DDIMScheduler, EulerAncestralDiscreteScheduler
+import torch
+import numpy as np
+import cv2
+from PIL import Image
+import spaces
 
-def greet(name):
-    return "Hello " + name + "!!"
+# 🌟 Auto-detect device (CPU/GPU)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+precision = torch.float16 if device == "cuda" else torch.float32
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+
+eulera_scheduler = EulerAncestralDiscreteScheduler.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="scheduler")
+
+# 🏗️ Load ControlNet model for Canny edge detection
+controlnet = ControlNetModel.from_pretrained(
+    "thepowerfuldeez/sd21-controlnet-canny",
+    torch_dtype=precision
+)
+
+# when test with other base model, you need to change the vae also.
+vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
+
+pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-base-1.0",
+    controlnet=controlnet,
+    vae=vae,
+    torch_dtype=torch.float16,
+    scheduler=eulera_scheduler,
+)
+
+pipe.to(device)
+
+# ✨ Enable attention slicing for CPU optimization
+#pipe.enable_attention_slicing()
+
+# 📸 Edge detection function using OpenCV (Canny)
+@spaces.GPU
+def apply_canny(image, low_threshold, high_threshold):
+    image = np.array(image.convert("RGB"))
+    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
+    edges = cv2.Canny(gray, low_threshold, high_threshold)
+    edges = np.stack([edges] * 3, axis=-1)  # Convert to RGB format
+    return Image.fromarray(edges)
+
+# 🎨 Image generation function
+@spaces.GPU
+def generate_image(prompt, input_image, low_threshold, high_threshold, strength, guidance):
+    # Apply edge detection
+    edge_detected = apply_canny(input_image, low_threshold, high_threshold)
+    
+    # Generate styled image using ControlNet
+    result = pipe(
+        prompt=prompt,
+        image=edge_detected,
+        num_inference_steps=30,
+        guidance_scale=guidance,
+        strength=strength
+    ).images[0]
+    
+    return edge_detected, result
+
+# 🖥️ Gradio UI
+with gr.Blocks() as demo:
+    gr.Markdown("# 🏗️ 3D Screenshot to Styled Render with ControlNet")
+
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(label="Upload 3D Screenshot", type="pil")
+            prompt = gr.Textbox(label="Style Prompt", placeholder="e.g., Futuristic building in sunset")
+            
+            low_threshold = gr.Slider(50, 150, value=100, label="Canny Edge Low Threshold")
+            high_threshold = gr.Slider(100, 200, value=150, label="Canny Edge High Threshold")
+            
+            strength = gr.Slider(0.1, 1.0, value=0.8, label="Denoising Strength")
+            guidance = gr.Slider(1, 20, value=7.5, label="Guidance Scale (Creativity)")
+            
+            generate_button = gr.Button("Generate Styled Image")
+
+        with gr.Column():
+            edge_output = gr.Image(label="Edge Detected Image")
+            result_output = gr.Image(label="Generated Styled Image")
+
+    # 🔗 Button Action
+    generate_button.click(
+        fn=generate_image,
+        inputs=[prompt, input_image, low_threshold, high_threshold, strength, guidance],
+        outputs=[edge_output, result_output]
+    )
+
+# 🚀 Launch the app
+demo.launch()