Update server1.py

server1.py

@@ -1,17 +1,18 @@
 # app.py
-# sin tildes / sin enye
+# comentarios sin tildes / sin enye
 
-import os, io, traceback, time
-from typing import Optional, List, Dict
+import os, io, traceback
+from typing import Optional, List, Tuple
 import torch
 from fastapi import FastAPI, File, UploadFile, Request
 from fastapi.responses import JSONResponse
 from PIL import Image, UnidentifiedImageError, ImageFile
 from torchvision import transforms as T
+from functools import lru_cache
 
 ImageFile.LOAD_TRUNCATED_IMAGES = True
 
-# ===== caches (ruta escribible) =====
+# ===== caches (usar ruta propia, escribible en runtime) =====
 CACHE_ROOT = os.environ.get("APP_CACHE", "/tmp/appcache")
 os.environ["XDG_CACHE_HOME"] = CACHE_ROOT
 os.environ["HF_HOME"] = os.path.join(CACHE_ROOT, "hf")
@@ -23,34 +24,38 @@ os.makedirs(os.environ["HF_HOME"], exist_ok=True)
 os.makedirs(os.environ["OPENCLIP_CACHE_DIR"], exist_ok=True)
 os.makedirs(os.environ["TORCH_HOME"], exist_ok=True)
 
-import open_clip  # importar tras setear caches
+import open_clip  # importar despues de ajustar caches
 
 # ===== limites basicos =====
-CPU_THREADS = int(os.environ.get("CPU_THREADS", max(1, min(8, os.cpu_count() or 1))))
-torch.set_num_threads(CPU_THREADS)
-os.environ["OMP_NUM_THREADS"] = str(CPU_THREADS)
-os.environ["MKL_NUM_THREADS"] = str(CPU_THREADS)
+# por defecto conservamos 1 hilo (tu baseline). Para probar mas:
+# export NUM_THREADS=4 (o el valor que quieras) sin tocar codigo
+NUM_THREADS = int(os.environ.get("NUM_THREADS", "1"))
+torch.set_num_threads(NUM_THREADS)
+os.environ["OMP_NUM_THREADS"] = str(NUM_THREADS)
+os.environ["MKL_NUM_THREADS"] = str(NUM_THREADS)
+try:
+    torch.set_num_interop_threads(1)
+except Exception:
+    pass
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 DTYPE  = torch.float16 if DEVICE == "cuda" else torch.float32
 if DEVICE == "cuda":
     torch.set_float32_matmul_precision("high")
 
-# ===== rutas a embeddings (compatibles con bigG/laion2b_s39b_b160k) =====
+# ===== rutas a embeddings =====
 MODEL_EMB_PATH = os.getenv("MODEL_EMB_PATH", "text_embeddings_modelos_bigg.pt")
 VERS_EMB_PATH  = os.getenv("VERS_EMB_PATH",  "text_embeddings_bigg.pt")
 
-# ===== modelo OpenCLIP bigG (el mismo con el que generaste los .pt) =====
-MODEL_NAME = "ViT-bigG-14"
-PRETRAINED = "laion2b_s39b_b160k"
+# ===== modelo PE bigG =====
+MODEL_NAME = "hf-hub:timm/PE-Core-bigG-14-448"
+PRETRAINED = None
 
-# Tamaño de entrada (por defecto el del preprocess = 448)
-FAST_SIZE = int(os.environ.get("FAST_SIZE", "448"))  # puedes dejar 448; no cambia la dimension de salida
+app = FastAPI(title="OpenCLIP PE bigG Vehicle API")
 
-app = FastAPI(title="OpenCLIP bigG Vehicle API (fast-safe)")
-
-# ===== carga modelo / preprocess =====
+# ===== modelo / preprocess =====
 _ret = open_clip.create_model_and_transforms(MODEL_NAME, pretrained=PRETRAINED)
+# versiones de open_clip devuelven (model, preprocess_train, preprocess_val)
 if isinstance(_ret, tuple) and len(_ret) == 3:
     clip_model, _preprocess_train, preprocess = _ret
 else:
@@ -60,16 +65,17 @@ clip_model = clip_model.to(device=DEVICE, dtype=DTYPE).eval()
 for p in clip_model.parameters():
     p.requires_grad = False
 
-# normalizacion & size desde preprocess (448). Permite bajar via FAST_SIZE sin tocar .pt
+# extraer normalizacion y size desde el preprocess devuelto
 normalize = next(t for t in getattr(preprocess, "transforms", []) if isinstance(t, T.Normalize))
-DEFAULT_SIZE = next((getattr(t, "size", None) for t in getattr(preprocess, "transforms", []) if hasattr(t, "size")), None)
-if isinstance(DEFAULT_SIZE, (tuple, list)):
-    DEFAULT_SIZE = max(DEFAULT_SIZE)
-SIZE = min(DEFAULT_SIZE or 448, FAST_SIZE)
+SIZE = next((getattr(t, "size", None) for t in getattr(preprocess, "transforms", []) if hasattr(t, "size")), None)
+if isinstance(SIZE, (tuple, list)):
+    SIZE = max(SIZE)
+if SIZE is None:
+    SIZE = 448  # PE bigG es 448; fallback
 
 transform = T.Compose([T.ToTensor(), T.Normalize(mean=normalize.mean, std=normalize.std)])
 
