import gradio as gr
from diffusers import StableDiffusionPipeline
import torch

# Check if CUDA is available, otherwise use CPU
device = "cuda" if torch.cuda.is_available() else "cpu"

# Load the Stable Diffusion model
pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16)
pipe = pipe.to(device)  # Move the model to the selected device (GPU or CPU)

def generate_image(prompt):
    # Generate the image using the model
    image = pipe(prompt).images[0]
    return image

# Set up the Gradio interface
interface = gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(label="Enter your prompt", placeholder="e.g., A futuristic cityscape at night"),
    outputs=gr.Image(label="Generated Image"),
    live=True,
    title="High-Quality Image Generator",
    description="This app generates high-quality images based on your text prompts using Stable Diffusion."
)

interface.launch()