Update app.py

app.py

@@ -268,10 +268,10 @@ def find_sharded_safetensors_files(repo_id, subfolder=None, hf_token=None, max_s
         
         # Filter for safetensors files in the subfolder
         if subfolder:
-            pattern = f"{subfolder}/"
+            pattern = f"{subfolder}/" if not subfolder.endswith("/") else subfolder
             safetensors_files = [f for f in repo_files if f.endswith('.safetensors') and f.startswith(pattern)]
             # Remove subfolder prefix
-            safetensors_files = [f[len(pattern):] for f in safetensors_files]
+            safetensors_files = [f[len(pattern):] for f in safetensors_files if len(f) > len(pattern)]
         else:
             safetensors_files = [f for f in repo_files if f.endswith('.safetensors')]
         
@@ -280,21 +280,28 @@ def find_sharded_safetensors_files(repo_id, subfolder=None, hf_token=None, max_s
         single_files = []
         
         for f in safetensors_files:
-            if "-of-" in f:
+            # Check for sharding pattern: model-XXXXX-of-YYYYY.safetensors
+            match = re.search(r'-\d{5}-of-\d{5}\.safetensors$', f)
+            if match:
                 sharded_files.append(f)
             else:
                 single_files.append(f)
         
-        # Return sharded files if found, otherwise single files
+        # If we have sharded files, return them sorted by shard number
         if sharded_files:
             # Sort by shard number for consistent ordering
-            sharded_files.sort(key=lambda x: int(re.search(r'-(\d+)-of-', x).group(1)))
+            def extract_shard_num(filename):
+                match = re.search(r'-(\d{5})-of-\d{5}\.safetensors$', filename)
+                return int(match.group(1)) if match else 0
+            
+            sharded_files.sort(key=extract_shard_num)
             # Limit number of shards to prevent accidental downloads of huge models
             if len(sharded_files) > max_shards:
                 raise ValueError(f"Too many shards found ({len(sharded_files)}). Maximum allowed is {max_shards}. "
                                f"Please specify a more specific pattern.")
             return sharded_files
         elif single_files:
+            # Return single files (non-sharded)
             return single_files
         else:
             return []
@@ -372,7 +379,7 @@ def download_model_files(source_type, repo_url, filename_pattern, model_format,
                     
                     return model_files, temp_dir
                 else:
-                    # Single file
+                    # SINGLE FILE SAFETENSORS - separate from shard discovery
                     progress(0.2, desc=f"Downloading {filename_pattern}...")
                     model_path = hf_hub_download(
                         repo_id=repo_id,