Update app.py

app.py

@@ -28,7 +28,7 @@ class Settings:
     AGGREGATOR_URL = os.getenv("AGGREGATOR_URL", "http://192.168.1.104:8002")
     
     # Model settings
-    MODEL_REPO = "https://huggingface.co/microsoft/florence-2-large"
+    MODEL_REPO = "https://huggingface.co/facebook/opt-125m"
     
     # Server settings
     TENSOR_SERVER_TIMEOUT = 30  # seconds
@@ -144,15 +144,36 @@ async def split_model_weights():
         import torch
         import math
         
+        # Install required packages if not present
+        try:
+            import safetensors
+        except ImportError:
+            print("[INFO] Installing required packages...")
+            import subprocess
+            subprocess.check_call(["pip", "install", "safetensors", "packaging"])
+        
         # Load the full model weights
         import torch
         from safetensors.torch import load_file as load_safetensors
 
-        # Try safetensors first, then fallback to pytorch
+        # Try safetensors first with chunked loading, then fallback to pytorch
         try:
             model_file = next(f for f in state.model_files.values() if f.endswith('.safetensors'))
             print(f"[INFO] Loading weights from safetensors file: {model_file}")
-            weights = load_safetensors(model_file)
+            try:
+                # Try direct loading first
+                weights = load_safetensors(model_file)
+            except Exception as e:
+                if "header too large" in str(e):
+                    print("[INFO] Large header detected, attempting chunked loading...")
+                    from safetensors import safe_open
+                    weights = {}
+                    with safe_open(model_file, framework="pt") as f:
+                        for key in f.keys():
+                            weights[key] = f.get_tensor(key)
+                    print("[INFO] Successfully loaded weights using chunked loading")
+                else:
+                    raise e
         except StopIteration:
             # No safetensors file found, try pytorch
             model_file = next(f for f in state.model_files.values() if f.endswith('.bin'))