Upload utils.py

utils.py

@@ -8,6 +8,7 @@ import re
 import urllib.parse
 import subprocess
 import time
+from typing import Any
 
 
 def get_token():
@@ -25,6 +26,17 @@ def set_token(token):
         print(f"Error: Failed to save token.")
 
 
+def get_state(state: dict, key: str):
+    if key in state.keys(): return state[key]
+    else:
+        print(f"State '{key}' not found.")
+        return None
+
+
+def set_state(state: dict, key: str, value: Any):
+    state[key] = value
+
+
 def get_user_agent():
     return 'Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:127.0) Gecko/20100101 Firefox/127.0'
 
@@ -113,31 +125,33 @@ def download_hf_file(directory, url, progress=gr.Progress(track_tqdm=True)):
 
 
 def download_thing(directory, url, civitai_api_key="", progress=gr.Progress(track_tqdm=True)): # requires aria2, gdown
-    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/")
-        download_hf_file(directory, url)
-    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}")
+    try:
+        url = url.strip()
+        if "drive.google.com" in url:
+            original_dir = os.getcwd()
+            os.chdir(directory)
+            subprocess.run(f"gdown --fuzzy {url}", shell=True)
+            os.chdir(original_dir)
+        elif "huggingface.co" in url:
+            url = url.replace("?download=true", "")
+            if "/blob/" in url: url = url.replace("/blob/", "/resolve/")
+            download_hf_file(directory, url)
+        elif "civitai.com" in url:
+            if civitai_api_key:
+                url = f"'{url}&token={civitai_api_key}'" if "?" in url else f"{url}?token={civitai_api_key}"
+                print(f"Downloading {url}")
+                subprocess.run(f"aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d {directory} {url}", shell=True)
+            else:
+                print("You need an API key to download Civitai models.")
         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}")
+            os.system(f"aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d {directory} {url}")
+    except Exception as e:
+        print(f"Failed to download: {e}")
 
 
 def get_local_file_list(dir_path):
     file_list = []
-    for file in Path(dir_path).glob("**/*.*"):
+    for file in Path(dir_path).glob("*/*.*"):
         if file.is_file():
             file_path = str(file)
             file_list.append(file_path)
@@ -145,30 +159,30 @@ def get_local_file_list(dir_path):
 
 
 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_file_list(temp_dir)
-        try:
+    try:
+        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_file_list(temp_dir)
             download_thing(temp_dir, url.strip(), civitai_key)
-        except Exception:
+            after = get_local_file_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 ""
-        after = get_local_file_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}")
+        print(f"Download completed: {url}")
+        return new_file
+    except Exception as e:
+        print(f"Download failed: {url} {e}")
         return ""
-    print(f"Download completed: {url}")
-    return new_file
 
 
 def download_repo(repo_id: str, dir_path: str, progress=gr.Progress(track_tqdm=True)): # for diffusers repo
@@ -183,17 +197,13 @@ def download_repo(repo_id: str, dir_path: str, progress=gr.Progress(track_tqdm=T
         return False
 
 
-def upload_repo(repo_id: str, dir_path: str, is_private: bool, progress=gr.Progress(track_tqdm=True)): # for diffusers repo
+def upload_repo(repo_id: str, dir_path: str, is_private: bool, is_pr: bool=False, progress=gr.Progress(track_tqdm=True)): # for diffusers repo
     hf_token = get_token()
     api = HfApi(token=hf_token)
     try:
         progress(0, desc="Start uploading...")
         api.create_repo(repo_id=repo_id, token=hf_token, private=is_private, exist_ok=True)
-        for path in Path(dir_path).glob("*"):
-            if path.is_dir():
-                api.upload_folder(repo_id=repo_id, folder_path=str(path), path_in_repo=path.name, token=hf_token)
-            elif path.is_file():
-                api.upload_file(repo_id=repo_id, path_or_fileobj=str(path), path_in_repo=path.name, token=hf_token)
+        api.upload_folder(repo_id=repo_id, folder_path=dir_path, path_in_repo="", create_pr=is_pr, token=hf_token)
         progress(1, desc="Uploaded.")
         return get_hf_url(repo_id, "model")
     except Exception as e:
@@ -201,6 +211,18 @@ def upload_repo(repo_id: str, dir_path: str, is_private: bool, progress=gr.Progr
         return ""
 
 
+def gate_repo(repo_id: str, gated_str: str, repo_type: str="model"):
+    hf_token = get_token()
+    api = HfApi(token=hf_token)
+    try:
+        if gated_str == "auto": gated = "auto"
+        elif gated_str == "manual": gated = "manual"
+        else: gated = False
+        api.update_repo_settings(repo_id=repo_id, gated=gated, repo_type=repo_type, token=hf_token)
+    except Exception as e:
+        print(f"Error: Failed to update settings {repo_id}. {e}")
+
+
 HF_SUBFOLDER_NAME = ["None", "user_repo"]