import gradio as gr
import torch


from diffusers import DiffusionPipeline

def generate_image(prompt):
    model_id = "runwayml/stable-diffusion-v1-5"
    pipeline = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
    pipeline = pipeline.to("cuda")
    num_inference_steps = 20
    
    generator = torch.Generator("cuda").manual_seed(0)
    image = pipeline(prompt, generator=generator, num_inference_steps=20).images[0]
    
    return image

iface = gr.Interface(
    fn=generate_image,
    inputs="text",
    outputs="image",
    title="Diffusion d'images",
    description="Générez des images à partir d'une description.",
)

iface.launch()