Upload 3 files

app.py

@@ -12,17 +12,24 @@ with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", css=css) as demo:
     """)
     with gr.Column(): 
         repo_id = gr.Textbox(label="Repo ID", placeholder="author/model", value="", lines=1)
-        is_half = gr.Checkbox(label="Half precision", value=True)
-        is_upload = gr.Checkbox(label="Upload safetensors to HF Repo", info="Fast download, but files will be public.", value=False)
-        uploaded_urls = gr.CheckboxGroup(visible=False, choices=[], value=None)
-        run_button = gr.Button(value="Convert")
+        with gr.Row():
+            dtype = gr.Radio(label="Output data type", choices=["fp16", "fp32", "bf16", "default"], value="fp16")
+            is_upload = gr.Checkbox(label="Upload safetensors to HF Repo", info="Fast download, but files will be public.", value=False)
+        with gr.Accordion("Advanced", open=False):
+            hf_token = gr.Textbox(label="Your HF write token", placeholder="hf_...", value="", max_lines=1)
+            with gr.Row():
+                newrepo_id = gr.Textbox(label="Upload repo name", placeholder="reponame", value="", max_lines=1)
+                newrepo_type = gr.Radio(label="Upload repo type", choices=["model", "dataset"], value="model")
+                is_private = gr.Checkbox(label="Create / Use private repo", value=True)
+        uploaded_urls = gr.CheckboxGroup(visible=False, choices=[], value=None) # hidden
+        run_button = gr.Button(value="Convert", variant="parimary")
         st_file = gr.Files(label="Output", interactive=False)
         st_md = gr.Markdown()
 
     gr.on(
         triggers=[repo_id.submit, run_button.click],
         fn=convert_repo_to_safetensors_multi,
-        inputs=[repo_id, st_file, is_upload, uploaded_urls, is_half],
+        inputs=[repo_id, hf_token, st_file, uploaded_urls, dtype, is_upload, newrepo_id, is_private],
         outputs=[st_file, uploaded_urls, st_md],
     )
 

convert_repo_to_safetensors.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):
+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")
@@ -325,14 +325,11 @@ def convert_diffusers_to_safetensors(model_path, checkpoint_path, half = True, u
     # Put together new checkpoint
     state_dict = {**unet_state_dict, **vae_state_dict, **text_enc_dict, **text_enc_2_dict}
 
-    if half:
-        state_dict = {k: v.half() for k, v in state_dict.items()}
+    if dtype == "fp16": state_dict = {k: v.half() for k, v in state_dict.items()}
+    elif dtype == "fp32": state_dict = {k: v.to(torch.float32) for k, v in state_dict.items()}
+    elif dtype == "bf16": state_dict = {k: v.to(torch.bfloat16) 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)
 
 
 def download_repo(repo_id, dir_path):
@@ -340,15 +337,15 @@ def download_repo(repo_id, dir_path):
     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
 
 
-def convert_repo_to_safetensors(repo_id):
+def convert_repo_to_safetensors(repo_id, dtype="fp16"):
     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)
-    convert_diffusers_to_safetensors(download_dir, output_filename)
+    convert_diffusers_to_safetensors(download_dir, output_filename, dtype)
     return output_filename
 
 
@@ -356,11 +353,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("--dtype", default="fp16", type=str, choices=["fp16", "fp32", "bf16", "default"], help='Output data type. (Default: "fp16")')
 
     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.dtype)
 
 
 # Usage: python convert_repo_to_safetensors.py --repo_id GraydientPlatformAPI/goodfit-pony41-xl

convert_repo_to_safetensors_gr.py

@@ -10,6 +10,26 @@ import torch
 from safetensors.torch import load_file, save_file
 import gradio as gr
 
+from huggingface_hub import HfApi, HfFolder, hf_hub_url, snapshot_download
+import os
+from pathlib import Path
+
+def get_token():
+    try:
+        token = HfFolder.get_token()
+    except Exception:
+        token = ""
+    return token
+
+def is_repo_exists(repo_id: str):
+    api = HfApi(token=get_token())
+    try:
+        if api.repo_exists(repo_id=repo_id): return True
+        else: return False
+    except Exception as e:
+        print(f"Error: Failed to connect {repo_id}. {e}")
+        return True # for safe
+
 # =================#
 # UNet Conversion #
 # =================#
@@ -269,7 +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, half = True, progress=gr.Progress(track_tqdm=True)):
+def convert_diffusers_to_safetensors(model_path, checkpoint_path, dtype="fp16", 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")
@@ -326,55 +346,56 @@ def convert_diffusers_to_safetensors(model_path, checkpoint_path, half = True, p
     # Put together new checkpoint
     state_dict = {**unet_state_dict, **vae_state_dict, **text_enc_dict, **text_enc_2_dict}
 
-    if half:
-        state_dict = {k: v.half() for k, v in state_dict.items()}
+    if dtype == "fp16": state_dict = {k: v.half() for k, v in state_dict.items()}
+    elif dtype == "fp32": state_dict = {k: v.to(torch.float32) for k, v in state_dict.items()}
+    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)):
-    from huggingface_hub import snapshot_download
     try:
-        snapshot_download(repo_id=repo_id, local_dir=dir_path)
+        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}")
         return
 
 
-def upload_safetensors_to_repo(filename, progress=gr.Progress(track_tqdm=True)):
-    from huggingface_hub import HfApi, hf_hub_url
-    import os
-    from pathlib import Path
+def upload_safetensors_to_repo(filename, repo_id, repo_type, is_private, progress=gr.Progress(track_tqdm=True)):
     output_filename = Path(filename).name
-    hf_token = os.environ.get("HF_TOKEN")
-    repo_id = os.environ.get("HF_OUTPUT_REPO")
-    api = HfApi()
+    api = HfApi(token=get_token())
     try:
+        if not is_repo_exists(repo_id): api.create_repo(repo_id, repo_type=repo_type, private=is_private)
         progress(0, desc="Start uploading...")
-        api.upload_file(path_or_fileobj=filename, path_in_repo=output_filename, repo_id=repo_id, token=hf_token)
+        api.upload_file(path_or_fileobj=filename, path_in_repo=output_filename, repo_type=repo_type, repo_id=repo_id)
         progress(1, desc="Uploaded.")
         url  = hf_hub_url(repo_id=repo_id, filename=output_filename)
     except Exception as e:
-        print(f"Error: Failed to upload to {repo_id}. ")
+        print(f"Error: Failed to upload to {repo_id}. {e}")
         return None
     return url
 
 
-def convert_repo_to_safetensors(repo_id, half=True, progress=gr.Progress(track_tqdm=True)):
+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)
-    convert_diffusers_to_safetensors(download_dir, output_filename, half)
+    convert_diffusers_to_safetensors(download_dir, output_filename, dtype)
     return output_filename
 
 
-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)
+def convert_repo_to_safetensors_multi(repo_id, hf_token, files, urls, dtype="fp16", is_upload=False,
+                                      newrepo_id="", repo_type="model", is_private=True, progress=gr.Progress(track_tqdm=True)):
+    if is_upload and newrepo_id and not hf_token: raise gr.Error("Invalid token.")
+    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)
     if not urls: urls = []
     url = ""
     if is_upload:
-        url = upload_safetensors_to_repo(file, half)
+        url = upload_safetensors_to_repo(file, newrepo_id, repo_type, is_private)
         if url: urls.append(url)
     md = ""
     for u in urls:
@@ -388,12 +409,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.")
+    parser.add_argument("--dtype", default="fp16", type=str, choices=["fp16", "fp32", "bf16", "default"], help='Output data type. (Default: "fp16")')
 
     args = parser.parse_args()
     assert args.repo_id is not None, "Must provide a Repo ID!"
 
-    convert_repo_to_safetensors(args.repo_id, args.half)
+    convert_repo_to_safetensors(args.repo_id, args.dtype)
 
 
 # Usage: python convert_repo_to_safetensors.py --repo_id GraydientPlatformAPI/goodfit-pony41-xl