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
@@ -145,17 +145,19 @@ async def split_model_weights():
         import math
         
         # Load the full model weights
-        model_file = next(f for f in state.model_files.values() if f.endswith('.safetensors') or f.endswith('.bin'))
-        loaded_weights = torch.load(model_file, weights_only=False)  # Explicitly allow non-weights for compatibility
-        
-        # For OPT models, weights are stored under model.decoder
-        if isinstance(loaded_weights, dict):
-            if 'model.decoder' in str(list(loaded_weights.keys())):
-                # Get all weights that are part of the decoder
-                weights = {k: v for k, v in loaded_weights.items() if k.startswith('model.decoder')}
-            else:
-                # Just use all weights if no decoder prefix found
-                weights = loaded_weights
+        import torch
+        from safetensors.torch import load_file as load_safetensors
+
+        # Try safetensors first, 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)
+        except StopIteration:
+            # No safetensors file found, try pytorch
+            model_file = next(f for f in state.model_files.values() if f.endswith('.bin'))
+            print(f"[INFO] Loading weights from PyTorch file: {model_file}")
+            weights = torch.load(model_file, map_location='cpu')
         
         # Calculate total model size and chunks
         total_size_bytes = sum(p.nelement() * p.element_size() for p in weights.values())