Upload convert_url_to_diffusers_sd_gr.py

convert_url_to_diffusers_sd_gr.py

@@ -16,7 +16,32 @@ def is_repo_name(s):
     return re.fullmatch(r'^[^/,\s\"\']+/[^/,\s\"\']+$', s)
 
 
-def download_thing(directory, url, civitai_api_key="", progress=gr.Progress(track_tqdm=True)):
+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])
+        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)):
+    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)
+        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()
@@ -27,9 +52,12 @@ def download_thing(directory, url, civitai_api_key="", progress=gr.Progress(trac
         url = url.replace("?download=true", "")
         if "/blob/" in url:
             url = url.replace("/blob/", "/resolve/")
-            os.system(f"aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 {url} -d {directory}  -o {url.split('/')[-1]}")
+        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]}")
+            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]
@@ -44,15 +72,15 @@ def download_thing(directory, url, civitai_api_key="", progress=gr.Progress(trac
 
 def get_local_model_list(dir_path):
     model_list = []
-    valid_extensions = ('.safetensors', '.ckpt', '.bin', '.pt', '.pth')
-    for file in Path(dir_path).glob("*"):
-        if file.suffix in valid_extensions:
-            file_path = str(Path(f"{dir_path}/{file.name}"))
+    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)):
+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
@@ -66,7 +94,7 @@ def get_download_file(temp_dir, url, civitai_key, progress=gr.Progress(track_tqd
         print(f"Start downloading: {url}")
         before = get_local_model_list(temp_dir)
         try:
-            download_thing(temp_dir, url.strip(), civitai_key)
+            download_thing(temp_dir, url.strip(), civitai_key, hf_token)
         except Exception:
             print(f"Download failed: {url}")
             return ""
@@ -215,11 +243,11 @@ def fuse_loras(pipe, lora_dict={}, temp_dir=".", civitai_key=""):
     pipe.unload_lora_weights()
 
 
-def convert_url_to_diffusers_sd(url, civitai_key="", is_upload_sf=False, half=True, vae=None, scheduler="Euler", lora_dict={},
+def convert_url_to_diffusers_sd(url, civitai_key="", hf_token="", is_upload_sf=False, half=True, vae=None, scheduler="Euler", lora_dict={},
                                   model_type="v1", sample_size=768, ema="ema", 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)
+    new_file = get_download_file(temp_dir, url, civitai_key, hf_token)
     if not new_file:
         print(f"Not found: {url}")
         return ""
@@ -283,23 +311,21 @@ def convert_url_to_diffusers_sd(url, civitai_key="", is_upload_sf=False, half=Tr
     return new_repo_name
 
 
-def is_repo_exists(repo_id):
+def is_repo_exists(repo_id, hf_token):
     from huggingface_hub import HfApi
-    api = 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(f"Error: Failed to connect {repo_id}. ")
         print(e)
+        print(f"Error: Failed to connect {repo_id}.")
         return True # for safe
 
 
-def create_diffusers_repo(new_repo_id, diffusers_folder, is_private, progress=gr.Progress(track_tqdm=True)):
+def create_diffusers_repo(new_repo_id, diffusers_folder, is_private, hf_token, progress=gr.Progress(track_tqdm=True)):
     from huggingface_hub import HfApi
-    import os
-    hf_token = os.environ.get("HF_TOKEN")
-    api = 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)
@@ -311,7 +337,7 @@ def create_diffusers_repo(new_repo_id, diffusers_folder, is_private, progress=gr
         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(f"Error: Failed to upload to {new_repo_id}.")
         print(e)
         return ""
     return url
@@ -328,7 +354,7 @@ def convert_url_to_diffusers_repo_sd(dl_url, hf_user, hf_repo, hf_token, civitai
         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")
     lora_dict = {lora1: lora1s, lora2: lora2s, lora3: lora3s, lora4: lora4s, lora5: lora5s}
-    new_path = convert_url_to_diffusers_sd(dl_url, civitai_key, is_upload_sf, half, vae, scheduler, lora_dict,
+    new_path = convert_url_to_diffusers_sd(dl_url, civitai_key, hf_token, is_upload_sf, half, vae, scheduler, lora_dict,
                                             model_type, sample_size, ema, False)
     if not new_path: return gr.update(value=repo_urls, choices=repo_urls), gr.update(value="")
     new_repo_id = f"{hf_user}/{Path(new_path).stem}"
@@ -337,11 +363,11 @@ def convert_url_to_diffusers_repo_sd(dl_url, hf_user, hf_repo, hf_token, civitai
         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):
+    if 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="")
-    repo_url = create_diffusers_repo(new_repo_id, new_path, is_private)
+    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)