-# ===== utils imagen =====
+# ===== utils imagen (sin cambios: letterbox + BICUBIC) =====
 def resize_letterbox(img: Image.Image, size: int) -> Image.Image:
     if img.mode != "RGB":
         img = img.convert("RGB")
@@ -83,7 +89,7 @@ def resize_letterbox(img: Image.Image, size: int) -> Image.Image:
     canvas.paste(img_resized, ((size-nw)//2, (size-nh)//2))
     return canvas
 
-# ===== cargar embeddings =====
+# ===== cargar embeddings (sin cambios) =====
 def _ensure_label_list(x):
     if isinstance(x, (list, tuple)):
         return list(x)
@@ -98,51 +104,37 @@ def _load_embeddings(path: str):
     embeds = embeds / embeds.norm(dim=-1, keepdim=True)
     return labels, embeds
 
-model_labels, model_embeddings_cpu     = _load_embeddings(MODEL_EMB_PATH)
-version_labels, version_embeddings_cpu = _load_embeddings(VERS_EMB_PATH)
-
-# mover a device una vez y cachear transpuestas
-model_embeddings_dev   = model_embeddings_cpu.to(device=DEVICE, dtype=DTYPE).contiguous()
-version_embeddings_dev = version_embeddings_cpu.to(device=DEVICE, dtype=DTYPE).contiguous()
-model_embeddings_T     = model_embeddings_dev.t().contiguous()
+model_labels, model_embeddings     = _load_embeddings(MODEL_EMB_PATH)
+version_labels, version_embeddings = _load_embeddings(VERS_EMB_PATH)
 
-# comprobar dimension de imagen vs textos
+# comprobar dimension (PE bigG mantiene 1280)
 with torch.inference_mode():
     dummy = torch.zeros(1, 3, SIZE, SIZE, device=DEVICE, dtype=DTYPE)
     img_dim = clip_model.encode_image(dummy).shape[-1]
-if model_embeddings_dev.shape[1] != img_dim or version_embeddings_dev.shape[1] != img_dim:
+if model_embeddings.shape[1] != img_dim or version_embeddings.shape[1] != img_dim:
     raise RuntimeError(
-        f"dimension mismatch: image={img_dim}, modelos={model_embeddings_dev.shape[1]}, "
-        f"versiones={version_embeddings_dev.shape[1]}. Recalcula embeddings con {MODEL_NAME}/{PRETRAINED}."
+        f"dimension mismatch: image={img_dim}, modelos={model_embeddings.shape[1]}, "
+        f"versiones={version_embeddings.shape[1]}. Recalcula embeddings con {MODEL_NAME}."
     )
 
-# ===== indice de versiones por modelo (para evitar startswith por request) =====
-from collections import defaultdict
-idx_by_model: Dict[str, List[int]] = defaultdict(list)
-
-# ordenar modelos por longitud para hacer match correcto de prefijo
-models_by_len = sorted(model_labels, key=len, reverse=True)
-for j, lab in enumerate(version_labels):
-    for m in models_by_len:
-        if lab.startswith(m):
-            idx_by_model[m].append(j)
-            break
-
-# pre-cachear sub-matrices transpuestas por modelo
-ver_index = {}
-for m, idxs in idx_by_model.items():
-    if idxs:
-        embT = version_embeddings_dev[idxs].t().contiguous()
-        labs = [version_labels[i] for i in idxs]
-        ver_index[m] = (embT, labs)
-    else:
-        ver_index[m] = (None, [])
-
-# ===== warm-up (reduce el primer request frio) =====
-with torch.inference_mode():
-    _ = clip_model.encode_image(torch.zeros(1,3,SIZE,SIZE, device=DEVICE, dtype=DTYPE))
-
-# ===== inferencia =====
+# ===== cache perezosa de sub-embeddings por modelo_full =====
+# no cambia precision; solo evita escanear version_labels en cada request
+_versions_cache: dict[str, Tuple[List[str], torch.Tensor]] = {}
+
+def _get_versions_subset(modelo_full: str) -> Tuple[List[str], Optional[torch.Tensor]]:
+    hit = _versions_cache.get(modelo_full)
+    if hit is not None:
+        return hit
+    idxs = [i for i, lab in enumerate(version_labels) if lab.startswith(modelo_full)]
+    if not idxs:
+        _versions_cache[modelo_full] = ([], None)
+        return _versions_cache[modelo_full]
+    labels_sub = [version_labels[i] for i in idxs]
+    embeds_sub = version_embeddings[idxs]  # copia de esas filas
+    _versions_cache[modelo_full] = (labels_sub, embeds_sub)
+    return _versions_cache[modelo_full]
+
+# ===== inferencia (sin cambios de logica/precision) =====
 @torch.inference_mode()
 def _encode_pil(img: Image.Image) -> torch.Tensor:
     img = resize_letterbox(img, SIZE)
@@ -152,62 +144,58 @@ def _encode_pil(img: Image.Image) -> torch.Tensor:
     feats = clip_model.encode_image(tensor)
     return feats / feats.norm(dim=-1, keepdim=True)
 
-def _top1(text_feats_T: torch.Tensor, img_feat: torch.Tensor):
-    sim = (img_feat @ text_feats_T)[0].float()
-    val, idx = torch.topk(sim, k=1)
-    conf = torch.softmax(val, dim=0)[0]
-    return int(idx), float(conf)*100.0
+def _topk_cosine(text_feats: torch.Tensor, text_labels: List[str], img_feat: torch.Tensor, k: int = 1):
+    sim = (img_feat.float() @ text_feats.to(img_feat.device).float().T)[0]
+    vals, idxs = torch.topk(sim, k=k)
+    conf = torch.softmax(vals, dim=0)
+    return [{"label": text_labels[int(i)], "confidence": round(float(c)*100.0, 2)} for i, c in zip(idxs, conf)]
 
 def process_image_bytes(front_bytes: bytes, back_bytes: Optional[bytes] = None):
     if not front_bytes or len(front_bytes) < 128:
         raise UnidentifiedImageError("imagen invalida")
 
     img_front = Image.open(io.BytesIO(front_bytes))
-    feat = _encode_pil(img_front)
+    feat_front = _encode_pil(img_front)
 
     if back_bytes:
         try:
             img_back = Image.open(io.BytesIO(back_bytes))
-            feat_b = _encode_pil(img_back)
-            feat = (feat + feat_b)
-            feat = feat / feat.norm(dim=-1, keepdim=True)
+            feat_back = _encode_pil(img_back)
+            img_feat = (feat_front + feat_back) / 2
+            img_feat = img_feat / img_feat.norm(dim=-1, keepdim=True)
         except Exception:
-            pass
-
-    # 1) modelo (top-1)
-    idx_m, _ = _top1(model_embeddings_T, feat)
-    modelo_full = model_labels[idx_m]
-
-    # 2) version (solo dentro del subconjunto indexado)
-    embT, labs = ver_index.get(modelo_full, (None, []))
-    version_out = ""
-    if embT is not None and len(labs) > 0:
-        idx_v, conf_v = _top1(embT, feat)
-        raw = labs[idx_v]
-        prefix = modelo_full + " "
-        ver = raw[len(prefix):] if raw.startswith(prefix) else raw
-        ver = ver.split(" ")[0]
-        if conf_v >= 30.0:
-            version_out = ver.title()
-
-    parts = modelo_full.split(" ", 1)
-    marca  = parts[0] if len(parts) >= 1 else ""
-    modelo = parts[1] if len(parts) == 2 else ""
-
-    return {"brand": marca.upper(), "model": modelo.title(), "version": version_out}
+            img_feat = feat_front
+    else:
+        img_feat = feat_front
+
+    # paso 1: modelo
+    top_model = _topk_cosine(model_embeddings, model_labels, img_feat, k=1)[0]
+    modelo_full = top_model["label"]
+
+    partes = modelo_full.split(" ", 1)
+    marca  = partes[0] if len(partes) >= 1 else ""
+    modelo = partes[1] if len(partes) == 2 else ""
+
+    # paso 2: versiones con cache (misma logica, sin bucle global cada vez)
+    labels_sub, embeds_sub = _get_versions_subset(modelo_full)
+    if not labels_sub:
+        return {"brand": marca.upper(), "model": modelo.title(), "version": ""}
+
+    # paso 3: version
+    top_ver = _topk_cosine(embeds_sub, labels_sub, img_feat, k=1)[0]
+    raw = top_ver["label"]
+    prefix = modelo_full + " "
+    ver = raw[len(prefix):] if raw.startswith(prefix) else raw
+    ver = ver.split(" ")[0]
+    if top_ver["confidence"] < 30.0:
+        ver = ""
+
+    return {"brand": marca.upper(), "model": modelo.title(), "version": ver.title() if ver else ""}
 
 # ===== endpoints =====
 @app.get("/")
 def root():
-    return {
-        "status": "ok",
-        "device": DEVICE,
-        "dtype": str(DTYPE),
-        "model": f"{MODEL_NAME}/{PRETRAINED}",
-        "size": SIZE,
-        "img_dim": int(model_embeddings_dev.shape[1]),
-        "threads": CPU_THREADS
-    }
+    return {"status": "ok", "device": DEVICE, "model": f"{MODEL_NAME}", "img_dim": int(model_embeddings.shape[1]), "threads": NUM_THREADS}
 
 @app.post("/predict/")
 async def predict(front: UploadFile = File(None), back: Optional[UploadFile] = File(None), request: Request = None):