Upload 12 files

README.md

@@ -1,5 +1,5 @@
 ---
-title: Convert diffusers SDXL repo to single Safetensors
+title: Convert HF Diffusers repo to single safetensors file V2 (for SDXL / SD 1.5 / LoRA)
 emoji: 🐶
 colorFrom: yellow
 colorTo: red
@@ -7,6 +7,7 @@ sdk: gradio
 sdk_version: 4.44.0
 app_file: app.py
 pinned: false
+license: mit
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -1,30 +1,47 @@
 import gradio as gr
 import os
-from convert_repo_to_safetensors_gr import convert_repo_to_safetensors_multi
+from convert_repo_to_safetensors_gr import convert_repo_to_safetensors_multi, clear_safetensors
 os.environ['HF_OUTPUT_REPO'] = 'John6666/safetensors_converting_test'
 
-css = """"""
+css = """
+.title { text-align: center; !important; }
+.footer { text-align: center; !important; }
+"""
 
-with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", css=css) as demo:
-    gr.Markdown(
-            f"""
-    - [A CLI version of this tool is available here](https://huggingface.co/spaces/John6666/convert_repo_to_safetensors/tree/main/local).
-    """)
+with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", fill_width=True, css=css, delete_cache=(60, 3600)) as demo:
+    gr.Markdown("# HF Diffusers repo to WebUI/ComfyUI single safetensors file converter (for SDXL / SD 1.5 / LoRA)", elem_classes="title")
     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)
+        with gr.Accordion("Advanced", open=False):
+            dtype = gr.Radio(label="Output data type", choices=["fp16", "fp32", "bf16", "fp8", "default"], value="fp16")
+            with gr.Accordion("Upload to your repo", open=True):
+                with gr.Row():
+                    hf_token = gr.Textbox(label="Your HF write token", placeholder="hf_...", value="", max_lines=1)
+                    gr.Markdown("Your token is available at [hf.co/settings/tokens](https://huggingface.co/settings/tokens).")
+                with gr.Row():
+                    newrepo_id = gr.Textbox(label="Upload repo ID", placeholder="yourid/newrepo", value="", max_lines=1)
+                    newrepo_type = gr.Radio(label="Upload repo type", choices=["model", "dataset"], value="model")
+                    is_private = gr.Checkbox(label="Create private repo", value=True)
+        uploaded_urls = gr.CheckboxGroup(visible=False, choices=[], value=None) # hidden
         run_button = gr.Button(value="Convert")
         st_file = gr.Files(label="Output", interactive=False)
         st_md = gr.Markdown()
+        delete_button = gr.Button(value="Delete Safetensors")
+    gr.DuplicateButton(value="Duplicate Space")
+    gr.Markdown(
+            f"""
+    - Thanks to [xi0v](https://huggingface.co/xi0v)
+    - [A CLI version of this tool is available here](https://huggingface.co/spaces/John6666/convert_repo_to_safetensors/tree/main/local).
+    """, elem_classes="footer")
 
     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, newrepo_type, is_private],
         outputs=[st_file, uploaded_urls, st_md],
     )
+    delete_button.click(clear_safetensors, None, [st_file], queue=False, show_api=False)
 
 demo.queue()
 demo.launch()

convert_repo_to_safetensors_gr.py

@@ -10,6 +10,15 @@ 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
+import shutil
+import gc
+from utils import get_token, set_token, is_repo_exists, get_model_type
+from convert_repo_to_safetensors_sd_gr import convert_repo_to_safetensors as convert_repo_to_safetensors_sd
+from convert_repo_to_safetensors_sdxl_lora_gr import convert_repo_to_safetensors_sdxl_lora
+
 # =================#
 # UNet Conversion #
 # =================#
@@ -269,8 +278,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)):
-    progress(0, desc="Start converting...")
+def convert_diffusers_to_safetensors(model_path, checkpoint_path, dtype="fp16", progress=gr.Progress(track_tqdm=True)):
     # 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")
@@ -326,74 +334,109 @@ 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()}
+    elif dtype == "fp8": state_dict = {k: v.to(torch.float8_e4m3fn) for k, v in state_dict.items()}
 
     save_file(state_dict, checkpoint_path)
-    progress(1, desc="Converted.")
 
 
+# https://huggingface.co/docs/huggingface_hub/v0.25.1/en/package_reference/file_download#huggingface_hub.snapshot_download
 def download_repo(repo_id, dir_path, progress=gr.Progress(track_tqdm=True)):
-    from huggingface_hub import snapshot_download
+    hf_token = get_token()
     try:
-        snapshot_download(repo_id=repo_id, local_dir=dir_path)
+        snapshot_download(repo_id=repo_id, local_dir=dir_path, token=hf_token, allow_patterns=["*.safetensors", "*.bin"],
+                          ignore_patterns=["*.fp16.*", "/*.safetensors", "/*.bin"], force_download=True)
     except Exception as e:
         print(f"Error: Failed to download {repo_id}. {e}")
