Prohect
+12
@@ -0,0 +1,12 @@
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*.egg-info/
|
||||
dist/
|
||||
build/
|
||||
.env
|
||||
.venv/
|
||||
*.swp
|
||||
*.swo
|
||||
.DS_Store
|
||||
|
||||
temp/
|
||||
@@ -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.
|
||||
|
||||
|
||||
|
||||
## 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
|
||||
|
||||
+14
@@ -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"]
|
||||
+136
@@ -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(), "<step_expression>", "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)
|
||||
@@ -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",
|
||||
}
|
||||
@@ -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 = []
|
||||
Reference in New Issue
Block a user