Upload 3 files

README.md

@@ -4,7 +4,7 @@ emoji: 🐶
 colorFrom: yellow
 colorTo: blue
 sdk: gradio
-sdk_version: 5.0.2
+sdk_version: 5.1.0
 app_file: app.py
 pinned: false
 license: mit

stkey.py

@@ -7,7 +7,7 @@ from safetensors.torch import load_file, save_file
 import torch
 
 
-SDXL_KEYS_FILE = "sdxl_keys.txt"
+SDXL_KEYS_FILE = "keys/sdxl_keys.txt"
 
 
 def list_uniq(l):

utils.py

@@ -6,6 +6,8 @@ import shutil
 import gc
 import re
 import urllib.parse
+import subprocess
+import time
 
 
 def get_token():
@@ -72,15 +74,24 @@ def list_sub(a, b):
 
 
 def is_repo_name(s):
-    return re.fullmatch(r'^[^/,\s\"\']+/[^/,\s\"\']+$', s)
+    return re.fullmatch(r'^[\w_\-\.]+/[\w_\-\.]+$', s)
+
+
+def get_hf_url(repo_id: str, repo_type: str="model"):
+    if repo_type == "dataset": url = f"https://huggingface.co/datasets/{repo_id}"
+    elif repo_type == "space": url = f"https://huggingface.co/spaces/{repo_id}"
+    else: url = f"https://huggingface.co/{repo_id}"
+    return url
 
 
 def split_hf_url(url: str):
     try:
-        s = list(re.findall(r'^(?:https?://huggingface.co/)(?:(datasets)/)?(.+?/.+?)/\w+?/.+?/(?:(.+)/)?(.+?.\w+)(?:\?download=true)?$', url)[0])
+        s = list(re.findall(r'^(?:https?://huggingface.co/)(?:(datasets|spaces)/)?(.+?/.+?)/\w+?/.+?/(?:(.+)/)?(.+?.\w+)(?:\?download=true)?$', url)[0])
         if len(s) < 4: return "", "", "", ""
         repo_id = s[1]
-        repo_type = "dataset" if s[0] == "datasets" else "model"
+        if s[0] == "datasets": repo_type = "dataset"
+        elif s[0] == "spaces": repo_type = "space"
+        else: repo_type = "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
@@ -92,14 +103,16 @@ 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:
-        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)
+        print(f"Downloading {url} to {directory}")
+        if subfolder is not None: path = hf_hub_download(repo_id=repo_id, filename=filename, subfolder=subfolder, repo_type=repo_type, local_dir=directory, token=hf_token)
+        else: path = hf_hub_download(repo_id=repo_id, filename=filename, repo_type=repo_type, local_dir=directory, token=hf_token)
+        return path
     except Exception as e:
         print(f"Failed to download: {e}")
+        return None
 
 
 def download_thing(directory, url, civitai_api_key="", progress=gr.Progress(track_tqdm=True)): # requires aria2, gdown
-    hf_token = get_token()
     url = url.strip()
     if "drive.google.com" in url:
         original_dir = os.getcwd()
