Update app.py

app.py

@@ -1,26 +1,33 @@
-from fastapi import FastAPI
-from pydantic import BaseModel
+# Installer les bibliothèques nécessaires
+!pip install diffusers[torch] transformers gradio
+
+# Importer les bibliothèques nécessaires
 from diffusers import DiffusionPipeline
 import torch
+import gradio as gr
 
-# Initialiser l'application FastAPI
-app = FastAPI()
-
-# Charger le modèle
+# Vérifier si un GPU est disponible et l'utiliser si possible
 device = "cuda" if torch.cuda.is_available() else "cpu"
-pipe = DiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4").to(device)
 
-# Définir le modèle de requête
-class Prompt(BaseModel):
-    text: str
+# Charger le modèle principal
+pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev").to(device)
+
+# Charger les poids LoRA
+pipe.load_lora_weights("gokaygokay/Flux-White-Background-LoRA")
+
+# Fonction pour générer l'image
+def generate_image(prompt):
+    image = pipe(prompt).images[0]
+    return image
 
-@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"}
+# Créer l'interface Gradio
+interface = gr.Interface(
+    fn=generate_image,
+    inputs="text",
+    outputs="image",
+    title="Générateur d'images FLUX",
+    description="Entrez un prompt pour générer une image avec le modèle FLUX.",
+)
 
-# Point d'entrée pour Uvicorn
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+# Lancer l'interface
+interface.launch()