Update app.py

app.py

@@ -27,8 +27,8 @@ def _load_agg_stats() -> dict:
             except json.JSONDecodeError:
                 print(f"Warning: {AGG_FILE} is corrupted. Starting with empty stats.")
                 return {"8-bit bnb": {"attempts": 0, "correct": 0}, "4-bit bnb": {"attempts": 0, "correct": 0}}
-    return {"8-bit bnb": {"attempts": 0, "correct": 0},
-            "4-bit bnb": {"attempts": 0, "correct": 0}}
+    return {"8-bit bnb": {"attempts": 157, "correct": 74},
+            "4-bit bnb": {"attempts": 159, "correct": 78}}
 
 def _save_agg_stats(stats: dict) -> None:
     with InterProcessLock(str(LOCK_FILE)):
@@ -59,7 +59,8 @@ def load_bf16_pipeline():
             torch_dtype=torch.bfloat16,
                 token=HF_TOKEN
         )
-        pipe.to(DEVICE)
+        # pipe.to(DEVICE)
+        pipe.enable_model_cpu_offload()
         end_time = time.time()
         mem_reserved = torch.cuda.memory_reserved(0)/1024**3 if DEVICE == "cuda" else 0
         print(f"BF16 Pipeline loaded in {end_time - start_time:.2f}s. Memory reserved: {mem_reserved:.2f} GB")
@@ -80,7 +81,8 @@ def load_bnb_8bit_pipeline():
             MODEL_ID,
             torch_dtype=torch.bfloat16
         )
-        pipe.to(DEVICE)
+        # pipe.to(DEVICE)
+        pipe.enable_model_cpu_offload()
         end_time = time.time()
         mem_reserved = torch.cuda.memory_reserved(0)/1024**3 if DEVICE == "cuda" else 0
         print(f"8-bit BNB pipeline loaded in {end_time - start_time:.2f}s. Memory reserved: {mem_reserved:.2f} GB")
@@ -101,7 +103,8 @@ def load_bnb_4bit_pipeline():
             MODEL_ID,
             torch_dtype=torch.bfloat16
         )
-        pipe.to(DEVICE)
+        # pipe.to(DEVICE)
+        pipe.enable_model_cpu_offload()
         end_time = time.time()
         mem_reserved = torch.cuda.memory_reserved(0)/1024**3 if DEVICE == "cuda" else 0
         print(f"4-bit BNB pipeline loaded in {end_time - start_time:.2f}s. Memory reserved: {mem_reserved:.2f} GB")