[Pfannkuchensack]

Summary

Every PiD decode node denormalized the latent before handing it to PiD (z / scale + shift, z * std + mean). PiD is trained on the normalized latent: each upstream tokenizer in nv-tlabs/PiD normalizes at encode (tokenizers/flux_vae.py, sdxl_vae.py, qwenimage_vae.py, flux2_vae.py), and the vendored pipeline_registry.py says the latent goes to PiD "directly — no extra denormalization". Fed the raw latent, PiD returned oversaturated, posterized images with visible banding.

• FLUX.1, Z-Image, SD3, SDXL, Qwen-Image decode and PiD Upscale now pass the stored latent through unchanged (Qwen-Image still drops the singleton frame dim).
• FLUX.2 is the exception: flux2_denoise stores the BN-denormalized, unpacked (B, 32, H/8, W/8) latent. The node now patchifies it with the (c ph pw) order and BN-normalizes it with the VAE's running statistics, which is exactly what upstream's FLUX.2 tokenizer encodes to (pack_flux2_latent_for_pid). The node's vae input is therefore required.
• The BN-stat lookup moves out of Flux2DenoiseInvocation._get_bn_stats into get_flux2_bn_stats (backend/flux2/sampling_utils.py), shared by denoise and PiD.
• The vae inputs on the Qwen-Image, SDXL and Z-Image decode nodes are no longer read. They stay as optional fields: existing workflows that connect a VAE still validate, and Generate/Canvas keep wiring the decode VAE there because Canvas img2img (findVaeSource, feat(canvas): restore PiD decoding and the Anima inpaint adapter on Canvas #9805) uses that edge to encode with the same VAE, including an SDXL VAE override or seamless.
• Node versions: flux2_pid_decode 1.1.0, the other six PiD nodes 1.0.1.

Related Issues / Discussions

• Builds on feat(canvas): restore PiD decoding and the Anima inpaint adapter on Canvas #9805 (PiD on Canvas), merged into this branch.

QA Instructions

A/B, old (main) vs this branch, RTX 4090. Four test images (512 px) per backbone, each encoded with that backbone's VAE, then decoded with the matching PiD node. PSNR of the PiD output, box-downscaled 4×, against the input image. VAE round trip shown as a ceiling. Same graphs, isolated root, use_cache: false, PiD seed 0.

Backbone (decoder)	VAE round trip	PiD main	PiD this PR
FLUX.1 (res2k)	32.96 dB	25.38 dB	29.34 dB
Z-Image (FLUX decoder)	32.96 dB	25.38 dB	29.34 dB
PiD Upscale (FLUX)	—	25.38 dB	29.34 dB
FLUX.2 (res2k)	35.50 dB	18.36 dB	26.63 dB
SD3 (res2k)	30.87 dB	22.08 dB	25.50 dB
SDXL (res2kto4k)	26.71 dB	16.59 dB	22.24 dB
Qwen-Image (res2kto4k)	31.57 dB	19.86 dB	25.19 dB

All 28 per-image comparisons improve. Visually, main is oversaturated with posterized skies and fur; this PR matches the VAE's colors and adds detail.

Browser (webv2 Generate): FLUX.2 Klein 4B + PiD 1.5 FLUX.2 and Z-Image Turbo + PiD 1.5 FLUX, Fit mode at 512 px. Both completed. The FLUX.2 graph carries the model_loader.vae → flux2_pid_decode.vae edge, and the Z-Image graph has no VAE edge on the decode. Output is clean at 2048 px.

Checks run

• uv run --no-sync pytest tests/backend/pid tests/app/invocations tests/backend/flux2 -n 4 (after merging main): 2921 passed.
• New tests/backend/pid/test_pid_latent_normalization.py:
  • the FLUX.2 pack against upstream's tokenizer recipe (a real eval-mode BatchNorm2d);
  • a round trip through flux2_denoise's output path (unpack_flux2 of the BN-denormalized latent);
  • rejection of non-FLUX.2 latents and of BN stats from another VAE.
• ruff check and ruff format --check on the changed files.
• New Canvas tests: SDXL img2img with PiD (fit and native) encodes with the VAE override / the seamless node, not the bundled VAE.
• webv2: vitest run --config vitest.config.mts src/features/generation/core (1044 passed), tsc --noEmit, lint:oxc, oxfmt --check, architecture:check.
• frontend/api: pnpm test.

Review

No material findings. Two small follow-ups came out of review:

• A clear error when the FLUX.2 decode gets BatchNorm statistics of another size (a wrong VAE), instead of a broadcast error.
• Two stale comments corrected.
• After merging main: Canvas img2img picks its encode VAE from the PiD decode's vae edge, so dropping that edge for SDXL would have silently ignored a VAE override or seamless. The edge stays on every PiD decode that takes one, and regression tests cover override and seamless.

Pre-existing and not changed here: sd3_i2l / sd3_l2i ignore the SD3 VAE's shift_factor (0.0609). PiD decodes from the SD3 denoiser are unaffected, but a latent from sd3_i2l carries a constant offset of about 0.093.

Compatibility / Rollout

• flux2_pid_decode 1.1.0: vae is now a required connection.
  • webv2-generated graphs always wired it already.
  • A saved workflow without that edge loads as an updatable node, and the editor reports the missing VAE connection before enqueue.
  • An API graph without it fails with RequiredConnectionException rather than producing a wrong image.
• Other PiD nodes: no field changes, patch bump only. Images from PiD decodes change (intentionally) for the same inputs.
• openapi.json and schema.ts are regenerated. No persisted-data changes.

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Meaningful regression coverage added / updated where needed; obsolete tests/code removed
• Persisted-state and API changes include required migrations / compatibility validation
• Relevant performance/efficiency opportunities considered; material claims have evidence
• Material review findings resolved and relevant checks rerun
• Documentation added / updated (if applicable)
• Updated What's New copy (if doing a release after this PR)