diff --git a/js/loraLoader.js b/js/loraLoader.js new file mode 100644 index 0000000..f609b0a --- /dev/null +++ b/js/loraLoader.js @@ -0,0 +1,90 @@ +import { app } from "/scripts/app.js"; +import { ComfyWidgets } from "/scripts/widgets.js"; +import { $el } from "/scripts/ui.js"; + +app.registerExtension({ + name: "pysssss.LoraLoader", + init() { + $el("style", { + textContent: ` + .litemenu-entry:hover .pysssss-combo-image { + display: block; + } + .pysssss-combo-image { + display: none; + position: absolute; + left: 0; + top: 0; + transform: translate(-100%, 0); + width: 256px; + height: 256px; + background-size: cover; + filter: brightness(65%); + } + `, + parent: document.body, + }); + + // Big ol' hack to get subclassing the context menu to work + // Replace the addItem function with our own that wraps the context of "this" with a proxy + // That proxy then replaces the constructor with another proxy + // That proxy then calls the custom ContextMenu that supports filters + const ctorProxy = new Proxy(LiteGraph.ContextMenu, { + construct(target, args) { + return new LiteGraph.ContextMenu(...args); + }, + }); + + const addItem = LiteGraph.ContextMenu.prototype.addItem; + LiteGraph.ContextMenu.prototype.addItem = function () { + const proxy = new Proxy(this, { + get(target, prop) { + if (prop === "constructor") { + return ctorProxy; + } + return target[prop]; + }, + }); + proxy.__target__ = this; + const el = addItem.apply(proxy, arguments); + if (arguments[1]?.image) { + el.textContent += " *"; + $el("div.pysssss-combo-image", { + parent: el, + style: { + backgroundImage: `url(/pysssss/view/${encodeURIComponent(arguments[1].image)})`, + }, + }); + } + return el; + }; + + // We also need to patch the ContextMenu constructor to unwrap the parent else it fails a LiteGraph type check + const ctxMenu = LiteGraph.ContextMenu; + LiteGraph.ContextMenu = function (values, options) { + ctxMenu.call(this, values, options); + }; + LiteGraph.ContextMenu.prototype = ctxMenu.prototype; + + const combo = ComfyWidgets["COMBO"]; + ComfyWidgets["COMBO"] = function (node, inputName, inputData) { + const type = inputData[0]; + const res = combo.apply(this, arguments); + if (type[0] && type[0].content) { + let value = res.widget.value; + Object.defineProperty(res.widget, "value", { + get() { + if (res.widget && new Error().stack.includes("drawNodeWidgets")) { + return value.content; + } + return value; + }, + set(v) { + value = v; + }, + }); + } + return res; + }; + }, +}); diff --git a/py/lora_loader.py b/py/lora_loader.py new file mode 100644 index 0000000..a497dd1 --- /dev/null +++ b/py/lora_loader.py @@ -0,0 +1,53 @@ +import os +from nodes import LoraLoader +import folder_paths +from server import PromptServer +from aiohttp import web + + +@PromptServer.instance.routes.get("/pysssss/view/{name}") +async def view_image(request): + name = request.match_info["name"] + pos = name.index("/") + type = name[0:pos] + name = name[pos+1:] + + image_path = folder_paths.get_full_path( + type, name) + if not image_path: + return web.Response(status=404) + + filename = os.path.basename(image_path) + return web.FileResponse(image_path, headers={"Content-Disposition": f"filename=\"{filename}\""}) + + +class LoraLoaderWithImages(LoraLoader): + @classmethod + def INPUT_TYPES(s): + types = super().INPUT_TYPES() + names = types["required"]["lora_name"][0] + + for idx, lora_name in enumerate(names): + lora_image = os.path.splitext(lora_name)[0] + ".png" + lora_image_path = folder_paths.get_full_path("loras", lora_image) + if not lora_image_path: + lora_image = os.path.splitext(lora_name)[0] + ".jpg" + lora_image_path = folder_paths.get_full_path( + "loras", lora_image) + + names[idx] = {"content": lora_name, + "image": f"loras/{lora_image}" if lora_image_path else None, } + + return types + + def load_lora(self, model, clip, lora_name, strength_model, strength_clip): + return super().load_lora(model, clip, lora_name["content"], strength_model, strength_clip) + + +NODE_CLASS_MAPPINGS = { + "LoraLoader|pysssss": LoraLoaderWithImages, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "LoraLoader|pysssss": "Lora Loader 🐍", +}