From b198f7fb258edee2fcae7bb5dbeaa2193c12b70d Mon Sep 17 00:00:00 2001 From: "kk297466058@gmail.com" Date: Fri, 8 Aug 2025 01:50:52 +0800 Subject: [PATCH] feat: add node loading function --- .gitignore | 161 +++++++++++++++++++++++++++++++ README.md | 218 ++++++++++-------------------------------- __init__.py | 51 ++++++---- js/latent_loader.js | 86 +++++++++++++++++ nodes/latent_nodes.py | 76 +++++++++++++++ requirements.txt | 2 + 6 files changed, 405 insertions(+), 189 deletions(-) create mode 100644 .gitignore create mode 100644 js/latent_loader.js create mode 100644 nodes/latent_nodes.py create mode 100644 requirements.txt diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..403ae4e --- /dev/null +++ b/.gitignore @@ -0,0 +1,161 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +.python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install different versions of packages depending +# on the platform. These differences will end up in the Pipfile.lock, which can be annoying +# to manage in a VCS. +# Pipfile.lock + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyderworkspace + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# ComfyUI specific +*.pt +*.pth +*.bin +*.onnx +*.safetensors +*.ckpt +*.jpeg +*.jpg +*.png +*.gif +*.latent +*.wav +*.mp3 +output/ +input/ +temp/ +/custom_nodes/* +!/custom_nodes/ComfyUI-Only +/models/ +/web/ +/comfy/ +.DS_Store + diff --git a/README.md b/README.md index 4264f48..789fa7b 100644 --- a/README.md +++ b/README.md @@ -1,209 +1,87 @@ # ComfyUI-Only -ComfyUI自定义节点集合，专注于workflow解析和图像处理功能。 +ComfyUI自定义节点集合，最初专注于workflow解析，现已扩展至提供高级文件加载功能。 ## 🌟 主要功能 -### 1. Workflow图片文件加载器 (WorkflowImageFileLoader) ⭐推荐 +### 1. 高级Latent加载器 (Load Latent (Advanced)) ⭐全新功能 + +- 📤 **从任意位置上传**：通过“Upload Latent”按钮，直接从您电脑的任何地方选择 `.latent` 文件。 +- ✨ **拖拽上传**：将 `.latent` 文件直接从您的文件管理器拖拽到节点上即可上传。 +- 🔄 **智能格式兼容**：自动识别并兼容标准的 `pickle` 格式和高速的 `safetensors` 格式。 +- 🤖 **智能结构解析**：自动解析多种 latent 内部结构，无论是标准的 `samples` 键，还是非标准的 `latent_tensor` 键，甚至是裸张量。 +- 🔢 **智能维度处理**：自动处理3D（图像）、4D（带批次的图像）和5D（视频）的 latent 张量，确保与下游节点（如VAEDecode）的兼容性。 + +### 2. Workflow图片文件加载器 (WorkflowImageFileLoader) + - 📁 直接从文件系统加载图片文件 - 🔍 自动解析图片中的workflow元数据信息 - 🎯 智能提取positive和negative提示词 -- 🏷️ **新增**：提取CheckpointLoaderSimple节点的模型名称 +- 🏷️ 提取CheckpointLoaderSimple节点的模型名称 - 📋 支持手动输入workflow JSON作为备选 - ✅ **解决了IMAGE类型不包含元数据的问题** -### 2. Workflow图片加载器 (WorkflowImageLoader) -- 📸 接收ComfyUI的IMAGE类型输入 -- ⚠️ **限制**：IMAGE类型不包含原始图片元数据 -- 📝 需要手动输入workflow JSON才能工作 -- 🏷️ **新增**：提取CheckpointLoaderSimple节点的模型名称 -- 🔄 主要用于连接其他节点的图片输出 +--- -### 3. Workflow JSON解析器 (WorkflowJSONParser) -- 📝 直接解析workflow JSON文本 -- 🔎 查找`cnr_id`为`comfyui_custom_nodes_alekpet`的节点 -- 🏷️ **新增**：查找`Node name for S&R`为`CheckpointLoaderSimple`的节点 -- 🧠 智能区分positive/negative提示词 -- 📊 提供详细的解析信息 +*（为保持简洁，其他旧节点说明已折叠，您可以在旧版本中查看它们）* ## 🚀 快速开始 ### 安装 -1. 将此项目克隆到ComfyUI的`custom_nodes`目录下 -2. 重启ComfyUI -3. 在节点菜单中找到`ComfyUI-Only`分类 +1. 将此项目克隆到ComfyUI的`custom_nodes`目录下。 +2. （可选）安装依赖：`pip install -r requirements.txt` +3. 重启ComfyUI。 +4. 在节点菜单中找到`ComfyUI-Only`分类。 ### 使用方法 -#### 方法一：使用图片文件加载器 ⭐推荐 -1. 添加`Workflow图片文件加载器`节点 -2. 在`image_file`下拉菜单中选择要解析的图片文件 -3. 或者点击上传按钮直接上传图片 -4. 自动获取输出的positive和negative提示词 +#### ⭐ 加载Latent文件 (推荐) +1. 在菜单中添加 `Load Latent (Advanced)` 节点。 +2. 点击 `Upload Latent` 按钮从您的电脑中选择一个 `.latent` 文件。 +3. 或者，直接将 `.latent` 文件拖拽到节点上。 +4. 节点会自动处理上传、解析和格式化，然后将 `LATENT` 输出连接到下一个节点（如 `VAEDecode`）。 -#### 方法二：使用图片加载器 + 手动JSON -1. 添加`Workflow图片加载器`节点 -2. 连接图片输入（来自其他节点） -3. **必须**在`workflow_json`字段手动输入JSON -4. 获取输出的positive和negative提示词 - -#### 方法三：使用JSON解析器 -1. 添加`Workflow JSON解析器`节点 -2. 将ComfyUI的workflow JSON粘贴到输入框 -3. 直接获取解析后的提示词 - -## 📋 节点输入输出 - -### WorkflowImageFileLoader ⭐推荐 -**输入:** -- `image_file` (选择列表) - 从文件系统选择图片文件，支持上传 -- `workflow_json` (STRING, 可选) - 手动输入的workflow JSON（备选） - -**输出:** -- `image` (IMAGE) - 加载的图片 -- `positive_prompt` (STRING) - 正面提示词 -- `negative_prompt` (STRING) - 负面提示词 -- `checkpoint_name` (STRING) - **新增**：CheckpointLoaderSimple节点的模型名称 -- `workflow_info` (STRING) - 解析状态信息 - -### WorkflowImageLoader -**输入:** -- `image` (IMAGE) - 要解析的图片（来自其他节点） -- `workflow_json` (STRING, **必需**) - 手动输入的workflow JSON - -**输出:** -- `image` (IMAGE) - 原始图片 -- `positive_prompt` (STRING) - 正面提示词 -- `negative_prompt` (STRING) - 负面提示词 -- `checkpoint_name` (STRING) - **新增**：CheckpointLoaderSimple节点的模型名称 -- `workflow_info` (STRING) - 解析状态信息 - -### WorkflowJSONParser -**输入:** -- `workflow_json` (STRING) - ComfyUI workflow JSON - -**输出:** -- `positive_prompt` (STRING) - 正面提示词 -- `negative_prompt` (STRING) - 负面提示词 -- `checkpoint_name` (STRING) - **新增**：CheckpointLoaderSimple节点的模型名称 -- `parse_info` (STRING) - 解析状态信息 - -## 🔧 工作原理 - -1. **节点识别**: 在workflow JSON中搜索两种类型的节点： - - `properties.cnr_id`为`comfyui_custom_nodes_alekpet`的节点（提示词） - - `properties["Node name for S&R"]`为`CheckpointLoaderSimple`的节点（模型） -2. **文本提取**: - - 从alekpet节点的`widgets_values[0]`中获取提示词内容 - - 从CheckpointLoaderSimple节点的`widgets_values[0]`中获取模型名称 -3. **智能分类**: 根据关键词分析，判断是positive还是negative提示词： - - **Positive**: 包含"masterpiece"、"best quality"、"best"等关键词 - - **Negative**: 包含"worst"、"bad"等关键词 - - **特征判断**: 以`, masterpiece, best quality..." - ] - }, - { - "id": 7, - "type": "CheckpointLoaderSimple", - "properties": { - "Node name for S&R": "CheckpointLoaderSimple" - }, - "widgets_values": [ - "uncannyValley_ilxl10Noob.safetensors" - ] - } - ] -} -``` - -**输出:** -- positive_prompt: `, masterpiece, best quality...` -- negative_prompt: `worst quality, bad anatomy, watermark...` -- checkpoint_name: `uncannyValley_ilxl10Noob.safetensors` - ## ⚙️ 依赖要求 - Python 3.8+ - ComfyUI - PyTorch -- Pillow -- NumPy +- **safetensors** (新) ## 📝 更新日志 +### v2.0.0 - 高级加载器版本 (开发代号：Phoenix) +- 🔥 **全新**：发布 `Load Latent (Advanced)` 节点，支持从任意位置上传和拖拽 latent 文件。 +- ✨ **新增**：为新节点编写了独立的前端JS扩展，实现了原生级别的UI体验。 +- 🤖 **增强**：加载器后端实现智能格式兼容（pickle/safetensors）、智能结构解析和智能维度处理。 +- 🧹 **重构**：项目结构调整，增加了 `js` 目录和 `requirements.txt`。 +- 📚 **文档**：全面重写 `README.md`，聚焦新功能，简化旧说明。 + ### v1.2.0 -- ✅ **新增**：CheckpointLoaderSimple节点的模型名称提取 +- ✅ 新增：CheckpointLoaderSimple节点的模型名称提取 - ✅ 所有节点新增checkpoint_name输出 -- ✅ 改进workflow解析算法，支持多种节点类型提取 -- ✅ 完善示例和文档说明 -### v1.1.0 -- ✅ 新增WorkflowImageFileLoader节点，支持直接读取图片文件 -- ✅ 改进workflow元数据提取算法 -- ✅ 完善错误处理和用户提示 - -### v1.0.0 -- ✅ 实现Workflow图片加载器 -- ✅ 实现Workflow JSON解析器 -- ✅ 智能提示词分类算法 -- ✅ 完整的错误处理和状态反馈 +*（更早的日志已省略）* ## 🤝 贡献 @@ -211,4 +89,4 @@ ComfyUI自定义节点集合，专注于workflow解析和图像处理功能。 ## 📄 许可证 -MIT License - 详见 [LICENSE](LICENSE) 文件 \ No newline at end of file +MIT License - 详见 [LICENSE](LICENSE) 文件 diff --git a/__init__.py b/__init__.py index 28f4811..25dff40 100644 --- a/__init__.py +++ b/__init__.py @@ -1,32 +1,45 @@ -""" -ComfyUI-Only 自定义节点包 -这是一个ComfyUI自定义节点集合，提供各种图像处理、文本工具和模型包装功能 -""" +import os +import importlib -# 导入节点模块 -from .nodes.image_processing_nodes import NODE_CLASS_MAPPINGS as IMAGE_NODES, NODE_DISPLAY_NAME_MAPPINGS as IMAGE_NAMES -# from .nodes.text_utility_nodes import NODE_CLASS_MAPPINGS as TEXT_NODES, NODE_DISPLAY_NAME_MAPPINGS as TEXT_NAMES -# from .nodes.model_wrapper_nodes import NODE_CLASS_MAPPINGS as MODEL_NODES, NODE_DISPLAY_NAME_MAPPINGS as MODEL_NAMES +# 告诉 ComfyUI 服务器我们的 Web 目录 +WEB_DIRECTORY = "./js" # ComfyUI 节点映射字典 NODE_CLASS_MAPPINGS = {} NODE_DISPLAY_NAME_MAPPINGS = {} -# 注册图像处理节点 -NODE_CLASS_MAPPINGS.update(IMAGE_NODES) -NODE_DISPLAY_NAME_MAPPINGS.update(IMAGE_NAMES) +# 动态加载所有节点 +# ====================================================================================================================== +# 获取当前文件所在目录 +current_dir = os.path.dirname(os.path.abspath(__file__)) +nodes_dir = os.path.join(current_dir, "nodes") -# 注册其他节点（暂时注释，等待实现） -# NODE_CLASS_MAPPINGS.update(TEXT_NODES) -# NODE_DISPLAY_NAME_MAPPINGS.update(TEXT_NAMES) -# NODE_CLASS_MAPPINGS.update(MODEL_NODES) -# NODE_DISPLAY_NAME_MAPPINGS.update(MODEL_NAMES) +# 遍历 nodes 目录下的所有 python 文件 +for filename in os.listdir(nodes_dir): + if filename.endswith(".py") and filename != "__init__.py": + module_name = f".nodes.{filename[:-3]}" + try: + # 动态导入模块 + module = importlib.import_module(module_name, package=__name__) + + # 从模块中获取节点映射 + if hasattr(module, "NODE_CLASS_MAPPINGS") and hasattr(module, "NODE_DISPLAY_NAME_MAPPINGS"): + NODE_CLASS_MAPPINGS.update(module.NODE_CLASS_MAPPINGS) + NODE_DISPLAY_NAME_MAPPINGS.update(module.NODE_DISPLAY_NAME_MAPPINGS) + print(f" - 成功加载模块: {module_name}") + else: + print(f" - 警告: 模块 {module_name} 缺少必要的映射变量") + + except Exception as e: + print(f" - 错误: 加载模块 {module_name} 失败 - {e}") +# ====================================================================================================================== # 导出ComfyUI需要的变量 -__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] print("🎨 ComfyUI-Only 自定义节点已加载") +print(f" - 前端扩展目录: {WEB_DIRECTORY}") print(f" 📦 已注册 {len(NODE_CLASS_MAPPINGS)} 个节点:") for name in NODE_CLASS_MAPPINGS.keys(): - print(f" - {name}") - + display_name = NODE_DISPLAY_NAME_MAPPINGS.get(name, name) + print(f" - {name} ({display_name})") \ No newline at end of file diff --git a/js/latent_loader.js b/js/latent_loader.js new file mode 100644 index 0000000..a3c47f9 --- /dev/null +++ b/js/latent_loader.js @@ -0,0 +1,86 @@ +import { app } from "/scripts/app.js"; +import { api } from "/scripts/api.js"; + +app.registerExtension({ + name: "Comfy.LatentLoader.Advanced.Final", // 使用新名称以避免浏览器缓存问题 + async beforeRegisterNodeDef(nodeType, nodeData) { + if (nodeData.name === "LatentLoaderAdvanced") { + + const onNodeCreated = nodeType.prototype.onNodeCreated; + nodeType.prototype.onNodeCreated = function () { + onNodeCreated?.apply(this, arguments); + + // 上传文件的共享函数 + const uploadFile = async (file) => { + try { + const body = new FormData(); + body.append("image", file); + body.append("overwrite", "true"); + const resp = await api.fetchApi("/upload/image", { + method: "POST", + body, + }); + + if (resp.status === 200) { + const data = await resp.json(); + const path = data.subfolder ? `${data.subfolder}/${data.name}` : data.name; + const textWidget = this.widgets.find((w) => w.name === "latent_file"); + if (textWidget) { + textWidget.value = path; + } + } else { + alert(`Upload Error: ${resp.status} - ${resp.statusText}`); + } + } catch (error) { + console.error("Upload failed:", error); + alert(`Upload failed: ${error}`); + } + }; + + // “上传”按钮的逻辑 + this.addWidget("button", "upload_latent", "Upload Latent", () => { + const inputEl = document.createElement("input"); + inputEl.type = "file"; + inputEl.accept = ".latent"; + document.body.appendChild(inputEl); + + const handleFileSelect = () => { + if (inputEl.files.length > 0) { + uploadFile(inputEl.files[0]); + } + inputEl.remove(); + }; + + inputEl.addEventListener("change", handleFileSelect); + inputEl.style.display = "none"; + inputEl.click(); + }); + + // 拖拽事件处理 + this.onDragOver = function(e) { + // 检查拖拽的是否是文件 + if (e.dataTransfer?.types.includes("Files")) { + e.preventDefault(); // 关键：阻止浏览器默认行为，以允许拖放 + return true; + } + return false; + }; + + this.onDragDrop = async function(e) { + e.preventDefault(); // 关键：阻止浏览器默认行为 + e.stopPropagation(); // 可选：阻止事件冒泡 + + let handled = false; + for (const file of e.dataTransfer.files) { + if (file.name.endsWith(".latent")) { + await uploadFile(file); // 等待上传完成 + handled = true; + break; // 只处理第一个有效文件 + } + } + return handled; + }; + }; + } + }, +}); diff --git a/nodes/latent_nodes.py b/nodes/latent_nodes.py new file mode 100644 index 0000000..a239bff --- /dev/null +++ b/nodes/latent_nodes.py @@ -0,0 +1,76 @@ +# -*- coding: utf-8 -*- + +import torch +import os +import folder_paths +import safetensors.torch + +class LatentLoaderAdvanced: + """ + 一个高级 Latent 加载器，通过自定义前端 UI 支持从外部拖拽或上传 .latent 文件。 + """ + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "latent_file": ("STRING", {"default": "", "multiline": False}), + }, + } + + RETURN_TYPES = ("LATENT",) + FUNCTION = "load_latent" + CATEGORY = "ComfyUI-Only/Latent" + + def load_latent(self, latent_file): + latent_path = folder_paths.get_annotated_filepath(latent_file) + + if not os.path.exists(latent_path): + raise FileNotFoundError(f"File not found at path: {latent_path}.") + + latent_data = None + try: + latent_data = safetensors.torch.load_file(latent_path, device="cpu") + except Exception: + try: + latent_data = torch.load(latent_path, map_location="cpu", weights_only=False) + except Exception as e: + raise RuntimeError(f"Failed to load file '{latent_file}'. It's not a valid safetensors or PyTorch file. Error: {e}") + + samples = None + if isinstance(latent_data, dict): + if 'samples' in latent_data and torch.is_tensor(latent_data['samples']): + samples = latent_data['samples'] + elif 'latent_tensor' in latent_data and torch.is_tensor(latent_data['latent_tensor']): + samples = latent_data['latent_tensor'] + else: + for key, value in latent_data.items(): + if torch.is_tensor(value) and value.numel() > 0: + samples = value + break + elif torch.is_tensor(latent_data): + samples = latent_data + + if samples is not None: + if samples.numel() == 0: + raise ValueError(f"The loaded latent file '{latent_file}' is empty or contains an empty tensor.") + + if samples.ndim == 3: + samples = samples.unsqueeze(0) + + if samples.ndim not in [4, 5]: + raise ValueError(f"Loaded latent tensor from '{latent_file}' has an unsupported shape: {samples.shape}. Expected a 3D, 4D or 5D tensor.") + + return ({"samples": samples},) + else: + raise ValueError(f"Could not extract a valid latent tensor from '{latent_file}'. The format may not be recognized.") + + +# 节点映射 +NODE_CLASS_MAPPINGS = { + "LatentLoaderAdvanced": LatentLoaderAdvanced, +} + +# 节点显示名称映射 +NODE_DISPLAY_NAME_MAPPINGS = { + "LatentLoaderAdvanced": "Load Latent (Advanced)", +} diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..01f8717 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,2 @@ +safetensors>=0.4.0 +