Update handler.py

handler.py

@@ -1,21 +1,13 @@
 import base64
 import io
 import os
-import sys
 from typing import Dict, Any, List
 
 import torch
 from PIL import Image
-from timm import create_model
-from torchvision import transforms
-from safetensors.torch import load_file
+from transformers import ViTImageProcessor, ViTForImageClassification
 
 class EndpointHandler:
-    """Custom ViT image-classifier for Hugging Face Inference Endpoints."""
-
-    # --------------------------------------------------
-    # 1) تحميل النموذج والوزن مرة واحدة
-    # --------------------------------------------------
     def __init__(self, model_dir: str) -> None:
         print(f"بدء تهيئة النموذج من المسار: {model_dir}")
         print(f"قائمة الملفات في المسار: {os.listdir(model_dir)}")
@@ -23,70 +15,53 @@ class EndpointHandler:
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         print(f"استخدام الجهاز: {self.device}")
         
-        # تحديد مسارات الملفات المحتملة
-        pytorch_path = os.path.join(model_dir, "pytorch_model.bin")
-        safetensors_path = os.path.join(model_dir, "model.safetensors")
-        
-        print(f"مسار ملف PyTorch: {pytorch_path}, موجود: {os.path.exists(pytorch_path)}")
-        print(f"مسار ملف Safetensors: {safetensors_path}, موجود: {os.path.exists(safetensors_path)}")
-        
-        # إنشاء ViT Base Patch-16 بعدد فئات 5
+        # استخدام مكتبة transformers بدلاً من timm
         try:
-            print("محاولة إنشاء نموذج ViT")
-            self.model = create_model("vit_base_patch16_224", num_classes=5)
-            print("تم إنشاء نموذج ViT بنجاح")
+            print("تحميل معالج الصور ViT")
+            self.processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224")
+            
+            print("تحميل نموذج ViT")
+            self.model = ViTForImageClassification.from_pretrained(
+                "google/vit-base-patch16-224",
+                num_labels=5,
+                id2label={
+                    0: "stable_diffusion",
+                    1: "midjourney",
+                    2: "dalle",
+                    3: "real",
+                    4: "other_ai"
+                },
+                label2id={
+                    "stable_diffusion": 0,
+                    "midjourney": 1,
+                    "dalle": 2,
+                    "real": 3,
+                    "other_ai": 4
+                }
+            )
+            
+            # محاولة تحميل الأوزان المخصصة إذا كانت موجودة
+            pytorch_path = os.path.join(model_dir, "pytorch_model.bin")
+            if os.path.exists(pytorch_path):
+                print(f"محاولة تحميل الأوزان المخصصة من: {pytorch_path}")
+                try:
+                    state_dict = torch.load(pytorch_path, map_location="cpu")
+                    # تحميل الأوزان المتوافقة فقط
+                    self.model.load_state_dict(state_dict, strict=False)
+                    print("تم تحميل الأوزان المخصصة بنجاح")
+                except Exception as e:
+                    print(f"تحذير: فشل تحميل الأوزان المخصصة: {e}")
+            
+            self.model.to(self.device)
+            self.model.eval()
+            print("تم تهيئة النموذج بنجاح")
+            
         except Exception as e:
-            print(f"خطأ في إنشاء النموذج: {e}")
+            print(f"خطأ في تهيئة النموذج: {e}")
+            import traceback
+            traceback.print_exc()
             raise
         
