Update handler.py

handler.py

@@ -3,52 +3,31 @@ 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
 
 
 class EndpointHandler:
-    """Custom pipeline for Hugging Face Inference Endpoints."""
+    """Custom pipeline for HF Inference Endpoints."""
 
     def __init__(self, model_dir: str):
-        # اختَر GPU إذا متاح
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-        # تحميل الوزن بصيغة safetensors
-        weights_path = os.path.join(model_dir, "model.safetensors")
-        weights = load_file(weights_path)
-
-        # إنشاء نموذج ViT مطابق لِما درّبتَه
+        weights = load_file(os.path.join(model_dir, "model.safetensors"))
         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.Image
-        • REST API: يستلم base64 فى data["inputs"] أو data["image"]
-        """
         img = None
         if isinstance(data, Image.Image):
             img = data
@@ -67,4 +46,3 @@ class EndpointHandler:
             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: correct handler indentation