+        gr.Warning(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()
+    hf_token = get_token()
+    api = HfApi(token=hf_token)
     try:
+        if not is_repo_exists(repo_id, repo_type): api.create_repo(repo_id=repo_id, repo_type=repo_type, token=hf_token, 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, revision="main", token=hf_token, repo_id=repo_id)
         progress(1, desc="Uploaded.")
-        url  = hf_hub_url(repo_id=repo_id, filename=output_filename)
+        url = hf_hub_url(repo_id=repo_id, repo_type=repo_type, 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}")
+        gr.Warning(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"
+    progress(0, desc="Start downloading...")
     download_repo(repo_id, download_dir)
-    convert_diffusers_to_safetensors(download_dir, output_filename, half)
+    progress(0, desc="Start converting...")
+    convert_diffusers_to_safetensors(download_dir, output_filename, dtype)
+    progress(1, desc="Converted.")
+    shutil.rmtree(download_dir)
     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 hf_token: set_token(hf_token)
+    else: set_token(os.environ.get("HF_TOKEN"))
+    if is_upload and newrepo_id and not hf_token: raise gr.Error("HF write token is required for this process.")
+    if not newrepo_id: newrepo_id = os.environ.get("HF_OUTPUT_REPO")
+    model_type = get_model_type(repo_id)
+    if model_type == "SDXL":
+        gr.Info(f"Converting {model_type} model.")
+        file = convert_repo_to_safetensors(repo_id, dtype)
+    elif model_type == "SD 1.5":
+        gr.Info(f"Converting {model_type} model.")
+        file = convert_repo_to_safetensors_sd(repo_id, dtype)
+    elif model_type == "LoRA":
+        gr.Info(f"Converting {model_type}.")
+        file = convert_repo_to_safetensors_sdxl_lora(repo_id)
+    else: raise gr.Error(f"Unsupported model type: {model_type}")
     if not urls: urls = []
+    if not files: files = []
     url = ""
     if is_upload:
-        url = upload_safetensors_to_repo(file, half)
-        if url: urls.append(url)
+        url = upload_safetensors_to_repo(file, newrepo_id, repo_type, is_private)
+        if url:
+            urls.append(url)
+            Path(file).unlink()
+        else: files.append(file)
+    else:
+        files.append(file)
+    progress(1, desc="Processing...")
     md = ""
     for u in urls:
         md += f"[Download {str(u).split('/')[-1]}]({str(u)})<br>"
-    if not files: files = []
-    files.append(file)
+    gc.collect()
     return gr.update(value=files), gr.update(value=urls, choices=urls), gr.update(value=md)
 
 
+def clear_safetensors():
+    for p in Path('.').glob('*.safetensors'):
+        p.unlink()
+        print("Deleted.")
+    gc.collect()
+    return gr.update(value=[])
+
+
 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", "fp8", "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

convert_repo_to_safetensors_sd_gr.py

@@ -0,0 +1,426 @@
+# Script for converting a HF Diffusers saved pipeline to a Stable Diffusion checkpoint.
+# *Only* converts the UNet, VAE, and Text Encoder.
+# Does not convert optimizer state or any other thing.
+
+import argparse
+import os.path as osp
+import re
+
+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
+import shutil
+import gc
+from utils import get_token, set_token, is_repo_exists
+
+# =================#
+# UNet Conversion #
+# =================#
+
+unet_conversion_map = [
+    # (stable-diffusion, HF Diffusers)
+    ("time_embed.0.weight", "time_embedding.linear_1.weight"),
+    ("time_embed.0.bias", "time_embedding.linear_1.bias"),
+    ("time_embed.2.weight", "time_embedding.linear_2.weight"),
+    ("time_embed.2.bias", "time_embedding.linear_2.bias"),
+    ("input_blocks.0.0.weight", "conv_in.weight"),
+    ("input_blocks.0.0.bias", "conv_in.bias"),
+    ("out.0.weight", "conv_norm_out.weight"),
+    ("out.0.bias", "conv_norm_out.bias"),
+    ("out.2.weight", "conv_out.weight"),
+    ("out.2.bias", "conv_out.bias"),
+]
+
+unet_conversion_map_resnet = [
+    # (stable-diffusion, HF Diffusers)
+    ("in_layers.0", "norm1"),
+    ("in_layers.2", "conv1"),
+    ("out_layers.0", "norm2"),
+    ("out_layers.3", "conv2"),
+    ("emb_layers.1", "time_emb_proj"),
+    ("skip_connection", "conv_shortcut"),
+]
+
+unet_conversion_map_layer = []
+# hardcoded number of downblocks and resnets/attentions...
+# would need smarter logic for other networks.
+for i in range(4):
+    # loop over downblocks/upblocks
+
+    for j in range(2):
+        # loop over resnets/attentions for downblocks
+        hf_down_res_prefix = f"down_blocks.{i}.resnets.{j}."
+        sd_down_res_prefix = f"input_blocks.{3*i + j + 1}.0."
+        unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix))
+
+        if i < 3:
+            # no attention layers in down_blocks.3
+            hf_down_atn_prefix = f"down_blocks.{i}.attentions.{j}."
+            sd_down_atn_prefix = f"input_blocks.{3*i + j + 1}.1."
+            unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix))
+
+    for j in range(3):
+        # loop over resnets/attentions for upblocks
+        hf_up_res_prefix = f"up_blocks.{i}.resnets.{j}."
+        sd_up_res_prefix = f"output_blocks.{3*i + j}.0."
+        unet_conversion_map_layer.append((sd_up_res_prefix, hf_up_res_prefix))
+
+        if i > 0:
+            # no attention layers in up_blocks.0
+            hf_up_atn_prefix = f"up_blocks.{i}.attentions.{j}."
+            sd_up_atn_prefix = f"output_blocks.{3*i + j}.1."
+            unet_conversion_map_layer.append((sd_up_atn_prefix, hf_up_atn_prefix))
+
+    if i < 3:
+        # no downsample in down_blocks.3
+        hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0.conv."
+        sd_downsample_prefix = f"input_blocks.{3*(i+1)}.0.op."
+        unet_conversion_map_layer.append((sd_downsample_prefix, hf_downsample_prefix))
+
+        # no upsample in up_blocks.3
+        hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
+        sd_upsample_prefix = f"output_blocks.{3*i + 2}.{1 if i == 0 else 2}."
+        unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix))
+
+hf_mid_atn_prefix = "mid_block.attentions.0."
+sd_mid_atn_prefix = "middle_block.1."
+unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))
+
+for j in range(2):
+    hf_mid_res_prefix = f"mid_block.resnets.{j}."
+    sd_mid_res_prefix = f"middle_block.{2*j}."
+    unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix))
+
+
+def convert_unet_state_dict(unet_state_dict):
+    # buyer beware: this is a *brittle* function,
+    # and correct output requires that all of these pieces interact in
+    # the exact order in which I have arranged them.
+    mapping = {k: k for k in unet_state_dict.keys()}
+    for sd_name, hf_name in unet_conversion_map:
+        mapping[hf_name] = sd_name
+    for k, v in mapping.items():
+        if "resnets" in k:
+            for sd_part, hf_part in unet_conversion_map_resnet:
+                v = v.replace(hf_part, sd_part)
+            mapping[k] = v
+    for k, v in mapping.items():
+        for sd_part, hf_part in unet_conversion_map_layer:
+            v = v.replace(hf_part, sd_part)
+        mapping[k] = v
+    new_state_dict = {v: unet_state_dict[k] for k, v in mapping.items()}
+    return new_state_dict
+
+
+# ================#
+# VAE Conversion #
+# ================#
+
+vae_conversion_map = [
+    # (stable-diffusion, HF Diffusers)
+    ("nin_shortcut", "conv_shortcut"),
+    ("norm_out", "conv_norm_out"),
+    ("mid.attn_1.", "mid_block.attentions.0."),
+]
+
+for i in range(4):
+    # down_blocks have two resnets
+    for j in range(2):
+        hf_down_prefix = f"encoder.down_blocks.{i}.resnets.{j}."
+        sd_down_prefix = f"encoder.down.{i}.block.{j}."
+        vae_conversion_map.append((sd_down_prefix, hf_down_prefix))
+
+    if i < 3:
+        hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0."
+        sd_downsample_prefix = f"down.{i}.downsample."
+        vae_conversion_map.append((sd_downsample_prefix, hf_downsample_prefix))
+
+        hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
+        sd_upsample_prefix = f"up.{3-i}.upsample."
+        vae_conversion_map.append((sd_upsample_prefix, hf_upsample_prefix))
+
+    # up_blocks have three resnets
+    # also, up blocks in hf are numbered in reverse from sd
+    for j in range(3):
+        hf_up_prefix = f"decoder.up_blocks.{i}.resnets.{j}."
+        sd_up_prefix = f"decoder.up.{3-i}.block.{j}."
+        vae_conversion_map.append((sd_up_prefix, hf_up_prefix))
+
+# this part accounts for mid blocks in both the encoder and the decoder
+for i in range(2):
+    hf_mid_res_prefix = f"mid_block.resnets.{i}."
+    sd_mid_res_prefix = f"mid.block_{i+1}."
+    vae_conversion_map.append((sd_mid_res_prefix, hf_mid_res_prefix))
+
+
+vae_conversion_map_attn = [
+    # (stable-diffusion, HF Diffusers)
+    ("norm.", "group_norm."),
+    ("q.", "query."),
+    ("k.", "key."),
+    ("v.", "value."),
+    ("proj_out.", "proj_attn."),
+]
+
+# This is probably not the most ideal solution, but it does work.
+vae_extra_conversion_map = [
+    ("to_q", "q"),
+    ("to_k", "k"),
+    ("to_v", "v"),
+    ("to_out.0", "proj_out"),
+]
+
+
+def reshape_weight_for_sd(w):
+    # convert HF linear weights to SD conv2d weights
+    if not w.ndim == 1:
+        return w.reshape(*w.shape, 1, 1)
+    else:
+        return w
+
+
+def convert_vae_state_dict(vae_state_dict):
+    mapping = {k: k for k in vae_state_dict.keys()}
+    for k, v in mapping.items():
+        for sd_part, hf_part in vae_conversion_map:
+            v = v.replace(hf_part, sd_part)
+        mapping[k] = v
+    for k, v in mapping.items():
+        if "attentions" in k:
+            for sd_part, hf_part in vae_conversion_map_attn:
+                v = v.replace(hf_part, sd_part)
+            mapping[k] = v
+    new_state_dict = {v: vae_state_dict[k] for k, v in mapping.items()}
+    weights_to_convert = ["q", "k", "v", "proj_out"]
+    keys_to_rename = {}
+    for k, v in new_state_dict.items():
+        for weight_name in weights_to_convert:
+            if f"mid.attn_1.{weight_name}.weight" in k:
+                print(f"Reshaping {k} for SD format")
+                new_state_dict[k] = reshape_weight_for_sd(v)
+        for weight_name, real_weight_name in vae_extra_conversion_map:
+            if f"mid.attn_1.{weight_name}.weight" in k or f"mid.attn_1.{weight_name}.bias" in k:
+                keys_to_rename[k] = k.replace(weight_name, real_weight_name)
+    for k, v in keys_to_rename.items():
+        if k in new_state_dict:
+            print(f"Renaming {k} to {v}")
+            new_state_dict[v] = reshape_weight_for_sd(new_state_dict[k])
+            del new_state_dict[k]
+    return new_state_dict
+
+
+# =========================#
+# Text Encoder Conversion #
+# =========================#
+
+
+textenc_conversion_lst = [
+    # (stable-diffusion, HF Diffusers)
+    ("resblocks.", "text_model.encoder.layers."),
+    ("ln_1", "layer_norm1"),
+    ("ln_2", "layer_norm2"),
+    (".c_fc.", ".fc1."),
+    (".c_proj.", ".fc2."),
+    (".attn", ".self_attn"),
+    ("ln_final.", "transformer.text_model.final_layer_norm."),
+    ("token_embedding.weight", "transformer.text_model.embeddings.token_embedding.weight"),
+    ("positional_embedding", "transformer.text_model.embeddings.position_embedding.weight"),
+]
+protected = {re.escape(x[1]): x[0] for x in textenc_conversion_lst}
+textenc_pattern = re.compile("|".join(protected.keys()))
+
+# Ordering is from https://github.com/pytorch/pytorch/blob/master/test/cpp/api/modules.cpp
+code2idx = {"q": 0, "k": 1, "v": 2}
+
+
+def convert_text_enc_state_dict_v20(text_enc_dict):
+    new_state_dict = {}
+    capture_qkv_weight = {}
+    capture_qkv_bias = {}
+    for k, v in text_enc_dict.items():
+        if (
+            k.endswith(".self_attn.q_proj.weight")
+            or k.endswith(".self_attn.k_proj.weight")
+            or k.endswith(".self_attn.v_proj.weight")
+        ):
+            k_pre = k[: -len(".q_proj.weight")]
+            k_code = k[-len("q_proj.weight")]
+            if k_pre not in capture_qkv_weight:
+                capture_qkv_weight[k_pre] = [None, None, None]
+            capture_qkv_weight[k_pre][code2idx[k_code]] = v
+            continue
+
+        if (
+            k.endswith(".self_attn.q_proj.bias")
+            or k.endswith(".self_attn.k_proj.bias")
+            or k.endswith(".self_attn.v_proj.bias")
+        ):
+            k_pre = k[: -len(".q_proj.bias")]
+            k_code = k[-len("q_proj.bias")]
+            if k_pre not in capture_qkv_bias:
+                capture_qkv_bias[k_pre] = [None, None, None]
+            capture_qkv_bias[k_pre][code2idx[k_code]] = v
+            continue
+
+        relabelled_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], k)
+        new_state_dict[relabelled_key] = v
+
+    for k_pre, tensors in capture_qkv_weight.items():
+        if None in tensors:
+            raise Exception("CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing")
+        relabelled_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], k_pre)
+        new_state_dict[relabelled_key + ".in_proj_weight"] = torch.cat(tensors)
+
+    for k_pre, tensors in capture_qkv_bias.items():
+        if None in tensors:
+            raise Exception("CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing")
+        relabelled_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], k_pre)
+        new_state_dict[relabelled_key + ".in_proj_bias"] = torch.cat(tensors)
+
+    return new_state_dict
+
+
+def convert_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)):
+    # 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")
+    text_enc_path = osp.join(model_path, "text_encoder", "model.safetensors")
+
+    # Load models from safetensors if it exists, if it doesn't pytorch
+    if osp.exists(unet_path):
+        unet_state_dict = load_file(unet_path, device="cpu")
+    else:
+        unet_path = osp.join(model_path, "unet", "diffusion_pytorch_model.bin")
+        unet_state_dict = torch.load(unet_path, map_location="cpu")
+
+    if osp.exists(vae_path):
+        vae_state_dict = load_file(vae_path, device="cpu")
+    else:
+        vae_path = osp.join(model_path, "vae", "diffusion_pytorch_model.bin")
+        vae_state_dict = torch.load(vae_path, map_location="cpu")
+
+    if osp.exists(text_enc_path):
+        text_enc_dict = load_file(text_enc_path, device="cpu")
+    else:
+        text_enc_path = osp.join(model_path, "text_encoder", "pytorch_model.bin")
+        text_enc_dict = torch.load(text_enc_path, map_location="cpu")
+
+    # Convert the UNet model
+    unet_state_dict = convert_unet_state_dict(unet_state_dict)
+    unet_state_dict = {"model.diffusion_model." + k: v for k, v in unet_state_dict.items()}
+
+    # Convert the VAE model
+    vae_state_dict = convert_vae_state_dict(vae_state_dict)
+    vae_state_dict = {"first_stage_model." + k: v for k, v in vae_state_dict.items()}
+
+    # Easiest way to identify v2.0 model seems to be that the text encoder (OpenCLIP) is deeper
+    is_v20_model = "text_model.encoder.layers.22.layer_norm2.bias" in text_enc_dict
+
+    if is_v20_model:
+        # Need to add the tag 'transformer' in advance so we can knock it out from the final layer-norm
+        text_enc_dict = {"transformer." + k: v for k, v in text_enc_dict.items()}
+        text_enc_dict = convert_text_enc_state_dict_v20(text_enc_dict)
+        text_enc_dict = {"cond_stage_model.model." + k: v for k, v in text_enc_dict.items()}
+    else:
+        text_enc_dict = convert_text_enc_state_dict(text_enc_dict)
+        text_enc_dict = {"cond_stage_model.transformer." + k: v for k, v in text_enc_dict.items()}
+
+    # Put together new checkpoint
+    state_dict = {**unet_state_dict, **vae_state_dict, **text_enc_dict}
+    
+    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()}
+    elif dtype == "fp8": state_dict = {k: v.to(torch.float8_e4m3fn) for k, v in state_dict.items()}
+
+    save_file(state_dict, checkpoint_path)
+
+
+# https://huggingface.co/docs/huggingface_hub/v0.25.1/en/package_reference/file_download#huggingface_hub.snapshot_download
+def download_repo(repo_id, dir_path):
+    hf_token = get_token()
+    try:
+        snapshot_download(repo_id=repo_id, local_dir=dir_path, token=hf_token, allow_patterns=["*.safetensors", "*.bin"],
+                          ignore_patterns=["*.fp16.*", "/*.safetensors", "/*.bin"])
+    except Exception as e:
+        print(f"Error: Failed to download {repo_id}. {e}")
+        gr.Warning(f"Error: Failed to download {repo_id}. {e}")
+        return
+
+
+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 = get_token()
+    api = HfApi(token=hf_token)
+    try:
+        if not is_repo_exists(repo_id, repo_type): api.create_repo(repo_id=repo_id, repo_type=repo_type, token=hf_token, private=is_private)
+        progress(0, desc="Start uploading...")
+        api.upload_file(path_or_fileobj=filename, path_in_repo=output_filename, repo_type=repo_type, revision="main", token=hf_token, repo_id=repo_id)
+        progress(1, desc="Uploaded.")
+        url = hf_hub_url(repo_id=repo_id, repo_type=repo_type, filename=output_filename)
+    except Exception as e:
+        print(f"Error: Failed to upload to {repo_id}. {e}")
+        gr.Warning(f"Error: Failed to upload to {repo_id}. {e}")
+        return None
+    return url
+
+
+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"
+    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.")
+    shutil.rmtree(download_dir)
+    return output_filename
+
+
+def convert_repo_to_safetensors_multi_sd(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 hf_token: set_token(hf_token)
+    else: set_token(os.environ.get("HF_TOKEN"))
+    if is_upload and newrepo_id and not hf_token: raise gr.Error("HF write token is required for this process.")
+    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, newrepo_id, repo_type, is_private)
+        if url: urls.append(url)
+    progress(1, desc="Processing...")
+    md = ""
+    for u in urls:
+        md += f"[Download {str(u).split('/')[-1]}]({str(u)})<br>"
+    if not files: files = []
+    files.append(file)
+    return gr.update(value=files), gr.update(value=urls, choices=urls), gr.update(value=md)
+
+
+def clear_safetensors():
+    for p in Path('.').glob('*.safetensors'):
+        p.unlink()
+        print("Deleted.")
+    gc.collect()
+    return gr.update(value=[])
+
+
+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", "fp8", "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.dtype)

