import torch
from diffusers import StableDiffusionPipeline
import gradio as gr

# Charger le modèle Stable Diffusion
device = "cuda" if torch.cuda.is_available() else "cpu"
model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16 if device == "cuda" else torch.float32)
pipe = pipe.to(device)

# Fonction pour générer une image
def generate_image(prompt, steps, guidance):
    image = pipe(prompt, num_inference_steps=steps, guidance_scale=guidance).images[0]
    return image

# Interface Gradio
iface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Prompt", placeholder="Entrez une description de l'image..."),
        gr.Slider(10, 50, value=25, step=1, label="Steps d'inférence"),
        gr.Slider(1.0, 20.0, value=7.5, step=0.1, label="Guidance Scale")
    ],
    outputs="image",
    title="Génération d'images avec Stable Diffusion"
)

iface.launch()