Update convert.py

convert.py

@@ -19,12 +19,24 @@ COMMIT_MESSAGE = " This PR adds weights in safetensors format"
 device = "cpu"
 
 def convert_pt_to_safetensors(model_path, safe_path):
-    model = torch.load(model_path, map_location="cpu", weights_only=False)
+    # Use weights_only=True for security unless you need to execute code from the pickle
+    # For model conversion, it's often safer. If it fails, you can revert to False.
+    try:
+        model = torch.load(model_path, map_location="cpu", weights_only=True)
+    except Exception:
+        # Fallback for models that can't be loaded with weights_only
+        model = torch.load(model_path, map_location="cpu")
+        
+    # If the loaded object is a state_dict, use it directly.
+    # If it's a full model object, get its state_dict.
+    if not isinstance(model, dict):
+        model = model.state_dict()
+        
     metadata = {"format":"pt"}
     save_file(model, safe_path, metadata)
 
 def convert_single(model_id: str, filename: str, folder: str, progress: Any, token: str):
-    progress(0, desc="Downloading model")
+    progress(0, desc=f"Downloading {filename}")
 
     local_file = os.path.join(model_id, filename)
     local_dir = os.path.dirname(local_file)
@@ -60,14 +72,34 @@ def convert(token: str, model_id: str, filename: str, your_model_id: str, progre
         new_pr = None
         try:
             converted_model = convert_single(model_id, filename, folder, progress, token)
-            progress(0.7, desc="Uploading to Hub")
-            new_pr  = api.upload_file(path_or_fileobj=converted_model, path_in_repo=filename, repo_id=your_model_id, repo_type="model", token=token, commit_message=pr_title, commit_description=COMMIT_MESSAGE.format(your_model_id), create_pr=True)
+
+            # Create the correct filename for the repository (e.g., "model.safetensors")
+            # os.path.splitext('path/to/model.pth') -> ('path/to/model', '.pth')
+            # We take the first part and add the new extension.
+            safetensors_filename = os.path.splitext(filename)[0] + ".safetensors"
+            
+            progress(0.7, desc=f"Uploading {safetensors_filename} to Hub")
+
+            # Use the new, correct filename in `path_in_repo`
+            new_pr  = api.upload_file(
+                path_or_fileobj=converted_model,
+                path_in_repo=safetensors_filename, # <-- CORRECTED
+                repo_id=your_model_id,
+                repo_type="model",
+                token=token,
+                commit_message=pr_title,
+                commit_description=COMMIT_MESSAGE.format(your_model_id),
+                create_pr=True
+            )
+            
             pr_number = new_pr.split("%2F")[-1].split("/")[0]
             link = f"Pr created at: {'https://huggingface.co/' + os.path.join(your_model_id, 'discussions', pr_number)}"
             progress(1, desc="Done")
         except Exception as e:
             raise gr.exceptions.Error(str(e))
         finally:
-            shutil.rmtree(folder)
+            # This check prevents an error if the folder was already removed or failed to create
+            if os.path.exists(folder):
+                shutil.rmtree(folder)
 
         return link