@@ -108,14 +121,8 @@ def download_thing(directory, url, civitai_api_key="", progress=gr.Progress(trac
         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 True or hf_token:
-            download_hf_file(directory, url)
-            #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]}")
+        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]
@@ -128,14 +135,13 @@ def download_thing(directory, url, civitai_api_key="", progress=gr.Progress(trac
         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', '.fp16.safetensors', '.sft')
+def get_local_file_list(dir_path):
+    file_list = []
     for file in Path(dir_path).glob("**/*.*"):
-        if file.is_file() and file.suffix in valid_extensions:
+        if file.is_file():
             file_path = str(file)
-            model_list.append(file_path)
-    return model_list
+            file_list.append(file_path)
+    return file_list
 
 
 def get_download_file(temp_dir, url, civitai_key, progress=gr.Progress(track_tqdm=True)):
@@ -150,13 +156,13 @@ def get_download_file(temp_dir, url, civitai_key, progress=gr.Progress(track_tqd
         new_file = f"{temp_dir}/{url.split('/')[-1]}"
     else:
         print(f"Start downloading: {url}")
-        before = get_local_model_list(temp_dir)
+        before = get_local_file_list(temp_dir)
         try:
             download_thing(temp_dir, url.strip(), civitai_key)
         except Exception:
             print(f"Download failed: {url}")
             return ""
-        after = get_local_model_list(temp_dir)
+        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}")
@@ -165,8 +171,7 @@ def get_download_file(temp_dir, url, civitai_key, progress=gr.Progress(track_tqd
     return new_file
 
 
-# https://huggingface.co/docs/huggingface_hub/v0.25.1/en/package_reference/file_download#huggingface_hub.snapshot_download
-def download_repo(repo_id, dir_path, progress=gr.Progress(track_tqdm=True)):
+def download_repo(repo_id: str, dir_path: str, progress=gr.Progress(track_tqdm=True)): # for diffusers repo
     hf_token = get_token()
     try:
         snapshot_download(repo_id=repo_id, local_dir=dir_path, token=hf_token, allow_patterns=["*.safetensors", "*.bin"],
@@ -178,20 +183,93 @@ def download_repo(repo_id, dir_path, progress=gr.Progress(track_tqdm=True)):
         return False
 
 
-def upload_repo(new_repo_id, dir_path, is_private, progress=gr.Progress(track_tqdm=True)):
+def upload_repo(repo_id: str, dir_path: str, is_private: bool, 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=new_repo_id, token=hf_token, private=is_private, exist_ok=True)
+        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=new_repo_id, folder_path=str(path), path_in_repo=path.name, token=hf_token)
+                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=new_repo_id, path_or_fileobj=str(path), path_in_repo=path.name, token=hf_token)
+                api.upload_file(repo_id=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}"
+        return get_hf_url(repo_id, "model")
     except Exception as e:
-        print(f"Error: Failed to upload to {new_repo_id}. {e}")
+        print(f"Error: Failed to upload to {repo_id}. {e}")
         return ""
-    return url
+
+
+HF_SUBFOLDER_NAME = ["None", "user_repo"]
+
+
+def duplicate_hf_repo(src_repo: str, dst_repo: str, src_repo_type: str, dst_repo_type: str,
+                      is_private: bool, subfolder_type: str=HF_SUBFOLDER_NAME[1], progress=gr.Progress(track_tqdm=True)):
+    hf_token = get_token()
+    api = HfApi(token=hf_token)
+    try:
+        if subfolder_type == "user_repo": subfolder = src_repo.replace("/", "_")
+        else: subfolder = ""
+        progress(0, desc="Start duplicating...")
+        api.create_repo(repo_id=dst_repo, repo_type=dst_repo_type, private=is_private, exist_ok=True, token=hf_token)
+        for path in api.list_repo_files(repo_id=src_repo, repo_type=src_repo_type, token=hf_token):
+            file = hf_hub_download(repo_id=src_repo, filename=path, repo_type=src_repo_type, token=hf_token)
+            if not Path(file).exists(): continue
+            if Path(file).is_dir(): # unused for now
+                api.upload_folder(repo_id=dst_repo, folder_path=file, path_in_repo=f"{subfolder}/{path}" if subfolder else path,
+                                  repo_type=dst_repo_type, token=hf_token)
+            elif Path(file).is_file():
+                api.upload_file(repo_id=dst_repo, path_or_fileobj=file, path_in_repo=f"{subfolder}/{path}" if subfolder else path,
+                                repo_type=dst_repo_type, token=hf_token)
+            if Path(file).exists(): Path(file).unlink()
+        progress(1, desc="Duplicated.")
+        return f"{get_hf_url(dst_repo, dst_repo_type)}/tree/main/{subfolder}" if subfolder else get_hf_url(dst_repo, dst_repo_type)
+    except Exception as e:
+        print(f"Error: Failed to duplicate repo {src_repo} to {dst_repo}. {e}")
+        return ""
+
+
+BASE_DIR = str(Path(__file__).resolve().parent.resolve())
+CIVITAI_API_KEY = os.environ.get("CIVITAI_API_KEY")
+
+
+def get_file(url: str, path: str): # requires aria2, gdown
+    print(f"Downloading {url} to {path}...")
+    get_download_file(path, url, CIVITAI_API_KEY)
+
+
+def git_clone(url: str, path: str, pip: bool=False, addcmd: str=""): # requires git
+    os.makedirs(str(Path(BASE_DIR, path)), exist_ok=True)
+    os.chdir(Path(BASE_DIR, path))
+    print(f"Cloning {url} to {path}...")
+    cmd = f'git clone {url}'
+    print(f'Running {cmd} at {Path.cwd()}')
+    i = subprocess.run(cmd, shell=True).returncode
+    if i != 0: print(f'Error occured at running {cmd}')
+    p = url.split("/")[-1]
+    if not Path(p).exists: return
+    if pip:
+        os.chdir(Path(BASE_DIR, path, p))
+        cmd = f'pip install -r requirements.txt'
+        print(f'Running {cmd} at {Path.cwd()}')
+        i = subprocess.run(cmd, shell=True).returncode
+        if i != 0: print(f'Error occured at running {cmd}')
+    if addcmd:
+        os.chdir(Path(BASE_DIR, path, p))
+        cmd = addcmd
+        print(f'Running {cmd} at {Path.cwd()}')
+        i = subprocess.run(cmd, shell=True).returncode
+        if i != 0: print(f'Error occured at running {cmd}')
+
+
+def run(cmd: str, timeout: float=0):
+    print(f'Running {cmd} at {Path.cwd()}')
+    if timeout == 0:
+        i = subprocess.run(cmd, shell=True).returncode
+        if i != 0: print(f'Error occured at running {cmd}')
+    else:
+        p = subprocess.Popen(cmd, shell=True)
+        time.sleep(timeout)
+        p.terminate()
+        print(f'Terminated in {timeout} seconds')