Update app.py

app.py

@@ -1,26 +1,29 @@
-# Installer les dépendances nécessaires
-!pip install torch torchvision transformers diffusers
-
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from diffusers import StableDiffusionPipeline
+from transformers 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
+# 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
 
-# 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}'")
+# 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)
 
-# Exemple d'utilisation
-prompt = "a cute animal in a forest"
-generate_animal_image(prompt)
+# Affichage de l'image générée
+for img in generated_images:
+    plt.imshow(img)
+    plt.axis('off')
+    plt.show()