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

# 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"  # Modèle à partir de Hugging Face
pipe = StableDiffusionPipeline.from_pretrained(model_id).to(device)

# Fonction pour générer des images
def generate_image(prompt, num_images=1):
    images = []
    for _ in range(num_images):
        with torch.no_grad():
            image = pipe(prompt).images[0]
            images.append(image)
    return images

# Génération d'images à partir d'un prompt
prompt = "a fantasy landscape with mountains and a river"
generated_images = generate_image(prompt, num_images=1)

# Affichage de l'image générée
for img in generated_images:
    plt.imshow(img)
    plt.axis('off')
    plt.show()