import gradio as gr
from diffusers import ControlNetModel, StableDiffusionControlNetPipeline, UniPCMultistepScheduler
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 ADE20k
controlnet = ControlNetModel.from_pretrained("thibaud/controlnet-sd21-ade20k-diffusers", torch_dtype=precision)

# Load the model from Hugging Face
pipe = StableDiffusionControlNetPipeline.from_pretrained(
    "stabilityai/stable-diffusion-2-1",
    torch_dtype=precision,
    controlnet=controlnet,
)

pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
pipe.to("cuda")  # Use GPU for faster inference


# 🎨 Image generation function
@spaces.GPU
def generate_image(prompt, input_image, strength, guidance, controlnet_conditioning_scale):
    
    # Generate styled image using ControlNet
    result = pipe(
        prompt=prompt,
        image=input_image,
        num_inference_steps=30,
        guidance_scale=guidance,
        controlnet_conditioning_scale=float(controlnet_conditioning_scale),
        strength=strength
    ).images[0]
    
    return result


# 🖥️ Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# 🏗️ 3D Screenshot to Styled Render with ADE20k Segmentation ControlNet")

    with gr.Row():
        with gr.Column():
            input_image = gr.Image(label="Upload a segmented image", type="pil")
            prompt = gr.Textbox(label="Style Prompt", placeholder="e.g., Futuristic building in sunset")
            
            strength = gr.Slider(0.1, 1.0, value=0.7, label="Denoising Strength")
            guidance = gr.Slider(1, 20, value=7.5, label="Guidance Scale (Creativity)")
            controlnet_conditioning_scale = gr.Slider(0, 2, value=0.5, step=0.01, label="ControlNet Conditioning Scale")
            
            generate_button = gr.Button("Generate Styled Image")

        with gr.Column():
            result_output = gr.Image(label="Generated Styled Image")

    # 🔗 Generate Button Action
    generate_button.click(
        fn=generate_image,
        inputs=[prompt, input_image, strength, guidance, controlnet_conditioning_scale],
        outputs=[result_output]
    )

# 🚀 Launch the app
demo.launch()