import gradio as gr
from diffusers import DiffusionPipeline
import torch

# Load the model
pipe = DiffusionPipeline.from_pretrained("dreamlike-art/dreamlike-photoreal-2.0")

# Check if GPU is available and move model to GPU if possible
if torch.cuda.is_available():
    pipe.to("cuda")
else:
    pipe.to("cpu")

# Define the image generation function
def generate_image(prompt):
    image = pipe(prompt).images[0]
    return image

# Set up the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("## Dreamlike Photoreal 2.0 Image Generator")

    prompt = gr.Textbox(
        label="Enter a creative prompt",
        placeholder="A futuristic city with flying cars"
    )
    
    image_output = gr.Image(label="Generated Image")

    generate_button = gr.Button("Generate Image")

    # Connect the button click to the image generation function
    generate_button.click(fn=generate_image, inputs=prompt, outputs=image_output)

# Launch the app
demo.launch()