Upload modutils.py

modutils.py

@@ -1,10 +1,12 @@
 import spaces
 import json
 import gradio as gr
-from huggingface_hub import HfApi
 import os
 from pathlib import Path
 from PIL import Image
+from huggingface_hub import HfApi, HfFolder, hf_hub_download, snapshot_download
+import urllib.parse
+import re
 
 
 from env import (HF_LORA_PRIVATE_REPOS1, HF_LORA_PRIVATE_REPOS2,
@@ -61,6 +63,48 @@ def get_local_model_list(dir_path):
     return model_list
 
 
+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.")
+
+
+set_token(HF_TOKEN)
+
+
+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, progress=gr.Progress(track_tqdm=True)):
+    hf_token = get_token()
+    repo_id, filename, subfolder, repo_type = split_hf_url(url)
+    try:
+        print(f"Downloading {url} to {directory}")
+        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_things(directory, url, hf_token="", civitai_api_key=""):
     url = url.strip()
     if "drive.google.com" in url:
@@ -73,11 +117,7 @@ def download_things(directory, url, hf_token="", civitai_api_key=""):
         # url = urllib.parse.quote(url, safe=':/')  # fix encoding
         if "/blob/" in url:
             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]}")
-        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]}")
+        download_hf_file(directory, url)
     elif "civitai.com" in url:
         if "?" in url:
             url = url.split("?")[0]
@@ -90,6 +130,33 @@ def download_things(directory, url, hf_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}")
 
 
+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_things(temp_dir, url.strip(), HF_TOKEN, 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
+
+
 def escape_lora_basename(basename: str):
     return basename.replace(".", "_").replace(" ", "_").replace(",", "")