Upload convert_url_to_diffusers_sdxl_gr.py

convert_url_to_diffusers_sdxl_gr.py

@@ -1,3 +1,13 @@
+import os
+if os.environ.get("SPACES_ZERO_GPU") is not None:
+    import spaces
+else:
+    class spaces:
+        @staticmethod
+        def GPU(func):
+            def wrapper(*args, **kwargs):
+                return func(*args, **kwargs)
+            return wrapper
 import argparse
 from pathlib import Path
 import os
@@ -5,108 +15,30 @@ 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
+from utils import (set_token, is_repo_exists, is_repo_name, get_download_file, upload_repo)
 
 
-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+?/.+?/(?:(.+)/)?(.+?.\w+)(?:\?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, hf_token="", progress=gr.Progress(track_tqdm=True)):
-    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="", hf_token="", progress=gr.Progress(track_tqdm=True)):
-    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, hf_token)
-            #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, hf_token, 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, hf_token)
-        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
+@spaces.GPU
+def fake_gpu():
+    pass
+
+
+TEMP_DIR = "."
+
+
+DTYPE_DICT = {
+    "fp16": torch.float16,
+    "bf16": torch.bfloat16,
+    "fp32": torch.float32,
+    "fp8": torch.float8_e4m3fn
+}
+
+
+def get_dtype(dtype: str):
+    return DTYPE_DICT.get(dtype, torch.float16)
 
 
 from diffusers import (
@@ -224,19 +156,19 @@ tags:
         f.write(md)
 
 
-def fuse_loras(pipe, lora_dict={}, temp_dir=".", civitai_key="", hf_token=""):
+def fuse_loras(pipe, lora_dict={}, temp_dir=TEMP_DIR, civitai_key=""):
     if not lora_dict or not isinstance(lora_dict, dict): return pipe
     a_list = []
     w_list = []
     for k, v in lora_dict.items():
         if not k: continue
-        new_lora_file = get_download_file(temp_dir, k, civitai_key, hf_token)
+        new_lora_file = get_download_file(temp_dir, k, civitai_key)
         if not new_lora_file or not Path(new_lora_file).exists():
             print(f"LoRA not found: {k}")
             continue
         w_name = Path(new_lora_file).name
         a_name = Path(new_lora_file).stem
-        pipe.load_lora_weights(new_lora_file, weight_name = w_name, adapter_name = a_name)
+        pipe.load_lora_weights(new_lora_file, weight_name=w_name, adapter_name=a_name)
         a_list.append(a_name)
         w_list.append(v)
     if not a_list: return pipe
@@ -246,121 +178,84 @@ 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, dtype="fp16", vae="",
-                                  scheduler="Euler a", lora_dict={}, is_local=True, clip="", progress=gr.Progress(track_tqdm=True)):
+def convert_url_to_diffusers_sdxl(url, civitai_key="", is_upload_sf=False, dtype="fp16", vae="", clip="",
+                                  scheduler="Euler a", lora_dict={}, is_local=True, 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)
+    temp_dir = TEMP_DIR
+    new_file = get_download_file(temp_dir, url, civitai_key)
     if not new_file:
         print(f"Not found: {url}")
         return ""
-    new_repo_name = Path(new_file).stem.replace(" ", "_").replace(",", "_").replace(".", "_") #
+    new_dir = Path(new_file).stem.replace(" ", "_").replace(",", "_").replace(".", "_") #
 
-    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
+    type_kwargs = {}
+    if dtype != "default": type_kwargs["torch_dtype"] = get_dtype(dtype)
 
     new_vae_file = ""
     if 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: my_vae = AutoencoderKL.from_single_file(new_vae_file, **type_kwargs)
-        kwargs["vae"] = my_vae
+            new_vae_file = get_download_file(temp_dir, vae, civitai_key)
+            my_vae = AutoencoderKL.from_single_file(new_vae_file, **type_kwargs) if new_vae_file else None
+        if my_vae: kwargs["vae"] = my_vae
 
     if clip:
         my_tokenizer = CLIPTokenizer.from_pretrained(clip)