convert_repo_to_safetensors_sdxl_lora_gr.py

@@ -0,0 +1,97 @@
+# Script for converting a Hugging Face Diffusers trained SDXL LoRAs to Kohya format
+# This means that you can input your diffusers-trained LoRAs and
+# Get the output to work with WebUIs such as AUTOMATIC1111, ComfyUI, SD.Next and others.
+
+# To get started you can find some cool `diffusers` trained LoRAs such as this cute Corgy
+# https://huggingface.co/ignasbud/corgy_dog_LoRA/, download its `pytorch_lora_weights.safetensors` file
+# and run the script:
+# python convert_diffusers_sdxl_lora_to_webui.py --input_lora pytorch_lora_weights.safetensors --output_lora corgy.safetensors
+# now you can use corgy.safetensors in your WebUI of choice!
+
+# To train your own, here are some diffusers training scripts and utils that you can use and then convert:
+# LoRA Ease - no code SDXL Dreambooth LoRA trainer: https://huggingface.co/spaces/multimodalart/lora-ease
+# Dreambooth Advanced Training Script - state of the art techniques such as pivotal tuning and prodigy optimizer:
+# - Script: https://github.com/huggingface/diffusers/blob/main/examples/advanced_diffusion_training/train_dreambooth_lora_sdxl_advanced.py
+# - Colab (only on Pro): https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/SDXL_Dreambooth_LoRA_advanced_example.ipynb
+# Canonical diffusers training scripts:
+# - Script: https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth_lora_sdxl.py
+# - Colab (runs on free tier): https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/SDXL_DreamBooth_LoRA_.ipynb
+
+import argparse
+import os
+
+from safetensors.torch import load_file, save_file
+from diffusers.utils import convert_all_state_dict_to_peft, convert_state_dict_to_kohya
+from pathlib import Path
+import gradio as gr
+
+from huggingface_hub import hf_hub_download, HfApi
+from huggingface_hub import HfApi, HfFolder, hf_hub_url, snapshot_download
+import os
+from pathlib import Path
+import shutil
+import gc
+from utils import get_token, set_token, is_repo_exists, get_model_type
+
+def convert_and_save(input_lora, output_lora=None):
+    if output_lora is None:
+        base_name = os.path.splitext(input_lora)[0]
+        output_lora = f"{base_name}_webui.safetensors"
+
+    diffusers_state_dict = load_file(input_lora)
+    try:
+        peft_state_dict = convert_all_state_dict_to_peft(diffusers_state_dict)
+        kohya_state_dict = convert_state_dict_to_kohya(peft_state_dict)
+    except Exception: # skipped
+        kohya_state_dict = diffusers_state_dict
+    save_file(kohya_state_dict, output_lora)
+
+
+def download_repo_lora(repo_id, local_file, progress=gr.Progress(track_tqdm=True)):
+    hf_token = get_token()
+    lora_filename = "pytorch_lora_weights.safetensors"
+    lora_path = Path(lora_filename)
+    api = HfApi(token=hf_token)
+    try:
+        if not api.file_exists(repo_id=repo_id, filename=lora_filename, token=hf_token):
+            print(f"Error: This repo isn't diffusers LoRA repo: {repo_id}.")
+            return None
+        if lora_path.exists():
+            print(f"Error: Download file already exists: {lora_filename}.")
+            return None
+        hf_hub_download(repo_id=repo_id, filename=lora_filename, local_dir=".")
+        if lora_path.exists(): lora_path.rename(Path(local_file))
+    except Exception as e:
+        print(f"Error: Failed to download from {repo_id}. {e}")
+        return local_file
+
+
+def convert_repo_to_safetensors_sdxl_lora(repo_id, progress=gr.Progress(track_tqdm=True)):
+    download_filename = f"{repo_id.split('/')[0]}_{repo_id.split('/')[-1]}_diffusers.safetensors"
+    output_filename = f"{repo_id.split('/')[0]}_{repo_id.split('/')[-1]}_webui.safetensors"
+    progress(0, desc="Start downloading...")
+    download_repo_lora(repo_id, download_filename)
+    progress(0, desc="Start converting...")
+    convert_and_save(download_filename, output_filename)
+    progress(1, desc="Converted.")
+    Path(download_filename).unlink()
+    return output_filename
+
+
+def convert_repo_to_safetensors_sdxl_lora_multi(repo_id, files, progress=gr.Progress(track_tqdm=True)):
+    file = convert_repo_to_safetensors_sdxl_lora(repo_id)
+    if not files: files = []
+    files.append(file)
+    return gr.update(value=files)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Convert LoRA model to PEFT and then to Kohya format from Repo.")
+    parser.add_argument("--repo_id", type=str, required=True, help="URL to the Repo of input LoRA model in the diffusers format.")
+
+    args = parser.parse_args()
+
+    convert_repo_to_safetensors_sdxl_lora(args.repo_id)
+
+
+# Usage: python convert_repo_to_safetensors_sdxl_lora.py --repo_id nroggendorff/zelda-lora

