diff --git a/README.md b/README.md index 74cf60c..ca5a3e3 100644 --- a/README.md +++ b/README.md @@ -28,7 +28,11 @@ I made this fork because the original repo was inactive and missing many things - Add `/` at the start of the search bar to see auto-complete suggestions. - Include models listed in ComfyUI's `extra_model_paths.yaml` or added in `ComfyUI/models`. - Sort for models (Date Created, Date Modified, Name). + +### Model Info + - View model metadata, including training tags and bucket resolutions. +- Delete or move a model. ### ComfyUI Node Graph @@ -62,7 +66,6 @@ I made this fork because the original repo was inactive and missing many things ### Model info window/panel (server load/send on demand) -- Move a model to a different directory. - Set preview image. - Optional (re)download `📥︎` model info from the internet and cache the text file locally. (requires checksum?) - Radio buttons to swap between downloaded and server view. diff --git a/__init__.py b/__init__.py index f0c91b9..416ee96 100644 --- a/__init__.py +++ b/__init__.py @@ -1,5 +1,6 @@ import os import pathlib +import shutil from datetime import datetime import sys import copy @@ -172,8 +173,8 @@ async def save_ui_settings(request): }) -@server.PromptServer.instance.routes.get("/model-manager/image/preview") -async def img_preview(request): +@server.PromptServer.instance.routes.get("/model-manager/preview/get") +async def get_model_preview(request): uri = request.query.get("uri") image_path = no_preview_image @@ -337,7 +338,7 @@ def download_file(url, filename, overwrite): api_key = server_settings["civitai_api_key"] if (api_key != ""): def_headers["Authorization"] = f"Bearer {api_key}" - url = url + f"?token={api_key}" + url = url + f"?token={api_key}" # TODO: Authorization didn't work in the header elif url.startswith("https://huggingface.co/"): api_key = server_settings["huggingface_api_key"] if api_key != "": @@ -445,6 +446,15 @@ async def get_model_info(request): info["Hash"] = metadata.get("sshs_model_hash", "") info["Output Name"] = metadata.get("ss_output_name", "") + file_name, _ = os.path.splitext(file) + txt_file = file_name + ".txt" + description = "" + if os.path.isfile(txt_file): + with open(txt_file, 'r', encoding="utf-8") as f: + description = f.read() + info["Description"] = description + + if metadata is not None: img_buckets = metadata.get("ss_bucket_info", "{}") if type(img_buckets) is str: img_buckets = json.loads(img_buckets) @@ -471,14 +481,6 @@ async def get_model_info(request): tags.sort(key=lambda x: x[1], reverse=True) info["Tags"] = tags - file_name, _ = os.path.splitext(file) - txt_file = file_name + ".txt" - description = "" - if os.path.isfile(txt_file): - with open(txt_file, 'r', encoding="utf-8") as f: - description = f.read() - info["Description"] = description - return web.json_response(info) @@ -531,16 +533,16 @@ async def download_model(request): image_uri = body.get("image") if image_uri is not None and image_uri != "": - # TODO: doesn't work for https://civitai.com/images/... - image_extension = None + image_extension = None # TODO: doesn't work for https://civitai.com/images/... for ext in image_extensions: if image_uri.endswith(ext): image_extension = ext break if image_extension is not None: + file_path_without_extension = name[:len(name) - len(model_extension)] image_name = os.path.join( directory, - (name[:len(name) - len(model_extension)]) + image_extension + file_path_without_extension + image_extension ) try: download_file(image_uri, image_name, overwrite) @@ -551,6 +553,48 @@ async def download_model(request): return web.json_response(result) +@server.PromptServer.instance.routes.post("/model-manager/model/move") +async def move_model(request): + body = await request.json() + model_type = body.get("type", None) + if model_type is None: + return web.json_response({ "success": False }) + + old_file = body.get("oldFile", None) + if old_file is None: + return web.json_response({ "success": False }) + old_file = search_path_to_system_path(old_file, model_type) + if not os.path.isfile(old_file): + return web.json_response({ "success": False }) + _, filename = os.path.split(old_file) + + new_path = body.get("newDirectory", None) + if new_path is None: + return web.json_response({ "success": False }) + new_path = search_path_to_system_path(new_path, model_type) + if not os.path.isdir(new_path): + return web.json_response({ "success": False }) + + new_file = os.path.join(new_path, filename) + try: + shutil.move(old_file, new_file) + except: + return web.json_response({ "success": False }) + + old_file_without_extension, _ = os.path.splitext(old_file) + new_file_without_extension, _ = os.path.splitext(new_file) + + for extension in image_extensions + (".txt",): + old_file = old_file_without_extension + extension + if os.path.isfile(old_file): + try: + shutil.move(old_file, new_file_without_extension + extension) + except Exception as e: + print(e, file=sys.stderr, flush=True) + + return web.json_response({ "success": True }) + + @server.PromptServer.instance.routes.post("/model-manager/model/delete") async def delete_model(request): result = { "success": False } diff --git a/web/model-manager.css b/web/model-manager.css index d89613a..d460bbb 100644 --- a/web/model-manager.css +++ b/web/model-manager.css @@ -42,7 +42,7 @@ } .comfy-tabs-body { - background-color: var(--comfy-input-bg); + background-color: var(--bg-color); border: 2px solid var(--border-color); border-top: none; padding: 16px 0px; @@ -254,13 +254,13 @@ } .model-manager ::-webkit-scrollbar-track { - background-color: #353535; + background-color: var(--comfy-input-bg); border-right: 1px solid var(--border-color); border-bottom: 1px solid var(--border-color); } .model-manager ::-webkit-scrollbar-thumb { - background-color: #a1a1a1; + background-color: var(--fg-color); border-radius: 3px; } @@ -331,20 +331,11 @@ color: var(--input-text); } -.model-manager .row { - position: relative; - padding-top: 2px; - margin-top: -2px; - padding-bottom: 18px; - margin-bottom: 1px; - top: -1px; - background-color: var(--comfy-input-bg); -} - .model-manager [data-name="Install"] .row, .model-manager [data-name="Models"] .row { position: sticky; z-index: 1; + top: 0; } .model-manager [data-name="Install"] input { @@ -359,7 +350,10 @@ } .model-manager .tab-header { - display: block; + display: flex; + padding: 8px 0; + flex-direction: column; + background-color: var(--bg-color); } .model-manager .tab-header-flex-block { @@ -393,6 +387,7 @@ position: absolute; background-color: var(--bg-color); border: 2px var(--border-color) solid; + color: var(--fg-color); max-height: 30vh; overflow: auto; border-radius: 10px; @@ -471,11 +466,16 @@ } .model-manager .model-info-view { - background-color: var(--comfy-menu-bg); - display: none; + background-color: var(--bg-color); + border: 2px solid var(--border-color); + box-sizing: border-box; + display: flex; + flex-direction: column; height: 100%; + margin-top: 40px; overflow-wrap: break-word; - width: 100%; + overflow-y: auto; + padding: 20px; } .model-manager .model-info-container { @@ -483,9 +483,7 @@ border-radius: 16px; color: var(--fg-color); margin-top: 8px; - max-height: 90%; padding: 16px; - overflow: auto; width: auto; } @@ -494,3 +492,7 @@ -ms-user-select: none; user-select: none; } + +.model-manager .download-model-infos { + padding: 16px 0; +} diff --git a/web/model-manager.js b/web/model-manager.js index 58bd38c..4d3ad3a 100644 --- a/web/model-manager.js +++ b/web/model-manager.js @@ -1094,7 +1094,7 @@ class ModelGrid { if (models.length > 0) { return models.map((item) => { const uri = item.post ?? "no-post"; - const imgUrl = `/model-manager/image/preview?uri=${uri}`; + const imgUrl = `/model-manager/preview/get?uri=${uri}`; const searchPath = item.path; const path = searchPathToSystemPath(searchPath, searchSeparator, systemSeparator); let buttons = []; @@ -1284,7 +1284,7 @@ class ModelManager extends ComfyDialog { /** @type {HTMLDivElement} */ modelGrid: null, /** @type {HTMLSelectElement} */ modelTypeSelect: null, /** @type {HTMLSelectElement} */ modelSortSelect: null, - /** @type {HTMLDivElement} */ searchDirectoryDropdown: null, + /** @type {HTMLDivElement} */ //searchDirectoryDropdown: null, /** @type {HTMLInputElement} */ modelContentFilter: null, /** @type {HTMLDivElement} */ sidebarButtons: null, @@ -1321,6 +1321,25 @@ class ModelManager extends ComfyDialog { constructor() { super(); + + const moveDestination = $el("input.search-text-area", { + placeholder: "/", + }); + let searchDropdown = null; + searchDropdown = new DirectoryDropdown( + moveDestination, + () => { + searchDropdown.update( + this.#data.modelDirectories, + this.#searchSeparator, + ); + }, + () => {}, + () => {}, + this.#searchSeparator, + true, + ); + this.element = $el( "div.comfy-modal.model-manager", { @@ -1330,6 +1349,7 @@ class ModelManager extends ComfyDialog { $el("div.comfy-modal-content", [ // TODO: settings.top_bar_left_to_right or settings.top_bar_right_to_left $el("div.model-info-view", { $: (el) => (this.#el.modelInfoView = el), + style: { display: "none" }, }, [ $el("div", { style: { @@ -1341,7 +1361,7 @@ class ModelManager extends ComfyDialog { textContent: "🗑︎", onclick: async(e) => { const affirmation = "delete"; - const confirmation = window.prompt("Type \"" + affirmation + "\" to delete the model PERMANENTLY?\n\nThis includes all image or text files."); + const confirmation = window.prompt("Type \"" + affirmation + "\" to delete the model PERMANENTLY.\n\nThis includes all image or text files."); let deleted = false; if (confirmation === affirmation) { const container = this.#el.modelInfoContainer; @@ -1370,6 +1390,50 @@ class ModelManager extends ComfyDialog { }, }), ]), + $el("div.row.tab-header", { + display: "block", + }, [ + $el("div.row.tab-header-flex-block", [ + $el("div.search-models", [ + moveDestination, + searchDropdown.element, + ]), + $el("button", { + textContent: "Move", + onclick: async(e) => { + const container = this.#el.modelInfoContainer; + const path = container.dataset.path; + const type = this.#el.modelTypeSelect.value; + const destination = moveDestination.value; + let moved = false; + await request( + `/model-manager/model/move`, + { + method: "POST", + body: JSON.stringify({ + "type": type, + "oldFile": path, + "newDirectory": destination, + }), + } + ) + .then((result) => { + if (result["success"]) + { + container.innerHTML = ""; + this.#el.modelInfoView.style.display = "none"; + this.#modelTab_updateModels(); + moved = true; + } + }) + .catch(err => {}); + if (!moved) { + buttonAlert(e.target, false); + } + }, + }), + ]), + ]), $el("div.model-info-container", { $: (el) => (this.#el.modelInfoContainer = el), "data-path": "", @@ -1597,7 +1661,8 @@ class ModelManager extends ComfyDialog { } infoHtml.append.apply(infoHtml, innerHtml); - this.#el.modelInfoView.style.display = "block"; + this.#el.modelInfoView.removeAttribute("style"); // remove "display: none" + // TODO: set default value of dropdown and value to model type? } /** @@ -2234,7 +2299,7 @@ class ModelManager extends ComfyDialog { * @returns {HTMLElement} */ #downloadTab_new() { - return $el("div", [ + return $el("div.tab-header", [ $el("div.row.tab-header-flex-block", [ $el("input.search-text-area", { $: (el) => (this.#el.modelInfoUrl = el), @@ -2252,7 +2317,7 @@ class ModelManager extends ComfyDialog { textContent: "🔍︎", }), ]), - $el("div", { + $el("div.download-model-infos", { $: (el) => (this.#el.modelInfos = el), }, [ $el("div", ["Input a URL to select a model to download."]),