Update convert.py

convert.py

@@ -19,24 +19,12 @@ COMMIT_MESSAGE = " This PR adds weights in safetensors format"
 device = "cpu"
 
 def convert_pt_to_safetensors(model_path, safe_path):
-    # 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()
-        
+    model = torch.load(model_path, map_location="cpu", weights_only=False)
     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=f"Downloading {filename}")
+    progress(0, desc="Downloading model")
 
     local_file = os.path.join(model_id, filename)
     local_dir = os.path.dirname(local_file)
@@ -71,34 +59,37 @@ 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)
-
-            # 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")
+            # --- START OF THE REQUESTED CHANGE ---
+            # Create the correct destination filename by replacing the extension
+            safetensors_filename = os.path.splitext(filename)[0] + ".safetensors"
+            # --- END OF THE REQUESTED CHANGE ---
 
-            # Use the new, correct filename in `path_in_repo`
+            progress(0.7, desc="Uploading to Hub")
+            
+            # Use the corrected 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),
+                path_or_fileobj=converted_model, 
+                path_in_repo=safetensors_filename, # <-- THE ONLY MODIFIED ARGUMENT
+                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:
-            # This check prevents an error if the folder was already removed or failed to create
-            if os.path.exists(folder):
-                shutil.rmtree(folder)
+            shutil.rmtree(folder)
+
+        finish_message(link)
+        return link
 
-        return link
+@spaces.GPU()
+def finish_message(link: str):
+    return print(link)