Skip to content

fix(pid): hand PiD the normalized latent it was trained on #9888

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Pfannkuchensack wants to merge 4 commits into
base: main
Choose a base branch
from
Open

fix(pid): hand PiD the normalized latent it was trained on #9888

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
98a11d7
fix(pid): hand PiD the normalized latent it was trained on
Pfannkuchensack Oct 16, 2026
644c22e
Merge branch 'main' into fix/pid-normalized-latent
Pfannkuchensack Oct 16, 2026
6cde691
Merge remote-tracking branch 'upstream/main' into fix/pid-normalized-…
Pfannkuchensack Oct 19, 2026
566d697
Merge branch 'main' into fix/pid-normalized-latent
Pfannkuchensack Oct 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions ...src/content/docs/Users Guide/04.Image Generation/Advanced Topics/pid-decode.mdx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,8 @@ When PiD mode is not *Off*, pick your **PiD Decoder** and **Gemma-2 Encoder** be

PiD is available in both the **Generate** tab (text-to-image) and on the **Canvas** (image-to-image), in both Fit and Native modes.

In the Workflow editor, each base has its own PiD decode node, such as *Latents to Image - Qwen-Image + PiD (4x SR)*. They take the denoised latent as the model produced it; only the **FLUX.2** node also needs the FLUX.2 VAE connected, because it reads the VAE's latent statistics to bring the latent into the form PiD was trained on.

## Running PiD on an existing image

Everything above describes PiD as a *decode* step: it needs the latent the model just produced, so it only exists inside a generation. If you want to push an image you already have through PiD, the Workflow editor has a separate **PiD Upscale (4x)** node.
Expand Down
40 changes: 3 additions & 37 deletions invokeai/app/invocations/flux2/flux2_denoise.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@
FLUX2_TXT2IMG_SCHEDULER_KWARGS,
compute_empirical_mu,
generate_img_ids_flux2,
get_flux2_bn_stats,
get_noise_flux2,
get_schedule_flux2,
pack_flux2,
Expand Down Expand Up @@ -175,44 +176,9 @@ class Flux2DenoiseInvocation(BaseInvocation):
)

def _get_bn_stats(self, context: InvocationContext) -> Optional[Tuple[torch.Tensor, torch.Tensor]]:
"""Extract BN statistics from the FLUX.2 VAE.

The FLUX.2 VAE uses batch normalization on the patchified 128-channel representation.
IMPORTANT: BFL FLUX.2 VAE uses affine=False, so there are NO learnable weight/bias.

BN formula (affine=False): y = (x - mean) / std
Inverse: x = y * std + mean

Returns:
Tuple of (bn_mean, bn_std) tensors of shape (128,), or None if BN layer not found.
"""
"""Extract the latent BN statistics (bn_mean, bn_std) from the wired FLUX.2 VAE, or None if it has none."""
with context.models.load(self.vae.vae).model_on_device() as (_, vae):
# Ensure VAE is in eval mode to prevent BN stats from being updated
vae.eval()

# Try to find the BN layer - it may be at different locations depending on model format
bn_layer = None
if hasattr(vae, "bn"):
bn_layer = vae.bn
elif hasattr(vae, "batch_norm"):
bn_layer = vae.batch_norm
elif hasattr(vae, "encoder") and hasattr(vae.encoder, "bn"):
bn_layer = vae.encoder.bn

if bn_layer is None:
return None

# Verify running statistics are initialized
if bn_layer.running_mean is None or bn_layer.running_var is None:
return None

# Get BN running statistics from VAE
bn_mean = bn_layer.running_mean.clone() # Shape: (128,)
bn_var = bn_layer.running_var.clone() # Shape: (128,)
bn_eps = bn_layer.eps if hasattr(bn_layer, "eps") else 1e-4 # BFL uses 1e-4
bn_std = torch.sqrt(bn_var + bn_eps)

return bn_mean, bn_std
return get_flux2_bn_stats(vae)

