From d64a145655f9a919fa279025f48dfbbf19983466 Mon Sep 17 00:00:00 2001 From: Prohect Date: Wed, 17 Jun 2026 07:57:43 +0800 Subject: [PATCH 1/2] Initial commit: comfyui-sampler-peek v0.1.0 - SamplerPeekAdvanced: drop-in replacement for SamplerCustomAdvanced with step-by-step latent decoding via VAE - StepExpression: evaluate math expressions in step context (step, n, total_steps) - PeekStepIndex: extract scalar step index from tensor - PeekConditionGate: conditional value pass-through for step-dependent routing - Safe expression evaluator with math functions (sin, cos, sqrt, clamp, lerp, etc.) --- .gitignore | 10 + README.md | 111 ++++++++++++ __init__.py | 10 + expressions.py | 137 ++++++++++++++ nodes.py | 481 +++++++++++++++++++++++++++++++++++++++++++++++++ pyproject.toml | 17 ++ 6 files changed, 766 insertions(+) create mode 100644 .gitignore create mode 100644 __init__.py create mode 100644 expressions.py create mode 100644 nodes.py create mode 100644 pyproject.toml diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..3d30e25 --- /dev/null +++ b/.gitignore @@ -0,0 +1,10 @@ +__pycache__/ +*.py[cod] +*.egg-info/ +dist/ +build/ +.env +.venv/ +*.swp +*.swo +.DS_Store diff --git a/README.md b/README.md index e69de29..30bbeaa 100644 --- a/README.md +++ b/README.md @@ -0,0 +1,111 @@ +# comfyui-sampler-peek + +**Peek into ComfyUI sampling loops.** Decode intermediate latents using a VAE +during `SamplerCustomAdvanced` execution, controlled by math expressions. + +![category: sampling](https://img.shields.io/badge/category-sampling-blue) + +## Overview + +Ever wanted to see what your image looks like *during* the sampling process? +This extension lets you do exactly that — and more. + +It replaces `SamplerCustomAdvanced` with `SamplerPeekAdvanced`, which wraps +the same sampling loop but adds a step-by-step callback. At steps matching a +configurable math expression, the current latent (x₀ prediction) is decoded +through a VAE and collected into an image batch. Step indices are recorded +alongside the images. + +## Nodes + +### SamplerPeekAdvanced (sampling/peek) +A drop-in replacement for `SamplerCustomAdvanced`. It accepts all the same +inputs, plus a VAE and decode configuration. + +| Input | Type | Description | +|----------------|----------|-----------------------------------------------| +| noise | NOISE | Noise source | +| guider | GUIDER | Guider (e.g., CFGGuider) | +| sampler | SAMPLER | Sampler to use | +| sigmas | SIGMAS | Sigmas schedule | +| latent_image | LATENT | Latent image to sample from | +| vae | VAE | VAE model for decoding intermediate latents | +| decode_expr | STRING | Expression for when to decode (see below) | +| start_step | INT | Only decode at ≥ this step (1-based) | +| end_step | INT | Only decode at ≤ this step (0 = no limit) | +| max_previews | INT | Max previews to store in RAM (0 = unlimited) | + +| Output | Type | Description | +|------------------|--------|------------------------------------------------| +| output | LATENT | Final sampled latent | +| denoised_output | LATENT | Final x₀ prediction | +| preview_images | IMAGE | Batch of decoded preview images | +| step_indices | INT | Step indices (1-based) for each preview image | + +### StepExpression (sampling/peek) +Evaluate a math expression with step context. Returns both the float value +and boolean truth. + +| Input | Type | Description | +|-------------|----------|--------------------------------------| +| step_index | INT | Current step index (1-based) | +| total_steps | INT | Total number of steps | +| expression | STRING | Math expression to evaluate | + +### PeekStepIndex (sampling/peek) +Extracts a scalar step index from the tensor output of `SamplerPeekAdvanced`. + +### PeekConditionGate (sampling/peek) +A conditional multiplexer: passes `true_value` when condition is True, +`false_value` otherwise. + +## Expression Syntax + +The `decode_expr` and `StepExpression` nodes use a safe math expression +evaluator. Available variables and functions: + +| Variable | Description | +|-------------|-----------------------| +| `step` | Current step (1..n) | +| `n` | Total steps | +| `total_steps` | Total steps (alias) | + +| Function | Description | +|-------------|------------------------| +| `sin`, `cos`, `tan` | Trig functions | +| `sqrt`, `log`, `exp` | Math functions | +| `ceil`, `floor` | Rounding | +| `clamp(x, lo, hi)` | Clamp value | +| `lerp(a, b, t)` | Linear interpolate | + +**Examples:** + +- `step % 5 == 0` — decode every 5th step +- `step >= 10 and step <= 20` — decode steps 10–20 +- `step == 1 or step == n` — decode first and last step +- `sin(step / n * pi)` — sine ramp for dynamic schedules +- `step % 2 == 0` — decode even-numbered steps + +## Installation + +1. Clone into ComfyUI's `custom_nodes` directory: + ```bash + cd ComfyUI/custom_nodes + git clone https://github.com/Prohect/comfyui-sampler-peek.git + ``` +2. Restart ComfyUI. + +No additional dependencies required — it uses only ComfyUI's built-in +libraries (PyTorch, etc.). + +## Usage + +1. Replace `SamplerCustomAdvanced` with `SamplerPeekAdvanced` in your workflow. +2. Connect a VAE to the `vae` input. +3. Set `decode_expr` to control which steps produce previews. +4. Connect `preview_images` to a `SaveImage` or `PreviewImage` node. +5. Queue and watch the sampling process unfold! + +## License + +MIT diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..dbe094b --- /dev/null +++ b/__init__.py @@ -0,0 +1,10 @@ +""" +comfyui-sampler-peek: Peek into the sampling process. + +Decode latents at intermediate steps during SamplerCustomAdvanced execution +using a given VAE, controlled by math expressions. +""" + +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/expressions.py b/expressions.py new file mode 100644 index 0000000..b81ca6b --- /dev/null +++ b/expressions.py @@ -0,0 +1,137 @@ +""" +Safe math expression evaluator for step-dependent expressions. + +Supports expressions with: + - step (int): current 1-indexed step number + - n / total_steps (int): total number of steps + - Standard math functions via the math module + - Standard Python operators: +, -, *, /, //, %, **, ==, !=, <, >, <=, >=, and, or, not + +Examples: + "step % 5 == 0" -> True every 5th step + "step >= 10 and step <= 20" -> True for steps 10-20 + "step == 1" -> True for the first step + "step == n" -> True for the last step + "sin(step / n * 3.14159)" -> sine ramp from 0 to pi +""" + +import math +import operator +from typing import Any, Callable + +# Allowed builtins: only safe ones +_ALLOWED_NAMES: dict[str, Any] = { + # Math constants + "pi": math.pi, + "e": math.e, + "tau": math.tau, + "inf": math.inf, + "nan": math.nan, + # Math functions + "abs": abs, + "round": round, + "min": min, + "max": max, + "sum": sum, + "len": len, + "int": int, + "float": float, + "bool": bool, + "sin": math.sin, + "cos": math.cos, + "tan": math.tan, + "asin": math.asin, + "acos": math.acos, + "atan": math.atan, + "atan2": math.atan2, + "sinh": math.sinh, + "cosh": math.cosh, + "tanh": math.tanh, + "exp": math.exp, + "log": math.log, + "log10": math.log10, + "log2": math.log2, + "sqrt": math.sqrt, + "pow": pow, + "ceil": math.ceil, + "floor": math.floor, + "trunc": math.trunc, + "fmod": math.fmod, + "fabs": math.fabs, + "gcd": math.gcd, + "degrees": math.degrees, + "radians": math.radians, + "clamp": lambda x, lo, hi: max(lo, min(x, hi)), + "lerp": lambda a, b, t: a + (b - a) * t, +} + + +class StepContext: + """Holds per-step context for expression evaluation.""" + + def __init__(self, step: int, total_steps: int): + self.step = step + self.n = total_steps + self.total_steps = total_steps + + def to_namespace(self) -> dict[str, Any]: + return {"step": self.step, "n": self.n, "total_steps": self.total_steps} + + +def evaluate_expression(expression: str, step: int, total_steps: int) -> Any: + """ + Evaluate a math expression in the context of a sampling step. + + Args: + expression: The math expression string. + step: Current step index (1-based). + total_steps: Total number of steps. + + Returns: + The evaluated result (bool, int, float, etc.) + + Raises: + SyntaxError: If the expression has invalid syntax. + ValueError: If the expression references disallowed names. + """ + if not expression or not expression.strip(): + return True + + ctx = StepContext(step, total_steps) + namespace = {**_ALLOWED_NAMES, **ctx.to_namespace()} + + # Compile for better error messages and slight speed improvement + try: + code = compile(expression.strip(), "", "eval") + except SyntaxError: + raise + + # Check that all names in the expression are allowed + for name in code.co_names: + if name not in namespace: + raise ValueError( + f"Expression uses disallowed name: '{name}'. " + f"Allowed names: {sorted(namespace.keys())}" + ) + + return eval(code, {"__builtins__": {}}, namespace) + + +def evaluate_condition(expression: str, step: int, total_steps: int) -> bool: + """ + Evaluate an expression and return its boolean truth value. + + This is the primary interface for step-filtering expressions. + """ + result = evaluate_expression(expression, step, total_steps) + return bool(result) + + +def evaluate_number(expression: str, step: int, total_steps: int) -> float: + """ + Evaluate an expression and return it as a float. + + This is for expressions that compute a dynamic value (e.g., CFG schedule). + """ + result = evaluate_expression(expression, step, total_steps) + return float(result) diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..15953b7 --- /dev/null +++ b/nodes.py @@ -0,0 +1,481 @@ +""" +Core nodes for comfyui-sampler-peek. + +Provides: + - SamplerPeekAdvanced: Wraps SamplerCustomAdvanced with step-by-step latent decoding. + - StepExpression: Evaluates math expressions in the context of a sampling step. + - PeekStepIndex: Outputs step index for use in downstream expression nodes. +""" + +from typing import Any + +import torch + +import comfy.latent_formats +import comfy.model_management +import comfy.sample +import comfy.samplers +import comfy.utils +import latent_preview +from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict + +from .expressions import evaluate_condition, evaluate_number + +# ────────────────────────────────────────────────────────────── +# SamplerPeekAdvanced +# ────────────────────────────────────────────────────────────── + + +class SamplerPeekAdvanced(ComfyNodeABC): + """ + A drop-in replacement for SamplerCustomAdvanced that can decode + intermediate latents during sampling using a VAE. + + At steps matching the decode_expression, the current latent (x0 + prediction) is decoded through the VAE and collected into an image + batch. Step indices are also recorded. + + Inputs: + noise: NOISE source + guider: GUIDER (e.g., CFGGuider) + sampler: SAMPLER + sigmas: SIGMAS schedule + latent_image: LATENT to sample from + vae: VAE model for decoding + decode_expr: STRING expression; when True, decode at this step. + Variables: step (1..n), n (total steps). + Default: all steps (empty string). + start_step: INT (1-based). Only decode at or after this step. + end_step: INT (1-based). Only decode at or before this step. + Set to 0 for no upper bound. Default: 0. + max_previews: INT. Maximum number of previews to store in RAM. + Default: 0 (unlimited). + + Outputs: + output: LATENT — final sampled latent + denoised_output: LATENT — final x0 prediction + preview_images: IMAGE — batch of decoded preview images + step_indices: INT — step indices corresponding to each preview + """ + + @classmethod + def INPUT_TYPES(cls) -> InputTypeDict: + return { + "required": { + "noise": (IO.NOISE, {"tooltip": "Noise source for sampling"}), + "guider": (IO.GUIDER, {"tooltip": "Guider (e.g., CFGGuider)"}), + "sampler": (IO.SAMPLER, {"tooltip": "Sampler to use"}), + "sigmas": (IO.SIGMAS, {"tooltip": "Sigmas schedule"}), + "latent_image": (IO.LATENT, {"tooltip": "Latent image to sample from"}), + "vae": ( + IO.VAE, + {"tooltip": "VAE model for decoding intermediate latents"}, + ), + "decode_expr": ( + IO.STRING, + { + "multiline": False, + "default": "", + "tooltip": ( + "Expression that determines when to decode. " + "Variables: step (1..n), n (total steps). " + "e.g., 'step % 5 == 0'. Empty = every step." + ), + }, + ), + "start_step": ( + IO.INT, + { + "default": 1, + "min": 1, + "max": 10000, + "step": 1, + "tooltip": "Only decode at or after this step (1-based)", + }, + ), + "end_step": ( + IO.INT, + { + "default": 0, + "min": 0, + "max": 10000, + "step": 1, + "tooltip": "Only decode at or before this step. 0 = no limit", + }, + ), + "max_previews": ( + IO.INT, + { + "default": 0, + "min": 0, + "max": 10000, + "step": 1, + "tooltip": "Max previews to store in RAM. 0 = unlimited", + }, + ), + } + } + + RETURN_TYPES = (IO.LATENT, IO.LATENT, IO.IMAGE, IO.INT) + RETURN_NAMES = ("output", "denoised_output", "preview_images", "step_indices") + OUTPUT_TOOLTIPS = ( + "Final sampled latent", + "Final x0 prediction", + "Batch of decoded preview images from intermediate steps", + "Step indices (1-based) for each preview image", + ) + FUNCTION = "sample" + CATEGORY = "sampling/peek" + DESCRIPTION = ( + "Extended SamplerCustomAdvanced that decodes intermediate latents " + "using a VAE at specified steps. Useful for visualizing the " + "sampling process or creating step-by-step animations." + ) + + def sample( + self, + noise: Any, + guider: Any, + sampler: Any, + sigmas: torch.Tensor, + latent_image: dict, + vae: Any, + decode_expr: str = "", + start_step: int = 1, + end_step: int = 0, + max_previews: int = 0, + ): + # ── prepare latent ────────────────────────────────── + latent = latent_image + latent_samples = latent["samples"] + latent = latent.copy() + latent_samples = comfy.sample.fix_empty_latent_channels( + guider.model_patcher, + latent_samples, + latent.get("downscale_ratio_spacial", None), + latent.get("downscale_ratio_temporal", None), + ) + latent["samples"] = latent_samples + + noise_mask = latent.get("noise_mask", None) + + # ── setup ──────────────────────────────────────────── + total_steps = sigmas.shape[-1] - 1 + if end_step <= 0: + effective_end = total_steps + else: + effective_end = min(end_step, total_steps) + effective_start = max(start_step, 1) + + # Accumulators for previews + preview_images: list[torch.Tensor] = [] + step_indices: list[int] = [] + + # Shared dict for x0 output from the original callback + x0_output: dict = {} + + # ── build chained callback ─────────────────────────── + original_callback = latent_preview.prepare_callback( + guider.model_patcher, total_steps, x0_output + ) + + def peek_callback(step: int, x0: torch.Tensor, x: torch.Tensor, total: int): + """Called at each sampling step. Decodes x0 when conditions met.""" + # Always forward to the original preview callback + if original_callback is not None: + original_callback(step, x0, x, total) + + current_step = step + 1 # 1-based + + # Check range bounds + if current_step < effective_start or current_step > effective_end: + return + + # Check max limit + if max_previews > 0 and len(preview_images) >= max_previews: + return + + # Check expression + try: + should_decode = evaluate_condition(decode_expr, current_step, total) + except Exception: + # If expression is invalid, skip this step + return + + if not should_decode: + return + + # Decode x0 through VAE + try: + # x0 is the predicted clean image at this step + # Move to appropriate device and decode + x0_decoded = vae.decode(x0[:1]) + preview_images.append(x0_decoded) + step_indices.append(current_step) + except Exception: + # If decoding fails, skip this step silently + pass + + # ── run sampling ───────────────────────────────────── + disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED + samples = guider.sample( + noise.generate_noise(latent), + latent_samples, + sampler, + sigmas, + denoise_mask=noise_mask, + callback=peek_callback, + disable_pbar=disable_pbar, + seed=noise.seed, + ) + samples = samples.to(comfy.model_management.intermediate_device()) + + # ── build outputs ──────────────────────────────────── + out = latent.copy() + out.pop("downscale_ratio_spacial", None) + out.pop("downscale_ratio_temporal", None) + out["samples"] = samples + + if "x0" in x0_output: + x0_out = guider.model_patcher.model.process_latent_out( + x0_output["x0"].cpu() + ) + if samples.is_nested: + from comfy import nested_tensor + + latent_shapes = [s.shape for s in samples.unbind()] + x0_out = nested_tensor.NestedTensor( + comfy.utils.unpack_latents(x0_out, latent_shapes) + ) + out_denoised = latent.copy() + out_denoised["samples"] = x0_out + else: + out_denoised = out + + # ── assemble preview image batch ───────────────────── + if len(preview_images) > 0: + preview_batch = torch.cat(preview_images, dim=0) + else: + # Return a 1x1 black image as placeholder + preview_batch = torch.zeros(1, 1, 1, 3, device="cpu", dtype=torch.float32) + + # Step indices as a 1D int tensor + step_tensor = torch.tensor( + step_indices if step_indices else [0], + dtype=torch.int32, + device="cpu", + ) + + return (out, out_denoised, preview_batch, step_tensor) + + +# ────────────────────────────────────────────────────────────── +# StepExpression +# ────────────────────────────────────────────────────────────── + + +class StepExpression(ComfyNodeABC): + """ + Evaluate a math expression in the context of a sampling step. + + Takes step index and total steps, evaluates an expression, + and returns a float value and boolean condition. + + Inputs: + step_index: INT — current step index (1-based) + total_steps: INT — total number of steps + expression: STRING — math expression. Variables: step, n. + + Outputs: + value: FLOAT — evaluated expression as float + condition: BOOLEAN — boolean truth of the expression + """ + + @classmethod + def INPUT_TYPES(cls) -> InputTypeDict: + return { + "required": { + "step_index": ( + IO.INT, + { + "default": 1, + "min": 1, + "max": 10000, + "tooltip": "Current step index (1-based)", + }, + ), + "total_steps": ( + IO.INT, + { + "default": 20, + "min": 1, + "max": 10000, + "tooltip": "Total number of sampling steps", + }, + ), + "expression": ( + IO.STRING, + { + "multiline": False, + "default": "step % 5 == 0", + "tooltip": ( + "Math expression to evaluate. " + "Variables: step (current), n (total steps). " + "Returns both float value and bool condition." + ), + }, + ), + } + } + + RETURN_TYPES = (IO.FLOAT, IO.BOOLEAN) + RETURN_NAMES = ("value", "condition") + OUTPUT_TOOLTIPS = ( + "Evaluated expression as a float", + "Boolean truth of the expression", + ) + FUNCTION = "evaluate" + CATEGORY = "sampling/peek" + DESCRIPTION = ( + "Evaluate a math expression in step context. " + "Useful for creating step-dependent schedules or conditions." + ) + + def evaluate(self, step_index: int, total_steps: int, expression: str): + from .expressions import evaluate_expression + + result = evaluate_expression(expression, step_index, total_steps) + return (float(result), bool(result)) + + +# ────────────────────────────────────────────────────────────── +# PeekStepIndex +# ────────────────────────────────────────────────────────────── + + +class PeekStepIndex(ComfyNodeABC): + """ + Simple pass-through node for step index from SamplerPeekAdvanced. + + Use this to route step indices to downstream nodes that need + step-dependent evaluation during sampling. + + Inputs: + step_indices: INT tensor from SamplerPeekAdvanced + + Outputs: + step_index: INT — the first step index value (for scalar use) + """ + + @classmethod + def INPUT_TYPES(cls) -> InputTypeDict: + return { + "required": { + "step_indices": ( + IO.INT, + { + "tooltip": "Step indices tensor from SamplerPeekAdvanced", + }, + ), + } + } + + RETURN_TYPES = (IO.INT,) + RETURN_NAMES = ("step_index",) + FUNCTION = "passthrough" + CATEGORY = "sampling/peek" + DESCRIPTION = "Extracts a scalar step index from a step indices tensor." + + def passthrough(self, step_indices): + if isinstance(step_indices, torch.Tensor): + if step_indices.numel() > 0: + val = int(step_indices.flatten()[0].item()) + else: + val = 0 + elif isinstance(step_indices, (list, tuple)): + val = int(step_indices[0]) if len(step_indices) > 0 else 0 + else: + val = int(step_indices) + return (val,) + + +# ────────────────────────────────────────────────────────────── +# PeekConditionGate +# ────────────────────────────────────────────────────────────── + + +class PeekConditionGate(ComfyNodeABC): + """ + A conditional gate: passes through inputs only when the condition + is true. Useful for step-dependent control flow in sampling. + + When condition is false, outputs a default/empty value. + + Inputs: + condition: BOOLEAN — gate condition + true_value: FLOAT — value to pass through when condition is True + false_value: FLOAT — value to pass through when condition is False + + Outputs: + value: FLOAT — the selected value + """ + + @classmethod + def INPUT_TYPES(cls) -> InputTypeDict: + return { + "required": { + "condition": ( + IO.BOOLEAN, + {"tooltip": "Gate condition"}, + ), + "true_value": ( + IO.FLOAT, + { + "default": 1.0, + "min": -10000.0, + "max": 10000.0, + "tooltip": "Value when condition is True", + }, + ), + "false_value": ( + IO.FLOAT, + { + "default": 0.0, + "min": -10000.0, + "max": 10000.0, + "tooltip": "Value when condition is False", + }, + ), + } + } + + RETURN_TYPES = (IO.FLOAT,) + RETURN_NAMES = ("value",) + FUNCTION = "gate" + CATEGORY = "sampling/peek" + DESCRIPTION = ( + "Passes through true_value when condition is True, " + "false_value otherwise. Useful for step-dependent routing." + ) + + def gate(self, condition: bool, true_value: float, false_value: float): + return (true_value if condition else false_value,) + + +# ────────────────────────────────────────────────────────────── +# Registry +# ────────────────────────────────────────────────────────────── + +NODE_CLASS_MAPPINGS = { + "SamplerPeekAdvanced": SamplerPeekAdvanced, + "StepExpression": StepExpression, + "PeekStepIndex": PeekStepIndex, + "PeekConditionGate": PeekConditionGate, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "SamplerPeekAdvanced": "Sampler Peek (Advanced)", + "StepExpression": "Step Expression", + "PeekStepIndex": "Peek Step Index", + "PeekConditionGate": "Peek Condition Gate", +} diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..535ef2a --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,17 @@ +[project] +name = "comfyui-sampler-peek" +version = "0.1.0" +description = "Peek into ComfyUI sampling loops - decode intermediate latents using a VAE during SamplerCustomAdvanced execution" +license = {text = "MIT"} +readme = "README.md" +requires-python = ">=3.10" + +[project.urls] +Homepage = "https://github.com/Prohect/comfyui-sampler-peek" + +[build-system] +requires = ["setuptools>=61.0"] +build-backend = "setuptools.build_meta" + +[tool.setuptools] +py-modules = [] From 39b752514f21a2aedcf388eaccae1a829a371404 Mon Sep 17 00:00:00 2001 From: Prohect Date: Wed, 17 Jun 2026 08:10:27 +0800 Subject: [PATCH 2/2] Add step-dependent CFG modulation and cleanup - Add optional cfg_expr input to SamplerPeekAdvanced for dynamic CFG - Dynamic CFG hooks predict_noise to evaluate expression at each step - Clean up unused imports in expressions.py - Update documentation and version to 0.2.0 --- .gitignore | 2 ++ README.md | 18 ++++++++++++++---- __init__.py | 6 +++++- expressions.py | 3 +-- nodes.py | 50 ++++++++++++++++++++++++++++++++++++++++++++------ pyproject.toml | 4 ++-- 6 files changed, 68 insertions(+), 15 deletions(-) diff --git a/.gitignore b/.gitignore index 3d30e25..7040b76 100644 --- a/.gitignore +++ b/.gitignore @@ -8,3 +8,5 @@ build/ *.swp *.swo .DS_Store + +temp/ diff --git a/README.md b/README.md index 30bbeaa..2073305 100644 --- a/README.md +++ b/README.md @@ -33,7 +33,8 @@ inputs, plus a VAE and decode configuration. | decode_expr | STRING | Expression for when to decode (see below) | | start_step | INT | Only decode at ≥ this step (1-based) | | end_step | INT | Only decode at ≤ this step (0 = no limit) | -| max_previews | INT | Max previews to store in RAM (0 = unlimited) | +| max_previews | INT | Max previews to store in RAM (0 = unlimited) | +| cfg_expr | STRING | (optional) Step-dependent CFG expression. e.g., `clamp(1 + step / n * 7, 1, 8)` for a linear ramp from 1 to 8 | | Output | Type | Description | |------------------|--------|------------------------------------------------| @@ -86,12 +87,20 @@ evaluator. Available variables and functions: - `sin(step / n * pi)` — sine ramp for dynamic schedules - `step % 2 == 0` — decode even-numbered steps +### Dynamic CFG + +Optionally provide a `cfg_expr` to modulate CFG at each sampling step: + +- `clamp(1 + step / n * 7, 1, 8)` — linear ramp from 1 to 8 +- `step < n / 2 ? 1 : 5` — use Python: `5 if step > n/2 else 1` +- `1 + sin(step / n * pi) * 4` — sinusoidal CFG swing + ## Installation 1. Clone into ComfyUI's `custom_nodes` directory: ```bash cd ComfyUI/custom_nodes - git clone https://github.com/Prohect/comfyui-sampler-peek.git + git clone https://github.com/aria4comfyui/comfyui-sampler-peek.git ``` 2. Restart ComfyUI. @@ -103,8 +112,9 @@ libraries (PyTorch, etc.). 1. Replace `SamplerCustomAdvanced` with `SamplerPeekAdvanced` in your workflow. 2. Connect a VAE to the `vae` input. 3. Set `decode_expr` to control which steps produce previews. -4. Connect `preview_images` to a `SaveImage` or `PreviewImage` node. -5. Queue and watch the sampling process unfold! +4. (Optional) Set `cfg_expr` for step-dependent CFG modulation. +5. Connect `preview_images` to a `SaveImage` or `PreviewImage` node. +6. Queue and watch the sampling process unfold! ## License diff --git a/__init__.py b/__init__.py index dbe094b..fcb9fc1 100644 --- a/__init__.py +++ b/__init__.py @@ -3,8 +3,12 @@ comfyui-sampler-peek: Peek into the sampling process. Decode latents at intermediate steps during SamplerCustomAdvanced execution using a given VAE, controlled by math expressions. + +Optionally supports step-dependent CFG modulation for dynamic guidance. """ from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS -__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] +VERSION = "0.2.0" + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "VERSION"] diff --git a/expressions.py b/expressions.py index b81ca6b..903ddef 100644 --- a/expressions.py +++ b/expressions.py @@ -16,8 +16,7 @@ Examples: """ import math -import operator -from typing import Any, Callable +from typing import Any # Allowed builtins: only safe ones _ALLOWED_NAMES: dict[str, Any] = { diff --git a/nodes.py b/nodes.py index 15953b7..b872b6a 100644 --- a/nodes.py +++ b/nodes.py @@ -9,14 +9,13 @@ Provides: from typing import Any -import torch - import comfy.latent_formats import comfy.model_management import comfy.sample import comfy.samplers import comfy.utils import latent_preview +import torch from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict from .expressions import evaluate_condition, evaluate_number @@ -35,6 +34,8 @@ class SamplerPeekAdvanced(ComfyNodeABC): prediction) is decoded through the VAE and collected into an image batch. Step indices are also recorded. + Optionally supports step-dependent CFG modulation via cfg_expr. + Inputs: noise: NOISE source guider: GUIDER (e.g., CFGGuider) @@ -50,6 +51,8 @@ class SamplerPeekAdvanced(ComfyNodeABC): Set to 0 for no upper bound. Default: 0. max_previews: INT. Maximum number of previews to store in RAM. Default: 0 (unlimited). + cfg_expr: STRING (optional). Step-dependent CFG expression. + e.g., "clamp(1 + step / n * 7, 1, 8)". Outputs: output: LATENT — final sampled latent @@ -113,13 +116,28 @@ class SamplerPeekAdvanced(ComfyNodeABC): "tooltip": "Max previews to store in RAM. 0 = unlimited", }, ), - } + }, + "optional": { + "cfg_expr": ( + IO.STRING, + { + "multiline": False, + "default": "", + "tooltip": ( + "Step-dependent CFG expression. Evaluated at each step " + "to dynamically set the guider CFG. Variables: step, n. " + "e.g., 'clamp(1 + step / n * 7, 1, 8)' for a linear ramp. " + "Empty = use the guider's default CFG." + ), + }, + ), + }, } RETURN_TYPES = (IO.LATENT, IO.LATENT, IO.IMAGE, IO.INT) RETURN_NAMES = ("output", "denoised_output", "preview_images", "step_indices") OUTPUT_TOOLTIPS = ( - "Final sampled latent", + "Final sampled latent (with optional step-dependent CFG applied)", "Final x0 prediction", "Batch of decoded preview images from intermediate steps", "Step indices (1-based) for each preview image", @@ -128,8 +146,8 @@ class SamplerPeekAdvanced(ComfyNodeABC): CATEGORY = "sampling/peek" DESCRIPTION = ( "Extended SamplerCustomAdvanced that decodes intermediate latents " - "using a VAE at specified steps. Useful for visualizing the " - "sampling process or creating step-by-step animations." + "using a VAE at specified steps, with optional step-dependent CFG. " + "Useful for visualizing the sampling process or creating step-by-step animations." ) def sample( @@ -144,6 +162,7 @@ class SamplerPeekAdvanced(ComfyNodeABC): start_step: int = 1, end_step: int = 0, max_previews: int = 0, + cfg_expr: str = "", ): # ── prepare latent ────────────────────────────────── latent = latent_image @@ -216,6 +235,25 @@ class SamplerPeekAdvanced(ComfyNodeABC): # If decoding fails, skip this step silently pass + # ── dynamic CFG (step-dependent) ───────────────────── + if cfg_expr: + step_counter = {"n": 1} # 1-based, mutated in closure + original_predict = guider.predict_noise + + def dynamic_predict(x, timestep, model_options=None, seed=None): + if model_options is None: + model_options = {} + current = step_counter["n"] + step_counter["n"] += 1 + try: + new_cfg = evaluate_number(cfg_expr, current, total_steps) + guider.set_cfg(new_cfg) + except Exception: + pass # keep current cfg on error + return original_predict(x, timestep, model_options, seed) + + guider.predict_noise = dynamic_predict + # ── run sampling ───────────────────────────────────── disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED samples = guider.sample( diff --git a/pyproject.toml b/pyproject.toml index 535ef2a..2de6b79 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,13 +1,13 @@ [project] name = "comfyui-sampler-peek" -version = "0.1.0" +version = "0.2.0" description = "Peek into ComfyUI sampling loops - decode intermediate latents using a VAE during SamplerCustomAdvanced execution" license = {text = "MIT"} readme = "README.md" requires-python = ">=3.10" [project.urls] -Homepage = "https://github.com/Prohect/comfyui-sampler-peek" +Homepage = "https://github.com/aria4comfyui/comfyui-sampler-peek" [build-system] requires = ["setuptools>=61.0"]