From 817cfa6e01df2cd15d820ab84b3ff46587873cf4 Mon Sep 17 00:00:00 2001 From: PALE Date: Sun, 10 May 2026 16:01:34 +0900 Subject: [PATCH 1/2] feat: Add UNET loader support with diffusion model preview MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add support for UNET (diffusion model) loader with image preview capabilities: - Python: * Import UNETLoader from nodes module * Add UNETLoaderWithImages class extending UNETLoader * Register UNETLoader node mapping as "Load Diffusion Model 🐍" - JavaScript: * Add DIFFUSION_MODEL_LOADER constant for UNET loader identification * Extend getType() to map UNET loader to "diffusion_models" type * Extend getWidgetName() to map "unet_name" widget for diffusion models * Add isImageLoader() helper function to consolidate loader type checks * Add hasExamples() helper function to explicitly identify nodes supporting examples Refactor and optimize existing code: * Replace hardcoded p1/p2 variables with dynamic modelTypes array iteration * Improve Promise handling for multiple model types in refreshComboInNodes() * Update UI setting name from "Lora & Checkpoint loader" to "Model loader" * Add null check for promptWidget to prevent runtime errors (fixes potential crash) This change enables users to preview diffusion models alongside existing Lora and Checkpoint model previews, maintaining full backward compatibility with existing functionality. - Existing Lora/Checkpoint features remain unchanged - All refactoring maintains identical behavior - Improves code maintainability and extensibility --- py/better_combos.py | 8 +++++++- web/js/betterCombos.js | 45 +++++++++++++++++++++++++++++------------- 2 files changed, 38 insertions(+), 15 deletions(-) diff --git a/py/better_combos.py b/py/better_combos.py index d31043b..d3c531a 100644 --- a/py/better_combos.py +++ b/py/better_combos.py @@ -1,6 +1,6 @@ import glob import os -from nodes import LoraLoader, CheckpointLoaderSimple +from nodes import LoraLoader, CheckpointLoaderSimple, UNETLoader import folder_paths from server import PromptServer from folder_paths import get_directory_by_type @@ -161,12 +161,18 @@ def load_checkpoint(self, **kwargs): return (*super().load_checkpoint(**kwargs), prompt) +class UNETLoaderWithImages(UNETLoader): + pass + + NODE_CLASS_MAPPINGS = { "LoraLoader|pysssss": LoraLoaderWithImages, "CheckpointLoader|pysssss": CheckpointLoaderSimpleWithImages, + "UNETLoader|pysssss": UNETLoaderWithImages, } NODE_DISPLAY_NAME_MAPPINGS = { "LoraLoader|pysssss": "Lora Loader 🐍", "CheckpointLoader|pysssss": "Checkpoint Loader 🐍", + "UNETLoader|pysssss": "Load Diffusion Model 🐍", } diff --git a/web/js/betterCombos.js b/web/js/betterCombos.js index cbf0a07..7c251d4 100644 --- a/web/js/betterCombos.js +++ b/web/js/betterCombos.js @@ -5,6 +5,7 @@ import { api } from "../../../scripts/api.js"; const CHECKPOINT_LOADER = "CheckpointLoader|pysssss"; const LORA_LOADER = "LoraLoader|pysssss"; +const DIFFUSION_MODEL_LOADER = "UNETLoader|pysssss"; const IMAGE_WIDTH = 384; const IMAGE_HEIGHT = 384; @@ -12,11 +13,28 @@ function getType(node) { if (node.comfyClass === CHECKPOINT_LOADER) { return "checkpoints"; } + if (node.comfyClass === DIFFUSION_MODEL_LOADER) { + return "diffusion_models"; + } return "loras"; } function getWidgetName(type) { - return type === "checkpoints" ? "ckpt_name" : "lora_name"; + if (type === "checkpoints") { + return "ckpt_name"; + } + if (type === "diffusion_models") { + return "unet_name"; + } + return "lora_name"; +} + +function isImageLoader(node) { + return node?.comfyClass === LORA_LOADER || node?.comfyClass === CHECKPOINT_LOADER || node?.comfyClass === DIFFUSION_MODEL_LOADER; +} + +function hasExamples(nodeData) { + return nodeData.name === LORA_LOADER || nodeData.name === CHECKPOINT_LOADER; } function encodeRFC3986URIComponent(str) { @@ -74,7 +92,7 @@ app.registerExtension({ const displayOptions = { "List (normal)": 0, "Tree (subfolders)": 1, "Thumbnails (grid)": 2 }; const displaySetting = app.ui.settings.addSetting({ id: "pysssss.Combo++.Submenu", - name: "🐍 Lora & Checkpoint loader display mode", + name: "🐍 Model loader display mode", defaultValue: 1, type: "combo", options: (value) => { @@ -157,15 +175,14 @@ app.registerExtension({ `, parent: document.body, }); - const p1 = loadImageList("checkpoints"); - const p2 = loadImageList("loras"); + const modelTypes = ["checkpoints", "loras", "diffusion_models"]; + const modelLists = Promise.all(modelTypes.map((type) => loadImageList(type))); const refreshComboInNodes = app.refreshComboInNodes; app.refreshComboInNodes = async function () { const r = await Promise.all([ refreshComboInNodes.apply(this, arguments), - loadImageList("checkpoints").catch(() => {}), - loadImageList("loras").catch(() => {}), + ...modelTypes.map((type) => loadImageList(type).catch(() => {})), ]); return r[0]; }; @@ -192,8 +209,7 @@ app.registerExtension({ const updateMenu = async (menu, type) => { try { - await p1; - await p2; + await modelLists; } catch (error) { console.error(error); console.error("Error loading pysssss.betterCombos data"); @@ -365,7 +381,7 @@ app.registerExtension({ const mutationObserver = new MutationObserver((mutations) => { const node = app.canvas.current_node; - if (!node || (node.comfyClass !== LORA_LOADER && node.comfyClass !== CHECKPOINT_LOADER)) { + if (!isImageLoader(node)) { return; } @@ -395,14 +411,15 @@ app.registerExtension({ mutationObserver.observe(document.body, { childList: true, subtree: false }); }, async beforeRegisterNodeDef(nodeType, nodeData, app) { - const isCkpt = nodeData.name === CHECKPOINT_LOADER; - const isLora = nodeData.name === LORA_LOADER; - if (isCkpt || isLora) { + if (hasExamples(nodeData)) { const onAdded = nodeType.prototype.onAdded; nodeType.prototype.onAdded = function () { onAdded?.apply(this, arguments); const { widget: exampleList } = ComfyWidgets["COMBO"](this, "example", [[""], {}], app); - this.widgets.find((w) => w.name === "prompt").computeSize = () => [0, -4]; + const promptWidget = this.widgets.find((w) => w.name === "prompt"); + if (promptWidget) { + promptWidget.computeSize = () => [0, -4]; + } let exampleWidget; const get = async (route, suffix) => { @@ -500,7 +517,7 @@ app.registerExtension({ img = this.imgs[this.overIndex]; } if (img) { - const nodes = app.graph._nodes.filter((n) => n.comfyClass === LORA_LOADER || n.comfyClass === CHECKPOINT_LOADER); + const nodes = app.graph._nodes.filter((n) => isImageLoader(n)); if (nodes.length) { options.unshift({ content: "Save as Preview", From dd1665bfd7ba6d85a20020e779eb7d49db28b66b Mon Sep 17 00:00:00 2001 From: PALE Date: Fri, 19 Jun 2026 20:04:21 +0900 Subject: [PATCH 2/2] Add diffusion model info menu --- web/js/modelInfo.js | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/web/js/modelInfo.js b/web/js/modelInfo.js index b5db20a..e5e96ae 100644 --- a/web/js/modelInfo.js +++ b/web/js/modelInfo.js @@ -300,6 +300,11 @@ class CheckpointInfoDialog extends ModelInfoDialog { } const lookups = {}; +const modelInfoTypes = { + Lora: { folderType: "loras", infoClass: LoraInfoDialog }, + Checkpoint: { folderType: "checkpoints", infoClass: CheckpointInfoDialog }, + "Diffusion Model": { folderType: "diffusion_models", infoClass: CheckpointInfoDialog }, +}; function addInfoOption(node, type, infoClass, widgetNamePattern, opts) { const widgets = widgetNamePattern @@ -326,12 +331,14 @@ function addInfoOption(node, type, infoClass, widgetNamePattern, opts) { } function addTypeOptions(node, typeName, options) { - const type = typeName.toLowerCase() + "s"; + const typeInfo = modelInfoTypes[typeName]; + if (!typeInfo) return; + + const { folderType: type, infoClass: cls } = typeInfo; const values = lookups[typeName][node.type]; if (!values) return; const widgets = Object.keys(values); - const cls = type === "loras" ? LoraInfoDialog : CheckpointInfoDialog; const opts = []; for (const w of widgets) { @@ -357,9 +364,9 @@ function addTypeOptions(node, typeName, options) { app.registerExtension({ name: "pysssss.ModelInfo", setup() { - const addSetting = (type, defaultValue) => { + const addSetting = (type, defaultValue, idType = type) => { app.ui.settings.addSetting({ - id: `pysssss.ModelInfo.${type}Nodes`, + id: `pysssss.ModelInfo.${idType}Nodes`, name: `🐍 Model Info - ${type} Nodes/Widgets`, type: "text", defaultValue, @@ -384,6 +391,7 @@ app.registerExtension({ "Checkpoint", ["CheckpointLoader.ckpt_name", "CheckpointLoaderSimple", "CheckpointLoader|pysssss", "Efficient Loader", "Eff. Loader SDXL"].join(",") ); + addSetting("Diffusion Model", ["UNETLoader.unet_name", "UNETLoader|pysssss"].join(","), "DiffusionModel"); app.ui.settings.addSetting({ id: `pysssss.ModelInfo.NsfwLevel`,