import streamlit as st
import torch
from diffusers import StableDiffusionPipeline
import matplotlib.pyplot as plt

# Configure l'app Streamlit
st.title("Générateur d'Images avec Stable Diffusion")

# Vérifiez si CUDA est disponible
device = "cuda" if torch.cuda.is_available() else "cpu"

# Chargez le modèle Stable Diffusion
model_id = "CompVis/stable-diffusion-v1-4"  # Vous pouvez également utiliser un autre modèle disponible
pipe = StableDiffusionPipeline.from_pretrained(model_id).to(device)

# Fonction pour générer des images
def generate_image(prompt):
    with torch.no_grad():
        image = pipe(prompt).images[0]
    return image

# Champ de texte pour le prompt
prompt = st.text_input("Entrez un prompt pour générer une image :", "a fantasy landscape with mountains and a river")

# Bouton pour générer l'image
if st.button("Générer l'Image"):
    if prompt:
        with st.spinner("Génération de l'image..."):
            generated_image = generate_image(prompt)

        # Affichage de l'image générée
        st.image(generated_image, caption=prompt, use_column_width=True)
    else:
        st.error("Veuillez entrer un prompt.")