Update app.py

app.py

@@ -1,20 +1,23 @@
+# app.py
+
+from fastapi import FastAPI
+from pydantic import BaseModel
 from diffusers import DiffusionPipeline
 import torch
-import matplotlib.pyplot as plt
 
-# Vérifier si un GPU est disponible et l'utiliser si possible
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# Initialiser l'application FastAPI
+app = FastAPI()
 
 # Charger le modèle
-pipe = DiffusionPipeline.from_pretrained("XLabs-AI/flux-controlnet-canny").to(device)
-
-# Définir le prompt
-prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+pipe = DiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4").to(device)
 
-# Générer l'image
-image = pipe(prompt).images[0]
+# Définir le modèle de requête
+class Prompt(BaseModel):
+    text: str
 
-# Afficher l'image
-plt.imshow(image)
-plt.axis("off")
-plt.show()
+@app.post("/generate")
+def generate_image(prompt: Prompt):
+    image = pipe(prompt.text).images[0]
+    image.save("output.png")  # Enregistrer l'image générée
+    return {"message": "Image generated", "image_path": "output.png"}