import gradio as gr
from diffusers import DiffusionPipeline
import torch

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

pipe = DiffusionPipeline.from_pretrained(
    model_id,
    torch_dtype=torch.float16
)

if torch.cuda.is_available():
    pipe = pipe.to("cuda")
else:
    pipe = pipe.to("cpu")


def generate_image(prompt):
    if not prompt.strip():
        return None

    image = pipe(prompt).images[0]
    return image


with gr.Blocks(title="AI Image Generator") as demo:

    gr.Markdown(
        """
        # 🎨 AI Image Generator
        Enter a prompt and let AI create an image for you!
        """
    )

    prompt = gr.Textbox(
        label="✨ Enter your prompt",
        placeholder="A cute cottage in a magical forest at sunset..."
    )

    generate_btn = gr.Button("🎨 Generate Image")

    output = gr.Image(label="Generated Image")

    generate_btn.click(
        fn=generate_image,
        inputs=prompt,
        outputs=output
    )


demo.launch()