Upload convert_url_to_diffusers_sdxl_gr.py

convert_url_to_diffusers_sdxl_gr.py

@@ -16,6 +16,31 @@ def is_repo_name(s):
     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])
+        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:
@@ -29,7 +54,8 @@ def download_thing(directory, url, civitai_api_key="", hf_token="", progress=gr.
             url = url.replace("/blob/", "/resolve/")
         user_header = f'"Authorization: Bearer {hf_token}"'
         if 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]}")
+            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:
@@ -47,9 +73,9 @@ def download_thing(directory, url, civitai_api_key="", hf_token="", progress=gr.
 def get_local_model_list(dir_path):
     model_list = []
     valid_extensions = ('.safetensors')
-    for file in Path(dir_path).glob("*"):
-        if file.suffix in valid_extensions:
-            file_path = str(Path(f"{dir_path}/{file.name}"))
+    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
 
@@ -287,11 +313,9 @@ def is_repo_exists(repo_id, hf_token):
         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)
@@ -303,7 +327,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
@@ -331,7 +355,7 @@ def convert_url_to_diffusers_repo(dl_url, hf_user, hf_repo, hf_token, civitai_ke
         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)