DivEye - PR (fix pickling)

app.py

@@ -51,7 +51,7 @@ def detect_ai_text(text):
     return message, round(ai_prob, 3), bar_data
 
 # Gradio app setup
-with gr.Blocks(title="DivEye", theme=theme) as demo:
+with gr.Blocks(title="DivEye") as demo:
     gr.HTML("""
     <div style="display: flex; justify-content: space-between; align-items: center; padding: 1.5rem; background: #f0f4f8; border-radius: 12px; margin-bottom: 1rem;">
         <div style="text-align: left; max-width: 70%;">

software.py

@@ -99,22 +99,42 @@ class BiScope:
 class Software:
     def __init__(self):
         self.token = os.getenv("HF_TOKEN") 
+        self.device_div = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        self.device_bi = self.device_div
 
-        self.div_tokenizer = AutoTokenizer.from_pretrained("tiiuae/falcon-7b", use_fast=False, trust_remote_code=True, use_auth_token=self.token)
-        self.div_model = AutoModelForCausalLM.from_pretrained(
-            "tiiuae/falcon-7b", torch_dtype=torch.float16, trust_remote_code=True, use_auth_token=self.token
-        )
-
-        self.bi_tokenizer = AutoTokenizer.from_pretrained("google/gemma-1.1-2b-it", use_fast=False, trust_remote_code=True, use_auth_token=self.token)
-        self.bi_model = AutoModelForCausalLM.from_pretrained(
-            "google/gemma-1.1-2b-it", torch_dtype=torch.float16, trust_remote_code=True, use_auth_token=self.token
-        )
-        
+        self.div_model = None
+        self.div_tokenizer = None
+        self.bi_model = None
+        self.bi_tokenizer = None        
         
         self.model_path = Path(__file__).parent / "model.json"
 
         self.model = xgb.XGBClassifier()
         self.model.load_model(self.model_path)
+
+    def _load_div_models(self):
+        if self.div_model is None or self.div_tokenizer is None:
+            self.div_tokenizer = AutoTokenizer.from_pretrained("tiiuae/falcon-7b", use_fast=False, trust_remote_code=True, use_auth_token=self.token)
+            self.div_model = AutoModelForCausalLM.from_pretrained(
+                "tiiuae/falcon-7b",
+                device_map="cuda",
+                torch_dtype=torch.float16,
+                trust_remote_code=True,
+                use_auth_token=self.token
+            )
+            self.div_model.to(self.device_div)
+
+    def _load_bi_models(self):
+        if self.bi_model is None or self.bi_tokenizer is None:
+            self.bi_tokenizer = AutoTokenizer.from_pretrained("google/gemma-1.1-2b-it", use_fast=False, trust_remote_code=True, use_auth_token=self.token)
+            self.bi_model = AutoModelForCausalLM.from_pretrained(
+                "google/gemma-1.1-2b-it",
+                device_map="cuda",
+                torch_dtype=torch.float16,
+                trust_remote_code=True,
+                use_auth_token=self.token
+            )
+            self.bi_model.to(self.device_bi)
         
     def load_data(self, jsonl_path):
         ids, texts = [], []
@@ -127,6 +147,9 @@ class Software:
     
     @spaces.GPU
     def evaluate(self, text):
+        self._load_div_models()
+        self._load_bi_models()
+        
         # Load models to GPUs.
         device_div = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
         if torch.cuda.device_count() > 1: