Upload 8 files

README.md

@@ -1,12 +1,13 @@
 ---
-title: Safetensors Key Checker
-emoji: 💻
-colorFrom: green
-colorTo: gray
+title: Safetensors file key checker
+emoji: 🐶
+colorFrom: yellow
+colorTo: blue
 sdk: gradio
 sdk_version: 5.0.2
 app_file: app.py
 pinned: false
+license: mit
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,48 @@
+import gradio as gr
+from stkey_gr import stkey_gr, KEYS_FILES
+
+css = """
+.title { font-size: 3em; align-items: center; text-align: center; }
+.info { align-items: center; text-align: center; }
+.block.result { margin: 1em 0; padding: 1em; box-shadow: 0 0 3px 3px #664422, 0 0 3px 2px #664422 inset; border-radius: 6px; background: #665544; }
+.desc [src$='#float'] { float: right; margin: 20px; }
+"""
+
+with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", fill_width=True, css=css, delete_cache=(60, 3600)) as demo:
+    with gr.Column():
+        gr.Markdown("# Safetensors file key checker.", elem_classes="title")
+        with gr.Group():
+            dl_url = gr.Textbox(label="Download URL(s)", placeholder="https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/sd_xl_base_1.0_0.9vae.safetensors\n...", value="", lines=2, max_lines=255)
+            with gr.Row():
+                is_validate = gr.Checkbox(label="Validate safetensors file keys", value=True)
+                rfile = gr.Radio(label="Reference file for validation", choices=KEYS_FILES, value=KEYS_FILES[0])
+            with gr.Accordion("Advanced", open=False):
+                with gr.Row():
+                    with gr.Column():
+                        hf_token = gr.Textbox(label="Your HF read token (Optional)", placeholder="hf_...", value="", max_lines=1)
+                        gr.Markdown("Your token is available at [hf.co/settings/tokens](https://huggingface.co/settings/tokens).", elem_classes="info")
+                    with gr.Column():
+                        civitai_key = gr.Textbox(label="Your Civitai Key (Optional)", value="", max_lines=1)
+                        gr.Markdown("Your Civitai API key is available at [https://civitai.com/user/account](https://civitai.com/user/account).", elem_classes="info")
+        run_button = gr.Button(value="Check", variant="primary")
+        uploaded_urls = gr.CheckboxGroup(visible=False, choices=[], value=[]) # hidden
+        urls_md = gr.Markdown("<br><br>", elem_classes="result")
+        out_files = gr.Files(label="Output", interactive=False, value=[])
+        with gr.Group():
+            with gr.Row():
+                missing = gr.JSON(value=[], label="Missing keys")
+                added = gr.JSON(value=[], label="Added keys")
+            with gr.Row():
+                keys = gr.JSON(value=[], label="All keys")
+                metadata = gr.JSON(value={}, label="Metadata")
+    gr.DuplicateButton(value="Duplicate Space")
+
+    gr.on(
+        triggers=[run_button.click],
+        fn=stkey_gr,
+        inputs=[dl_url, civitai_key, hf_token, uploaded_urls, out_files, is_validate, rfile],
+        outputs=[uploaded_urls, out_files, urls_md, metadata, keys, missing, added],
+    )
+
+demo.queue()
+demo.launch()

packages.txt

@@ -0,0 +1 @@
+git-lfs aria2

requirements.txt

@@ -0,0 +1,4 @@
+huggingface-hub
+gdown
+safetensors
+torch

stkey.py

@@ -0,0 +1,121 @@
+import argparse
+from pathlib import Path
+import json
+import re
+import gc
+from safetensors.torch import load_file, save_file
+import torch
+
+
+SDXL_KEYS_FILE = "sdxl_keys.txt"
+
+
+def list_uniq(l):
+    return sorted(set(l), key=l.index)
+
+
+def read_safetensors_metadata(path: str):
+    with open(path, 'rb') as f:
+        header_size = int.from_bytes(f.read(8), 'little')
+        header_json = f.read(header_size).decode('utf-8')
+        header = json.loads(header_json)
+        metadata = header.get('__metadata__', {})
+        return metadata
+
+
+def keys_from_file(path: str):
+    keys = []
+    try:
+        with open(str(Path(path)), encoding='utf-8', mode='r') as f:
+            lines = f.readlines()
+        for line in lines:
+            keys.append(line.strip())
+    except Exception as e:
+        print(e)
+    finally:
+        return keys
+
+
+def validate_keys(keys: list[str], rfile: str=SDXL_KEYS_FILE):
+    missing = []
+    added = []
+    try:
+        rkeys = keys_from_file(rfile)
+        all_keys = list_uniq(keys + rkeys)
+        for key in all_keys:
+            if key in set(rkeys) and key not in set(keys): missing.append(key)
+            if key in set(keys) and key not in set(rkeys): added.append(key)
+    except Exception as e:
+        print(e)
+    finally:
+        return missing, added
+
+
+def read_safetensors_key(path: str):
+    try:
+        keys = []
+        state_dict = load_file(str(Path(path)))
+        for k, v in state_dict.items():
+            keys.append(k)
+    except Exception as e:
+        print(e)
+    finally:
+        del state_dict
+        torch.cuda.empty_cache()
+        gc.collect()
+        return keys
+
+
+def write_safetensors_key(keys: list[str], path: str, is_validate: bool=True, rpath: str=SDXL_KEYS_FILE):
+    if len(keys) == 0: return False
+    try:
+        with open(str(Path(path)), encoding='utf-8', mode='w') as f:
+            f.write("\n".join(keys))
+        if is_validate:
+            missing, added = validate_keys(keys, rpath)
+            with open(str(Path(path).stem + "_missing.txt"), encoding='utf-8', mode='w') as f:
+                f.write("\n".join(missing))
+            with open(str(Path(path).stem + "_added.txt"), encoding='utf-8', mode='w') as f:
+                f.write("\n".join(added))
+        return True
+    except Exception as e:
+        print(e)
+        return False
+
+
+def stkey(input: str, out_filename: str="", is_validate: bool=True, rfile: str=SDXL_KEYS_FILE):
+    keys = read_safetensors_key(input)
+    if len(keys) != 0 and out_filename: write_safetensors_key(keys, out_filename, is_validate, rfile)
+    if len(keys) != 0:
+        print("Metadata:")
+        print(read_safetensors_metadata(input))
+        print("\nKeys:")
+        print("\n".join(keys))
+        if is_validate:
+            missing, added = validate_keys(keys, rfile)
+            print("\nMissing Keys:")
+            print("\n".join(missing))
+            print("\nAdded Keys:")
+            print("\n".join(added))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("input", type=str, help="Input safetensors file.")
+    parser.add_argument("-s", "--save", action="store_true", default=False, help="Output to text file.")
+    parser.add_argument("-o", "--output", default="", type=str, help="Output to specific text file.")
+    parser.add_argument("-v", "--val", action="store_false", default=True, help="Disable key validation.")
+    parser.add_argument("-r", "--rfile", default=SDXL_KEYS_FILE, type=str, help="Specify reference file to validate keys.")
+
+    args = parser.parse_args()
+
+    if args.save: out_filename = Path(args.input).stem + ".txt"
+    out_filename = args.output if args.output else out_filename
+
+    stkey(args.input, out_filename, args.val, args.rfile)
+
+
+# Usage:
+# python stkey.py sd_xl_base_1.0_0.9vae.safetensors
+# python stkey.py sd_xl_base_1.0_0.9vae.safetensors -s
+# python stkey.py sd_xl_base_1.0_0.9vae.safetensors -o key.txt

stkey_gr.py