def _bn_normalize(
self,
Expand Down
98 changes: 25 additions & 73 deletions invokeai/app/invocations/pid/flux2_pid_decode.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,28 +6,15 @@
4B and 9B FLUX.2 Klein variants share the same 32-channel VAE, so this one node
covers both.

Latent layout (the important difference from the FLUX.1 node):
Latent layout and normalization (the important difference from the FLUX.1 node):

* ``flux2_denoise`` stores an *unpacked* ``(B, 32, H/8, W/8)`` latent that is
already **BN-denormalized** (``x * bn_std + bn_mean`` is applied before the
unpack, see ``flux2_denoise.py``). That is exactly the raw latent the FLUX.2
VAE's conv decoder consumes.
* PiD's FLUX.2 backbone expects the **packed** ``(B, 128, H/16, W/16)``
representation (``lq_latent_channels=128``, ``latent_spatial_down_factor=16``
in ``backend/pid/decode.py``). We therefore patchify the stored latent
(2x2 spatial patches folded into channels: 32*4 = 128) *before* handing it to
PiD - mirroring ``pack_flux2`` but keeping a spatial ``(B, C, h, w)`` layout
instead of the transformer's ``(B, seq, C)`` sequence layout.

Denormalization: unlike FLUX.1 (single ``scale``/``shift``) and Z-Image
(checkpoint-specific ``scaling_factor``/``shift_factor``), the FLUX.2 VAE
(``AutoencoderKLFlux2``) exposes **no** scalar ``scaling_factor``/``shift_factor``
at all - its only normalization is the per-channel BatchNorm applied/inverted
*outside* the VAE in ``flux2_denoise``. So the packed latent is already in PiD's
expected raw space and no further scaling is needed (identity fallbacks below).
We still accept an optional ``vae`` input and read the constants at runtime (like
the Z-Image node) so any future FLUX.2 VAE variant that does expose scalar
constants is honored automatically.
already **BN-denormalized** - the raw latent the FLUX.2 VAE's conv decoder
consumes.
* PiD's FLUX.2 backbone was trained on what the upstream FLUX.2 tokenizer encodes
to: the 2x2-patchified ``(B, 128, H/16, W/16)`` latent, **BN-normalized** with
the VAE's running statistics. ``pack_flux2_latent_for_pid`` restores that, so
this node needs the FLUX.2 VAE for its BatchNorm statistics.
"""

from contextlib import ExitStack
Expand All @@ -50,6 +37,7 @@
from invokeai.app.invocations.model import Gemma2EncoderField, PiDDecoderField, VAEField
from invokeai.app.invocations.primitives import ImageOutput
from invokeai.app.services.shared.invocation_context import InvocationContext
from invokeai.backend.flux2.sampling_utils import get_flux2_bn_stats
from invokeai.backend.model_manager.taxonomy import BaseModelType
from invokeai.backend.pid._src.networks.pid_net import PidNet
from invokeai.backend.pid.decode import (
Expand All @@ -58,33 +46,25 @@
assert_pid_decoder_matches_base,
encode_caption_for_pid,
estimate_pid_decode_working_memory,
pack_flux2_latent_for_pid,
)
from invokeai.backend.util.devices import TorchDevice

# FLUX.2 uses per-channel BatchNorm (affine=False) for latent normalization, and
# that BN is already inverted in flux2_denoise before the latent is stored. The
# FLUX.2 VAE (AutoencoderKLFlux2) has no scalar scaling_factor/shift_factor, so
# the identity transform below is the correct default: the stored (packed) latent
# is already the raw representation PiD was trained on.
_FLUX2_VAE_SCALING_FACTOR_FALLBACK: float = 1.0
_FLUX2_VAE_SHIFT_FACTOR_FALLBACK: float = 0.0


@invocation(
"flux2_pid_decode",
title="Latents to Image - FLUX.2 + PiD (4x SR)",
tags=["latents", "image", "pid", "flux2", "klein", "upscale"],
category="latents",
version="1.0.0",
version="1.1.0",
classification=Classification.Prototype,
)
class Flux2PiDDecodeInvocation(BaseInvocation, WithMetadata, WithBoard):
"""Decode a FLUX.2 Klein latent with the PiD pixel-diffusion decoder.

Produces a 4x super-resolved image in a single pass. The stored FLUX.2 latent
is patchified from ``(B, 32, H/8, W/8)`` to the ``(B, 128, H/16, W/16)`` layout
PiD's FLUX.2 backbone expects, then decoded directly (it is already in raw,
BN-denormalized space; see the module docstring).
is patchified from ``(B, 32, H/8, W/8)`` to ``(B, 128, H/16, W/16)`` and
BN-normalized with the VAE's statistics, the form PiD's FLUX.2 backbone expects.
"""

latents: LatentsField = InputField(description=FieldDescriptions.latents, input=Input.Connection)
Expand All @@ -102,12 +82,9 @@ class Flux2PiDDecodeInvocation(BaseInvocation, WithMetadata, WithBoard):
description="PiD FLUX.2 decoder checkpoint.",
input=Input.Connection,
)
vae: VAEField | None = InputField(
default=None,
vae: VAEField = InputField(
title="VAE",
description="FLUX.2 VAE, used only to read a scalar scaling_factor / shift_factor if one exists. "
"FLUX.2 normalises latents with BatchNorm (already inverted in flux2_denoise), so this is "
"normally an identity transform and the input can be left unconnected.",
description="FLUX.2 VAE. Its BatchNorm statistics normalize the latent the way PiD was trained on.",
input=Input.Connection,
)
num_inference_steps: int = InputField(
Expand All @@ -131,39 +108,21 @@ def invoke(self, context: InvocationContext) -> ImageOutput:

latents = context.tensors.load(self.latents.latents_name)

# 1) Patchify the stored FLUX.2 latent into PiD's expected layout.
# flux2_denoise stores an unpacked (B, 32, H/8, W/8) latent; PiD's
# FLUX.2 backbone wants the packed (B, 128, H/16, W/16) form (32*4=128
# channels, spatial halved). This mirrors pack_flux2's 2x2 patchify but
# keeps a spatial (B, C, h, w) layout rather than a (B, seq, C) sequence.
if latents.shape[-3] != 32:
# 1) Patchify and BN-normalize the stored latent into the form PiD's FLUX.2 backbone was trained on.
with context.models.load(self.vae.vae).model_on_device() as (_, vae):
bn_stats = get_flux2_bn_stats(vae)
if bn_stats is None:
raise ValueError(
f"FLUX.2 PiD decode expected a 32-channel latent from flux2_denoise, got shape "
f"{tuple(latents.shape)}. The upstream node must output the unpacked FLUX.2 latent."
"FLUX.2 PiD decode needs the FLUX.2 VAE's BatchNorm statistics to normalize the latent, "
"but the connected VAE has none. Connect the FLUX.2 VAE."
)
packed = rearrange(latents, "b c (h ph) (w pw) -> b (c ph pw) h w", ph=2, pw=2)
packed = pack_flux2_latent_for_pid(latents, *bn_stats)
context.logger.info(
f"FLUX.2 PiD decode: stored latent shape={tuple(latents.shape)} -> packed for PiD "
f"shape={tuple(packed.shape)} (expect [B, 128, H/16, W/16]) dtype={packed.dtype}"
)

# 2) Resolve the scalar scaling/shift (identity for current FLUX.2 VAEs).
scaling_factor = _FLUX2_VAE_SCALING_FACTOR_FALLBACK
shift_factor = _FLUX2_VAE_SHIFT_FACTOR_FALLBACK
if self.vae is not None:
vae_info = context.models.load(self.vae.vae)
with vae_info.model_on_device() as (_, vae):
config = getattr(vae, "config", None)
if config is not None and hasattr(config, "scaling_factor"):
scaling_factor = float(config.scaling_factor)
shift_factor = float(getattr(config, "shift_factor", None) or 0.0)
else:
scaling_factor = float(getattr(vae, "scale_factor", scaling_factor))
shift_factor = float(getattr(vae, "shift_factor", shift_factor))
del vae_info
TorchDevice.empty_cache()

# 3) Encode caption with Gemma-2.
# 2) Encode caption with Gemma-2.
gemma_text_encoder_info = context.models.load(self.gemma2_encoder.text_encoder)
gemma_tokenizer_info = context.models.load(self.gemma2_encoder.tokenizer)
with ExitStack() as stack:
Expand Down Expand Up @@ -199,7 +158,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
context.models.offload_from_vram(self.gemma2_encoder.text_encoder)
TorchDevice.empty_cache()

# 4) Run PiD decode (the loader already returns a live PidNet).
# 3) Run PiD decode (the loader already returns a live PidNet).
pid_info = context.models.load(self.pid_decoder.decoder)
# The working-memory estimate scales with the OUTPUT pixel count, so it must see the PACKED latent
# (spatial H/16), not the unpacked one - otherwise it over-reserves by 4x.
Expand All @@ -218,20 +177,13 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
device = TorchDevice.choose_torch_device()
dtype = next(iter(pid_net.parameters())).dtype

# The packed latent is already BN-denormalized (raw VAE-input space); the scalar transform below is
# identity for current FLUX.2 VAEs and only bites if a VAE ever exposes real scalar constants.
denorm_latent = packed.to(device=device, dtype=dtype) / scaling_factor + shift_factor
context.logger.info(
f"FLUX.2 PiD denorm_latent stats[min={denorm_latent.min().item():.3f} "
f"max={denorm_latent.max().item():.3f} mean={denorm_latent.mean().item():.3f}] "
f"using scale={scaling_factor:.4f} shift={shift_factor:.4f}"
)
latent = packed.to(device=device, dtype=dtype)
caption_embs = caption_embs.to(device=device, dtype=dtype)

context.util.signal_progress("Running PiD decoder")
decoder = PiDDecoder(pid_net, backbone=BaseModelType.Flux2)
x0 = decoder.decode(
latent=denorm_latent,
latent=latent,
caption_embs=caption_embs,
caption_mask=caption_mask,
config=PiDDecodeConfig(
Expand Down
20 changes: 9 additions & 11 deletions invokeai/app/invocations/pid/flux_pid_decode.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,10 @@
Replaces the regular FLUX VAE decode with the PiD pixel-diffusion super-res
decoder (``PiD_res2k_sr4x_official_flux_distill_4step``). Produces a 4x
super-resolved image from a FLUX latent in a single 4-step distill pass.

The stored latent goes to PiD unchanged: PiD was trained on the FLUX-normalized latent
(``scale * (raw - shift)``, what the upstream tokenizer encodes to), not the raw latent the
AutoEncoder's conv decoder consumes.
"""

