commit 02a0b7082cca84228ea1dcc43bb62ee05d1faece Author: XinPeng76 Date: Wed Jul 8 14:05:13 2026 +0800 initial public release diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..4d809b9 --- /dev/null +++ b/.gitignore @@ -0,0 +1,10 @@ +__pycache__/ +*.py[cod] +*$py.class + +.env +.venv/ +venv/ + +.DS_Store +Thumbs.db diff --git a/README.md b/README.md new file mode 100644 index 0000000..fa05521 --- /dev/null +++ b/README.md @@ -0,0 +1,155 @@ +# ComfyUI_TJS - Truncated Jump Sampling + +中文 | [English](README_EN.md) + +ComfyUI_TJS 是一个用于 ComfyUI 的自定义节点插件，用来实验 +**TJS（Truncated Jump Sampling，截断跳步采样）**。 + +该插件用于验证一种基于 endpoint / denoised 预测的训练无关采样加速思路。 + +## 功能简介 + +TJS 的核心思想是：不一定要把扩散/流匹配采样轨迹完整跑完。模型在中间步已经能给出一个 +`x0` / `denoised` 端点预测，因此可以先运行一部分采样步骤，再直接输出端点估计，从而减少 NFE。 + +简化流程如下： + +```text +k* = ceil(gamma * K) +run sampler from sigma[0] to sigma[k*] +return x0_hat(xt, sigma[k*]) with one endpoint model call +``` + +其中： + +- `K` 是原始完整采样步数。 +- `gamma` 是早退比例，例如 `0.6` 表示只跑大约 60% 的轨迹。 +- `k*` 是实际早退步。 +- `x0_hat` 是模型在中间 latent 上预测的干净端点。 + +在 ComfyUI 中，endpoint estimate 对应模型包装器给出的原生 `denoised` latent，因此理论上可以兼容 diffusion、flow matching、Flux/SD3 类模型包装。 + +当 `gamma = 1.0` 时，节点会运行完整采样 schedule，并跳过额外 endpoint call，等价于普通 KSampler 的边界情况。 + +## 安装 + +把本文件夹复制到 ComfyUI 的自定义节点目录： + +```text +ComfyUI/custom_nodes/ComfyUI_TJS/ +``` + +然后重启 ComfyUI。 + +## 节点说明 + +### TJS Sampler (Truncated Jump Sampling) + +这是主要可用节点，可以作为普通文生图采样器的替代节点使用。 + +输入： + +| Input | 中文说明 | +|---|---| +| `model` | ComfyUI 加载的模型 | +| `total_steps` | 原始完整采样步数 `K`，例如 `30` | +| `early_exit_gamma` | 早退比例，例如 `0.6` | +| `cfg` | CFG scale | +| `sampler_name` | 截断采样阶段使用的 ComfyUI sampler | +| `scheduler` | ComfyUI sigma scheduler | +| `positive` / `negative` | 正向/反向条件 | +| `latent_image` | 空 latent 或输入 latent | +| `seed` | 随机种子 | +| `model_type` | 信息标记：`auto`、`diffusion` 或 `flow` | +| `denoise` | 可选 denoise 强度 | + +输出： + +| Output | 中文说明 | +|---|---| +| `latent_x0` | endpoint decoded latent，接到 VAE Decode | +| `latent_xt` | 早退时刻的中间 noisy latent | +| `k_star` | 实际早退步 | +| `nfe_used` | 早退时为 `k_star + 1`；`gamma = 1.0` 时为 `K` | +| `nfe_saving_pct` | 相对完整步数节省的 NFE 百分比 | +| `sigma_at_exit` | endpoint decode 使用的 sigma | + +### TJS Decode (Endpoint / Advanced KSampler) + +状态：实验性 / TODO。 + +这个节点目前还没有完全写对，暂时不要把它当成可靠正式功能。它的目标用法是配合 +KSampler Advanced：先让 KSampler Advanced 跑到某个中间 `end_at_step`，保留 leftover noisy latent， +再用 TJS Decode 把这个中间状态转换成 endpoint `x0` 预测。 + +后续需要继续修复和验证的问题： + +- 如何严格对齐 KSampler Advanced 的中间 latent 与对应 sigma。 +- 如何保证不同 sampler / scheduler 下 endpoint decode 的输入缩放完全正确。 +- 如何在工作流里可靠传递中间步信息，而不是让用户手动猜连续时间 `t`。 + +当前节点设计为：用户不需要知道连续时间 `t` 或 sigma，只需要填入和 KSampler Advanced 相同的 +`steps`、`sampler_name`、`scheduler`、`end_at_step`，节点内部尝试重建： + +```text +sigma* = sigmas[end_at_step] +``` + +注意：在 KSampler Advanced 中，需要设置： + +```text +return_with_leftover_noise = enable +``` + +如果关闭 leftover noise，高级采样器会强制采到 sigma=0，输出已经是完整去噪 latent，此时就没有可供 TJS decode 的中间 noisy endpoint state。 + +### TJS Decode (Manual Sigma) + +调试节点。只有当你已经明确知道某个中间 latent 对应的精确 sigma 时才建议使用。 + +普通 Advanced KSampler 工作流不要优先用这个节点。 + +## 使用示例 + +### 直接使用 TJS Sampler + +例如设置： + +```text +total_steps = 30 +early_exit_gamma = 0.6 +``` + +节点会运行到： + +```text +k* = ceil(0.6 * 30) = 18 +``` + +然后执行一次 endpoint call。总成本约为： + +```text +NFE = 18 + 1 = 19 +``` + +相对 30 步完整采样，约节省 `36.7%` NFE。 + +### 计划中的 Advanced KSampler 用法 + +这部分仍是 TODO，当前 TJS Decode 节点还需要修复。 + +目标工作流如下： + +1. 使用 KSampler Advanced。 +2. 设置 `steps = 30`，`end_at_step = 18`。 +3. 设置 `return_with_leftover_noise = enable`。 +4. 把 KSampler Advanced 输出的中间 latent 接到 `TJS Decode (Endpoint / Advanced KSampler)`。 +5. TJS Decode 使用相同的 `steps`、`end_at_step`、`sampler_name`、`scheduler`。 +6. 把 `latent_x0` 接到 VAE Decode。 + +## 当前限制 + +- `TJS Sampler` 是当前主要实现。 +- `TJS Decode (Endpoint / Advanced KSampler)` 仍是实验节点，需要后续修复。 +- 不同模型、sampler、scheduler 对 latent 缩放和 denoised 语义可能不同，实际效果需要逐模型验证。 +- 当前插件更适合作为算法验证和 ComfyUI 原型实验，不建议直接用于严肃生产流程。 diff --git a/README_EN.md b/README_EN.md new file mode 100644 index 0000000..c233633 --- /dev/null +++ b/README_EN.md @@ -0,0 +1,171 @@ +# ComfyUI_TJS - Truncated Jump Sampling + +[中文](README.md) | English + +ComfyUI_TJS is a custom node plugin for ComfyUI, built to experiment with +**TJS (Truncated Jump Sampling)**. + +This plugin explores a training-free sampling acceleration idea based on +endpoint / denoised prediction. + +## What It Does + +TJS is based on a simple idea: diffusion and flow-matching samplers may not +need to traverse the full trajectory. At an intermediate step, the model +already produces an `x0` / `denoised` endpoint estimate. TJS runs only part of +the sampling trajectory, then uses that endpoint prediction as the final latent. + +Simplified algorithm: + +```text +k* = ceil(gamma * K) +run sampler from sigma[0] to sigma[k*] +return x0_hat(xt, sigma[k*]) with one endpoint model call +``` + +Where: + +- `K` is the full sampling step budget. +- `gamma` is the early-exit ratio, for example `0.6`. +- `k*` is the actual early-exit step. +- `x0_hat` is the clean endpoint estimate predicted from the intermediate latent. + +In ComfyUI, this endpoint estimate corresponds to the model wrapper's native +`denoised` latent. In principle, this makes the approach usable with diffusion, +flow-matching, Flux/SD3-style model wrappers, and related models. + +When `gamma = 1.0`, the node runs the full sampling schedule and skips the +extra endpoint call, matching the ordinary KSampler boundary case. + +## Installation + +Copy this folder into ComfyUI's custom node directory: + +```text +ComfyUI/custom_nodes/ComfyUI_TJS/ +``` + +Then restart ComfyUI. + +## Nodes + +### TJS Sampler (Truncated Jump Sampling) + +This is the main usable node. It can be used as an experimental replacement for +a normal text-to-image sampler. + +Inputs: + +| Input | Description | +|---|---| +| `model` | Loaded ComfyUI model | +| `total_steps` | Full sampling step budget `K`, for example `30` | +| `early_exit_gamma` | Early-exit ratio, for example `0.6` | +| `cfg` | CFG scale | +| `sampler_name` | ComfyUI sampler used for the truncated trajectory | +| `scheduler` | ComfyUI sigma scheduler | +| `positive` / `negative` | Positive / negative conditioning | +| `latent_image` | Empty latent or input latent | +| `seed` | Random seed | +| `model_type` | Informational selector: `auto`, `diffusion`, or `flow` | +| `denoise` | Optional denoise strength | + +Outputs: + +| Output | Description | +|---|---| +| `latent_x0` | Endpoint decoded latent; connect this to VAE Decode | +| `latent_xt` | Intermediate noisy latent at the early-exit step | +| `k_star` | Actual early-exit step | +| `nfe_used` | `k_star + 1` for early exit, or `K` when `gamma = 1.0` | +| `nfe_saving_pct` | NFE saving percentage relative to full sampling | +| `sigma_at_exit` | Sigma used by the endpoint decode | + +### TJS Decode (Endpoint / Advanced KSampler) + +Status: experimental / TODO. + +This node is not fully correct yet. Do not treat it as a reliable production +feature. The intended future workflow is to pair it with KSampler Advanced: +run KSampler Advanced to an intermediate `end_at_step`, keep the leftover noisy +latent, then use TJS Decode to convert that intermediate state into an endpoint +`x0` prediction. + +Open issues to fix: + +- Strictly align the KSampler Advanced intermediate latent with the correct sigma. +- Ensure the input scaling for endpoint decode is correct across samplers and schedulers. +- Pass intermediate step information reliably through the workflow instead of asking users to guess continuous time `t`. + +The current node attempts to infer: + +```text +sigma* = sigmas[end_at_step] +``` + +from the same `steps`, `sampler_name`, `scheduler`, and `end_at_step` used by +KSampler Advanced. + +Important: in KSampler Advanced, use: + +```text +return_with_leftover_noise = enable +``` + +If leftover noise is disabled, KSampler Advanced forces the final sigma to zero, +so the output is already fully denoised and there is no intermediate noisy state +left for TJS Decode. + +### TJS Decode (Manual Sigma) + +Debug-only node. Use it only when you already know the exact sigma for an +intermediate latent. + +For normal Advanced KSampler workflows, prefer the planned +`TJS Decode (Endpoint / Advanced KSampler)` workflow after it is fixed. + +## Example + +### Direct TJS Sampler Usage + +Example settings: + +```text +total_steps = 30 +early_exit_gamma = 0.6 +``` + +The node runs to: + +```text +k* = ceil(0.6 * 30) = 18 +``` + +Then it performs one endpoint call: + +```text +NFE = 18 + 1 = 19 +``` + +Compared with full 30-step sampling, this saves about `36.7%` NFE. + +### Planned Advanced KSampler Workflow + +This part is still TODO because the current TJS Decode node needs further +repair and validation. + +Target workflow: + +1. Use KSampler Advanced. +2. Set `steps = 30` and `end_at_step = 18`. +3. Set `return_with_leftover_noise = enable`. +4. Feed the intermediate latent into `TJS Decode (Endpoint / Advanced KSampler)`. +5. Use the same `steps`, `end_at_step`, `sampler_name`, and `scheduler` in TJS Decode. +6. Connect `latent_x0` to VAE Decode. + +## Current Limitations + +- `TJS Sampler` is the main implemented path. +- `TJS Decode (Endpoint / Advanced KSampler)` is still experimental and needs future repair. +- Latent scaling and `denoised` semantics may differ across models, samplers, and schedulers. +- This plugin is intended for research prototyping and ComfyUI experiments, not production workflows. diff --git a/SDXL-test.ipynb b/SDXL-test.ipynb new file mode 100644 index 0000000..ab8b8be --- /dev/null +++ b/SDXL-test.ipynb @@ -0,0 +1,310 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "id": "0fafbffb", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "2026-06-23 23:56:09.848939: I tensorflow/core/util/port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n", + "2026-06-23 23:56:09.873201: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n", + "To enable the following instructions: AVX2 AVX_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n" + ] + }, + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "04b0363565e64194a37b97ca8983f4d0", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Loading pipeline components...: 0%| | 0/7 [00:00" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Saved to figs/x0_visual_comparison.pdf and .png\n" + ] + } + ], + "source": [ + "import matplotlib.pyplot as plt\n", + "import matplotlib.gridspec as gridspec\n", + "\n", + "n_prompts = len(PROMPTS)\n", + "n_steps = len(CAPTURE_STEPS)\n", + "\n", + "# 核心逻辑：每 1 个提示词占用 2 行：第 1 行是图像序列，第 2 行是纯文字标签\n", + "# 总行数为 n_prompts * 2\n", + "fig = plt.figure(figsize=(n_steps * 2.2, n_prompts * 2.8))\n", + "gs = gridspec.GridSpec(n_prompts * 2, n_steps, figure=fig, \n", + " height_ratios=[1, 0.2] * n_prompts, # 图像高，文字区矮\n", + " hspace=0.05, wspace=0.05, \n", + " bottom=0.05, top=0.95, left=0.05, right=0.95)\n", + "\n", + "for ri, prompt in enumerate(PROMPTS):\n", + " images = all_x0[prompt]\n", + " \n", + " # 1. 绘制图像序列 (占据每一组的第 0 行)\n", + " for ci in range(n_steps):\n", + " ax = plt.subplot(gs[ri * 2, ci])\n", + " ax.imshow(images[ci])\n", + " ax.set_xticks([])\n", + " ax.set_yticks([])\n", + " \n", + " # 顶部 NFE 标识\n", + " if ri == 0:\n", + " k = CAPTURE_STEPS[ci]\n", + " ax.set_title(f\"NFE={k+1}\", fontsize=12, fontweight='bold', pad=10)\n", + " \n", + " # 2. 绘制提示词 (占据每一组的第 1 行)\n", + " # 跨越所有列 (colspan=n_steps)\n", + " ax_text = plt.subplot(gs[ri * 2 + 1, :])\n", + " ax_text.text(0.5, 0.5, f\"Prompt: {prompt}\", \n", + " fontsize=12, fontweight='bold', \n", + " ha='center', va='center')\n", + " ax_text.axis('off') # 隐藏坐标轴\n", + "\n", + "plt.show()\n", + "#fig.savefig(\"figs/x0_visual_comparison.pdf\", dpi=300, bbox_inches=\"tight\", pad_inches=0.05)\n", + "#fig.savefig(\"figs/x0_visual_comparison.png\", dpi=200, bbox_inches=\"tight\", pad_inches=0.05)\n", + "print(\"Saved to figs/x0_visual_comparison.pdf and .png\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "70661b0e", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "torch2.8", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.12" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/TJS.json b/TJS.json new file mode 100644 index 0000000..eb087d3 --- /dev/null +++ b/TJS.json @@ -0,0 +1 @@ +{"id":"28031df2-54e9-405b-a770-d935b5fff509","revision":0,"last_node_id":30,"last_link_id":51,"nodes":[{"id":20,"type":"VAEDecode","pos":[2106.274488284699,1073.2045011376285],"size":[210,46],"flags":{},"order":9,"mode":0,"inputs":[{"label":"Latent","localized_name":"Latent","name":"samples","type":"LATENT","link":26},{"label":"VAE","localized_name":"vae","name":"vae","type":"VAE","link":27}],"outputs":[{"label":"图像","localized_name":"图像","name":"IMAGE","type":"IMAGE","slot_index":0,"links":[28]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.75","Node name for S&R":"VAEDecode","widget_ue_connectable":{}},"widgets_values":[]},{"id":12,"type":"SaveImage","pos":[1419.1406290017885,315.19472595192394],"size":[210,270],"flags":{},"order":10,"mode":0,"inputs":[{"label":"图像","localized_name":"图片","name":"images","type":"IMAGE","link":15},{"label":"filename_prefix","localized_name":"文件名前缀","name":"filename_prefix","type":"STRING","widget":{"name":"filename_prefix"},"link":null}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.3.75","Node name for S&R":"SaveImage","widget_ue_connectable":{}},"widgets_values":["ComfyUI"]},{"id":11,"type":"VAEDecode","pos":[1195.8677790253187,322.6816177414051],"size":[210,46],"flags":{},"order":7,"mode":0,"inputs":[{"label":"Latent","localized_name":"Latent","name":"samples","type":"LATENT","link":14},{"label":"VAE","localized_name":"vae","name":"vae","type":"VAE","link":17}],"outputs":[{"label":"图像","localized_name":"图像","name":"IMAGE","type":"IMAGE","slot_index":0,"links":[15]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.75","Node name for S&R":"VAEDecode","widget_ue_connectable":{}},"widgets_values":[]},{"id":9,"type":"SaveImage","pos":[1639.4469633436274,318.5350523011988],"size":[210,270.00000000000006],"flags":{},"order":11,"mode":0,"inputs":[{"label":"图像","localized_name":"图片","name":"images","type":"IMAGE","link":9},{"label":"filename_prefix","localized_name":"文件名前缀","name":"filename_prefix","type":"STRING","widget":{"name":"filename_prefix"},"link":null}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.3.75","Node name for S&R":"SaveImage","widget_ue_connectable":{}},"widgets_values":["ComfyUI"]},{"id":3,"type":"KSampler","pos":[925.5265539168148,201.27608189006332],"size":[315,474],"flags":{},"order":4,"mode":0,"inputs":[{"label":"模型","localized_name":"模型","name":"model","type":"MODEL","link":1},{"label":"正面条件","localized_name":"正面条件","name":"positive","type":"CONDITIONING","link":4},{"label":"负面条件","localized_name":"负面条件","name":"negative","type":"CONDITIONING","link":6},{"label":"Latent","localized_name":"Latent图像","name":"latent_image","type":"LATENT","link":2},{"label":"seed","localized_name":"随机种","name":"seed","type":"INT","widget":{"name":"seed"},"link":null},{"label":"steps","localized_name":"步数","name":"steps","type":"INT","widget":{"name":"steps"},"link":null},{"label":"cfg","localized_name":"CFG","name":"cfg","type":"FLOAT","widget":{"name":"cfg"},"link":null},{"label":"sampler_name","localized_name":"采样器","name":"sampler_name","type":"COMBO","widget":{"name":"sampler_name"},"link":null},{"label":"scheduler","localized_name":"调度器","name":"scheduler","type":"COMBO","widget":{"name":"scheduler"},"link":null},{"label":"denoise","localized_name":"降噪","name":"denoise","type":"FLOAT","widget":{"name":"denoise"},"link":null}],"outputs":[{"label":"Latent","localized_name":"Latent","name":"LATENT","type":"LATENT","slot_index":0,"links":[14]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.75","Node name for S&R":"KSampler","widget_ue_connectable":{}},"widgets_values":[156680208700287,"fixed",20,7.5,"euler","karras",1]},{"id":6,"type":"CLIPTextEncode","pos":[333.1020376965134,167.33970479161064],"size":[422.84503173828125,164.31304931640625],"flags":{},"order":2,"mode":0,"inputs":[{"label":"CLIP","localized_name":"clip","name":"clip","type":"CLIP","link":3},{"label":"text","localized_name":"文本","name":"text","type":"STRING","widget":{"name":"text"},"link":null}],"outputs":[{"label":"条件","localized_name":"条件","name":"CONDITIONING","type":"CONDITIONING","slot_index":0,"links":[4,10,23]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.75","Node name for S&R":"CLIPTextEncode","widget_ue_connectable":{}},"widgets_values":["a golden retriever playing in a field of sunflowers, sunny day, 4k\n",[false,true]]},{"id":7,"type":"CLIPTextEncode","pos":[413,389],"size":[425.27801513671875,180.6060791015625],"flags":{},"order":3,"mode":0,"inputs":[{"label":"CLIP","localized_name":"clip","name":"clip","type":"CLIP","link":5},{"label":"text","localized_name":"文本","name":"text","type":"STRING","widget":{"name":"text"},"link":null}],"outputs":[{"label":"条件","localized_name":"条件","name":"CONDITIONING","type":"CONDITIONING","slot_index":0,"links":[6,11,24]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.75","Node name for S&R":"CLIPTextEncode","widget_ue_connectable":{}},"widgets_values":["text, watermark",[false,true]]},{"id":8,"type":"VAEDecode","pos":[1374.1153394196851,747.6203684212926],"size":[210,46],"flags":{},"order":8,"mode":0,"inputs":[{"label":"Latent","localized_name":"Latent","name":"samples","type":"LATENT","link":13},{"label":"VAE","localized_name":"vae","name":"vae","type":"VAE","link":8}],"outputs":[{"label":"图像","localized_name":"图像","name":"IMAGE","type":"IMAGE","slot_index":0,"links":[9]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.75","Node name for S&R":"VAEDecode","widget_ue_connectable":{}},"widgets_values":[]},{"id":21,"type":"SaveImage","pos":[1863.736971362057,313.6466554311809],"size":[210,270.00000000000006],"flags":{},"order":12,"mode":0,"inputs":[{"label":"图像","localized_name":"图片","name":"images","type":"IMAGE","link":28},{"label":"filename_prefix","localized_name":"文件名前缀","name":"filename_prefix","type":"STRING","widget":{"name":"filename_prefix"},"link":null}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.3.75","Node name for S&R":"SaveImage","widget_ue_connectable":{}},"widgets_values":["ComfyUI"]},{"id":5,"type":"EmptyLatentImage","pos":[81.23746893053843,700.9537880546733],"size":[315,106],"flags":{},"order":0,"mode":0,"inputs":[{"label":"width","localized_name":"宽度","name":"width","type":"INT","widget":{"name":"width"},"link":null},{"label":"height","localized_name":"高度","name":"height","type":"INT","widget":{"name":"height"},"link":null},{"label":"batch_size","localized_name":"批次大小","name":"batch_size","type":"INT","widget":{"name":"batch_size"},"link":null}],"outputs":[{"label":"Latent","localized_name":"Latent","name":"LATENT","type":"LATENT","slot_index":0,"links":[2,12,25]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.75","Node name for S&R":"EmptyLatentImage","widget_ue_connectable":{}},"widgets_values":[1024,1024,1]},{"id":10,"type":"TJSSampler","pos":[884.7394908592985,914.0882997056509],"size":[356.917578125,562.0000000000001],"flags":{},"order":5,"mode":0,"inputs":[{"label":"model","localized_name":"model","name":"model","type":"MODEL","link":16},{"label":"positive","localized_name":"positive","name":"positive","type":"CONDITIONING","link":10},{"label":"negative","localized_name":"negative","name":"negative","type":"CONDITIONING","link":11},{"label":"latent_image","localized_name":"latent_image","name":"latent_image","type":"LATENT","link":12},{"label":"total_steps","localized_name":"total_steps","name":"total_steps","type":"INT","widget":{"name":"total_steps"},"link":null},{"label":"early_exit_gamma","localized_name":"early_exit_gamma","name":"early_exit_gamma","type":"FLOAT","widget":{"name":"early_exit_gamma"},"link":null},{"label":"cfg","localized_name":"cfg","name":"cfg","type":"FLOAT","widget":{"name":"cfg"},"link":null},{"label":"sampler_name","localized_name":"sampler_name","name":"sampler_name","type":"COMBO","widget":{"name":"sampler_name"},"link":null},{"label":"scheduler","localized_name":"scheduler","name":"scheduler","type":"COMBO","widget":{"name":"scheduler"},"link":null},{"label":"seed","localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null},{"label":"model_type","localized_name":"model_type","name":"model_type","type":"COMBO","widget":{"name":"model_type"},"link":null},{"label":"denoise","localized_name":"denoise","name":"denoise","shape":7,"type":"FLOAT","widget":{"name":"denoise"},"link":null}],"outputs":[{"label":"latent_x0","localized_name":"latent_x0","name":"latent_x0","type":"LATENT","links":[13]},{"label":"latent_xt","localized_name":"latent_xt","name":"latent_xt","type":"LATENT","links":[]},{"label":"k_star","localized_name":"k_star","name":"k_star","type":"INT","links":null},{"label":"nfe_used","localized_name":"nfe_used","name":"nfe_used","type":"INT","links":null},{"label":"nfe_saving_pct","localized_name":"nfe_saving_pct","name":"nfe_saving_pct","type":"FLOAT","links":null},{"label":"sigma_at_exit","localized_name":"sigma_at_exit","name":"sigma_at_exit","type":"FLOAT","links":null}],"properties":{"Node name for S&R":"TJSSampler","widget_ue_connectable":{}},"widgets_values":[20,0.71,7.5,"euler","karras",156680208700287,"fixed","diffusion",1]},{"id":4,"type":"CheckpointLoaderSimple","pos":[26,474],"size":[315,122],"flags":{},"order":1,"mode":0,"inputs":[{"label":"ckpt_name","localized_name":"Checkpoint名称","name":"ckpt_name","type":"COMBO","widget":{"name":"ckpt_name"},"link":null}],"outputs":[{"label":"模型","localized_name":"模型","name":"MODEL","type":"MODEL","slot_index":0,"links":[1,16,22]},{"label":"CLIP","localized_name":"CLIP","name":"CLIP","type":"CLIP","slot_index":1,"links":[3,5]},{"label":"VAE","localized_name":"VAE","name":"VAE","type":"VAE","slot_index":2,"links":[8,17,27]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.75","Node name for S&R":"CheckpointLoaderSimple","widget_ue_connectable":{}},"widgets_values":["realisticVisionV60B1_v60B1VAE.safetensors",null]},{"id":19,"type":"TJSSampler","pos":[1383.324590045547,855.0022826646398],"size":[356.917578125,562],"flags":{},"order":6,"mode":0,"inputs":[{"label":"model","localized_name":"model","name":"model","type":"MODEL","link":22},{"label":"positive","localized_name":"positive","name":"positive","type":"CONDITIONING","link":23},{"label":"negative","localized_name":"negative","name":"negative","type":"CONDITIONING","link":24},{"label":"latent_image","localized_name":"latent_image","name":"latent_image","type":"LATENT","link":25},{"label":"total_steps","localized_name":"total_steps","name":"total_steps","type":"INT","widget":{"name":"total_steps"},"link":null},{"label":"early_exit_gamma","localized_name":"early_exit_gamma","name":"early_exit_gamma","type":"FLOAT","widget":{"name":"early_exit_gamma"},"link":null},{"label":"cfg","localized_name":"cfg","name":"cfg","type":"FLOAT","widget":{"name":"cfg"},"link":null},{"label":"sampler_name","localized_name":"sampler_name","name":"sampler_name","type":"COMBO","widget":{"name":"sampler_name"},"link":null},{"label":"scheduler","localized_name":"scheduler","name":"scheduler","type":"COMBO","widget":{"name":"scheduler"},"link":null},{"label":"seed","localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null},{"label":"model_type","localized_name":"model_type","name":"model_type","type":"COMBO","widget":{"name":"model_type"},"link":null},{"label":"denoise","localized_name":"denoise","name":"denoise","shape":7,"type":"FLOAT","widget":{"name":"denoise"},"link":null}],"outputs":[{"label":"latent_x0","localized_name":"latent_x0","name":"latent_x0","type":"LATENT","links":[26]},{"label":"latent_xt","localized_name":"latent_xt","name":"latent_xt","type":"LATENT","links":[]},{"label":"k_star","localized_name":"k_star","name":"k_star","type":"INT","links":null},{"label":"nfe_used","localized_name":"nfe_used","name":"nfe_used","type":"INT","links":null},{"label":"nfe_saving_pct","localized_name":"nfe_saving_pct","name":"nfe_saving_pct","type":"FLOAT","links":null},{"label":"sigma_at_exit","localized_name":"sigma_at_exit","name":"sigma_at_exit","type":"FLOAT","links":null}],"properties":{"Node name for S&R":"TJSSampler","widget_ue_connectable":{}},"widgets_values":[20,1,7.5,"euler","karras",156680208700287,"fixed","diffusion",1]}],"links":[[1,4,0,3,0,"MODEL"],[2,5,0,3,3,"LATENT"],[3,4,1,6,0,"CLIP"],[4,6,0,3,1,"CONDITIONING"],[5,4,1,7,0,"CLIP"],[6,7,0,3,2,"CONDITIONING"],[8,4,2,8,1,"VAE"],[9,8,0,9,0,"IMAGE"],[10,6,0,10,1,"CONDITIONING"],[11,7,0,10,2,"CONDITIONING"],[12,5,0,10,3,"LATENT"],[13,10,0,8,0,"LATENT"],[14,3,0,11,0,"LATENT"],[15,11,0,12,0,"IMAGE"],[16,4,0,10,0,"MODEL"],[17,4,2,11,1,"VAE"],[22,4,0,19,0,"MODEL"],[23,6,0,19,1,"CONDITIONING"],[24,7,0,19,2,"CONDITIONING"],[25,5,0,19,3,"LATENT"],[26,19,0,20,0,"LATENT"],[27,4,2,20,1,"VAE"],[28,20,0,21,0,"IMAGE"]],"groups":[],"config":{},"extra":{"ds":{"scale":1.2839025177495058,"offset":[-712.2695660853183,-370.4487874559189]},"workflowRendererVersion":"LG","ue_links":[],"links_added_by_ue":[]},"version":0.4} \ No newline at end of file diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..bb760e5 --- /dev/null +++ b/__init__.py @@ -0,0 +1,5 @@ +"""ComfyUI_TJS custom node package.""" + +from .tjs_nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..a11444a --- /dev/null +++ b/nodes.py @@ -0,0 +1,31 @@ +"""Compatibility wrapper for older imports. + +ComfyUI loads this custom node through __init__.py, which imports tjs_nodes. +This file re-exports the same mappings in case a local workflow imports +ComfyUI_TJS.nodes directly. +""" + +try: + from .tjs_nodes import ( + NODE_CLASS_MAPPINGS, + NODE_DISPLAY_NAME_MAPPINGS, + TJSDecode, + TJSDecodeManualSigma, + TJSSampler, + ) +except ImportError: + from tjs_nodes import ( # type: ignore + NODE_CLASS_MAPPINGS, + NODE_DISPLAY_NAME_MAPPINGS, + TJSDecode, + TJSDecodeManualSigma, + TJSSampler, + ) + +__all__ = [ + "NODE_CLASS_MAPPINGS", + "NODE_DISPLAY_NAME_MAPPINGS", + "TJSDecode", + "TJSDecodeManualSigma", + "TJSSampler", +] diff --git a/tjs_nodes.py b/tjs_nodes.py new file mode 100644 index 0000000..42d8a5b --- /dev/null +++ b/tjs_nodes.py @@ -0,0 +1,481 @@ +""" +ComfyUI nodes for Truncated Jump Sampling (TJS). + +TJS runs the ordinary sampler for k* early steps, then asks the model for the +endpoint prediction at the early-exit state. In ComfyUI/k-diffusion terms that +endpoint prediction is the model's `denoised` latent at sigma*. +""" + +import math + +import torch + +import comfy.model_management +import comfy.sample +import comfy.samplers + +try: + import comfy.utils +except Exception: # pragma: no cover - older ComfyUI builds may differ. + pass + +import latent_preview + + +def _progress_enabled(): + utils = getattr(comfy, "utils", None) + return bool(getattr(utils, "PROGRESS_BAR_ENABLED", True)) + + +def _sampler_object(name): + if hasattr(comfy.samplers, "sampler_object"): + return comfy.samplers.sampler_object(name) + return comfy.samplers.ksampler(name) + + +def _calculate_sigmas(model, steps, sampler_name, scheduler, denoise): + device = comfy.model_management.get_torch_device() + if hasattr(comfy.samplers, "KSampler"): + sampler = comfy.samplers.KSampler( + model, + steps, + device, + sampler=sampler_name, + scheduler=scheduler, + denoise=denoise, + ) + return sampler.sigmas.to(device) + + from comfy_extras.nodes_custom_sampler import BasicScheduler + + return BasicScheduler().get_sigmas(model, scheduler, steps, denoise)[0].to(device) + + +def _fix_latent(model, latent): + fixed = latent.copy() + fixed["samples"] = comfy.sample.fix_empty_latent_channels(model, latent["samples"]) + return fixed + + +def _prepare_noise(latent, seed): + samples = latent["samples"] + batch_inds = latent.get("batch_index") + if hasattr(comfy.sample, "prepare_noise"): + return comfy.sample.prepare_noise(samples, seed, batch_inds) + + generator = torch.Generator("cpu").manual_seed(seed) + return torch.randn( + samples.size(), + dtype=samples.dtype, + layout=samples.layout, + generator=generator, + device="cpu", + ) + + +def _zeros_like_noise(latent_samples): + return torch.zeros( + latent_samples.size(), + dtype=latent_samples.dtype, + layout=latent_samples.layout, + device="cpu", + ) + + +def _sample_custom( + model, + noise, + cfg, + sampler, + sigmas, + positive, + negative, + latent_image, + noise_mask=None, + callback=None, + seed=0, +): + disable_pbar = not _progress_enabled() + + if hasattr(comfy.sample, "sample_custom"): + try: + return comfy.sample.sample_custom( + model, + noise, + cfg, + sampler, + sigmas, + positive, + negative, + latent_image, + noise_mask=noise_mask, + callback=callback, + disable_pbar=disable_pbar, + seed=seed, + ) + except TypeError as exc: + if "noise_mask" not in str(exc): + raise + return comfy.sample.sample_custom( + model, + noise, + cfg, + sampler, + sigmas, + positive, + negative, + latent_image, + denoise_mask=noise_mask, + callback=callback, + disable_pbar=disable_pbar, + seed=seed, + ) + + device = comfy.model_management.get_torch_device() + return comfy.samplers.sample( + model, + noise, + positive, + negative, + cfg, + device, + sampler, + sigmas, + latent_image=latent_image, + denoise_mask=noise_mask, + callback=callback, + disable_pbar=disable_pbar, + seed=seed, + ) + + +def _callback_x0_to_latent(model, x0_output, fallback_samples): + if "x0" not in x0_output: + return fallback_samples + return model.model.process_latent_out(x0_output["x0"].detach().cpu()) + + +def _endpoint_decode(model, latent_xt, sigma, cfg, positive, negative, seed=0): + """One model evaluation at sigma, returning ComfyUI's denoised endpoint.""" + sigma = float(sigma) + if sigma <= 0.0: + return latent_xt + + device = comfy.model_management.get_torch_device() + sigmas = torch.tensor([sigma, 0.0], dtype=torch.float32, device=device) + sampler = _sampler_object("euler") + noise = _zeros_like_noise(latent_xt) + x0_output = {} + callback = latent_preview.prepare_callback(model, 1, x0_output) + + samples = _sample_custom( + model, + noise, + cfg, + sampler, + sigmas, + positive, + negative, + latent_xt, + noise_mask=None, + callback=callback, + seed=seed, + ) + return _callback_x0_to_latent(model, x0_output, samples) + + +class TJSSampler: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model": ("MODEL",), + "total_steps": ("INT", {"default": 30, "min": 2, "max": 10000}), + "early_exit_gamma": ( + "FLOAT", + { + "default": 0.6, + "min": 0.05, + "max": 1.0, + "step": 0.05, + "display": "slider", + }, + ), + "cfg": ( + "FLOAT", + {"default": 7.5, "min": 0.0, "max": 100.0, "step": 0.1}, + ), + "sampler_name": (comfy.samplers.KSampler.SAMPLERS,), + "scheduler": (comfy.samplers.KSampler.SCHEDULERS,), + "positive": ("CONDITIONING",), + "negative": ("CONDITIONING",), + "latent_image": ("LATENT",), + "seed": ( + "INT", + {"default": 0, "min": 0, "max": 0xffffffffffffffff}, + ), + "model_type": ( + ["auto", "diffusion", "flow"], + {"default": "auto"}, + ), + }, + "optional": { + "denoise": ( + "FLOAT", + {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}, + ), + }, + } + + RETURN_TYPES = ("LATENT", "LATENT", "INT", "INT", "FLOAT", "FLOAT") + RETURN_NAMES = ( + "latent_x0", + "latent_xt", + "k_star", + "nfe_used", + "nfe_saving_pct", + "sigma_at_exit", + ) + FUNCTION = "sample" + CATEGORY = "sampling/TJS" + + def sample( + self, + model, + total_steps, + early_exit_gamma, + cfg, + sampler_name, + scheduler, + positive, + negative, + latent_image, + seed, + model_type="auto", + denoise=1.0, + ): + latent = _fix_latent(model, latent_image) + latent_samples = latent["samples"] + noise_mask = latent.get("noise_mask") + + sigmas = _calculate_sigmas(model, total_steps, sampler_name, scheduler, denoise) + actual_steps = max(1, int(sigmas.shape[-1]) - 1) + + gamma = float(early_exit_gamma) + k_star = max(1, math.ceil(gamma * total_steps)) + k_star = min(k_star, actual_steps) + sigma_star = float(sigmas[k_star].detach().cpu()) + + trunc_sigmas = sigmas[: k_star + 1].clone() + noise = _prepare_noise(latent, seed) + sampler = _sampler_object(sampler_name) + callback = latent_preview.prepare_callback(model, k_star) + + latent_xt = _sample_custom( + model, + noise, + cfg, + sampler, + trunc_sigmas, + positive, + negative, + latent_samples, + noise_mask=noise_mask, + callback=callback, + seed=seed, + ) + + if sigma_star <= 0.0 or k_star >= actual_steps: + latent_x0 = latent_xt + nfe_used = k_star + else: + latent_x0 = _endpoint_decode( + model, + latent_xt, + sigma_star, + cfg, + positive, + negative, + seed=seed, + ) + nfe_used = k_star + 1 + nfe_saving = max(0.0, (1.0 - nfe_used / float(total_steps)) * 100.0) + + out_x0 = latent.copy() + out_x0["samples"] = latent_x0 + + out_xt = latent.copy() + out_xt["samples"] = latent_xt + + print( + "[ComfyUI_TJS] " + f"K={total_steps} gamma={early_exit_gamma:.3f} " + f"k*={k_star} sigma*={sigma_star:.6g} " + f"NFE={nfe_used} saving={nfe_saving:.1f}% mode={model_type}" + ) + return (out_x0, out_xt, k_star, nfe_used, round(nfe_saving, 1), sigma_star) + + +class TJSDecode: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model": ("MODEL",), + "latent": ("LATENT",), + "steps": ("INT", {"default": 30, "min": 1, "max": 10000}), + "end_at_step": ( + "INT", + { + "default": 18, + "min": 0, + "max": 10000, + "tooltip": "Use the same end_at_step as KSampler Advanced.", + }, + ), + "cfg": ( + "FLOAT", + {"default": 7.5, "min": 0.0, "max": 100.0, "step": 0.1}, + ), + "sampler_name": (comfy.samplers.KSampler.SAMPLERS,), + "scheduler": (comfy.samplers.KSampler.SCHEDULERS,), + "positive": ("CONDITIONING",), + "negative": ("CONDITIONING",), + "model_type": ( + ["auto", "diffusion", "flow"], + {"default": "auto"}, + ), + }, + "optional": { + "seed": ( + "INT", + {"default": 0, "min": 0, "max": 0xffffffffffffffff}, + ), + "denoise": ( + "FLOAT", + {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}, + ), + }, + } + + RETURN_TYPES = ("LATENT", "FLOAT", "INT") + RETURN_NAMES = ("latent_x0", "sigma_at_decode", "decode_step") + FUNCTION = "decode" + CATEGORY = "sampling/TJS" + + def decode( + self, + model, + latent, + steps, + end_at_step, + cfg, + sampler_name, + scheduler, + positive, + negative, + model_type="auto", + seed=0, + denoise=1.0, + ): + fixed = _fix_latent(model, latent) + sigmas = _calculate_sigmas(model, steps, sampler_name, scheduler, denoise) + actual_steps = max(1, int(sigmas.shape[-1]) - 1) + decode_step = max(0, min(int(end_at_step), actual_steps)) + sigma = float(sigmas[decode_step].detach().cpu()) + + latent_x0 = _endpoint_decode( + model, + fixed["samples"], + sigma, + cfg, + positive, + negative, + seed=seed, + ) + out = fixed.copy() + out["samples"] = latent_x0 + print( + "[ComfyUI_TJS] advanced decode " + f"steps={steps} end_at_step={end_at_step} " + f"decode_step={decode_step} sigma={sigma:.6g} mode={model_type}" + ) + return (out, sigma, decode_step) + + +class TJSDecodeManualSigma: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model": ("MODEL",), + "latent": ("LATENT",), + "sigma": ( + "FLOAT", + {"default": 1.0, "min": 0.0, "max": 1000.0, "step": 0.01}, + ), + "cfg": ( + "FLOAT", + {"default": 7.5, "min": 0.0, "max": 100.0, "step": 0.1}, + ), + "positive": ("CONDITIONING",), + "negative": ("CONDITIONING",), + "model_type": ( + ["auto", "diffusion", "flow"], + {"default": "auto"}, + ), + }, + "optional": { + "seed": ( + "INT", + {"default": 0, "min": 0, "max": 0xffffffffffffffff}, + ), + }, + } + + RETURN_TYPES = ("LATENT",) + RETURN_NAMES = ("latent_x0",) + FUNCTION = "decode" + CATEGORY = "sampling/TJS" + + def decode( + self, + model, + latent, + sigma, + cfg, + positive, + negative, + model_type="auto", + seed=0, + ): + fixed = _fix_latent(model, latent) + latent_x0 = _endpoint_decode( + model, + fixed["samples"], + sigma, + cfg, + positive, + negative, + seed=seed, + ) + out = fixed.copy() + out["samples"] = latent_x0 + print( + f"[ComfyUI_TJS] manual decode sigma={float(sigma):.6g} " + f"mode={model_type}" + ) + return (out,) + + +NODE_CLASS_MAPPINGS = { + "TJSSampler": TJSSampler, + "TJSDecode": TJSDecode, + "TJSDecodeManualSigma": TJSDecodeManualSigma, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "TJSSampler": "TJS Sampler (Truncated Jump Sampling)", + "TJSDecode": "TJS Decode (Endpoint / Advanced KSampler)", + "TJSDecodeManualSigma": "TJS Decode (Manual Sigma)", +}