Update handler.py

handler.py

@@ -3,39 +3,61 @@ import torch
 from torchvision import transforms
 from PIL import Image
 from safetensors.torch import load_file
-from timm import create_model   # timm ضرورى للتعامل مع ViT
+from timm import create_model  # timm ضروري لتشغيل ViT
 
 
-class EndpointHandler:      # اسم الفئة مهم جداً
+class EndpointHandler:
+    """Custom pipeline for Hugging Face Inference Endpoints."""
+
     def __init__(self, model_dir: str):
+        # اختَر GPU إذا متاح
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
         # تحميل الوزن بصيغة safetensors
-        weights = load_file(os.path.join(model_dir, "model.safetensors"))
+        weights_path = os.path.join(model_dir, "model.safetensors")
+        weights = load_file(weights_path)
+
+        # إنشاء نموذج ViT مطابق لِما درّبتَه
         self.model = create_model("vit_base_patch16_224", num_classes=5)
         self.model.load_state_dict(weights)
         self.model.eval().to(self.device)
 
-        self.transform = transforms.Compose([
-            transforms.Resize((224, 224), interpolation=Image.BICUBIC),
-            transforms.ToTensor(),
-        ])
+        # تحويـلات الصورة
+        self.transform = 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",
+        ]
 
+    # ---------- helpers ----------
     def _prep(self, img: Image.Image):
         return self.transform(img.convert("RGB")).unsqueeze(0).to(self.device)
 
+    # ---------- main entry ----------
     def __call__(self, data):
-        # يدعم: Widget (PIL) أو REST (base64)
+        """
+        يدعم:
+        • Widget: يستلم PIL.Image
+        • REST API: يستلم base64 فى data["inputs"] أو data["image"]
+        """
         img = None
         if isinstance(data, Image.Image):
             img = data
         elif isinstance(data, dict):
-            b = data.get("inputs") or data.get("image")
-            if isinstance(b, (str, bytes)):
-                b = b.encode() if isinstance(b, str) else b
-                img = Image.open(io.BytesIO(base64.b64decode(b)))
+            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)))
 
         if img is None:
             return {"error": "No image provided"}