import gradio as gr
import torch
from diffusers import StableDiffusionDepth2ImgPipeline

# Load model
pipe = StableDiffusionDepth2ImgPipeline.from_pretrained(
    "stabilityai/stable-diffusion-2-depth",
    torch_dtype=torch.float16,
    variant="fp16",
    use_safetensors=True
).to("cuda" if torch.cuda.is_available() else "cpu")

# Define the function
def decorate_room(image, prompt, strength=0.7, guidance_scale=7.5):
    output = pipe(prompt=prompt, image=image, strength=strength, guidance_scale=guidance_scale)
    return output.images[0]

# Create Gradio UI
demo = gr.Interface(
    fn=decorate_room,
    inputs=[
        gr.Image(type="pil"),
        gr.Textbox(label="Prompt"),
        gr.Slider(0, 1, value=0.7, label="Denoising Strength"),
        gr.Slider(0, 20, value=7.5, label="Guidance Scale")
    ],
    outputs=gr.Image(type="pil"),
    title="Parentsphere Room Decorator",
    description="Upload a room photo and describe how you want it redecorated!"
)

demo.launch()