Update app.py

updates with quantization

app.py

@@ -20,7 +20,7 @@ logger = logging.getLogger(__name__)
 
 # ── Config ────────────────────────────────────────────────────────────────────
 API_KEY = os.environ.get("API_KEY", "changeme")
-DEVICE  = "cuda" if torch.cuda.is_available() else "cpu"
+DEVICE  = "cpu"  # quantized models only run on CPU
 
 # ── API Key Auth ──────────────────────────────────────────────────────────────
 api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
@@ -30,14 +30,19 @@ def verify_key(key: str = Security(api_key_header)):
         raise HTTPException(status_code=401, detail="Invalid or missing API key")
     return key
 
-# ── Model — loaded once from local folder ─────────────────────────────────────
+# ── Model ─────────────────────────────────────────────────────────────────────
 @lru_cache(maxsize=1)
 def load_model():
-    logger.info(f"Loading model from ./model on {DEVICE}...")
+    logger.info("Loading quantized model from ./model...")
     processor = TrOCRProcessor.from_pretrained('./model')
-    model = VisionEncoderDecoderModel.from_pretrained('./model').to(DEVICE)
+    model = VisionEncoderDecoderModel.from_pretrained('./model')
+    model = torch.quantization.quantize_dynamic(
+        model,
+        {torch.nn.Linear},
+        dtype=torch.qint8
+    )
     model.eval()
-    logger.info("✅ Model ready")
+    logger.info("✅ Quantized model ready")
     return processor, model
 
 # ── Schemas ───────────────────────────────────────────────────────────────────
@@ -60,7 +65,7 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# ── Pre-warm on startup ───────────────────────────────────────────────────────
+# ── Startup ───────────────────────────────────────────────────────────────────
 @app.on_event("startup")
 async def startup():
     logger.info("Pre-warming model...")
@@ -74,31 +79,37 @@ def health():
         "status": "ok",
         "device": DEVICE,
         "model": "anuashok/ocr-captcha-v3 (local)",
-        "quantized": False
+        "quantized": True
     }
 
 @app.post("/solve-captcha-base64", response_model=SolveResponse)
 def solve(req: SolveRequest, _: str = Depends(verify_key)):
     start = time.time()
     try:
-        # Decode image
         raw = req.image_base64
         if "," in raw:
             raw = raw.split(",", 1)[1]
 
         image = Image.open(io.BytesIO(base64.b64decode(raw))).convert("RGB")
 
-        # Run inference
         processor, model = load_model()
         with torch.no_grad():
-            pixel_values = processor(image, return_tensors="pt").pixel_values.to(DEVICE)
+            pixel_values = processor(
+                image,
+                return_tensors="pt"
+            ).pixel_values
+
             generated_ids = model.generate(
                 pixel_values,
                 max_new_tokens=10,
                 num_beams=1,
                 do_sample=False
             )
-            text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+
+            text = processor.batch_decode(
+                generated_ids,
+                skip_special_tokens=True
+            )[0]
 
         elapsed = time.time() - start
         logger.info(f"✅ Solved: '{text}' in {elapsed:.2f}s")
@@ -117,6 +128,5 @@ def solve(req: SolveRequest, _: str = Depends(verify_key)):
             processing_time=time.time() - start
         )
 
-# ── Launch ────────────────────────────────────────────────────────────────────
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)