Update handler.py

handler.py

@@ -5,25 +5,24 @@ from typing import Dict, Any
 
 import torch
 from PIL import Image
-from safetensors.torch import load_file
 from timm import create_model
 from torchvision import transforms
 
 
 class EndpointHandler:
-    """Custom image-classification pipeline for Hugging Face Inference Endpoints."""
+    """Custom ViT image-classifier for Hugging Face Inference Endpoints."""
 
     # --------------------------------------------------
-    # 1) تحميل النموذج والوزن مرة واحدة عند تشغيل الـ Endpoint
+    # 1) تحميل النموذج والوزن مرة واحدة
     # --------------------------------------------------
     def __init__(self, model_dir: str) -> None:
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-        # وزن محفوظ بصيغة safetensors
-        weights_path = os.path.join(model_dir, "model.safetensors")
-        state_dict = load_file(weights_path)
+        # ─── تحميل وزن ViT الصحيح ───
+        weights_path = os.path.join(model_dir, "pytorch_model.bin")  # موجود بالفعل
+        state_dict = torch.load(weights_path, map_location="cpu")
 
-        # أنشئ نفس معماريّة ViT التى درّبتها (num_classes = 5)
+        # إنشاء ViT Base Patch-16 بعدد فئات 5
         self.model = create_model("vit_base_patch16_224", num_classes=5)
         self.model.load_state_dict(state_dict)
         self.model.eval().to(self.device)
@@ -47,25 +46,23 @@ class EndpointHandler:
     # --------------------------------------------------
     # 2) دوال مساعدة
     # --------------------------------------------------
-    def _image_from_bytes(self, b: bytes) -> Image.Image:
-        """decode base64 → PIL"""
-        return Image.open(io.BytesIO(base64.b64decode(b)))
-
-    def _to_tensor(self, img: Image.Image) -> torch.Tensor:
-        """PIL → tensor (1 × 3 × 224 × 224) على نفس الجهاز"""
+    def _img_to_tensor(self, img: Image.Image) -> torch.Tensor:
         return self.preprocess(img.convert("RGB")).unsqueeze(0).to(self.device)
 
+    def _decode_b64(self, b: bytes) -> Image.Image:
+        return Image.open(io.BytesIO(base64.b64decode(b)))
+
     # --------------------------------------------------
-    # 3) الدالة الرئيسة التى تستدعيها المنصّة لكل طلب
+    # 3) الدالة الرئيسة
     # --------------------------------------------------
     def __call__(self, data: Any) -> Dict[str, float]:
         """
         يدعم:
-        • Widget — يمرّر PIL.Image مباشرةً
-        • REST — يمرّر dict وفيه مفتاح "inputs" أو "image" (base64)
+        • Widget (PIL.Image)
+        • REST  (base64 فى data["inputs"] أو data["image"])
         """
-        # — الحصول على صورة PIL —
         img: Image.Image | None = None
+
         if isinstance(data, Image.Image):
             img = data
         elif isinstance(data, dict):
@@ -73,14 +70,13 @@ class EndpointHandler:
             if isinstance(payload, (str, bytes)):
                 if isinstance(payload, str):
                     payload = payload.encode()
-                img = self._image_from_bytes(payload)
+                img = self._decode_b64(payload)
 
         if img is None:
             return {"error": "No image provided"}
 
-        # — الاستدلال —
         with torch.no_grad():
-            logits = self.model(self._to_tensor(img))
+            logits = self.model(self._img_to_tensor(img))
             probs = torch.nn.functional.softmax(logits.squeeze(0), dim=0)
 
         return {self.labels[i]: float(probs[i]) for i in range(len(self.labels))}