app.py

import gradio as gr
from diffusers import ControlNetModel, StableDiffusionXLControlNetPipeline, AutoencoderKL
import torch
import numpy as np
import cv2
from PIL import Image
import spaces

# 🌟 Auto-detect device (CPU/GPU)
device = "cuda"
precision = torch.float16


# 🏗️ Load ControlNet model for Canny edge detection
controlnet = ControlNetModel.from_pretrained(
    "diffusers/controlnet-canny-sdxl-1.0",
    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=precision)

pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    controlnet=controlnet,
    vae=vae,
    torch_dtype=precision,
)

pipe.enable_model_cpu_offload()
pipe.to(device)

## to move in the function
controlnet_conditioning_scale = 0.5  # recommended for good generalization

# 📸 Edge detection function using OpenCV (Canny)
@spaces.GPU
def apply_canny(image, low_threshold, high_threshold):
    image = np.array(image)
    image = cv2.Canny(image, low_threshold, high_threshold)
    image = image[:, :, None]
    image = np.concatenate([image, image, image], axis=2)
    return Image.fromarray(image)

# 🎨 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,
        controlnet_conditioning_scale=controlnet_conditioning_scale,
        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()