Update app.py

app.py

@@ -148,16 +148,21 @@ async def split_model_weights():
         import shutil
         from pathlib import Path
         
-        # Find model file (safetensors or pytorch)
-        try:
-            model_file = next(f for f in state.model_files.values() if f.endswith('.safetensors'))
-            print(f"[INFO] Found safetensors file: {model_file}")
-        except StopIteration:
-            try:
-                model_file = next(f for f in state.model_files.values() if f.endswith('.bin'))
-                print(f"[INFO] Found PyTorch file: {model_file}")
-            except StopIteration:
-                raise Exception("No model weight files found")
+        # Find model weight files (safetensors or pytorch)
+        weight_files = [f for f in state.model_files.values() if f.endswith(('.safetensors', '.bin'))]
+        
+        if not weight_files:
+            raise Exception("No model weight files found")
+            
+        # The current splitting logic only supports splitting a single file.
+        # If there are multiple files, we assume they are sharded and need a different approach.
+        # For now, we will select the largest file to split, or the first one if all are small.
+        model_file = max(weight_files, key=os.path.getsize) if len(weight_files) > 1 else weight_files[0]
+        
+        if len(weight_files) > 1:
+            print(f"[WARN] Found multiple weight files. Selecting the largest one for splitting: {model_file}")
+        else:
+            print(f"[INFO] Found model weight file: {model_file}")
 
         # Get file size and calculate chunks
         try:
@@ -690,35 +695,18 @@ async def download_model_files():
             try:
                 print(f"[INFO] Downloading model files from {repo_id}...")
                 
-                # First download config and other small files
-                config_files = ["config.json", "tokenizer_config.json", "vocab.json", "generation_config.json"]
-                for filename in config_files:
-                    try:
-                        file_path = hf_hub_download(
-                            repo_id=repo_id,
-                            filename=filename,
-                            local_dir=model_path,
-                            force_download=True
-                        )
-                        print(f"[INFO] Downloaded {filename}")
-                    except Exception as e:
-                        print(f"[WARN] Could not download {filename}: {str(e)}")
+                # Use snapshot_download to get all necessary files at once, which supports all weight file names
+                print("[INFO] Downloading all model files (this may take a while)...")
                 
-                # Then download the model weights
-                print("[INFO] Downloading model weights (this may take a while)...")
-                for weight_file in ["pytorch_model.bin", "model.safetensors"]:
-                    try:
-                        file_path = hf_hub_download(
-                            repo_id=repo_id,
-                            filename=weight_file,
-                            local_dir=model_path,
-                            force_download=True
-                        )
-                        print(f"[INFO] Successfully downloaded {weight_file}")
-                        break  # Stop after first successful weight file download
-                    except Exception as e:
-                        print(f"[WARN] Could not download {weight_file}: {str(e)}")
-                        continue
+                # snapshot_download is the most robust way to get all files matching patterns
+                # This addresses the user's request to download model files that are not just "pytorch.bin"
+                model_path = snapshot_download(
+                    repo_id=repo_id,
+                    local_dir=model_path,
+                    allow_patterns=["*.bin", "*.safetensors", "*.json", "*.txt", "tokenizer.model"],
+                    ignore_patterns=["*.msgpack", "*.onnx"], # Ignore non-PyTorch/safetensors formats
+                    force_download=True
+                )
                 
                 print(f"[INFO] All files downloaded to {model_path}")
                 state.is_model_loaded = True