+        if my_tokenizer: kwargs["tokenizer"] = my_tokenizer
         my_text_encoder = CLIPTextModel.from_pretrained(clip, **type_kwargs)
-        kwargs["tokenizer"] = my_tokenizer
-        kwargs["text_encoder"] = my_text_encoder
+        if my_text_encoder: 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)
+    pipe = fuse_loras(pipe, lora_dict, temp_dir, civitai_key)
 
     sconf = get_scheduler_config(scheduler)
     pipe.scheduler = sconf[0].from_config(pipe.scheduler.config, **sconf[1])
 
-    pipe.save_pretrained(new_repo_name, safe_serialization=True, use_safetensors=True)
+    pipe.save_pretrained(new_dir, safe_serialization=True, use_safetensors=True)
 
-    if Path(new_repo_name).exists(): save_readme_md(new_repo_name, url)
+    if Path(new_dir).exists(): save_readme_md(new_dir, url)
 
     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()))
+        if not is_repo_name(new_file) and is_upload_sf: shutil.move(str(Path(new_file).resolve()), str(Path(new_dir, 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):
-    api = HfApi(token=hf_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
-
-
-def create_diffusers_repo(new_repo_id, diffusers_folder, is_private, hf_token, progress=gr.Progress(track_tqdm=True)):
-    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, 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)
-            elif path.is_file():
-                api.upload_file(repo_id=new_repo_id, path_or_fileobj=str(path), path_in_repo=path.name, token=hf_token)
-        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}. {e}")
-        return ""
-    return url
+    return new_dir
 
 
 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",
+                                  urls=[], dtype="fp16", vae="", 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)):
+    is_local = False
     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(visible=True)
+    if not hf_user: raise gr.Error(f"Invalid user name: {hf_user}")
+    if not hf_repo and os.environ.get("HF_REPO"): hf_repo = os.environ.get("HF_REPO") # default reponame
+    set_token(hf_token)
     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, dtype, vae, scheduler, lora_dict, False, clip)
+    new_path = convert_url_to_diffusers_sdxl(dl_url, civitai_key, is_upload_sf, dtype, vae, clip, scheduler, lora_dict, is_local)
     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(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(visible=True)
-    repo_url = create_diffusers_repo(new_repo_id, new_path, is_private, hf_token)
+    if not is_repo_name(new_repo_id): raise gr.Error(f"Invalid repo name: {new_repo_id}")
+    if not is_overwrite and is_repo_exists(new_repo_id): raise gr.Error(f"Repo already exists: {new_repo_id}")
+    repo_url = upload_repo(new_repo_id, new_path, is_private)
     shutil.rmtree(new_path)
-    if not repo_urls: repo_urls = []
-    repo_urls.append(repo_url)
+    if not urls: urls = []
+    urls.append(repo_url)
     md = "### Your new repo:\n"
-    for u in repo_urls:
+    for u in 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)
+    return gr.update(value=urls, choices=urls), gr.update(value=md)
 
 
 if __name__ == "__main__":
@@ -386,10 +281,11 @@ if __name__ == "__main__":
     args = parser.parse_args()
     assert args.url is not None, "Must provide a URL!"
 
+    is_local = True
     lora_dict = {args.lora1: args.lora1s, args.lora2: args.lora2s, args.lora3: args.lora3s, args.lora4: args.lora4s, args.lora5: args.lora5s}
-
     if args.loras and Path(args.loras).exists():
         for p in Path(args.loras).glob('**/*.safetensors'):
             lora_dict[str(p)] = 1.0
+    clip = ""
 
-    convert_url_to_diffusers_sdxl(args.url, args.civitai_key, args.dtype, args.vae, args.scheduler, lora_dict, True)
+    convert_url_to_diffusers_sdxl(args.url, args.civitai_key, args.dtype, args.vae, clip, args.scheduler, lora_dict, is_local)