LucidFlux PiD Decode (latent -> 4K)

Here's the strange bit about the 4K half of LucidFlux: it doesn't use a VAE. LucidFlux PiD Decode (latent -> 4K) takes the sampler's latent and hands you back an image four times wider and taller, skipping VAE decode entirely.

The mechanism is NVIDIA's PiD - a pixel-space DiT distilled down to a four-step student. It isn't an upscaler in the "add a bit of detail" sense. It regenerates the image at target resolution from the latent plus a caption embedding, which is why the pack ships a gemma_prompt_embedding.pt for it: PiD wants text conditioning even though the LucidFlux stage never did. (You don't load Gemma-2. The embeddings are precomputed, same trick as the main pipeline.)

The resolution math, plainly

PiD multiplies by 4. Target size is latent × 8 × 4, so a 128×128 latent lands at 4096×4096 and a 64² latent gives you 2048². That's the only relationship you need. If you got 1024 instead of 4096, your latent was 32² and you should go look at your EmptySD3LatentImage size - or your VAELoader's empty latent, if you're generating the latent some other way.

The node also converts the latent into the space PiD expects before running: ComfyUI's sampler returns raw VAE latents, and PiD wants the diffusers-normalised version, (z − 0.1159) × 0.3611. It does that internally, so what you wire in is just the plain LATENT output of the sampler.

The inputs worth touching

latent - straight from LucidFlux Sampler. Don't VAE-decode first; that's the point.

pid_model - defaults to auto, which searches models/diffusion_models/ and models/LucidFlux/ for files with "pid" in the name, preferring anything containing "2kto4k". Grab PiD_v1pt5_res2kto4k_sr4x_official_flux_distill_4step.pth from nvidia/PiD and the auto path just works. none is a legacy alias for auto, so setting it changes nothing.

caption_embeddings - gemma_prompt_embedding.pt, in models/LucidFlux/. Note this dropdown filters on "gemma", so the null_caption_embs.pt file that's also in that HuggingFace repo will not appear here.

num_steps - default 4, max 8. The distilled model's schedule is [0.999, 0.866, 0.634, 0.342, 0]; anything above 4 gets subsampled from that table, not interpolated fresh. Leave it at 4.

unload_models - default true, and this is the setting that matters. It calls ComfyUI's unload_all_models() before decoding. PiD's network is moved onto the GPU by the pack directly, outside ComfyUI's memory manager, so ComfyUI won't make room for it - and the gentler free_memory(need) call only partially unloads the Flux model. The source comments are blunt about this being the direct cause of 4K OOMs with 11GB of Flux still resident. Turn it false only if you're doing something clever like reusing Flux afterwards in the same run.

streaming_prefetch - 0 keeps the whole PiD net on the card, 1–2 streams it layer by layer. Worth knowing: streaming shaves weight memory, but the 4K peak is dominated by activations, not weights, so it won't rescue a 12GB card at 4096². The tooltip notes the official default is 1 while the node default is 0; if you're tight on VRAM, set it to 1 and leave it.

Then seed, degrade_sigma (0 = treat the input as clean, which it is coming from a sampler), and dtype. There's also an optional reference + color_strength pair for a colour-alignment pass against DiffBIR's ldr, same idea as LucidFluxColorAlign but done inside the decode.

The VRAM reality

The pack's own measurements, bf16 with streaming on: 1024² ≈ 1.5GB, 2048² ≈ 5.2GB, 3072² ≈ 11.2GB, 4096² ≈ 18.6GB. Peak scales with output pixels, roughly 1.1e-6 × pixels + 0.35 GB.

So: 4096² isn't a 12GB job. It will run on Windows by spilling into shared system memory, and you'll see 20GB+ in Task Manager and wait a long time for it. The honest workflow on a mid-range card is to land the 2K image through the VAE path and use 2048² from PiD when you want the higher-fidelity variant.

Install

cd ComfyUI/custom_nodes
git clone https://github.com/smthemex/ComfyUI_LucidFlux

Weights: the PiD checkpoint into ComfyUI/models/diffusion_models/, and gemma_prompt_embedding.pt plus (optionally) null_caption_embs.pt into ComfyUI/models/LucidFlux/. Both from the HF links in the README. No new Python dependencies - torch and einops cover the whole pack.

What goes wrong

The weight isn't found. The filename must contain "pid". PiD_...pth is fine, decoder_final.pth is not, and auto picks the first match it finds sorted with 2kto4k first.

OOM at 4K with a big card. Check unload_models is still true and drop your latent size. If you've got two branches in the graph both holding latents, ComfyUI may have models resident that PiD can't evict.

Blurry or "regenerated" faces. That's PiD doing its job - it's a generative decoder, not a resampler, so at 4× it's synthesising plausible detail. For photo restoration the community's current favourite for detail-preserving work is SeedVR2, which invents less; PiD's argument is that it's one latent away with no extra graph. Try color_strength with a reference wired from DiffBIR's ldr before you blame the model.