Upload convert_repo_to_safetensors_gr.py

convert_repo_to_safetensors_gr.py

@@ -289,8 +289,7 @@ def convert_openai_text_enc_state_dict(text_enc_dict):
     return text_enc_dict
 
 
-def convert_diffusers_to_safetensors(model_path, checkpoint_path, dtype="fp16", progress=gr.Progress(track_tqdm=True)):
-    progress(0, desc="Start converting...")
+def convert_diffusers_to_safetensors(model_path, checkpoint_path, dtype="fp16"):
     # Path for safetensors
     unet_path = osp.join(model_path, "unet", "diffusion_pytorch_model.safetensors")
     vae_path = osp.join(model_path, "vae", "diffusion_pytorch_model.safetensors")
@@ -351,12 +350,10 @@ def convert_diffusers_to_safetensors(model_path, checkpoint_path, dtype="fp16",
     elif dtype == "bf16": state_dict = {k: v.to(torch.bfloat16) for k, v in state_dict.items()}
 
     save_file(state_dict, checkpoint_path)
-    progress(1, desc="Converted.")
 
 
-def download_repo(repo_id, dir_path, progress=gr.Progress(track_tqdm=True)):
+def download_repo(repo_id, dir_path):
     try:
-        progress(0, desc="Start downloading...")
         snapshot_download(repo_id=repo_id, local_dir=dir_path, token=get_token())
     except Exception as e:
         print(f"Error: Failed to download {repo_id}. {e}")
@@ -383,8 +380,11 @@ def upload_safetensors_to_repo(filename, repo_id, repo_type, is_private, progres
 def convert_repo_to_safetensors(repo_id, dtype="fp16", 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, progress)
-    convert_diffusers_to_safetensors(download_dir, output_filename, dtype, progress)
+    progress(0, desc="Start downloading...")
+    download_repo(repo_id, download_dir)
+    progress(0, desc="Start converting...")
+    convert_diffusers_to_safetensors(download_dir, output_filename, dtype)
+    progress(1, desc="Converted.")
     return output_filename
 
 
@@ -394,11 +394,11 @@ def convert_repo_to_safetensors_multi(repo_id, hf_token, files, urls, dtype="fp1
     if hf_token: HfFolder.save_token(hf_token)
     else: HfFolder.save_token(os.environ.get("HF_TOKEN"))
     if not newrepo_id: newrepo_id = os.environ.get("HF_OUTPUT_REPO")
-    file = convert_repo_to_safetensors(repo_id, dtype, progress)
+    file = convert_repo_to_safetensors(repo_id, dtype)
     if not urls: urls = []
     url = ""
     if is_upload:
-        url = upload_safetensors_to_repo(file, newrepo_id, repo_type, is_private, progress)
+        url = upload_safetensors_to_repo(file, newrepo_id, repo_type, is_private)
         if url: urls.append(url)
     md = ""
     for u in urls: