Upload 4 files

convert_url_to_diffusers_sdxl_gr.py

@@ -3,7 +3,13 @@ from pathlib import Path
 import os
 import torch
 from diffusers import StableDiffusionXLPipeline, AutoencoderKL
+from transformers import CLIPTokenizer, CLIPTextModel
 import gradio as gr
+from huggingface_hub import hf_hub_download, HfApi
+import urllib.parse
+import re
+import shutil
+import gc
 # also requires aria, gdown, peft, huggingface_hub, safetensors, transformers, accelerate, pytorch_lightning
 
 
@@ -12,15 +18,12 @@ def list_sub(a, b):
 
 
 def is_repo_name(s):
-    import re
     return re.fullmatch(r'^[^/,\s\"\']+/[^/,\s\"\']+$', s)
 
 
 def split_hf_url(url: str):
-    import re
-    import urllib.parse
     try:
-        s = list(re.findall(r'^(?:https?://huggingface.co/)(?:(datasets)/)?(.+?/.+?)/\w+?/.+?/(?:(.+)/)?(.+?.safetensors)(?:\?download=true)?$', url)[0])
+        s = list(re.findall(r'^(?:https?://huggingface.co/)(?:(datasets)/)?(.+?/.+?)/\w+?/.+?/(?:(.+)/)?(.+?.\w+)(?:\?download=true)?$', url)[0])
         if len(s) < 4: return "", "", "", ""
         repo_id = s[1]
         repo_type = "dataset" if s[0] == "datasets" else "model"
@@ -32,7 +35,6 @@ def split_hf_url(url: str):
 
 
 def download_hf_file(directory, url, hf_token="", progress=gr.Progress(track_tqdm=True)):
-    from huggingface_hub import hf_hub_download
     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)
@@ -244,8 +246,8 @@ def fuse_loras(pipe, lora_dict={}, temp_dir=".", civitai_key="", hf_token=""):
     return pipe
 
 
-def convert_url_to_diffusers_sdxl(url, civitai_key="", hf_token="", is_upload_sf=False, half=True, vae=None,
-                                  scheduler="Euler a", lora_dict={}, is_local=True, progress=gr.Progress(track_tqdm=True)):
+def convert_url_to_diffusers_sdxl(url, civitai_key="", hf_token="", is_upload_sf=False, dtype="fp16", vae="",
+                                  scheduler="Euler a", lora_dict={}, is_local=True, clip="", progress=gr.Progress(track_tqdm=True)):
     progress(0, desc="Start converting...")
     temp_dir = "."
     new_file = get_download_file(temp_dir, url, civitai_key, hf_token)
@@ -254,72 +256,66 @@ def convert_url_to_diffusers_sdxl(url, civitai_key="", hf_token="", is_upload_sf
         return ""
     new_repo_name = Path(new_file).stem.replace(" ", "_").replace(",", "_").replace(".", "_") #
 
-    pipe = None
-    if is_repo_name(url):
-        if half:
-            pipe = StableDiffusionXLPipeline.from_pretrained(new_file, use_safetensors=True, torch_dtype=torch.float16)
-        else:
-            pipe = StableDiffusionXLPipeline.from_pretrained(new_file, use_safetensors=True)
-    else:
-        if half:
-            pipe = StableDiffusionXLPipeline.from_single_file(new_file, use_safetensors=True, torch_dtype=torch.float16)
-        else:
-            pipe = StableDiffusionXLPipeline.from_single_file(new_file, use_safetensors=True)
+    type_kwargs = {}
+    kwargs = {}
+    if dtype == "fp16": type_kwargs["torch_dtype"] = torch.float16
+    elif dtype == "fp32": type_kwargs["torch_dtype"] = torch.float32
+    elif dtype == "bf16": type_kwargs["torch_dtype"] = torch.bfloat16
+    elif dtype == "fp8": type_kwargs["torch_dtype"] = torch.float8_e4m3fn
 
     new_vae_file = ""
     if vae:
-        if is_repo_name(vae):
-            if half:
-                pipe.vae = AutoencoderKL.from_pretrained(vae, torch_dtype=torch.float16)
-            else:
-                pipe.vae = AutoencoderKL.from_pretrained(vae)
+        if is_repo_name(vae): my_vae = AutoencoderKL.from_pretrained(vae, **type_kwargs)
         else:
             new_vae_file = get_download_file(temp_dir, vae, civitai_key, hf_token)
-            if new_vae_file and half:
-                pipe.vae = AutoencoderKL.from_single_file(new_vae_file, torch_dtype=torch.float16)
-            elif new_vae_file:
-                pipe.vae = AutoencoderKL.from_single_file(new_vae_file)
+            if new_vae_file: my_vae = AutoencoderKL.from_single_file(new_vae_file, **type_kwargs)
+        kwargs["vae"] = my_vae
+
+    if clip:
+        my_tokenizer = CLIPTokenizer.from_pretrained(clip)
+        my_text_encoder = CLIPTextModel.from_pretrained(clip, **type_kwargs)
+        kwargs["tokenizer"] = my_tokenizer
+        kwargs["text_encoder"] = my_text_encoder
+
+    pipe = None
+    if is_repo_name(url): pipe = StableDiffusionXLPipeline.from_pretrained(new_file, use_safetensors=True, **kwargs, **type_kwargs)
+    else: pipe = StableDiffusionXLPipeline.from_single_file(new_file, use_safetensors=True, **kwargs, **type_kwargs)
 
     pipe = fuse_loras(pipe, lora_dict, temp_dir, civitai_key, hf_token)
 
     sconf = get_scheduler_config(scheduler)
     pipe.scheduler = sconf[0].from_config(pipe.scheduler.config, **sconf[1])
 
-    if half:
-        pipe.save_pretrained(new_repo_name, safe_serialization=True, use_safetensors=True)
-    else:
-        pipe.save_pretrained(new_repo_name, safe_serialization=True, use_safetensors=True)
+    pipe.save_pretrained(new_repo_name, safe_serialization=True, use_safetensors=True)
 
-    if Path(new_repo_name).exists():
-        save_readme_md(new_repo_name, url)
+    if Path(new_repo_name).exists(): save_readme_md(new_repo_name, url)
 
-    if not is_repo_name(new_file) and is_upload_sf:
-        import shutil
-        shutil.move(str(Path(new_file).resolve()), str(Path(new_repo_name, Path(new_file).name).resolve()))
-    elif not is_local: os.remove(new_file)
+    if not is_local:
+        if not is_repo_name(new_file) and is_upload_sf: shutil.move(str(Path(new_file).resolve()), str(Path(new_repo_name, Path(new_file).name).resolve()))
+        else: os.remove(new_file)
+    del pipe
+    torch.cuda.empty_cache()
+    gc.collect()
 
     progress(1, desc="Converted.")
     return new_repo_name
 
 
 def is_repo_exists(repo_id, hf_token):
-    from huggingface_hub import HfApi
     api = HfApi(token=hf_token)
     try:
         if api.repo_exists(repo_id=repo_id): return True
         else: return False
     except Exception as e:
-        print(e)
-        print(f"Error: Failed to connect {repo_id}.")
+        print(f"Error: Failed to connect {repo_id}. {e}")
         return True # for safe
 
 
 def create_diffusers_repo(new_repo_id, diffusers_folder, is_private, hf_token, progress=gr.Progress(track_tqdm=True)):
-    from huggingface_hub import HfApi
     api = HfApi(token=hf_token)
     try:
         progress(0, desc="Start uploading...")
-        api.create_repo(repo_id=new_repo_id, token=hf_token, private=is_private)
+        api.create_repo(repo_id=new_repo_id, token=hf_token, private=is_private, exist_ok=True)
         for path in Path(diffusers_folder).glob("*"):
             if path.is_dir():
                 api.upload_folder(repo_id=new_repo_id, folder_path=str(path), path_in_repo=path.name, token=hf_token)
@@ -328,39 +324,40 @@ def create_diffusers_repo(new_repo_id, diffusers_folder, is_private, hf_token, p
         progress(1, desc="Uploaded.")
         url = f"https://huggingface.co/{new_repo_id}"
     except Exception as e:
-        print(f"Error: Failed to upload to {new_repo_id}.")
-        print(e)
+        print(f"Error: Failed to upload to {new_repo_id}. {e}")
         return ""
     return url
 
 
-def convert_url_to_diffusers_repo(dl_url, hf_user, hf_repo, hf_token, civitai_key="", is_private=True, is_upload_sf=False, repo_urls=[], half=True, vae=None,
-                                  scheduler="Euler a", lora1=None, lora1s=1.0, lora2=None, lora2s=1.0, lora3=None, lora3s=1.0,
+def convert_url_to_diffusers_repo(dl_url, hf_user, hf_repo, hf_token, civitai_key="", is_private=True, is_overwrite=False, is_upload_sf=False,
+                                  repo_urls=[], dtype="fp16", vae=None, clip="", scheduler="Euler a",
+                                  lora1=None, lora1s=1.0, lora2=None, lora2s=1.0, lora3=None, lora3s=1.0,
                                   lora4=None, lora4s=1.0, lora5=None, lora5s=1.0, progress=gr.Progress(track_tqdm=True)):
-    import shutil
+    if not civitai_key and os.environ.get("CIVITAI_API_KEY"): civitai_key = os.environ.get("CIVITAI_API_KEY") # default Civitai API key
+    if not hf_token and os.environ.get("HF_TOKEN"): hf_token = os.environ.get("HF_TOKEN") # default HF write token
+    if not hf_user and os.environ.get("HF_USER"): hf_user = os.environ.get("HF_USER") # default username
     if not hf_user:
         print(f"Invalid user name: {hf_user}")
         progress(1, desc=f"Invalid user name: {hf_user}")
-        return gr.update(value=repo_urls, choices=repo_urls), gr.update(value="")
-    if not civitai_key and os.environ.get("CIVITAI_API_KEY"): civitai_key = os.environ.get("CIVITAI_API_KEY")
+        return gr.update(value=repo_urls, choices=repo_urls), gr.update(visible=True)
     lora_dict = {lora1: lora1s, lora2: lora2s, lora3: lora3s, lora4: lora4s, lora5: lora5s}
-    new_path = convert_url_to_diffusers_sdxl(dl_url, civitai_key, hf_token, is_upload_sf, half, vae, scheduler, lora_dict, False)
+    new_path = convert_url_to_diffusers_sdxl(dl_url, civitai_key, hf_token, is_upload_sf, dtype, vae, scheduler, lora_dict, False, clip)
     if not new_path: return ""
     new_repo_id = f"{hf_user}/{Path(new_path).stem}"
     if hf_repo != "": new_repo_id = f"{hf_user}/{hf_repo}"
     if not is_repo_name(new_repo_id):
         print(f"Invalid repo name: {new_repo_id}")
         progress(1, desc=f"Invalid repo name: {new_repo_id}")
-        return gr.update(value=repo_urls, choices=repo_urls), gr.update(value="")
-    if is_repo_exists(new_repo_id, hf_token):
+        return gr.update(value=repo_urls, choices=repo_urls), gr.update(visible=True)
+    if not is_overwrite and is_repo_exists(new_repo_id, hf_token):
         print(f"Repo already exists: {new_repo_id}")
         progress(1, desc=f"Repo already exists: {new_repo_id}")
-        return gr.update(value=repo_urls, choices=repo_urls), gr.update(value="")
+        return gr.update(value=repo_urls, choices=repo_urls), gr.update(visible=True)
     repo_url = create_diffusers_repo(new_repo_id, new_path, is_private, hf_token)
     shutil.rmtree(new_path)
     if not repo_urls: repo_urls = []
     repo_urls.append(repo_url)
-    md = "Your new repo:<br>"
+    md = "### Your new repo:\n"
     for u in repo_urls:
         md += f"[{str(u).split('/')[-2]}/{str(u).split('/')[-1]}]({str(u)})<br>"
     return gr.update(value=repo_urls, choices=repo_urls), gr.update(value=md)
@@ -370,7 +367,7 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser()
 
     parser.add_argument("--url", default=None, type=str, required=True, help="URL 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")')
     parser.add_argument("--scheduler", default="Euler a", type=str, choices=list(SCHEDULER_CONFIG_MAP.keys()), required=False, help="Scheduler name to use.")
     parser.add_argument("--vae", default=None, type=str, required=False, help="URL of the VAE to use.")
     parser.add_argument("--civitai_key", default=None, type=str, required=False, help="Civitai API Key (If you want to download file from Civitai).")
@@ -395,4 +392,4 @@ if __name__ == "__main__":
         for p in Path(args.loras).glob('**/*.safetensors'):
             lora_dict[str(p)] = 1.0
 
-    convert_url_to_diffusers_sdxl(args.url, args.civitai_key, args.half, args.vae, args.scheduler, lora_dict, True)
+    convert_url_to_diffusers_sdxl(args.url, args.civitai_key, args.dtype, args.vae, args.scheduler, lora_dict, True)