local/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):
+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,8 +325,9 @@ def convert_diffusers_to_safetensors(model_path, checkpoint_path, half = True):
     # 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)
 
@@ -336,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, half=True):
+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, half)
+    convert_diffusers_to_safetensors(download_dir, output_filename, dtype)
     return output_filename
 
 
@@ -352,12 +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("--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

local/convert_repo_to_safetensors_sd.py

@@ -0,0 +1,365 @@
+# Script for converting a HF Diffusers saved pipeline to a Stable Diffusion checkpoint.
+# *Only* converts the UNet, VAE, and Text Encoder.
+# Does not convert optimizer state or any other thing.
+
+import argparse
+import os.path as osp
+import re
+
+import torch
+from safetensors.torch import load_file, save_file
+
+
+# =================#
+# UNet Conversion #
+# =================#
+
+unet_conversion_map = [
+    # (stable-diffusion, HF Diffusers)
+    ("time_embed.0.weight", "time_embedding.linear_1.weight"),
+    ("time_embed.0.bias", "time_embedding.linear_1.bias"),
+    ("time_embed.2.weight", "time_embedding.linear_2.weight"),
+    ("time_embed.2.bias", "time_embedding.linear_2.bias"),
+    ("input_blocks.0.0.weight", "conv_in.weight"),
+    ("input_blocks.0.0.bias", "conv_in.bias"),
+    ("out.0.weight", "conv_norm_out.weight"),
+    ("out.0.bias", "conv_norm_out.bias"),
+    ("out.2.weight", "conv_out.weight"),
+    ("out.2.bias", "conv_out.bias"),
+]
+
+unet_conversion_map_resnet = [
+    # (stable-diffusion, HF Diffusers)
+    ("in_layers.0", "norm1"),
+    ("in_layers.2", "conv1"),
+    ("out_layers.0", "norm2"),
+    ("out_layers.3", "conv2"),
+    ("emb_layers.1", "time_emb_proj"),
+    ("skip_connection", "conv_shortcut"),
+]
+
+unet_conversion_map_layer = []
+# hardcoded number of downblocks and resnets/attentions...
+# would need smarter logic for other networks.
+for i in range(4):
+    # loop over downblocks/upblocks
+
+    for j in range(2):
+        # loop over resnets/attentions for downblocks
+        hf_down_res_prefix = f"down_blocks.{i}.resnets.{j}."
+        sd_down_res_prefix = f"input_blocks.{3*i + j + 1}.0."
+        unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix))
+
+        if i < 3:
+            # no attention layers in down_blocks.3
+            hf_down_atn_prefix = f"down_blocks.{i}.attentions.{j}."
+            sd_down_atn_prefix = f"input_blocks.{3*i + j + 1}.1."
+            unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix))
+
+    for j in range(3):
+        # loop over resnets/attentions for upblocks
+        hf_up_res_prefix = f"up_blocks.{i}.resnets.{j}."
+        sd_up_res_prefix = f"output_blocks.{3*i + j}.0."
+        unet_conversion_map_layer.append((sd_up_res_prefix, hf_up_res_prefix))
+
+        if i > 0:
+            # no attention layers in up_blocks.0
+            hf_up_atn_prefix = f"up_blocks.{i}.attentions.{j}."
+            sd_up_atn_prefix = f"output_blocks.{3*i + j}.1."
+            unet_conversion_map_layer.append((sd_up_atn_prefix, hf_up_atn_prefix))
+
+    if i < 3:
+        # no downsample in down_blocks.3
+        hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0.conv."
+        sd_downsample_prefix = f"input_blocks.{3*(i+1)}.0.op."
+        unet_conversion_map_layer.append((sd_downsample_prefix, hf_downsample_prefix))
+
+        # no upsample in up_blocks.3
+        hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
+        sd_upsample_prefix = f"output_blocks.{3*i + 2}.{1 if i == 0 else 2}."
+        unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix))
+
+hf_mid_atn_prefix = "mid_block.attentions.0."
+sd_mid_atn_prefix = "middle_block.1."
+unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))
+
+for j in range(2):
+    hf_mid_res_prefix = f"mid_block.resnets.{j}."
+    sd_mid_res_prefix = f"middle_block.{2*j}."
+    unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix))
+
+
+def convert_unet_state_dict(unet_state_dict):
+    # buyer beware: this is a *brittle* function,
+    # and correct output requires that all of these pieces interact in
+    # the exact order in which I have arranged them.
+    mapping = {k: k for k in unet_state_dict.keys()}
+    for sd_name, hf_name in unet_conversion_map:
+        mapping[hf_name] = sd_name
+    for k, v in mapping.items():
+        if "resnets" in k:
+            for sd_part, hf_part in unet_conversion_map_resnet:
+                v = v.replace(hf_part, sd_part)
+            mapping[k] = v
+    for k, v in mapping.items():
+        for sd_part, hf_part in unet_conversion_map_layer:
+            v = v.replace(hf_part, sd_part)
+        mapping[k] = v
+    new_state_dict = {v: unet_state_dict[k] for k, v in mapping.items()}
+    return new_state_dict
+
+
+# ================#
+# VAE Conversion #
+# ================#
+
+vae_conversion_map = [
+    # (stable-diffusion, HF Diffusers)
+    ("nin_shortcut", "conv_shortcut"),
+    ("norm_out", "conv_norm_out"),
+    ("mid.attn_1.", "mid_block.attentions.0."),
+]
+
+for i in range(4):
+    # down_blocks have two resnets
+    for j in range(2):
+        hf_down_prefix = f"encoder.down_blocks.{i}.resnets.{j}."
+        sd_down_prefix = f"encoder.down.{i}.block.{j}."
+        vae_conversion_map.append((sd_down_prefix, hf_down_prefix))
+
+    if i < 3:
+        hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0."
+        sd_downsample_prefix = f"down.{i}.downsample."
+        vae_conversion_map.append((sd_downsample_prefix, hf_downsample_prefix))
+
+        hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
+        sd_upsample_prefix = f"up.{3-i}.upsample."
+        vae_conversion_map.append((sd_upsample_prefix, hf_upsample_prefix))
+
+    # up_blocks have three resnets
+    # also, up blocks in hf are numbered in reverse from sd
+    for j in range(3):
+        hf_up_prefix = f"decoder.up_blocks.{i}.resnets.{j}."
+        sd_up_prefix = f"decoder.up.{3-i}.block.{j}."
+        vae_conversion_map.append((sd_up_prefix, hf_up_prefix))
+
+# this part accounts for mid blocks in both the encoder and the decoder
+for i in range(2):
+    hf_mid_res_prefix = f"mid_block.resnets.{i}."
+    sd_mid_res_prefix = f"mid.block_{i+1}."
+    vae_conversion_map.append((sd_mid_res_prefix, hf_mid_res_prefix))
+
+
+vae_conversion_map_attn = [
+    # (stable-diffusion, HF Diffusers)
+    ("norm.", "group_norm."),
+    ("q.", "query."),
+    ("k.", "key."),
+    ("v.", "value."),
+    ("proj_out.", "proj_attn."),
+]
+
+# This is probably not the most ideal solution, but it does work.
+vae_extra_conversion_map = [
+    ("to_q", "q"),
+    ("to_k", "k"),
+    ("to_v", "v"),
+    ("to_out.0", "proj_out"),
+]
+
+
+def reshape_weight_for_sd(w):
+    # convert HF linear weights to SD conv2d weights
+    if not w.ndim == 1:
+        return w.reshape(*w.shape, 1, 1)
+    else:
+        return w
+
+
+def convert_vae_state_dict(vae_state_dict):
+    mapping = {k: k for k in vae_state_dict.keys()}
+    for k, v in mapping.items():
+        for sd_part, hf_part in vae_conversion_map:
+            v = v.replace(hf_part, sd_part)
+        mapping[k] = v
+    for k, v in mapping.items():
+        if "attentions" in k:
+            for sd_part, hf_part in vae_conversion_map_attn:
+                v = v.replace(hf_part, sd_part)
+            mapping[k] = v
+    new_state_dict = {v: vae_state_dict[k] for k, v in mapping.items()}
+    weights_to_convert = ["q", "k", "v", "proj_out"]
+    keys_to_rename = {}
+    for k, v in new_state_dict.items():
+        for weight_name in weights_to_convert:
+            if f"mid.attn_1.{weight_name}.weight" in k:
+                print(f"Reshaping {k} for SD format")
+                new_state_dict[k] = reshape_weight_for_sd(v)
+        for weight_name, real_weight_name in vae_extra_conversion_map:
+            if f"mid.attn_1.{weight_name}.weight" in k or f"mid.attn_1.{weight_name}.bias" in k:
+                keys_to_rename[k] = k.replace(weight_name, real_weight_name)
+    for k, v in keys_to_rename.items():
+        if k in new_state_dict:
+            print(f"Renaming {k} to {v}")
+            new_state_dict[v] = reshape_weight_for_sd(new_state_dict[k])
+            del new_state_dict[k]
+    return new_state_dict
+
+
+# =========================#
+# Text Encoder Conversion #
+# =========================#
+
+
+textenc_conversion_lst = [
+    # (stable-diffusion, HF Diffusers)
+    ("resblocks.", "text_model.encoder.layers."),
+    ("ln_1", "layer_norm1"),
+    ("ln_2", "layer_norm2"),
+    (".c_fc.", ".fc1."),
+    (".c_proj.", ".fc2."),
+    (".attn", ".self_attn"),
+    ("ln_final.", "transformer.text_model.final_layer_norm."),
+    ("token_embedding.weight", "transformer.text_model.embeddings.token_embedding.weight"),
+    ("positional_embedding", "transformer.text_model.embeddings.position_embedding.weight"),
+]
+protected = {re.escape(x[1]): x[0] for x in textenc_conversion_lst}
+textenc_pattern = re.compile("|".join(protected.keys()))
+
+# Ordering is from https://github.com/pytorch/pytorch/blob/master/test/cpp/api/modules.cpp
+code2idx = {"q": 0, "k": 1, "v": 2}
+
+
+def convert_text_enc_state_dict_v20(text_enc_dict):
+    new_state_dict = {}
+    capture_qkv_weight = {}
+    capture_qkv_bias = {}
+    for k, v in text_enc_dict.items():
+        if (
+            k.endswith(".self_attn.q_proj.weight")
+            or k.endswith(".self_attn.k_proj.weight")
+            or k.endswith(".self_attn.v_proj.weight")
+        ):
+            k_pre = k[: -len(".q_proj.weight")]
+            k_code = k[-len("q_proj.weight")]
+            if k_pre not in capture_qkv_weight:
+                capture_qkv_weight[k_pre] = [None, None, None]
+            capture_qkv_weight[k_pre][code2idx[k_code]] = v
+            continue
+
+        if (
+            k.endswith(".self_attn.q_proj.bias")
+            or k.endswith(".self_attn.k_proj.bias")
+            or k.endswith(".self_attn.v_proj.bias")
+        ):
+            k_pre = k[: -len(".q_proj.bias")]
+            k_code = k[-len("q_proj.bias")]
+            if k_pre not in capture_qkv_bias:
+                capture_qkv_bias[k_pre] = [None, None, None]
+            capture_qkv_bias[k_pre][code2idx[k_code]] = v
+            continue
+
+        relabelled_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], k)
+        new_state_dict[relabelled_key] = v
+
+    for k_pre, tensors in capture_qkv_weight.items():
+        if None in tensors:
+            raise Exception("CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing")
+        relabelled_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], k_pre)
+        new_state_dict[relabelled_key + ".in_proj_weight"] = torch.cat(tensors)
+
+    for k_pre, tensors in capture_qkv_bias.items():
+        if None in tensors:
+            raise Exception("CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing")
+        relabelled_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], k_pre)
+        new_state_dict[relabelled_key + ".in_proj_bias"] = torch.cat(tensors)
+
+    return new_state_dict
+
+
+def convert_text_enc_state_dict(text_enc_dict):
+    return text_enc_dict
+
+
+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")
+    text_enc_path = osp.join(model_path, "text_encoder", "model.safetensors")
+
+    # Load models from safetensors if it exists, if it doesn't pytorch
+    if osp.exists(unet_path):
+        unet_state_dict = load_file(unet_path, device="cpu")
+    else:
+        unet_path = osp.join(model_path, "unet", "diffusion_pytorch_model.bin")
+        unet_state_dict = torch.load(unet_path, map_location="cpu")
+
+    if osp.exists(vae_path):
+        vae_state_dict = load_file(vae_path, device="cpu")
+    else:
+        vae_path = osp.join(model_path, "vae", "diffusion_pytorch_model.bin")
+        vae_state_dict = torch.load(vae_path, map_location="cpu")
+
+    if osp.exists(text_enc_path):
+        text_enc_dict = load_file(text_enc_path, device="cpu")
+    else:
+        text_enc_path = osp.join(model_path, "text_encoder", "pytorch_model.bin")
+        text_enc_dict = torch.load(text_enc_path, map_location="cpu")
+
+    # Convert the UNet model
+    unet_state_dict = convert_unet_state_dict(unet_state_dict)
+    unet_state_dict = {"model.diffusion_model." + k: v for k, v in unet_state_dict.items()}
+
+    # Convert the VAE model
+    vae_state_dict = convert_vae_state_dict(vae_state_dict)
+    vae_state_dict = {"first_stage_model." + k: v for k, v in vae_state_dict.items()}
+
+    # Easiest way to identify v2.0 model seems to be that the text encoder (OpenCLIP) is deeper
+    is_v20_model = "text_model.encoder.layers.22.layer_norm2.bias" in text_enc_dict
+
+    if is_v20_model:
+        # Need to add the tag 'transformer' in advance so we can knock it out from the final layer-norm
+        text_enc_dict = {"transformer." + k: v for k, v in text_enc_dict.items()}
+        text_enc_dict = convert_text_enc_state_dict_v20(text_enc_dict)
+        text_enc_dict = {"cond_stage_model.model." + k: v for k, v in text_enc_dict.items()}
+    else:
+        text_enc_dict = convert_text_enc_state_dict(text_enc_dict)
+        text_enc_dict = {"cond_stage_model.transformer." + k: v for k, v in text_enc_dict.items()}
+
+    # Put together new checkpoint
+    state_dict = {**unet_state_dict, **vae_state_dict, **text_enc_dict}
+    
+    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)
+
+
+def download_repo(repo_id, dir_path):
+    from huggingface_hub import snapshot_download
+    try:
+        snapshot_download(repo_id=repo_id, local_dir=dir_path)
+    except Exception as e:
+        print(f"Error: Failed to download {repo_id}. ")
+        return
+
+
+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, dtype)
+    return output_filename
+
+
+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, args.dtype)

