EnoderBOT / app.py
Enoder's picture
Create app.py
adce9cf verified
Raw
History Blame
989 Bytes
# 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)