Update handler.py

handler.py

@@ -1,48 +1,87 @@
-import base64, io, os
+import base64
+import io
+import os
+from typing import Dict, Any
+
 import torch
-from torchvision import transforms
 from PIL import Image
 from safetensors.torch import load_file
 from timm import create_model
+from torchvision import transforms
 
 
 class EndpointHandler:
-    """Custom pipeline for HF Inference Endpoints."""
+    """Custom image-classification pipeline for Hugging Face Inference Endpoints."""
 
-    def __init__(self, model_dir: str):
+    # --------------------------------------------------
+    # 1) تحميل النموذج والوزن مرة واحدة عند تشغيل الـ Endpoint
+    # --------------------------------------------------
+    def __init__(self, model_dir: str) -> None:
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-        weights = load_file(os.path.join(model_dir, "model.safetensors"))
+        # وزن محفوظ بصيغة safetensors
+        weights_path = os.path.join(model_dir, "model.safetensors")
+        state_dict = load_file(weights_path)
+
+        # أنشئ نفس معماريّة ViT التى درّبتها (num_classes = 5)
         self.model = create_model("vit_base_patch16_224", num_classes=5)
-        self.model.load_state_dict(weights)
+        self.model.load_state_dict(state_dict)
         self.model.eval().to(self.device)
 
-        self.transform = transforms.Compose([
-            transforms.Resize((224, 224), interpolation=Image.BICUBIC),
-            transforms.ToTensor(),
-        ])
+        # محوّلات التحضير
+        self.preprocess = transforms.Compose(
+            [
+                transforms.Resize((224, 224), interpolation=Image.BICUBIC),
+                transforms.ToTensor(),
+            ]
+        )
+
+        self.labels = [
+            "stable_diffusion",
+            "midjourney",
+            "dalle",
+            "real",
+            "other_ai",
+        ]
 
-        self.labels = ["stable_diffusion", "midjourney", "dalle", "real", "other_ai"]
+    # --------------------------------------------------
+    # 2) دوال مساعدة
+    # --------------------------------------------------
+    def _image_from_bytes(self, b: bytes) -> Image.Image:
+        """decode base64 → PIL"""
+        return Image.open(io.BytesIO(base64.b64decode(b)))
 
-    def _prep(self, img: Image.Image):
-        return self.transform(img.convert("RGB")).unsqueeze(0).to(self.device)
+    def _to_tensor(self, img: Image.Image) -> torch.Tensor:
+        """PIL → tensor (1 × 3 × 224 × 224) على نفس الجهاز"""
+        return self.preprocess(img.convert("RGB")).unsqueeze(0).to(self.device)
 
-    def __call__(self, data):
-        img = None
+    # --------------------------------------------------
+    # 3) الدالة الرئيسة التى تستدعيها المنصّة لكل طلب
+    # --------------------------------------------------
+    def __call__(self, data: Any) -> Dict[str, float]:
+        """
+        يدعم:
+        • Widget — يمرّر PIL.Image مباشرةً
+        • REST — يمرّر dict وفيه مفتاح "inputs" أو "image" (base64)
+        """
+        # — الحصول على صورة PIL —
+        img: Image.Image | None = None
         if isinstance(data, Image.Image):
             img = data
         elif isinstance(data, dict):
-            b64 = data.get("inputs") or data.get("image")
-            if isinstance(b64, (str, bytes)):
-                if isinstance(b64, str):
-                    b64 = b64.encode()
-                img = Image.open(io.BytesIO(base64.b64decode(b64)))
+            payload = data.get("inputs") or data.get("image")
+            if isinstance(payload, (str, bytes)):
+                if isinstance(payload, str):
+                    payload = payload.encode()
+                img = self._image_from_bytes(payload)
 
         if img is None:
             return {"error": "No image provided"}
 
+        # — الاستدلال —
         with torch.no_grad():
-            logits = self.model(self._prep(img))
+            logits = self.model(self._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))}
+fix: proper indentation for handler