local/convert_repo_to_safetensors_sdxl_lora.py

@@ -0,0 +1,77 @@
+# Script for converting a Hugging Face Diffusers trained SDXL LoRAs to Kohya format
+# This means that you can input your diffusers-trained LoRAs and
+# Get the output to work with WebUIs such as AUTOMATIC1111, ComfyUI, SD.Next and others.
+
+# To get started you can find some cool `diffusers` trained LoRAs such as this cute Corgy
+# https://huggingface.co/ignasbud/corgy_dog_LoRA/, download its `pytorch_lora_weights.safetensors` file
+# and run the script:
+# python convert_diffusers_sdxl_lora_to_webui.py --input_lora pytorch_lora_weights.safetensors --output_lora corgy.safetensors
+# now you can use corgy.safetensors in your WebUI of choice!
+
+# To train your own, here are some diffusers training scripts and utils that you can use and then convert:
+# LoRA Ease - no code SDXL Dreambooth LoRA trainer: https://huggingface.co/spaces/multimodalart/lora-ease
+# Dreambooth Advanced Training Script - state of the art techniques such as pivotal tuning and prodigy optimizer:
+# - Script: https://github.com/huggingface/diffusers/blob/main/examples/advanced_diffusion_training/train_dreambooth_lora_sdxl_advanced.py
+# - Colab (only on Pro): https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/SDXL_Dreambooth_LoRA_advanced_example.ipynb
+# Canonical diffusers training scripts:
+# - Script: https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth_lora_sdxl.py
+# - Colab (runs on free tier): https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/SDXL_DreamBooth_LoRA_.ipynb
+
+import argparse
+import os
+
+from safetensors.torch import load_file, save_file
+from diffusers.utils import convert_all_state_dict_to_peft, convert_state_dict_to_kohya
+from pathlib import Path
+
+def convert_and_save(input_lora, output_lora=None):
+    if output_lora is None:
+        base_name = os.path.splitext(input_lora)[0]
+        output_lora = f"{base_name}_webui.safetensors"
+
+    diffusers_state_dict = load_file(input_lora)
+    try:
+        peft_state_dict = convert_all_state_dict_to_peft(diffusers_state_dict)
+        kohya_state_dict = convert_state_dict_to_kohya(peft_state_dict)
+    except Exception: # skipped
+        kohya_state_dict = diffusers_state_dict
+    save_file(kohya_state_dict, output_lora)
+
+
+def download_repo_lora(repo_id, local_file):
+    from huggingface_hub import hf_hub_download, HfApi
+    lora_filename = "pytorch_lora_weights.safetensors"
+    lora_path = Path(lora_filename)
+    api = HfApi()
+    try:
+        if not api.file_exists(repo_id=repo_id, filename=lora_filename):
+            print(f"Error: This repo isn't diffusers LoRA repo: {repo_id}. ")
+            return None
+        if lora_path.exists():
+            print(f"Error: Download file already exists: {lora_filename}. ")
+            return None
+        hf_hub_download(repo_id=repo_id, filename=lora_filename, local_dir=".")
+        if lora_path.exists(): lora_path.rename(Path(local_file))
+    except Exception as e:
+        print(f"Error: Failed to download from {repo_id}. {e}")
+        return local_file
+
+
+def convert_repo_to_safetensors_sdxl_lora(repo_id):
+    download_filename = f"{repo_id.split('/')[0]}_{repo_id.split('/')[-1]}_diffusers.safetensors"
+    output_filename = f"{repo_id.split('/')[0]}_{repo_id.split('/')[-1]}_webui.safetensors"
+    download_repo_lora(repo_id, download_filename)
+    convert_and_save(download_filename, output_filename)
+    return output_filename
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Convert LoRA model to PEFT and then to Kohya format from Repo.")
+    parser.add_argument("--repo_id", type=str, required=True, help="URL to the Repo of input LoRA model in the diffusers format.")
+
+    args = parser.parse_args()
+
+    convert_repo_to_safetensors_sdxl_lora(args.repo_id)
+
+
+# Usage: python convert_repo_to_safetensors_sdxl_lora.py --repo_id nroggendorff/zelda-lora

