initial public release

.gitignore

@@ -0,0 +1,10 @@
+__pycache__/
+*.py[cod]
+*$py.class
+
+.env
+.venv/
+venv/
+
+.DS_Store
+Thumbs.db

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 原型实验，不建议直接用于严肃生产流程。

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.

__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"]

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",
+]

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)",
+}