@@ -0,0 +1,104 @@
+import gradio as gr
+from huggingface_hub import HfApi, hf_hub_url
+import os
+from pathlib import Path
+import gc
+import re
+import json
+from utils import get_token, set_token, is_repo_exists, get_user_agent, get_download_file, list_uniq
+from stkey import read_safetensors_key, read_safetensors_metadata, validate_keys, write_safetensors_key
+
+
+TEMP_DIR = "."
+KEYS_FILES = ["sdxl_keys.txt"]
+
+
+def parse_urls(s):
+    url_pattern = "https?://[\\w/:%#\\$&\\?\\(\\)~\\.=\\+\\-]+"
+    try:
+        urls = re.findall(url_pattern, s)
+        return list(urls)
+    except Exception:
+        return []
+
+
+def to_urls(l: list[str]):
+    return "\n".join(l)
+
+
+def uniq_urls(s):
+    return to_urls(list_uniq(parse_urls(s)))
+
+
+def upload_safetensors_to_repo(filename, repo_id, repo_type, is_private, progress=gr.Progress(track_tqdm=True)):
+    output_filename = Path(filename).name
+    hf_token = get_token()
+    api = HfApi(token=hf_token)
+    try:
+        if not is_repo_exists(repo_id, repo_type): api.create_repo(repo_id=repo_id, repo_type=repo_type, token=hf_token, private=is_private)
+        progress(0, desc=f"Start uploading... {filename} to {repo_id}")
+        api.upload_file(path_or_fileobj=filename, path_in_repo=output_filename, repo_type=repo_type, revision="main", token=hf_token, repo_id=repo_id)
+        progress(1, desc="Uploaded.")
+        url = hf_hub_url(repo_id=repo_id, repo_type=repo_type, filename=output_filename)
+    except Exception as e:
+        print(f"Error: Failed to upload to {repo_id}. {e}")
+        gr.Warning(f"Error: Failed to upload to {repo_id}. {e}")
+        return None
+    finally:
+        Path(filename).unlink()
+    return url
+
+
+def download_file(dl_url, civitai_key, progress=gr.Progress(track_tqdm=True)):
+    download_dir = TEMP_DIR
+    progress(0, desc=f"Start downloading... {dl_url}")
+    output_filename = get_download_file(download_dir, dl_url, civitai_key)
+    return output_filename
+
+
+def get_stkey(filename: str, is_validate: bool=True, rfile: str=KEYS_FILES[0], progress=gr.Progress(track_tqdm=True)):
+    paths = []
+    metadata = {}
+    keys = []
+    missing = []
+    added = []
+    try:
+        progress(0, desc=f"Loading keys... {filename}")
+        keys = read_safetensors_key(filename)
+        if len(keys) == 0: raise Exception("No keys found.")
+        progress(0.5, desc=f"Checking keys... {filename}")
+        if write_safetensors_key(keys, str(Path(filename).stem + ".txt"), is_validate, rfile):
+            paths.append(str(Path(filename).stem + ".txt"))
+            paths.append(str(Path(filename).stem + "_missing.txt"))
+            paths.append(str(Path(filename).stem + "_added.txt"))
+        missing, added = validate_keys(keys, rfile)
+        metadata = read_safetensors_metadata(filename)
+    except Exception as e:
+        print(f"Error: Failed check {filename}. {e}")
+        gr.Warning(f"Error: Failed check {filename}. {e}")
+    finally:
+        Path(filename).unlink()
+        return paths, metadata, keys, missing, added
+
+
+def stkey_gr(dl_url: str, civitai_key: str, hf_token: str, urls: list[str], files: list[str], 
+             is_validate=True, rfile=KEYS_FILES[0], progress=gr.Progress(track_tqdm=True)):
+    if hf_token: set_token(hf_token)
+    else: set_token(os.environ.get("HF_TOKEN")) # default huggingface token
+    if not civitai_key: civitai_key = os.environ.get("CIVITAI_API_KEY") # default Civitai API key
+    dl_urls = parse_urls(dl_url)
+    if not urls: urls = []
+    if not files: files = []
+    metadata = {}
+    keys = []
+    missing = []
+    added = []
+    for u in dl_urls:
+        file = download_file(u, civitai_key)
+        if not Path(file).exists() or not Path(file).is_file(): continue
+        paths, metadata, keys, missing, added = get_stkey(file, is_validate, rfile)
+        if len(paths) != 0: files.extend(paths)
+    progress(1, desc="Processing...")
+    gc.collect()
+    return gr.update(value=urls, choices=urls), gr.update(value=files), gr.update(visible=True), metadata, keys, missing, added
+

utils.py

@@ -0,0 +1,165 @@
+import gradio as gr
+from huggingface_hub import HfApi, HfFolder, hf_hub_download
+import os
+from pathlib import Path
+import shutil
+import gc
+import re
+import urllib.parse
+
+
+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.")
+
+
+def get_user_agent():
+    return 'Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:127.0) Gecko/20100101 Firefox/127.0'
+
+
+def is_repo_exists(repo_id: str, repo_type: str="model"):
+    hf_token = get_token()
+    api = HfApi(token=hf_token)
+    try:
+        if api.repo_exists(repo_id=repo_id, repo_type=repo_type, token=hf_token): return True
+        else: return False
+    except Exception as e:
+        print(f"Error: Failed to connect {repo_id} ({repo_type}). {e}")
+        return True # for safe
+
+
+MODEL_TYPE_CLASS = {
+    "diffusers:StableDiffusionPipeline": "SD 1.5",
+    "diffusers:StableDiffusionXLPipeline": "SDXL",
+    "diffusers:FluxPipeline": "FLUX",
+}
+
+
+def get_model_type(repo_id: str):
+    hf_token = get_token()
+    api = HfApi(token=hf_token)
+    lora_filename = "pytorch_lora_weights.safetensors"
+    diffusers_filename = "model_index.json"
+    default = "SDXL"
+    try:
+        if api.file_exists(repo_id=repo_id, filename=lora_filename, token=hf_token): return "LoRA"
+        if not api.file_exists(repo_id=repo_id, filename=diffusers_filename, token=hf_token): return "None"
+        model = api.model_info(repo_id=repo_id, token=hf_token)
+        tags = model.tags
+        for tag in tags:
+            if tag in MODEL_TYPE_CLASS.keys(): return MODEL_TYPE_CLASS.get(tag, default)
+    except Exception:
+        return default
+    return default
+
+
+def list_uniq(l):
+        return sorted(set(l), key=l.index)
+
+
+def list_sub(a, b):
+    return [e for e in a if e not in b]
+
+
+def is_repo_name(s):
+    return re.fullmatch(r'^[^/,\s\"\']+/[^/,\s\"\']+$', s)
+
+
+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:
+        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="", 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()
+        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/")
+        #user_header = f'"Authorization: Bearer {hf_token}"'
+        if 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]}")
+    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}")
+        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}")
+
+
+def get_local_model_list(dir_path):
+    model_list = []
+    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)):
+    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_thing(temp_dir, url.strip(), 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