from contextlib import ExitStack
Expand All @@ -25,7 +29,6 @@
from invokeai.app.invocations.model import Gemma2EncoderField, PiDDecoderField
from invokeai.app.invocations.primitives import ImageOutput
from invokeai.app.services.shared.invocation_context import InvocationContext
from invokeai.backend.flux.util import get_flux_ae_params
from invokeai.backend.model_manager.taxonomy import BaseModelType
from invokeai.backend.pid._src.networks.pid_net import PidNet
from invokeai.backend.pid.decode import (
Expand All @@ -43,15 +46,14 @@
title="Latents to Image - FLUX + PiD (4x SR)",
tags=["latents", "image", "pid", "flux", "upscale"],
category="latents",
version="1.0.0",
version="1.0.1",
classification=Classification.Prototype,
)
class FluxPiDDecodeInvocation(BaseInvocation, WithMetadata, WithBoard):
"""Decode a FLUX latent with the PiD pixel-diffusion decoder.

The FLUX AutoEncoder usually denormalises the stored latent internally
before its conv decoder runs (`z / scale + shift`); we apply the same
transform manually here so PiD sees the raw latent it was trained on.
Produces a 4x super-resolved image in a single pass. PiD takes the stored,
FLUX-normalized latent as is.
"""

latents: LatentsField = InputField(description=FieldDescriptions.latents, input=Input.Connection)
Expand Down Expand Up @@ -144,17 +146,13 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
device = TorchDevice.choose_torch_device()
dtype = next(iter(pid_net.parameters())).dtype

# FLUX latent is stored in normalised form (matching FluxAutoEncoder
# state); denormalise so PiD sees the same representation it
# consumed during training.
ae = get_flux_ae_params()
denorm_latent = latents.to(device=device, dtype=dtype) / ae.scale_factor + ae.shift_factor
latent = latents.to(device=device, dtype=dtype)
caption_embs = caption_embs.to(device=device, dtype=dtype)

context.util.signal_progress("Running PiD decoder")
decoder = PiDDecoder(pid_net, backbone=BaseModelType.Flux)
x0 = decoder.decode(
latent=denorm_latent,
latent=latent,
caption_embs=caption_embs,
caption_mask=caption_mask,
config=PiDDecodeConfig(
Expand Down
Loading
Loading