import gradio as gr
import torch
from diffusers import StableDiffusionPipeline

model_name = "runwayml/stable-diffusion-v1-5"
device = "cuda" if torch.cuda.is_available() else "cpu"

if device == "cuda":
    pipe = StableDiffusionPipeline.from_pretrained(model_name, torch_dtype=torch.float16).to(device)
else:
    pipe = StableDiffusionPipeline.from_pretrained(model_name).to(device)

def generate_image(prompt):
    with torch.no_grad():
        image = pipe(prompt, num_inference_steps=50, guidance_scale=9.5).images[0]

    return image

iface = gr.Interface(
    fn=generate_image,
    inputs="text",
    outputs="image",
    title="BL's T2I Generator",
    description="Enter a prompt to generate an image."
)

iface.launch()