Update server1.py

server1.py

@@ -1,16 +1,16 @@
 # app.py
 # comentarios sin tildes / sin enye
 
-import os, io
+import os, io, traceback
 from typing import Optional
 import torch
-from fastapi import FastAPI, File, UploadFile
+from fastapi import FastAPI, File, UploadFile, Request
 from fastapi.responses import JSONResponse
 from PIL import Image, UnidentifiedImageError
 import open_clip
-from fastapi import FastAPI, File, UploadFile, Request 
-import traceback                                      
-import os
+from torchvision import transforms as T
+
+# caches locales
 os.environ.setdefault("HF_HOME", "/app/cache")
 os.environ.setdefault("XDG_CACHE_HOME", "/app/cache")
 os.environ.setdefault("HUGGINGFACE_HUB_CACHE", "/app/cache/huggingface")
@@ -18,8 +18,6 @@ os.environ.setdefault("TRANSFORMERS_CACHE", "/app/cache/huggingface")
 os.environ.setdefault("TORCH_HOME", "/app/cache/torch")
 os.makedirs("/app/cache", exist_ok=True)
 
-from torchvision import transforms as T
-
 # limites basicos
 torch.set_num_threads(1)
 os.environ["OMP_NUM_THREADS"] = "1"
@@ -28,7 +26,6 @@ os.environ["MKL_NUM_THREADS"] = "1"
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 DTYPE  = torch.float16 if DEVICE == "cuda" else torch.float32
 
-# nombres de ficheros (en el mismo repo)
 MODEL_EMB_PATH = os.getenv("MODEL_EMB_PATH", "text_embeddings_modelos_h14.pt")
 VERS_EMB_PATH  = os.getenv("VERS_EMB_PATH",  "text_embeddings_h14.pt")
 
@@ -60,12 +57,12 @@ def _ensure_label_list(x):
 def _load_embeddings(path: str):
     ckpt = torch.load(path, map_location="cpu")
     labels = _ensure_label_list(ckpt["labels"])
-    embeds = ckpt["embeddings"].to("cpu", dtype=torch.float16)
+    embeds = ckpt["embeddings"].to("cpu")   # guardados como fp16
     embeds = embeds / embeds.norm(dim=-1, keepdim=True)
     return labels, embeds
 
-model_labels, model_embeddings   = _load_embeddings(MODEL_EMB_PATH)   # "Marca Modelo"
-version_labels, version_embeddings = _load_embeddings(VERS_EMB_PATH)  # "Marca Modelo Version"
+model_labels, model_embeddings     = _load_embeddings(MODEL_EMB_PATH)
+version_labels, version_embeddings = _load_embeddings(VERS_EMB_PATH)
 
 # ============== inferencia ==============
 @torch.inference_mode()
@@ -79,6 +76,8 @@ def _encode_image(img_tensor: torch.Tensor) -> torch.Tensor:
 
 def _predict_top(text_feats_dev: torch.Tensor, text_labels: list[str], image_tensor: torch.Tensor, topk: int = 1):
     img_f = _encode_image(image_tensor)
+    # casteamos embeddings al mismo dtype que la imagen
+    text_feats_dev = text_feats_dev.to(device=img_f.device, dtype=img_f.dtype)
     sim = (100.0 * img_f @ text_feats_dev.T).softmax(dim=-1)[0]
     vals, idxs = torch.topk(sim, k=topk)
     return [{"label": text_labels[i], "confidence": round(float(v)*100.0, 2)} for v, i in zip(vals, idxs)]
@@ -94,7 +93,7 @@ def process_image_bytes(image_bytes: bytes):
     img_tensor = transform(img).unsqueeze(0).to(device=DEVICE, dtype=DTYPE)
 
     # paso 1: top-1 modelo
-    model_feats_dev = model_embeddings.to(DEVICE) if DEVICE == "cuda" else model_embeddings
+    model_feats_dev = model_embeddings.to(device=DEVICE, dtype=DTYPE)
     top_model = _predict_top(model_feats_dev, model_labels, img_tensor, topk=1)[0]
     modelo_full = top_model["label"]; conf_m = top_model["confidence"]
 
@@ -112,7 +111,7 @@ def process_image_bytes(image_bytes: bytes):
 
     idxs = [i for _, i in matches]
     labels_sub = [lab for lab, _ in matches]
-    embeds_sub = version_embeddings[idxs].to(DEVICE) if DEVICE == "cuda" else version_embeddings[idxs]
+    embeds_sub = version_embeddings[idxs].to(device=DEVICE, dtype=DTYPE)
 
     # paso 3: top-1 version
     top_ver = _predict_top(embeds_sub, labels_sub, img_tensor, topk=1)[0]
@@ -137,23 +136,16 @@ def process_image_bytes(image_bytes: bytes):
 def root():
     return {"status": "ok", "device": DEVICE}
 
-
 @app.post("/predict")
 async def predict(front: UploadFile = File(None), back: Optional[UploadFile] = File(None), request: Request = None):
     try:
-        # log de cabeceras y tipos
         if request:
             print("headers:", dict(request.headers))
         if front is None:
-            print("no llego 'front'")
             return JSONResponse(content={"code": 400, "error": "faltan archivos: 'front' es obligatorio"}, status_code=200)
 
-        print("front filename:", front.filename, "content_type:", front.content_type)
         front_bytes = await front.read()
-        print("front size:", len(front_bytes))
-
         if back is not None:
-            print("back filename:", back.filename, "content_type:", back.content_type)
             _ = await back.read()
 
         result = process_image_bytes(front_bytes)
@@ -163,3 +155,4 @@ async def predict(front: UploadFile = File(None), back: Optional[UploadFile] = F
         print("EXCEPTION:", repr(e))
         traceback.print_exc()
         return JSONResponse(content={"code": 404, "data": {}, "error": str(e)}, status_code=200)
+