Update handler.py

handler.py

@@ -11,39 +11,30 @@ class EndpointHandler:
         print(f"بدء تهيئة النموذج من المسار: {model_dir}")
         print(f"قائمة الملفات في المسار: {os.listdir(model_dir)}")
 
-        # تحديد الجهاز المستخدم (CPU في معظم بيئات Edge Function)
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         print(f"استخدام الجهاز: {self.device}")
 
         try:
-            # استخدام pipeline مباشرة من Hugging Face مع تحديد خيارات تحسين الذاكرة
             print("تحميل النموذج باستخدام pipeline")
-            
-            # تحميل النموذج مباشرة من Hugging Face مع تفعيل خيارات تحسين الذاكرة
+
             self.classifier = pipeline(
                 task="image-classification",
                 model="yaya36095/ai-source-detector",
-                device=self.device,
-                torch_dtype=torch.float16,  # استخدام دقة أقل لتوفير الذاكرة
-                low_cpu_mem_usage=True      # تقليل استخدام ذاكرة CPU
+                device=0 if torch.cuda.is_available() else -1,
+                torch_dtype=torch.float16  # يُفضل فقط لو الجهاز يدعم
             )
-            
+
             print("تم تحميل النموذج بنجاح")
-            
+            self.fallback_mode = False
+
         except Exception as e:
             print(f"خطأ أثناء تهيئة النموذج: {e}")
-            # محاولة بديلة باستخدام تكوين مخصص إذا فشلت الطريقة الأولى
             try:
                 print("محاولة تحميل النموذج بطريقة بديلة...")
-                
-                # تحميل التكوين فقط (ملف صغير) بدلاً من النموذج الكامل
                 config = AutoConfig.from_pretrained("yaya36095/ai-source-detector")
-                
-                # إنشاء وظيفة محاكاة بسيطة للتصنيف
                 self.fallback_mode = True
                 self.config = config
                 print("تم التحويل إلى وضع المحاكاة البسيطة")
-                
             except Exception as e2:
                 print(f"فشلت المحاولة البديلة أيضًا: {e2}")
                 raise
@@ -61,7 +52,6 @@ class EndpointHandler:
 
         img = None
         try:
-            # استخراج الصورة من البيانات المدخلة
             if isinstance(data, Image.Image):
                 img = data
             elif isinstance(data, dict):
@@ -75,13 +65,8 @@ class EndpointHandler:
                 print("لم يتم العثور على صورة صالحة")
                 return [{"label": "error", "score": 1.0}]
 
-            # التحقق من وجود وضع المحاكاة البسيطة
-            if hasattr(self, 'fallback_mode') and self.fallback_mode:
+            if self.fallback_mode:
                 print("استخدام وضع المحاكاة البسيطة")
-                # تحليل بسيط للصورة واستخدام قيم افتراضية
-                # يمكن تحسين هذا الجزء بإضافة تحليل بسيط للصورة
-                
-                # استخدام قيم افتراضية متوازنة
                 results = [
                     {"label": "real", "score": 0.5},
                     {"label": "stable_diffusion", "score": 0.2},
@@ -89,26 +74,18 @@ class EndpointHandler:
                     {"label": "dalle", "score": 0.1},
                     {"label": "other_ai", "score": 0.05}
                 ]
-                
-                # ترتيب النتائج تنازليًا حسب النتيجة
                 results.sort(key=lambda x: x["score"], reverse=True)
-                best = results[0]
-                print(f"أفضل نتيجة (محاكاة): {best}")
-                return [best]
-            
-            # استخدام النموذج الكامل إذا كان متاحًا
-            print("تصنيف الصورة باستخدام النموذج")
+                return [results[0]]
+
+            print("تصنيف الصورة باستخدام النموذج الحقيقي")
             results = self.classifier(img)
-            
+
             if isinstance(results, list) and len(results) > 0:
-                best = results[0]
-                print(f"أفضل نتيجة: {best}")
-                return [best]
+                return [results[0]]
             else:
                 print("لم يتم الحصول على نتائج صالحة من النموذج")
                 return [{"label": "error", "score": 1.0}]
 
         except Exception as e:
             print(f"حدث استثناء: {e}")
-            # في حالة حدوث خطأ، نعود بنتيجة محايدة بدلاً من خطأ
             return [{"label": "real", "score": 0.5}]