Update handler.py

handler.py

@@ -17,35 +17,50 @@ class EndpointHandler:
     # --------------------------------------------------
     def __init__(self, model_dir: str) -> None:
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-        # ─── البحث عن ملف النموذج المتاح ─── 
-        # أولاً نبحث عن ملف safetensors
-        safetensors_path = os.path.join(model_dir, "model.safetensors")
+        
+        # تحديد مسارات الملفات المحتملة
         pytorch_path = os.path.join(model_dir, "pytorch_model.bin")
+        safetensors_path = os.path.join(model_dir, "model.safetensors")
         
-        # التحقق من وجود الملفات وتحميل الملف المناسب
-        if os.path.exists(safetensors_path):
-            print(f"Loading model from safetensors: {safetensors_path}")
-            state_dict = load_file(safetensors_path, device="cpu")
-        elif os.path.exists(pytorch_path):
-            print(f"Loading model from PyTorch bin: {pytorch_path}")
-            state_dict = torch.load(pytorch_path, map_location="cpu")
-        else:
-            raise FileNotFoundError(f"No model file found at {safetensors_path} or {pytorch_path}")
-
         # إنشاء ViT Base Patch-16 بعدد فئات 5
         self.model = create_model("vit_base_patch16_224", num_classes=5)
-        self.model.load_state_dict(state_dict)
+        
+        # محاولة تحميل النموذج من 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")
+                self.model.load_state_dict(state_dict)
+                print("تم تحميل النموذج بنجاح من pytorch_model.bin")
+                model_loaded = True
+            except Exception as e:
+                print(f"خطأ في تحميل pytorch_model.bin: {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)
-
+        
         # محوّلات التحضير
-        self.preprocess = 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",