Update app.py

app.py

@@ -1,12 +1,18 @@
-from fastapi import FastAPI, File, UploadFile
-from fastapi.responses import JSONResponse
+# app.py
+from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse
+from pydantic import BaseModel
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
-from PIL import Image
-from transformers import AutoImageProcessor, AutoBackbone
-import io
 
-app = FastAPI(title="FAST TextNet API")
+# Modèle HF Flan-T5
+MODEL_NAME = "google/flan-t5-base"
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
+model.eval()
+
+app = FastAPI(title="Flan-T5 Service")
 
 app.add_middleware(
     CORSMiddleware,
@@ -15,44 +21,20 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-print("Chargement du modèle...")
-processor = AutoImageProcessor.from_pretrained("czczup/textnet-base")
-model = AutoBackbone.from_pretrained("czczup/textnet-base")
-model.eval()
-print("Modèle prêt !")
+class ImproveRequest(BaseModel):
+    text: str
 
 @app.get("/")
 def health():
-    return {"status": "ok", "model": "czczup/textnet-base"}
+    return {"status": "ok", "model": MODEL_NAME}
 
-@app.post("/detect")
-async def detect_text(file: UploadFile = File(...)):
+@app.post("/improve")
+def improve_text(req: ImproveRequest):
     try:
-        image_bytes = await file.read()
-        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-
-        inputs = processor(image, return_tensors="pt")
-
+        inputs = tokenizer(req.text, return_tensors="pt")
         with torch.no_grad():
-            outputs = model(**inputs)
-
-        # Récupère tous les feature maps
-        feature_maps = []
-        for fm in outputs.feature_maps:
-            feature_maps.append({
-                "shape": list(fm.shape),
-                "mean": float(fm.mean()),
-                "std": float(fm.std()),
-                "min": float(fm.min()),
-                "max": float(fm.max()),
-            })
-
-        return JSONResponse({
-            "success": True,
-            "image_size": list(image.size),
-            "feature_maps": feature_maps,
-            "num_stages": len(feature_maps),
-        })
-
+            outputs = model.generate(**inputs, max_new_tokens=512)
+        improved = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return JSONResponse({"success": True, "improved_text": improved})
     except Exception as e:
-        return JSONResponse({"success": False, "error": str(e)}, status_code=500)
+        return JSONResponse({"success": False, "error": str(e)}, status_code=500)