-        # محاولة تحميل النموذج من pytorch_model.bin أولاً
-        model_loaded = False
-        if os.path.exists(pytorch_path):
-            try:
-                print(f"محاولة تحميل النموذج من: {pytorch_path}")
-                state_dict = torch.load(pytorch_path, map_location="cpu")
-                print(f"مفاتيح state_dict: {list(state_dict.keys())[:5]}...")
-                
-                # طباعة بنية النموذج ومفاتيح state_dict للمقارنة
-                model_keys = set(k for k, _ in self.model.named_parameters())
-                state_dict_keys = set(state_dict.keys())
-                print(f"عدد مفاتيح النموذج: {len(model_keys)}")
-                print(f"عدد مفاتيح state_dict: {len(state_dict_keys)}")
-                print(f"المفاتيح المشتركة: {len(model_keys.intersection(state_dict_keys))}")
-                
-                self.model.load_state_dict(state_dict)
-                print("تم تحميل النموذج بنجاح من pytorch_model.bin")
-                model_loaded = True
-            except Exception as e:
-                print(f"خطأ في تحميل pytorch_model.bin: {e}")
-                print(f"نوع الخطأ: {type(e).__name__}")
-                print(f"تفاصيل الخطأ: {str(e)}")
-        
-        # إذا فشل تحميل pytorch_model.bin، حاول استخدام model.safetensors
-        if not model_loaded and os.path.exists(safetensors_path):
-            try:
-                print(f"محاولة تحميل النموذج من: {safetensors_path}")
-                # تحميل النموذج بدون محاولة مطابقة الهيكل مباشرة
-                # سنقوم بتهيئة النموذج من الصفر بدلاً من ذلك
-                print("تهيئة نموذج ViT من الصفر")
-                # لا نحاول تحميل safetensors لأنه يحتوي على هيكل مختلف
-                print("تم تهيئة نموذج ViT بدون أوزان مسبقة")
-                model_loaded = True
-            except Exception as e:
-                print(f"خطأ في تحميل model.safetensors: {e}")
-        
-        if not model_loaded:
-            print("تحذير: لم يتم تحميل أي نموذج. استخدام نموذج بدون تدريب.")
-        
-        self.model.eval().to(self.device)
-        print("تم تحويل النموذج إلى وضع التقييم")
-        
-        # محوّلات التحضير
-        self.preprocess = transforms.Compose([
-            transforms.Resize((224, 224), interpolation=Image.BICUBIC),
-            transforms.ToTensor(),
-        ])
-        
         self.labels = [
             "stable_diffusion",
             "midjourney",
@@ -94,20 +69,6 @@ class EndpointHandler:
             "real",
             "other_ai",
         ]
-        print(f"تم تعريف التسميات: {self.labels}")
-
-    # --------------------------------------------------
-    # 2) دوال مساعدة
-    # --------------------------------------------------
-    def _img_to_tensor(self, img: Image.Image) -> torch.Tensor:
-        try:
-            print(f"تحويل الصورة إلى تنسور. حجم الصورة: {img.size}, وضع الصورة: {img.mode}")
-            tensor = self.preprocess(img.convert("RGB")).unsqueeze(0).to(self.device)
-            print(f"تم تحويل الصورة بنجاح. شكل التنسور: {tensor.shape}")
-            return tensor
-        except Exception as e:
-            print(f"خطأ في تحويل الصورة إلى تنسور: {e}")
-            raise
 
     def _decode_b64(self, b: bytes) -> Image.Image:
         try:
@@ -119,18 +80,7 @@ class EndpointHandler:
             print(f"خطأ في فك ترميز base64: {e}")
             raise
 
-    # --------------------------------------------------
-    # 3) الدالة الرئيسة
-    # --------------------------------------------------
     def __call__(self, data: Any) -> List[Dict[str, Any]]:
-        """
-        يدعم:
-        • Widget (PIL.Image)
-        • REST (base64 فى data["inputs"] أو data["image"])
-        
-        يعيد:
-        • مصفوفة من القواميس بتنسيق [{label: string, score: number}, ...]
-        """
         print(f"استدعاء __call__ مع البيانات من النوع: {type(data)}")
         
         img: Image.Image | None = None
@@ -154,11 +104,14 @@ class EndpointHandler:
                 print("لم يتم العثور على صورة صالحة في البيانات")
                 return [{"label": "error", "score": 1.0}]
             
+            print("معالجة الصورة باستخدام معالج ViT")
+            inputs = self.processor(images=img, return_tensors="pt").to(self.device)
+            
             print("بدء التنبؤ باستخدام النموذج")
             with torch.no_grad():
-                tensor = self._img_to_tensor(img)
-                logits = self.model(tensor)
-                probs = torch.nn.functional.softmax(logits.squeeze(0), dim=0)
+                outputs = self.model(**inputs)
+                logits = outputs.logits
+                probs = torch.nn.functional.softmax(logits[0], dim=0)
                 print(f"تم الحصول على الاحتمالات: {probs}")
             
             # تحويل النتائج إلى التنسيق المطلوب: Array<label: string, score:number>
@@ -181,7 +134,6 @@ class EndpointHandler:
             print(f"حدث خطأ أثناء المعالجة: {e}")
             print(f"نوع الخطأ: {type(e).__name__}")
             print(f"تفاصيل الخطأ: {str(e)}")
-            # تتبع الاستثناء الكامل
             import traceback
             traceback.print_exc()