Klein Edit Scheduler

The Klein Edit Scheduler is this pack's showpiece, and it's the most opinionated node of the four. Where the other schedulers hand you a slider or two, this one gives you a live graph widget, a draggable XY pad, and a draw mode that turns your sigma schedule into a row of dots you can move by hand. It's built for editing workflows on flow-matching models - the author's stated targets are Z-Image Turbo and Flux 2 Klein (hence the name; Klein's distilled line natively runs 4 steps, which is the default here).

Why you'd reach for it: editing is where "pick a preset" stops cutting it. denoise is your edit strength - how much of the source image survives - and how you spend the remaining steps decides whether you get a subtle relight or a full redraw. shift and curve redistribute those steps. When the built-in schedulers don't give you the control you want, this is the node where you stop guessing and start shaping.

How it works

In parametric mode (the default), it builds the schedule in the standard flow-matching way: t = linspace(0, 1, steps + 1), then a power curve t^curve (below 1 bunches steps at the start, above 1 at the end), then a timestep-shift transform t / (t + shift·(1 − t)) that pushes steps toward high sigma (shift > 1 means more early denoising), and finally blends sigmas = denoise·(1 − t) + sigma_min·t, pinning the endpoints to denoise and sigma_min.

The interactive widget rides on top of that. Top panel shows the curve; the bottom panel is an XY pad - drag left/right for curve, up/down for shift, and the curve redraws live. Click the toggle and you're in draw mode: one draggable dot per sigma step. Dots move vertically to set sigma values directly, middle dots also drag horizontally to tighten or loosen step density in a region, and monotonicity is enforced automatically so you can't draw a schedule that goes backwards. The first and last dots (teal) stay synced to the denoise and sigma_min sliders. Drawn sigmas get saved into the workflow and handed to the sampler as-is, bypassing the formula entirely.

Inputs and outputs

Six required inputs, two optional:

• model - your checkpoint, passed straight through
• steps - default 4, 1–100 ("Klein native = 4")
• denoise - 0.001 to 2.0, "Starting sigma / edit strength." The editing lever
• sigma_min - 0 to 1. Raise to 0.01–0.03 to soften the final step
• shift - 0.01 to 20, ">1 = more steps at high sigma"
• curve - 0.01 to 10, "<1 = bunch at start, >1 = bunch at end"
• draw_mode - parametric or draw (optional, defaults to parametric)
• custom_sigmas - a JSON array of sigma values used in draw mode (optional)

Outputs are MODEL (passthrough - wire it onward) and SIGMAS, which goes into the sigmas input of SamplerCustomAdvanced or SamplerCustom. This is a scheduling node, not a sampler - it shapes sigmas, it doesn't run the denoise loop.

Install

Same pack, same two routes:

cd ComfyUI/custom_nodes
git clone https://github.com/capitan01R/ComfyUI-CapitanZiT-Scheduler

restart, or install via ComfyUI Manager (search "ComfyUI-CapitanZiT-Scheduler"). The interactive widget lives in JavaScript, so give the browser a full reload after installing before you expect the graph to appear.

Where people get burned

• Default 4 steps is a Klein thing, not a Z-Image thing. The tooltip says it plainly: "Klein native = 4." On Turbo you usually want 8–9, so don't blindly leave the default when you switch models.
• For reference preservation in Klein edits, keep the denoise low. That's the classic "how much of the original survives" knob - the most-asked question in the release thread was about exactly this, and the answer is the edit-strength slider, not the curve.
• The scheduler only schedules. Wire SIGMAS into a custom sampler; forgetting the MODEL passthrough output is the classic footgun.
• One thing this node actually avoids: because shift and curve live inside the node, you don't hit the classic trap of a scheduler that silently ignores your external shift. No bong_tangent surprises here.