Update bin2safetensors/convert.py (#2)

- Update bin2safetensors/convert.py (62bc6924b6d614a9e1b4ad5f3f1dd90b06573eab)


Co-authored-by: None <[email protected]>

bin2safetensors/convert.py

@@ -312,7 +312,7 @@ def convert(api: "HfApi", model_id: str, force: bool = False) -> Tuple["CommitIn
         return new_pr, errors
 
 
-def main(input_directory, output_directory):
+def main(input_directory, output_directory, delete_files, delete_input_directory):
     # Get a list of all files in the input directory
     files = os.listdir(input_directory)
     
@@ -360,11 +360,29 @@ def main(input_directory, output_directory):
         output_filename = os.path.join(output_directory, f"model-{i:05d}-of-{yyyyy:05d}.safetensors")
         convert_file(input_filename, output_filename)
         print(f"Converted {input_filename} to {output_filename}")
+        
+        # Delete the pytorch_model file if the delete_files flag or delete_input_directory flag are set
+        if delete_files or delete_input_directory:
+            os.remove(input_filename)
+            print(f"Deleted {input_filename}")
+    
+    # Check if there are any remaining pytorch_model files in the input directory
+    remaining_model_files = [file for file in os.listdir(input_directory) if re.match(r'pytorch_model-\d{5}-of-\d{5}\.bin', file)]
+    
+    if len(remaining_model_files) == 0:
+        # Delete the input directory if all files have been converted successfully
+        if delete_input_directory:
+            shutil.rmtree(input_directory)
+            print(f"Deleted input directory {input_directory}")
+    else:
+        print("Warning: Input directory still contains pytorch_model files and won't be deleted.")
     
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Convert pytorch_model model to safetensor and copy JSON and .model files.")
     parser.add_argument("input_directory", help="Path to the input directory containing pytorch_model files")
     parser.add_argument("output_directory", help="Path to the output directory for converted safetensor files")
+    parser.add_argument("-d", "--delete", action="store_true", help="Delete pytorch_model files after conversion")
+    parser.add_argument("-D", "--delete-input", action="store_true", help="Delete pytorch_model files after conversion as well as the input directory after all files are converted")
     args = parser.parse_args()
     
-    main(args.input_directory, args.output_directory)
+    main(args.input_directory, args.output_directory, args.delete, args.delete_input)