Update server1.py

server1.py

@@ -7,7 +7,7 @@ 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
+import multiprocessing as mp
 
 ImageFile.LOAD_TRUNCATED_IMAGES = True
 
@@ -25,10 +25,12 @@ os.makedirs(os.environ["TORCH_HOME"], exist_ok=True)
 
 import open_clip  # importar despues de ajustar caches
 
-# ===== limites basicos =====
-torch.set_num_threads(1)
-os.environ["OMP_NUM_THREADS"] = "1"
-os.environ["MKL_NUM_THREADS"] = "1"
+# ===== limites basicos (usar todos los nucleos) =====
+NTHREADS = max(1, mp.cpu_count())
+torch.set_num_threads(NTHREADS)
+os.environ["OMP_NUM_THREADS"] = str(NTHREADS)
+os.environ["MKL_NUM_THREADS"] = str(NTHREADS)
+# opcional: torch.set_num_interop_threads(1)
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 DTYPE  = torch.float16 if DEVICE == "cuda" else torch.float32
@@ -40,7 +42,6 @@ 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 PE bigG =====
-# usamos HF Hub; no pasamos 'pretrained' clasico
 MODEL_NAME = "hf-hub:timm/PE-Core-bigG-14-448"
 PRETRAINED = None
 
@@ -52,36 +53,22 @@ _ret = open_clip.create_model_and_transforms(MODEL_NAME, pretrained=PRETRAINED)
 if isinstance(_ret, tuple) and len(_ret) == 3:
     clip_model, _preprocess_train, preprocess = _ret
 else:
-    # fallback por si alguna version devuelve solo 2
     clip_model, preprocess = _ret
 
 clip_model = clip_model.to(device=DEVICE, dtype=DTYPE).eval()
 for p in clip_model.parameters():
     p.requires_grad = False
 
-# extraer normalizacion y size desde el preprocess devuelto
-normalize = next(t for t in getattr(preprocess, "transforms", []) if isinstance(t, T.Normalize))
+# opcional GPU: formato canales para mejorar rendimiento
+if DEVICE == "cuda":
+    clip_model = clip_model.to(memory_format=torch.channels_last)
+
+# obtener SIZE desde preprocess solo para chequeos
 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)  # por si viene como (H,W)
+    SIZE = max(SIZE)
 if SIZE is None:
-    SIZE = 448  # PE bigG es 448; fallback por seguridad
-
-transform = T.Compose([T.ToTensor(), T.Normalize(mean=normalize.mean, std=normalize.std)])
-
-# ===== utils imagen =====
-def resize_letterbox(img: Image.Image, size: int) -> Image.Image:
-    if img.mode != "RGB":
-        img = img.convert("RGB")
-    w, h = img.size
-    if w == 0 or h == 0:
-        raise UnidentifiedImageError("imagen invalida")
-    scale = size / max(w, h)
-    nw, nh = max(1, int(w*scale)), max(1, int(h*scale))
-    img_resized = img.resize((nw, nh), Image.BICUBIC)
-    canvas = Image.new("RGB", (size, size), (0, 0, 0))
-    canvas.paste(img_resized, ((size-nw)//2, (size-nh)//2))
-    return canvas
+    SIZE = 448  # fallback por seguridad
 
 # ===== cargar embeddings =====
 def _ensure_label_list(x):
@@ -103,8 +90,12 @@ version_labels, version_embeddings = _load_embeddings(VERS_EMB_PATH)
 
 # 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]
+    dummy = Image.new("RGB", (SIZE, SIZE), (0, 0, 0))
+    tensor = preprocess(dummy).unsqueeze(0).to(device=DEVICE)
+    if DEVICE == "cuda":
+        tensor = tensor.to(dtype=DTYPE)
+    img_dim = clip_model.encode_image(tensor).shape[-1]
+
 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.shape[1]}, "
@@ -114,15 +105,17 @@ if model_embeddings.shape[1] != img_dim or version_embeddings.shape[1] != img_di
 # ===== inferencia =====
 @torch.inference_mode()
 def _encode_pil(img: Image.Image) -> torch.Tensor:
-    img = resize_letterbox(img, SIZE)
-    tensor = transform(img).unsqueeze(0).to(device=DEVICE)
+    if img.mode != "RGB":
+        img = img.convert("RGB")
+    x = preprocess(img).unsqueeze(0).to(device=DEVICE)
     if DEVICE == "cuda":
-        tensor = tensor.to(dtype=DTYPE)
-    feats = clip_model.encode_image(tensor)
+        x = x.to(dtype=DTYPE)
+    feats = clip_model.encode_image(x)
     return feats / feats.norm(dim=-1, keepdim=True)
 
 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]
+    tf = text_feats.to(img_feat.device, non_blocking=True)
+    sim = (img_feat @ tf.T)[0]  # tensores ya normalizados
     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)]