Upload 2 files

convert_repo_to_safetensors_gr.py

@@ -269,7 +269,7 @@ def convert_openai_text_enc_state_dict(text_enc_dict):
     return text_enc_dict
 
 
-def convert_diffusers_to_safetensors(model_path, checkpoint_path, half = True, use_safetensors = True, progress=gr.Progress(track_tqdm=True)):
+def convert_diffusers_to_safetensors(model_path, checkpoint_path, half = True, progress=gr.Progress(track_tqdm=True)):
     progress(0, desc="Start converting...")
     # Path for safetensors
     unet_path = osp.join(model_path, "unet", "diffusion_pytorch_model.safetensors")
@@ -329,12 +329,7 @@ def convert_diffusers_to_safetensors(model_path, checkpoint_path, half = True, u
     if half:
         state_dict = {k: v.half() for k, v in state_dict.items()}
 
-    if use_safetensors:
-        save_file(state_dict, checkpoint_path)
-    else:
-        state_dict = {"state_dict": state_dict}
-        torch.save(state_dict, checkpoint_path)
-    
+    save_file(state_dict, checkpoint_path)
     progress(1, desc="Converted.")
 
 
@@ -343,7 +338,7 @@ def download_repo(repo_id, dir_path, progress=gr.Progress(track_tqdm=True)):
     try:
         snapshot_download(repo_id=repo_id, local_dir=dir_path)
     except Exception as e:
-        print(f"Error: Failed to download {repo_id}. ")
+        print(f"Error: Failed to download {repo_id}. {e}")
         return
 
 
@@ -366,7 +361,7 @@ def upload_safetensors_to_repo(filename, progress=gr.Progress(track_tqdm=True)):
     return url
 
 
-def convert_repo_to_safetensors(repo_id, progress=gr.Progress(track_tqdm=True)):
+def convert_repo_to_safetensors(repo_id, half=True, progress=gr.Progress(track_tqdm=True)):
     download_dir = f"{repo_id.split('/')[0]}_{repo_id.split('/')[-1]}"
     output_filename = f"{repo_id.split('/')[0]}_{repo_id.split('/')[-1]}.safetensors"
     download_repo(repo_id, download_dir)
@@ -374,7 +369,7 @@ def convert_repo_to_safetensors(repo_id, progress=gr.Progress(track_tqdm=True)):
     return output_filename
 
 
-def convert_repo_to_safetensors_multi(repo_id, files, is_upload, urls, progress=gr.Progress(track_tqdm=True)):
+def convert_repo_to_safetensors_multi(repo_id, files, is_upload, urls, half=True, progress=gr.Progress(track_tqdm=True)):
     file = convert_repo_to_safetensors(repo_id)
     if not urls: urls = []
     url = ""
@@ -393,11 +388,12 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser()
 
     parser.add_argument("--repo_id", default=None, type=str, required=True, help="HF Repo ID of the model to convert.")
+    parser.add_argument("--half", default=True, help="Save weights in half precision.")
 
     args = parser.parse_args()
     assert args.repo_id is not None, "Must provide a Repo ID!"
 
-    convert_repo_to_safetensors(args.repo_id)
+    convert_repo_to_safetensors(args.repo_id, args.half)
 
 
 # Usage: python convert_repo_to_safetensors.py --repo_id GraydientPlatformAPI/goodfit-pony41-xl