From 529115b0a414e1e4523f7eec98e5f370243ab585 Mon Sep 17 00:00:00 2001
From: Christian Bastian <bastian.43@osu.edu>
Date: Mon, 12 Feb 2024 21:46:03 -0500
Subject: [PATCH] Set API keys for Civitai & HuggingFace in
 `server_settings.yaml`.

---
 .gitignore           |  1 +
 README.md            |  5 ++---
 __init__.py          | 34 ++++++++++++++++++++++------------
 web/model-manager.js |  1 +
 4 files changed, 26 insertions(+), 15 deletions(-)

diff --git a/.gitignore b/.gitignore
index c91d84d..32b09a4 100644
--- a/.gitignore
+++ b/.gitignore
@@ -159,3 +159,4 @@ cython_debug/
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
 ui_settings.yaml
+server_settings.yaml
diff --git a/README.md b/README.md
index 50830af..6e180b0 100644
--- a/README.md
+++ b/README.md
@@ -27,7 +27,8 @@ Currently it is still missing some features it should have.
 - Increased supported preview image types.
 - Correctly change colors using ComfyUI's theme colors.
 - Simplified UI.
-- Settings tab and config file.
+- Civitai and HuggingFace API token configurable in `server_settings.yaml`.
+- Settings tab saved in `ui_settings.yaml`.
   - Hide/Show 'add' and 'copy-to-clipboard' buttons.
   - Text to always search.
   - Show/Hide add embedding extension.
@@ -65,8 +66,6 @@ Currently it is still missing some features it should have.
   - &#9744; Check search code is optimized to avoid recalculation on every minor input change
 - &#9744; Directory dropdown
   - &#9744; Use always filter to filter directory content auto-suggest dropdown
-    - &#9744; Generalize model list filtering code to reuse approach
-  - &#9744; Generalize search dropdown for download location selection
 - &#9744; Filters dropdown
   - &#9744; Stable Diffusion model version/Clip/Upscale/?
   - &#9744; Favorites
diff --git a/__init__.py b/__init__.py
index 1ab9c17..6c36e61 100644
--- a/__init__.py
+++ b/__init__.py
@@ -28,6 +28,7 @@ index_uri = os.path.join(extension_uri, "index.json")
 #checksum_cache_uri = os.path.join(extension_uri, "checksum_cache.txt")
 no_preview_image = os.path.join(extension_uri, "no-preview.png")
 ui_settings_uri = os.path.join(extension_uri, "ui_settings.yaml")
+server_settings_uri = os.path.join(extension_uri, "server_settings.yaml")
 
 fallback_model_extensions = set([".bin", ".ckpt", ".onnx", ".pt", ".pth", ".safetensors"]) # TODO: magic values
 image_extensions = (".apng", ".gif", ".jpeg", ".jpg", ".png", ".webp")
@@ -122,13 +123,15 @@ def ui_rules():
     ]
 
 
-#def server_rules():
-#    Rule = config_loader.Rule
-#    return [
-#        Rule("model_extension_download_whitelist", [".safetensors"], list),
-#        Rule("civitai_api_key", "", str),
-#    ]
-
+def server_rules():
+    Rule = config_loader.Rule
+    return [
+        #Rule("model_extension_download_whitelist", [".safetensors"], list),
+        Rule("civitai_api_key", "", str),
+        Rule("huggingface_api_key", "", str),
+    ]
+server_settings = config_loader.yaml_load(server_settings_uri, server_rules())
+config_loader.yaml_save(server_settings_uri, server_rules(), server_settings)
 
 @server.PromptServer.instance.routes.get("/model-manager/settings/load")
 async def load_ui_settings(request):
@@ -384,11 +387,6 @@ async def directory_list(request):
     return web.json_response(dir_list)
 
 
-def_headers = {
-    "User-Agent": "Mozilla/5.0 (iPad; CPU OS 12_2 like Mac OS X) AppleWebKit/605.1.15 (KHTML, like Gecko) Mobile/15E148"
-}
-
-
 def download_file(url, filename, overwrite):
     if not overwrite and os.path.isfile(filename):
         raise Exception("File already exists!")
@@ -396,6 +394,18 @@ def download_file(url, filename, overwrite):
     # TODO: clear any previous failed partial download file
     dl_filename = filename + ".download"
 
+    def_headers = {
+        "User-Agent": "Mozilla/5.0 (iPad; CPU OS 12_2 like Mac OS X) AppleWebKit/605.1.15 (KHTML, like Gecko) Mobile/15E148",
+    }
+    if url.startswith("https://civitai.com/"):
+        api_key = server_settings["civitai_api_key"]
+        if (api_key != ""):
+            def_headers["Authorization"] = f"Bearer {api_key}"
+    elif url.startswith("https://huggingface.co/"):
+        api_key = server_settings["huggingface_api_key"]
+        if api_key != "":
+            def_headers["Authorization"] = f"Bearer {api_key}"
+
     rh = requests.get(url=url, stream=True, verify=False, headers=def_headers, proxies=None, allow_redirects=False)
     if not rh.ok:
         raise Exception("Unable to download")
diff --git a/web/model-manager.js b/web/model-manager.js
index 34b8003..8b8a553 100644
--- a/web/model-manager.js
+++ b/web/model-manager.js
@@ -2009,6 +2009,7 @@ class ModelManager extends ComfyDialog {
                             record["type"] = els.modelTypeSelect.value;
                             if (record["type"] === "") { return; } // TODO: notify user in app
                             record["path"] = els.saveDirectoryPath.value;
+                            if (record["path"] === "/") { return; } // TODO: notify user in app
                             record["name"] = (() => {
                                 const filename = info["fileName"];
                                 const name = els.filename.value;