feat: add EndpointHandler

handler.py

@@ -0,0 +1,47 @@
+import base64, io, os
+import torch
+from torchvision import transforms
+from PIL import Image
+from safetensors.torch import load_file
+from timm import create_model   # timm ضرورى للتعامل مع ViT
+
+
+class EndpointHandler:      # اسم الفئة مهم جداً
+    def __init__(self, model_dir: str):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        # تحميل الوزن بصيغة safetensors
+        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.labels = ['stable_diffusion', 'midjourney', 'dalle', 'real', 'other_ai']
+
+    def _prep(self, img: Image.Image):
+        return self.transform(img.convert("RGB")).unsqueeze(0).to(self.device)
+
+    def __call__(self, data):
+        # يدعم: Widget (PIL) أو REST (base64)
+        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)))
+
+        if img is None:
+            return {"error": "No image provided"}
+
+        with torch.no_grad():
+            logits = self.model(self._prep(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))}