change to api

Dockerfile

@@ -2,8 +2,7 @@ FROM python:3.10
 
 WORKDIR /app
 
-RUN pip install --no-cache-dir torch torchvision torchaudio diffusers gradio 
-RUN pip install -U diffusers
+RUN pip install --no-cache-dir torch torchvision torchaudio diffusers fastapi uvicorn
 RUN pip install git+https://github.com/huggingface/diffusers.git
 RUN pip install transformers
 RUN pip install accelerate
@@ -14,4 +13,4 @@ COPY . /app
 
 EXPOSE 7860
 
-CMD ["python", "app.py"]
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -1,26 +1,40 @@
 import os
-from diffusers import FluxPipeline # type: ignore
-import gradio as gr # type: ignore
-from huggingface_hub import login, InferenceClient
+import io
+import torch
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import StreamingResponse
+from diffusers import FluxPipeline  # type: ignore
 
-token = os.getenv("HF_TOKEN")
+app = FastAPI()
 
-client = InferenceClient(
-	provider="together",
-	api_key=token,
-    model="black-forest-labs/FLUX.1-dev")
+# Récupération du token et authentification
+token = os.getenv("HF_TOKEN")
+login(token=token)
 
-def generate_image(prompt):
-    image = client.text_to_image(prompt)
+def generate_image(prompt: str):
+    image = pipe(
+        prompt,
+        height=1024,
+        width=1024,
+        guidance_scale=3.5,
+        num_inference_steps=50,
+        max_sequence_length=512,
+        generator=torch.Generator("cpu").manual_seed(0)
+    ).images[0]
     return image
 
-gradio_app = gr.Interface(
-    fn=generate_image,
-    inputs=gr.Textbox(label="Entrez une description"),
-    outputs=gr.Image(label="Image générée"),
-    title="Générateur d'images IA",
-    description="Entrez une description et générez une image correspondante."
-)
+@app.get("/generate")
+def generate(prompt: str):
+    try:
+        image = generate_image(prompt)
+        # On sauvegarde l'image dans un buffer en mémoire
+        buf = io.BytesIO()
+        image.save(buf, format="PNG")
+        buf.seek(0)
+        return StreamingResponse(buf, media_type="image/png")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
 
 if __name__ == "__main__":
-    gradio_app.launch()
+    import uvicorn # type: ignore
+    uvicorn.run(app, host="0.0.0.0", port=7860)