PiD KSampler Capture

PiD is picky about its inputs. It doesn't want your finished, fully-denoised latent - it wants the latent at a specific intermediate step of denoising, plus the sigma (noise level) at that moment. That's exactly what this node exists to hand you. PiD KSampler Capture is a KSampler-compatible sampler that runs your normal generation and, at a step you choose, quietly copies the latent and its sigma to CPU and passes them out alongside the usual final result.

Think of it as a KSampler with a dashcam. You get three outputs: final_latent (your normal result, wire it to VAE Decode if you still want the conventional image), pid_latent (the captured intermediate), and pid_sigma (the noise level at capture). The last two go straight into PiD Prepare or PiD Decode.

How it works

Under the hood it's a normal KSampler run (comfy.sample), with a callback that watches the step counter. NVIDIA's PiD counts "K" as denoising passes completed, so when capture_step equals the current step, the node snapshots the latent x and reads sigmas[step]. Set capture_step to 0 and you capture the initial noisy latent; set it at or beyond the total steps and you get the final clean latent with sigma = 0 (which PiD treats as no-detail-regeneration).

The pack also adds a flowmatch_euler_discrete scheduler to the menu - a mirror of diffusers' FlowMatchEulerDiscreteScheduler with a configurable shift - which is what the flow-matching backbones (Z-Image, Flux, Flux2, Qwen-Image) want. flowmatch_shift defaults to 3.0, matching the Z-Image scheduler config.

Inputs worth setting

• capture_step - the star of the show. The README's recommended capture settings are the fastest route to good results: Flux/SD3 at 28 steps capture at 24; Flux2 at 50 capture at 46; Flux2-Klein-4b/9b at just 4 steps capture at 4; Qwen-Image 50/44; Z-Image 50/46 with flowmatch_shift=3.0; Z-Image-Turbo 9/9. Capturing too late in the run means PiD has less to regenerate; too early means it's fighting raw noise.
• steps, cfg, sampler_name, scheduler, denoise - same semantics as KSampler. Defaults are euler + flowmatch_euler_discrete, which covers most backbones.
• flowmatch_shift - a timestep shift knob for flow-matching schedules. This is a "tune this instead of hunting for a scheduler" thing; the Z-Image family wants 3.0.

The workflow

PiD Text Prompt -> CLIP Text Encode -> PiD KSampler Capture
                                     ├── final_latent -> VAE Decode -> (optional normal image)
                                     ├── pid_latent  -> PiD Prepare (latent)
                                     └── pid_sigma   -> PiD Prepare (sigma)

PiD Prepare reads the captured sigma off the latent automatically if you leave sigma at zero, so in practice you just wire pid_latent across and you're done.

Setup and gotchas

Install once for the pack: ComfyUI Manager → "ComfyUI-PiD", or git clone https://github.com/Merserk/ComfyUI-PiD.git into ComfyUI/custom_nodes, pip install -r requirements.txt, restart. Needs ComfyUI 0.28.0+ and a recent ComfyUI generally - the node imports comfy.samplers.KSampler internals, so it tracks the core's pace.

Two things that actually bite:

• Qwen-Image + fp8 loader. Loading Qwen-Image with weight_dtype='fp8_e4m3fn_fast' produces speckled/artifact latents during capture. The node raises a clear error telling you to use default weight dtype - listen to it.
• Capture at the right step. Capture a fully denoised latent (capture_step ≥ total steps) and PiD gets sigma 0 - which is fine, it just becomes a straight decode with no detail regeneration. Capture at step 0 on a 50-step run and you're feeding PiD pure noise. Both are valid setups, but neither is the interesting one.