# Installer les dépendances nécessaires
!pip install torch torchvision transformers diffusers

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from diffusers import StableDiffusionPipeline
import matplotlib.pyplot as plt

# Chargement du modèle de génération d'images Stable Diffusion depuis Hugging Face
model_id = "CompVis/stable-diffusion-v1-4"  # Stable Diffusion v1.4
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe = pipe.to("cuda")  # Utiliser le GPU si disponible

# Fonction pour générer et sauvegarder une image d'animal
def generate_animal_image(prompt, output_path="animal_image.png"):
    with torch.no_grad():
        image = pipe(prompt).images[0]
        image.save(output_path)
        plt.imshow(image)
        plt.axis('off')
        plt.show()
        print(f"Generated image saved as '{output_path}'")

# Exemple d'utilisation
prompt = "a cute animal in a forest"
generate_animal_image(prompt)