local/requirements.txt

@@ -1,3 +1,7 @@
 torch
 safetensors
-huggingface-hub
+huggingface-hub
+accelerate
+diffusers
+transformers
+peft

packages.txt

@@ -0,0 +1 @@
+git-lfs aria2

requirements.txt

@@ -1,3 +1,7 @@
 torch
 safetensors
-huggingface-hub
+huggingface-hub
+accelerate
+diffusers
+transformers
+peft

utils.py

@@ -0,0 +1,161 @@
+import gradio as gr
+from huggingface_hub import HfApi, HfFolder, hf_hub_download
+import os
+from pathlib import Path
+import shutil
+import gc
+import re
+import urllib.parse
+
+
+def get_token():
+    try:
+        token = HfFolder.get_token()
+    except Exception:
+        token = ""
+    return token
+
+
+def set_token(token):
+    try:
+        HfFolder.save_token(token)
+    except Exception:
+        print(f"Error: Failed to save token.")
+
+
+def get_user_agent():
+    return 'Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:127.0) Gecko/20100101 Firefox/127.0'
+
+
+def is_repo_exists(repo_id: str, repo_type: str="model"):
+    hf_token = get_token()
+    api = HfApi(token=hf_token)
+    try:
+        if api.repo_exists(repo_id=repo_id, repo_type=repo_type, token=hf_token): return True
+        else: return False
+    except Exception as e:
+        print(f"Error: Failed to connect {repo_id} ({repo_type}). {e}")
+        return True # for safe
+
+
+MODEL_TYPE_CLASS = {
+    "diffusers:StableDiffusionPipeline": "SD 1.5",
+    "diffusers:StableDiffusionXLPipeline": "SDXL",
+    "diffusers:FluxPipeline": "FLUX",
+}
+
+
+def get_model_type(repo_id: str):
+    hf_token = get_token()
+    api = HfApi(token=hf_token)
+    lora_filename = "pytorch_lora_weights.safetensors"
+    diffusers_filename = "model_index.json"
+    default = "SDXL"
+    try:
+        if api.file_exists(repo_id=repo_id, filename=lora_filename, token=hf_token): return "LoRA"
+        if not api.file_exists(repo_id=repo_id, filename=diffusers_filename, token=hf_token): return "None"
+        model = api.model_info(repo_id=repo_id, token=hf_token)
+        tags = model.tags
+        for tag in tags:
+            if tag in MODEL_TYPE_CLASS.keys(): return MODEL_TYPE_CLASS.get(tag, default)
+    except Exception:
+        return default
+    return default
+
+
+def list_sub(a, b):
+    return [e for e in a if e not in b]
+
+
+def is_repo_name(s):
+    return re.fullmatch(r'^[^/,\s\"\']+/[^/,\s\"\']+$', s)
+
+
+def split_hf_url(url: str):
+    try:
+        s = list(re.findall(r'^(?:https?://huggingface.co/)(?:(datasets)/)?(.+?/.+?)/\w+?/.+?/(?:(.+)/)?(.+?.safetensors)(?:\?download=true)?$', url)[0])
+        if len(s) < 4: return "", "", "", ""
+        repo_id = s[1]
+        repo_type = "dataset" if s[0] == "datasets" else "model"
+        subfolder = urllib.parse.unquote(s[2]) if s[2] else None
+        filename = urllib.parse.unquote(s[3])
+        return repo_id, filename, subfolder, repo_type
+    except Exception as e:
+        print(e)
+
+
+def download_hf_file(directory, url, progress=gr.Progress(track_tqdm=True)):
+    hf_token = get_token()
+    repo_id, filename, subfolder, repo_type = split_hf_url(url)
+    try:
+        if subfolder is not None: hf_hub_download(repo_id=repo_id, filename=filename, subfolder=subfolder, repo_type=repo_type, local_dir=directory, token=hf_token)
+        else: hf_hub_download(repo_id=repo_id, filename=filename, repo_type=repo_type, local_dir=directory, token=hf_token)
+    except Exception as e:
+        print(f"Failed to download: {e}")
+
+
+def download_thing(directory, url, civitai_api_key="", progress=gr.Progress(track_tqdm=True)): # requires aria2, gdown
+    hf_token = get_token()
+    url = url.strip()
+    if "drive.google.com" in url:
+        original_dir = os.getcwd()
+        os.chdir(directory)
+        os.system(f"gdown --fuzzy {url}")
+        os.chdir(original_dir)
+    elif "huggingface.co" in url:
+        url = url.replace("?download=true", "")
+        if "/blob/" in url:
+            url = url.replace("/blob/", "/resolve/")
+        #user_header = f'"Authorization: Bearer {hf_token}"'
+        if hf_token:
+            download_hf_file(directory, url)
+            #os.system(f"aria2c --console-log-level=error --summary-interval=10 --header={user_header} -c -x 16 -k 1M -s 16 {url} -d {directory}  -o {url.split('/')[-1]}")
+        else:
+            os.system(f"aria2c --optimize-concurrent-downloads --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 {url} -d {directory}  -o {url.split('/')[-1]}")
+    elif "civitai.com" in url:
+        if "?" in url:
+            url = url.split("?")[0]
+        if civitai_api_key:
+            url = url + f"?token={civitai_api_key}"
+            os.system(f"aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d {directory} {url}")
+        else:
+            print("You need an API key to download Civitai models.")
+    else:
+        os.system(f"aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d {directory} {url}")
+
+
+def get_local_model_list(dir_path):
+    model_list = []
+    valid_extensions = ('.safetensors')
+    for file in Path(dir_path).glob("**/*.*"):
+        if file.is_file() and file.suffix in valid_extensions:
+            file_path = str(file)
+            model_list.append(file_path)
+    return model_list
+
+
+def get_download_file(temp_dir, url, civitai_key, progress=gr.Progress(track_tqdm=True)):
+    if not "http" in url and is_repo_name(url) and not Path(url).exists():
+        print(f"Use HF Repo: {url}")
+        new_file = url
+    elif not "http" in url and Path(url).exists():
+        print(f"Use local file: {url}")
+        new_file = url
+    elif Path(f"{temp_dir}/{url.split('/')[-1]}").exists():
+        print(f"File to download alreday exists: {url}")
+        new_file = f"{temp_dir}/{url.split('/')[-1]}"
+    else:
+        print(f"Start downloading: {url}")
+        before = get_local_model_list(temp_dir)
+        try:
+            download_thing(temp_dir, url.strip(), civitai_key)
+        except Exception:
+            print(f"Download failed: {url}")
+            return ""
+        after = get_local_model_list(temp_dir)
+        new_file = list_sub(after, before)[0] if list_sub(after, before) else ""
+    if not new_file:
+        print(f"Download failed: {url}")
+        return ""
+    print(f"Download completed: {url}")
+    return new_file