import torch
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
import gradio as gr

# Check if CUDA is available
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

if device == "cuda":
    torch.cuda.empty_cache()

model_id = "stabilityai/stable-diffusion-2-1"

# Use appropriate dtype based on device
if device == "cuda":
    pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
else:
    pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)

pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe = pipe.to(device)

def generate_image(prompt, width, height):
    image = pipe(prompt, width=int(width), height=int(height)).images[0]
    return image

iface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Prompt", value="a house in front of the ocean and a dog is running in the field"),
        gr.Number(label="Width", value=1000),
        gr.Number(label="Height", value=1000)
    ],
    outputs=gr.Image(type="pil"),
    title="Stable Diffusion Image Generator"
)

if __name__ == "__main__":
    iface.launch()