Add progress bar.

convert.py

@@ -20,9 +20,10 @@ from transformers import CONFIG_MAPPING
 COMMIT_MESSAGE = " This PR adds fp32 and fp16 weights in PyTorch and safetensors format to {}"
 
 
-def convert_single(model_id: str, filename: str, model_type: str, sample_size: int, scheduler_type: str, extract_ema: bool, folder: str):
+def convert_single(model_id: str, filename: str, model_type: str, sample_size: int, scheduler_type: str, extract_ema: bool, folder: str, progress):
     from_safetensors = filename.endswith(".safetensors")
 
+    progress(0, desc="Downloading model")
     local_file = os.path.join(model_id, filename)
     ckpt_file = local_file if os.path.isfile(local_file) else hf_hub_download(repo_id=model_id, filename=filename)
 
@@ -40,9 +41,11 @@ def convert_single(model_id: str, filename: str, model_type: str, sample_size: i
     config_file = BytesIO(requests.get(config_url).content)
 
     if model_type == "ControlNet":
+        progress(0.2, desc="Converting ControlNet Model")
         pipeline = download_controlnet_from_original_ckpt(ckpt_file, config_file, image_size=sample_size, from_safetensors=from_safetensors, extract_ema=extract_ema)
         to_args = {"dtype": torch.float16}
     else:
+        progress(0.1, desc="Converting Model")
         pipeline = download_from_original_stable_diffusion_ckpt(ckpt_file, config_file, image_size=sample_size, scheduler_type=scheduler_type, from_safetensors=from_safetensors, extract_ema=extract_ema)
         to_args = {"torch_dtype": torch.float16}
 
@@ -68,7 +71,7 @@ def previous_pr(api: "HfApi", model_id: str, pr_title: str) -> Optional["Discuss
                 return discussion
 
 
-def convert(token: str, model_id: str, filename: str, model_type: str, sample_size: int = 512, scheduler_type: str = "pndm", extract_ema: bool = True):
+def convert(token: str, model_id: str, filename: str, model_type: str, sample_size: int = 512, scheduler_type: str = "pndm", extract_ema: bool = True, progress=gr.Progress()):
     api = HfApi()
 
     pr_title = "Adding `diffusers` weights of this model"
@@ -78,10 +81,12 @@ def convert(token: str, model_id: str, filename: str, model_type: str, sample_si
         os.makedirs(folder)
         new_pr = None
         try:
-            folder = convert_single(model_id, filename, model_type, sample_size, scheduler_type, extract_ema, folder)
+            folder = convert_single(model_id, filename, model_type, sample_size, scheduler_type, extract_ema, folder, progress)
+            progress(0.7, desc="Uploading to Hub")
             new_pr  = api.upload_folder(folder_path=folder, path_in_repo="./", repo_id=model_id, repo_type="model", token=token, commit_message=pr_title, commit_description=COMMIT_MESSAGE.format(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(model_id, 'discussions', pr_number)}"
+            progress(1, desc="Done")
         except Exception as e:
             raise gr.exceptions.Error(str(e))
         finally: