diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..7040b76 --- /dev/null +++ b/.gitignore @@ -0,0 +1,12 @@ +__pycache__/ +*.py[cod] +*.egg-info/ +dist/ +build/ +.env +.venv/ +*.swp +*.swo +.DS_Store + +temp/ diff --git a/README.md b/README.md index e69de29..2073305 100644 --- a/README.md +++ b/README.md @@ -0,0 +1,121 @@ +# 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) | +| 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 | +|------------------|--------|------------------------------------------------| +| 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 + +### 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/aria4comfyui/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. (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 + +MIT diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..fcb9fc1 --- /dev/null +++ b/__init__.py @@ -0,0 +1,14 @@ +""" +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 + +VERSION = "0.2.0" + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "VERSION"] diff --git a/expressions.py b/expressions.py new file mode 100644 index 0000000..903ddef --- /dev/null +++ b/expressions.py @@ -0,0 +1,136 @@ +""" +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 +from typing import Any + +# 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..b872b6a --- /dev/null +++ b/nodes.py @@ -0,0 +1,519 @@ +""" +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 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 + +# ────────────────────────────────────────────────────────────── +# 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. + + Optionally supports step-dependent CFG modulation via cfg_expr. + + 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). + cfg_expr: STRING (optional). Step-dependent CFG expression. + e.g., "clamp(1 + step / n * 7, 1, 8)". + + 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", + }, + ), + }, + "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 (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", + ) + FUNCTION = "sample" + CATEGORY = "sampling/peek" + DESCRIPTION = ( + "Extended SamplerCustomAdvanced that decodes intermediate latents " + "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( + 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, + cfg_expr: str = "", + ): + # ── 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 + + # ── 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( + 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..2de6b79 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,17 @@ +[project] +name = "comfyui-sampler-peek" +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/aria4comfyui/comfyui-sampler-peek" + +[build-system] +requires = ["setuptools>=61.0"] +build-backend = "setuptools.build_meta" + +[tool.setuptools] +py-modules = []