diff --git a/docs/src/content/docs/Users Guide/04.Image Generation/Advanced Topics/pid-decode.mdx b/docs/src/content/docs/Users Guide/04.Image Generation/Advanced Topics/pid-decode.mdx index 1fe55efa95e..15debf6cffd 100644 --- a/docs/src/content/docs/Users Guide/04.Image Generation/Advanced Topics/pid-decode.mdx +++ b/docs/src/content/docs/Users Guide/04.Image Generation/Advanced Topics/pid-decode.mdx @@ -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. diff --git a/invokeai/app/invocations/flux2/flux2_denoise.py b/invokeai/app/invocations/flux2/flux2_denoise.py index 928d2e1d117..a8c189e1c00 100644 --- a/invokeai/app/invocations/flux2/flux2_denoise.py +++ b/invokeai/app/invocations/flux2/flux2_denoise.py @@ -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, @@ -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, diff --git a/invokeai/app/invocations/pid/flux2_pid_decode.py b/invokeai/app/invocations/pid/flux2_pid_decode.py index 449eaaeb46d..ceb4fbfe793 100644 --- a/invokeai/app/invocations/pid/flux2_pid_decode.py +++ b/invokeai/app/invocations/pid/flux2_pid_decode.py @@ -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 @@ -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 ( @@ -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) @@ -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( @@ -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: @@ -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. @@ -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( diff --git a/invokeai/app/invocations/pid/flux_pid_decode.py b/invokeai/app/invocations/pid/flux_pid_decode.py index a084caeceaa..50976ae1168 100644 --- a/invokeai/app/invocations/pid/flux_pid_decode.py +++ b/invokeai/app/invocations/pid/flux_pid_decode.py @@ -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 @@ -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 ( @@ -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) @@ -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( diff --git a/invokeai/app/invocations/pid/pid_upscale.py b/invokeai/app/invocations/pid/pid_upscale.py index 9db7864e721..d0f8d974fc3 100644 --- a/invokeai/app/invocations/pid/pid_upscale.py +++ b/invokeai/app/invocations/pid/pid_upscale.py @@ -4,7 +4,7 @@ latent. Pipeline:: image - -> FLUX VAE encode (denormalised back to raw) + -> FLUX VAE encode (FLUX-normalized latent, the form PiD was trained on) -> Gemma-2 caption encode -> PiD decoder (4x SR) -> image (4x linear) @@ -12,7 +12,7 @@ This is the PiD analogue of ESRGAN / SUPIR: a one-shot, end-to-end pixel upscaler. The FLUX VAE is also valid for Z-Image inputs (they share the same 16-channel encoder). SD3 / FLUX.2 upscale paths would each need their -own invocation with the matching VAE encode and latent denormalisation; +own invocation with the matching VAE encode and latent normalisation; they are deferred until we have the matching PiD checkpoints to validate against. """ @@ -37,7 +37,7 @@ from invokeai.app.invocations.primitives import ImageOutput from invokeai.app.invocations.vae.flux_vae_encode import FluxVaeEncodeInvocation from invokeai.app.services.shared.invocation_context import InvocationContext -from invokeai.backend.flux.util import get_flux_ae_params, is_flux_family_vae +from invokeai.backend.flux.util import is_flux_family_vae from invokeai.backend.model_manager.taxonomy import BaseModelType from invokeai.backend.pid._src.networks.pid_net import PidNet from invokeai.backend.pid.decode import ( @@ -56,7 +56,7 @@ title="PiD Upscale (4x) - FLUX VAE", tags=["upscale", "image", "pid", "super-resolution", "flux"], category="image", - version="1.0.0", + version="1.0.1", classification=Classification.Prototype, ) class PiDUpscaleInvocation(BaseInvocation, WithMetadata, WithBoard): @@ -71,8 +71,8 @@ class PiDUpscaleInvocation(BaseInvocation, WithMetadata, WithBoard): image: ImageField = InputField(description="Image to upscale.") vae: VAEField = InputField( - description="The FLUX.1 VAE, which Z-Image also uses. PiD upscale runs the FLUX backbone " - "and undoes FLUX VAE scaling, so a VAE of another family is not accepted.", + description="The FLUX.1 VAE, which Z-Image also uses. PiD's FLUX decoder was trained on latents " + "normalized with its constants, so a VAE of another family is not accepted.", input=Input.Connection, ) gemma2_encoder: Gemma2EncoderField = InputField( @@ -109,7 +109,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput: node_title="PiD Upscale", ) - # 1) Encode the source image into a FLUX raw latent. + # 1) Encode the source image into a FLUX-normalized latent. pil_image = context.images.get_pil(self.image.image_name).convert("RGB") image_tensor = image_resized_to_grid_as_tensor(pil_image) if image_tensor.dim() == 3: @@ -119,7 +119,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput: # Reject a foreign VAE up front with a clear error rather than failing on a stripped assert # inside `vae_encode` under `python -O`. The check is on the latent normalisation rather than # on the class: the FLUX.1 autoencoder is an `AutoencoderKL` now, and so are SD, SD 3.5 and - # CogView 4 -- and the scaling this node undoes below is precisely what differs between them. + # CogView 4 -- and the normalisation PiD's FLUX decoder was trained on is precisely what differs. if not is_flux_family_vae(vae_info.model): raise ValueError( "PiD Upscale requires the FLUX.1 autoencoder (Z-Image uses the same one). The VAE " @@ -127,12 +127,9 @@ def invoke(self, context: InvocationContext) -> ImageOutput: "so its latents do not mean what the PiD decoder expects." ) context.util.signal_progress("Running VAE encode") - normalised_latent = FluxVaeEncodeInvocation.vae_encode(vae_info=vae_info, image_tensor=image_tensor) - # The FLUX VAE encode emits `scale * (raw - shift)`. PiD expects raw, so undo it. - ae = get_flux_ae_params() - raw_latent = normalised_latent / ae.scale_factor + ae.shift_factor - raw_latent = raw_latent.to("cpu") # park while we swap to Gemma - del normalised_latent + # The FLUX VAE encode emits `scale * (raw - shift)`, the latent PiD's FLUX decoder was trained on. + latent = FluxVaeEncodeInvocation.vae_encode(vae_info=vae_info, image_tensor=image_tensor) + latent = latent.to("cpu") # park while we swap to Gemma TorchDevice.empty_cache() # 2) Encode the caption with Gemma-2. @@ -177,7 +174,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput: # peak that will not happen (or too little for one that will). pid_memory_optimization = context.config.get().pid_memory_optimization estimated_working_memory = estimate_pid_decode_working_memory( - raw_latent, + latent, BaseModelType.Flux, pid_memory_optimization, pid_net=pid_info.model, @@ -188,7 +185,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput: device = TorchDevice.choose_torch_device() dtype = next(iter(pid_net.parameters())).dtype - latent_on_device = raw_latent.to(device=device, dtype=dtype) + latent_on_device = latent.to(device=device, dtype=dtype) caption_embs = caption_embs.to(device=device, dtype=dtype) context.util.signal_progress("Running PiD decoder") diff --git a/invokeai/app/invocations/pid/qwen_image_pid_decode.py b/invokeai/app/invocations/pid/qwen_image_pid_decode.py index 83f5375ac08..8ae3953d949 100644 --- a/invokeai/app/invocations/pid/qwen_image_pid_decode.py +++ b/invokeai/app/invocations/pid/qwen_image_pid_decode.py @@ -7,19 +7,13 @@ Qwen-Image is 16-channel at an 8x spatial down-factor (``_PER_BACKBONE[QwenImage]`` in ``backend/pid/decode.py``: ``lq_latent_channels=16``, ``latent_spatial_down_factor=8``), -so no packing is needed. Two Qwen-specific wrinkles, both handled below and both -verified against the existing ``qwen_image_l2i`` node: - -1. **5D latent.** The denoiser stores a 5D ``(B, 16, num_frames, H, W)`` latent - (Qwen's VAE is a video-style autoencoder). PiD is a 2D image decoder, so we - drop the singleton temporal dim before decoding. -2. **Per-channel normalization.** Unlike FLUX / Z-Image / SDXL (a scalar - ``scaling_factor`` / ``shift``), the Qwen VAE normalizes each of the 16 latent - channels by its own ``latents_mean`` / ``latents_std`` vector. Denormalization - is therefore ``z_raw = z_norm * latents_std + latents_mean`` per channel - - exactly the transform ``qwen_image_l2i`` applies before ``vae.decode``, so PiD - (which replaces that decode) sees the same raw latent. We read the vectors from - the VAE config when a ``vae`` is wired, with the diffusers defaults as fallback. +so no packing is needed. The denoiser stores a 5D ``(B, 16, num_frames, H, W)`` +latent (Qwen's VAE is a video-style autoencoder); PiD is a 2D image decoder, so we +drop the singleton temporal dim before decoding. + +The latent stays in the Qwen VAE's per-channel normalized space +(``(raw - latents_mean) / latents_std``, what the upstream tokenizer encodes to): +PiD was trained on that, not on the raw latent ``qwen_image_l2i`` hands the VAE. """ from contextlib import ExitStack @@ -53,31 +47,20 @@ ) from invokeai.backend.util.devices import TorchDevice -# Per-channel Qwen-Image VAE normalization constants (diffusers AutoencoderKLQwenImage defaults, z_dim=16). Used -# only as a fallback when no `vae` is wired; prefer the wired VAE config's latents_mean / latents_std at runtime. -_QWEN_VAE_LATENTS_MEAN_FALLBACK: list[float] = [ - -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508, - 0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921, -] # fmt: skip -_QWEN_VAE_LATENTS_STD_FALLBACK: list[float] = [ - 2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743, - 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160, -] # fmt: skip - @invocation( "qwen_image_pid_decode", title="Latents to Image - Qwen-Image + PiD (4x SR)", tags=["latents", "image", "pid", "qwen-image", "upscale"], category="latents", - version="1.0.0", + version="1.0.1", classification=Classification.Prototype, ) class QwenImagePiDDecodeInvocation(BaseInvocation, WithMetadata, WithBoard): """Decode a Qwen-Image latent with the PiD pixel-diffusion decoder. Produces a 4x super-resolved image in a single pass. The 5D Qwen latent is - reduced to 2D and per-channel denormalized (``z * std + mean``) before PiD. + reduced to 2D and otherwise handed to PiD in its normalized form. """ latents: LatentsField = InputField(description=FieldDescriptions.latents, input=Input.Connection) @@ -98,8 +81,8 @@ class QwenImagePiDDecodeInvocation(BaseInvocation, WithMetadata, WithBoard): vae: VAEField | None = InputField( default=None, title="VAE", - description="Qwen-Image VAE, used to read the per-channel latents_mean / latents_std. " - "If omitted, the diffusers default Qwen-Image constants are used.", + description="Not used: PiD decodes the normalized latent directly. Kept so existing workflows " + "that connect a VAE stay valid.", input=Input.Connection, ) num_inference_steps: int = InputField( @@ -134,30 +117,12 @@ def invoke(self, context: InvocationContext) -> ImageOutput: if latents.ndim != 4 or latents.shape[-3] != 16: raise ValueError(f"Qwen-Image PiD decode expected a 16-channel latent, got shape {tuple(latents.shape)}.") - # 2) Resolve the per-channel latents_mean / latents_std used to denormalise the stored latent. - latents_mean = list(_QWEN_VAE_LATENTS_MEAN_FALLBACK) - latents_std = list(_QWEN_VAE_LATENTS_STD_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) - cfg_mean = getattr(config, "latents_mean", None) if config is not None else None - cfg_std = getattr(config, "latents_std", None) if config is not None else None - if cfg_mean is not None and cfg_std is not None: - latents_mean = [float(x) for x in cfg_mean] - latents_std = [float(x) for x in cfg_std] - del vae_info - TorchDevice.empty_cache() - if len(latents_mean) != 16 or len(latents_std) != 16: - raise ValueError( - f"Qwen-Image VAE latents_mean/latents_std must have 16 entries, got {len(latents_mean)}/{len(latents_std)}." - ) context.logger.info( f"Qwen-Image PiD decode: latent shape={tuple(latents.shape)} (expect [B, 16, H/8, W/8]) " - f"dtype={latents.dtype} per-channel denorm (mean/std from {'VAE config' if self.vae else 'fallback'})" + f"dtype={latents.dtype}" ) - # 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: @@ -193,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) # Read once: the estimate and the decode must agree, or the cache reserves headroom for a # peak that will not happen (or too little for one that will). @@ -210,17 +175,13 @@ def invoke(self, context: InvocationContext) -> ImageOutput: device = TorchDevice.choose_torch_device() dtype = next(iter(pid_net.parameters())).dtype - # Per-channel denormalise: z_raw = z_norm * std + mean (the transform qwen_image_l2i applies before - # vae.decode). mean/std are (16,) -> (1, 16, 1, 1) to broadcast over the (B, 16, H, W) latent. - mean_t = torch.tensor(latents_mean, device=device, dtype=dtype).view(1, 16, 1, 1) - std_t = torch.tensor(latents_std, device=device, dtype=dtype).view(1, 16, 1, 1) - denorm_latent = latents.to(device=device, dtype=dtype) * std_t + mean_t + 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.QwenImage) x0 = decoder.decode( - latent=denorm_latent, + latent=latent, caption_embs=caption_embs, caption_mask=caption_mask, config=PiDDecodeConfig( diff --git a/invokeai/app/invocations/pid/sd3_pid_decode.py b/invokeai/app/invocations/pid/sd3_pid_decode.py index 59eae48a61a..e9b61d3062a 100644 --- a/invokeai/app/invocations/pid/sd3_pid_decode.py +++ b/invokeai/app/invocations/pid/sd3_pid_decode.py @@ -3,6 +3,10 @@ Replaces SD3's AutoencoderKL decode with the PiD pixel-diffusion super-res decoder (``PiD_res2k_sr4x_official_sd3_distill_4step``). Produces a 4x super-resolved image from an SD3 latent in a 4-step distill pass. + +The stored latent goes to PiD unchanged: PiD was trained on the SD3-normalized latent +(``scale * (raw - shift)``, what the upstream tokenizer encodes to), not the raw latent the +VAE's conv decoder consumes. """ from contextlib import ExitStack @@ -34,7 +38,6 @@ encode_caption_for_pid, estimate_pid_decode_working_memory, ) -from invokeai.backend.sd3.vae import SD3_VAE_SCALING_FACTOR, SD3_VAE_SHIFT_FACTOR from invokeai.backend.util.devices import TorchDevice @@ -43,7 +46,7 @@ title="Latents to Image - SD3 + PiD (4x SR)", tags=["latents", "image", "pid", "sd3", "upscale"], category="latents", - version="1.0.0", + version="1.0.1", classification=Classification.Prototype, ) class SD3PiDDecodeInvocation(BaseInvocation, WithMetadata, WithBoard): @@ -137,13 +140,13 @@ def invoke(self, context: InvocationContext) -> ImageOutput: device = TorchDevice.choose_torch_device() dtype = next(iter(pid_net.parameters())).dtype - denorm_latent = latents.to(device=device, dtype=dtype) / SD3_VAE_SCALING_FACTOR + SD3_VAE_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.StableDiffusion3) x0 = decoder.decode( - latent=denorm_latent, + latent=latent, caption_embs=caption_embs, caption_mask=caption_mask, config=PiDDecodeConfig( diff --git a/invokeai/app/invocations/pid/sdxl_pid_decode.py b/invokeai/app/invocations/pid/sdxl_pid_decode.py index fa34744f052..57f950d62fb 100644 --- a/invokeai/app/invocations/pid/sdxl_pid_decode.py +++ b/invokeai/app/invocations/pid/sdxl_pid_decode.py @@ -6,14 +6,9 @@ SDXL latents are 4-channel at an 8x spatial down-factor (``_PER_BACKBONE[SDXL]`` in ``backend/pid/decode.py``: ``lq_latent_channels=4``, ``latent_spatial_down_factor=8``), -so - unlike FLUX.2 - no patchify/pack is needed; the stored latent goes straight -to PiD after denormalization. - -Denormalization: SDXL's VAE (``AutoencoderKL``) exposes a scalar -``scaling_factor`` (0.13025) and no shift, so the stored latent is denormalized -as ``z / scaling_factor + shift`` (matching the FLUX / Z-Image nodes). We read -the constants from the VAE config at runtime when a ``vae`` is wired, falling -back to the documented SDXL constants otherwise. +so - unlike FLUX.2 - no patchify/pack is needed. The stored latent goes to PiD +unchanged: PiD was trained on the scaled latent (``0.13025 * raw``, what the upstream +tokenizer encodes to), not the raw latent the VAE's conv decoder consumes. """ from contextlib import ExitStack @@ -47,26 +42,20 @@ ) from invokeai.backend.util.devices import TorchDevice -# SDXL VAE constants (diffusers `stabilityai/sdxl-vae` config: scaling_factor=0.13025, no shift). Prefer reading -# scaling_factor / shift_factor from the wired VAE config at runtime; use these only as a fallback. -_SDXL_VAE_SCALING_FACTOR_FALLBACK: float = 0.13025 -_SDXL_VAE_SHIFT_FACTOR_FALLBACK: float = 0.0 - @invocation( "sdxl_pid_decode", title="Latents to Image - SDXL + PiD (4x SR)", tags=["latents", "image", "pid", "sdxl", "upscale"], category="latents", - version="1.0.0", + version="1.0.1", classification=Classification.Prototype, ) class SDXLPiDDecodeInvocation(BaseInvocation, WithMetadata, WithBoard): """Decode an SDXL latent with the PiD pixel-diffusion decoder. - Produces a 4x super-resolved image in a single pass. The SDXL latent is - 4-channel at an 8x down-factor, so it is denormalized (``z / scaling_factor``) - and handed straight to PiD - no packing needed. + Produces a 4x super-resolved image in a single pass. The stored, scaled SDXL + latent goes to PiD as is. """ latents: LatentsField = InputField(description=FieldDescriptions.latents, input=Input.Connection) @@ -87,8 +76,8 @@ class SDXLPiDDecodeInvocation(BaseInvocation, WithMetadata, WithBoard): vae: VAEField | None = InputField( default=None, title="VAE", - description="SDXL VAE, used to read scaling_factor / shift_factor. " - "If omitted, the SDXL fallback constants (0.13025 / 0.0) are used.", + description="Not used: PiD decodes the normalized latent directly. Kept so existing workflows " + "that connect a VAE stay valid.", input=Input.Connection, ) num_inference_steps: int = InputField( @@ -112,28 +101,11 @@ def invoke(self, context: InvocationContext) -> ImageOutput: latents = context.tensors.load(self.latents.latents_name) - # 1) Resolve the VAE scaling/shift used to denormalise the stored SDXL latent. Prefer the VAE config; fall - # back to the documented SDXL constants (0.13025 / 0.0). - scaling_factor = _SDXL_VAE_SCALING_FACTOR_FALLBACK - shift_factor = _SDXL_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() context.logger.info( - f"SDXL PiD decode: latent shape={tuple(latents.shape)} (expect [B, 4, H/8, W/8]) dtype={latents.dtype} " - f"using scale={scaling_factor:.5f} shift={shift_factor:.5f}" + f"SDXL PiD decode: latent shape={tuple(latents.shape)} (expect [B, 4, H/8, W/8]) dtype={latents.dtype}" ) - # 2) Encode caption with Gemma-2. + # 1) 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: @@ -169,7 +141,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput: context.models.offload_from_vram(self.gemma2_encoder.text_encoder) TorchDevice.empty_cache() - # 3) Run PiD decode (the loader already returns a live PidNet). + # 2) Run PiD decode (the loader already returns a live PidNet). pid_info = context.models.load(self.pid_decoder.decoder) # Read once: the estimate and the decode must agree, or the cache reserves headroom for a # peak that will not happen (or too little for one that will). @@ -186,14 +158,13 @@ def invoke(self, context: InvocationContext) -> ImageOutput: device = TorchDevice.choose_torch_device() dtype = next(iter(pid_net.parameters())).dtype - # SDXL latents come out of the LDM in the VAE-normalized space; denormalise so PiD sees the raw latent. - denorm_latent = latents.to(device=device, dtype=dtype) / scaling_factor + 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.StableDiffusionXL) x0 = decoder.decode( - latent=denorm_latent, + latent=latent, caption_embs=caption_embs, caption_mask=caption_mask, config=PiDDecodeConfig( diff --git a/invokeai/app/invocations/pid/z_image_pid_decode.py b/invokeai/app/invocations/pid/z_image_pid_decode.py index 7d835ce4b12..cb982deeaaa 100644 --- a/invokeai/app/invocations/pid/z_image_pid_decode.py +++ b/invokeai/app/invocations/pid/z_image_pid_decode.py @@ -5,6 +5,10 @@ Z-Image latents. This node replaces the regular Z-Image VAE decode with a PiD super-resolution decode (4x scale, ~256×256 latent → 2048×2048 image by default). + +Z-Image normalizes its latents with the FLUX VAE constants, so the stored latent is already in the +FLUX-normalized space the FLUX decoder was trained on (``scale * (raw - shift)``, what the upstream +tokenizer encodes to) and goes to PiD unchanged. """ from contextlib import ExitStack @@ -38,20 +42,13 @@ ) from invokeai.backend.util.devices import TorchDevice -# Fallback Z-Image VAE constants. PiD's pipeline_registry.py explicitly notes -# the exact values depend on the pretrained checkpoint, so prefer reading them -# from the VAE config at runtime (see `vae` input below) and use these only as -# a last resort. -_ZIMAGE_VAE_SCALING_FACTOR_FALLBACK: float = 0.3611 -_ZIMAGE_VAE_SHIFT_FACTOR_FALLBACK: float = 0.1159 - @invocation( "z_image_pid_decode", title="Latents to Image - Z-Image + PiD (4x SR)", tags=["latents", "image", "pid", "z-image", "upscale"], category="latents", - version="1.0.0", + version="1.0.1", classification=Classification.Prototype, ) class ZImagePiDDecodeInvocation(BaseInvocation, WithMetadata, WithBoard): @@ -80,8 +77,8 @@ class ZImagePiDDecodeInvocation(BaseInvocation, WithMetadata, WithBoard): vae: VAEField | None = InputField( default=None, title="VAE", - description="Z-Image VAE used to read scaling_factor / shift_factor. " - "If omitted, the FLUX.1 fallback constants (0.3611 / 0.1159) are used.", + description="Not used: PiD decodes the normalized latent directly. Kept so existing workflows " + "that connect a VAE stay valid.", input=Input.Connection, ) num_inference_steps: int = InputField( @@ -105,32 +102,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput: latents = context.tensors.load(self.latents.latents_name) - # 1) Resolve the VAE scaling/shift used to denormalise the stored - # Z-Image latent. PiD's pipeline_registry says these are - # checkpoint-specific for Z-Image, so prefer the VAE config when - # available and fall back to the FLUX values otherwise. - scaling_factor = _ZIMAGE_VAE_SCALING_FACTOR_FALLBACK - shift_factor = _ZIMAGE_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: - # FluxAutoEncoder stores the constants directly on the module. - scaling_factor = float(getattr(vae, "scale_factor", scaling_factor)) - shift_factor = float(getattr(vae, "shift_factor", shift_factor)) - del vae_info - TorchDevice.empty_cache() - context.logger.info( - f"Z-Image PiD decode: latent shape={tuple(latents.shape)} dtype={latents.dtype} " - f"stats[min={latents.min().item():.3f} max={latents.max().item():.3f} " - f"mean={latents.mean().item():.3f}] using scale={scaling_factor:.4f} shift={shift_factor:.4f}" - ) - - # 2) Encode caption with Gemma-2. + # 1) 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: @@ -187,23 +159,13 @@ def invoke(self, context: InvocationContext) -> ImageOutput: device = TorchDevice.choose_torch_device() dtype = next(iter(pid_net.parameters())).dtype - # Z-Image latents come out of the diffusers pipeline normalised - # by the VAE constants. PiD expects the raw latent. - denorm_latent = latents.to(device=device, dtype=dtype) / scaling_factor + shift_factor - context.logger.info( - f"denorm_latent stats[min={denorm_latent.min().item():.3f} " - f"max={denorm_latent.max().item():.3f} mean={denorm_latent.mean().item():.3f} " - f"std={denorm_latent.float().std().item():.3f}]; " - f"caption_embs shape={tuple(caption_embs.shape)} " - f"stats[min={caption_embs.min().item():.3f} max={caption_embs.max().item():.3f} " - f"mean={caption_embs.mean().item():.3f} std={caption_embs.float().std().item():.3f}]" - ) + 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( diff --git a/invokeai/backend/flux/util.py b/invokeai/backend/flux/util.py index 78fdeb56981..b3900b3ccaf 100644 --- a/invokeai/backend/flux/util.py +++ b/invokeai/backend/flux/util.py @@ -120,11 +120,11 @@ def is_flux_family_vae(vae: Any) -> bool: latent width, so the channel count alone does not separate them either. What does separate them is the normalisation, and that is also exactly what a caller of this - depends on: `pid_upscale` undoes `scale * (raw - shift)` with the constants from - `get_flux_ae_params()`, and a VAE that used different ones hands it a latent that means something - else. So the test is on the constants themselves -- Z-Image passes, because it is the same - autoencoder (`_name_or_path: "flux-dev"` in its published config); SD 3.5 does not, because its - `scaling_factor` is 1.5305. + depends on: `pid_upscale` hands the encoded latent to a PiD decoder trained on + `scale * (raw - shift)` with the constants from `get_flux_ae_params()`, and a VAE that used + different ones hands it a latent that means something else. So the test is on the constants + themselves -- Z-Image passes, because it is the same autoencoder (`_name_or_path: "flux-dev"` in + its published config); SD 3.5 does not, because its `scaling_factor` is 1.5305. """ config = getattr(vae, "config", None) return config is not None and is_flux_family_vae_config(config) diff --git a/invokeai/backend/flux2/sampling_utils.py b/invokeai/backend/flux2/sampling_utils.py index 88649603fb7..c7564ef5de4 100644 --- a/invokeai/backend/flux2/sampling_utils.py +++ b/invokeai/backend/flux2/sampling_utils.py @@ -57,6 +57,36 @@ def get_noise_flux2( ).to(device=device, dtype=dtype) +def get_flux2_bn_stats(vae: torch.nn.Module) -> tuple[torch.Tensor, torch.Tensor] | None: + """Extract the latent BatchNorm statistics from a FLUX.2 VAE. + + The FLUX.2 VAE normalizes the patchified 128-channel latent with a BatchNorm that has + affine=False, so there are no learnable weight/bias: y = (x - mean) / std, x = y * std + mean. + + Returns: + (bn_mean, bn_std) of shape (128,), or None if the VAE has no initialized BN layer. + """ + # Eval mode keeps a forward pass from updating the running statistics. + vae.eval() + + # The BN layer sits at different locations depending on the 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 or bn_layer.running_mean is None or bn_layer.running_var is None: + return None + + bn_mean = bn_layer.running_mean.clone() + bn_eps = bn_layer.eps if hasattr(bn_layer, "eps") else 1e-4 # BFL uses 1e-4 + bn_std = torch.sqrt(bn_layer.running_var.clone() + bn_eps) + return bn_mean, bn_std + + def pack_flux2(x: torch.Tensor) -> torch.Tensor: """Pack latent image to flattened array of patch embeddings for FLUX.2. diff --git a/invokeai/backend/pid/decode.py b/invokeai/backend/pid/decode.py index 8e4b74bb047..895af18aa53 100644 --- a/invokeai/backend/pid/decode.py +++ b/invokeai/backend/pid/decode.py @@ -27,6 +27,7 @@ from typing import Any, Optional import torch +from einops import rearrange from torch import Tensor from invokeai.backend.model_manager.taxonomy import BaseModelType @@ -594,7 +595,9 @@ def decode( Args: latent: ``[B, C_lat, H_lat, W_lat]`` LQ latent (the LDM's x0 - output, scaled per the backbone's VAE convention). + output) in the backbone VAE's *normalized* space — what the + upstream tokenizers' ``encode`` returns and the transformer + denoises in. Never the raw latent a VAE's conv decoder consumes. caption_embs: ``[B, T, 2304]`` Gemma-2-2b-it caption embeddings (output of `_encode_text_raw` upstream — InvokeAI callers produce this via `Gemma2EncoderLoader`). @@ -754,6 +757,29 @@ def assert_pid_decoder_matches_base(decoder_base: BaseModelType, node_base: Base ) +def pack_flux2_latent_for_pid(latents: Tensor, bn_mean: Tensor, bn_std: Tensor) -> Tensor: + """Bring a stored FLUX.2 latent into the layout and normalization PiD's FLUX.2 backbone was trained on. + + ``flux2_denoise`` stores the raw ``(B, 32, H/8, W/8)`` latent: it inverts the VAE's BatchNorm before + unpacking. PiD instead takes the latent the upstream FLUX.2 tokenizer encodes to (``flux2_vae.py``): + 2x2-patchified to ``(B, 128, H/16, W/16)`` with channel order ``(c ph pw)`` — the order ``pack_flux2`` + uses, so the BN statistics index the same channels — and then BN-normalized. + """ + if latents.ndim != 4 or latents.shape[1] != 32: + raise ValueError( + f"FLUX.2 PiD decode expected an unpacked 32-channel FLUX.2 latent, got shape {tuple(latents.shape)}." + ) + packed = rearrange(latents, "b c (h ph) (w pw) -> b (c ph pw) h w", ph=2, pw=2) + if bn_mean.numel() != packed.shape[1] or bn_std.numel() != packed.shape[1]: + raise ValueError( + f"FLUX.2 PiD decode needs BatchNorm statistics for {packed.shape[1]} channels, got " + f"{bn_mean.numel()}. Connect the FLUX.2 VAE." + ) + bn_mean = bn_mean.to(device=packed.device, dtype=packed.dtype).view(1, -1, 1, 1) + bn_std = bn_std.to(device=packed.device, dtype=packed.dtype).view(1, -1, 1, 1) + return (packed - bn_mean) / bn_std + + __all__ = [ "BACKBONE_DISCRIMINATOR_KEY", "PID_CHI_PROMPT", @@ -765,5 +791,6 @@ def assert_pid_decoder_matches_base(decoder_base: BaseModelType, node_base: Base "build_pid_net", "encode_caption_for_pid", "load_pid_decoder", + "pack_flux2_latent_for_pid", "required_pid_net_shapes", ] diff --git a/invokeai/frontend/api/openapi.json b/invokeai/frontend/api/openapi.json index 537d0836d51..029225b40db 100644 --- a/invokeai/frontend/api/openapi.json +++ b/invokeai/frontend/api/openapi.json @@ -37653,7 +37653,7 @@ "category": "latents", "class": "invocation", "classification": "prototype", - "description": "Decode a FLUX.2 Klein latent with the PiD pixel-diffusion decoder.\n\nProduces a 4x super-resolved image in a single pass. The stored FLUX.2 latent\nis patchified from ``(B, 32, H/8, W/8)`` to the ``(B, 128, H/16, W/16)`` layout\nPiD's FLUX.2 backbone expects, then decoded directly (it is already in raw,\nBN-denormalized space; see the module docstring).", + "description": "Decode a FLUX.2 Klein latent with the PiD pixel-diffusion decoder.\n\nProduces a 4x super-resolved image in a single pass. The stored FLUX.2 latent\nis patchified from ``(B, 32, H/8, W/8)`` to ``(B, 128, H/16, W/16)`` and\nBN-normalized with the VAE's statistics, the form PiD's FLUX.2 backbone expects.", "node_pack": "invokeai", "properties": { "board": { @@ -37786,11 +37786,10 @@ } ], "default": null, - "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.", "field_kind": "input", "input": "connection", - "orig_default": null, - "orig_required": false, + "orig_required": true, "title": "VAE" }, "num_inference_steps": { @@ -37827,7 +37826,7 @@ "tags": ["latents", "image", "pid", "flux2", "klein", "upscale"], "title": "Latents to Image - FLUX.2 + PiD (4x SR)", "type": "object", - "version": "1.0.0", + "version": "1.1.0", "output": { "$ref": "#/components/schemas/ImageOutput" } @@ -40340,7 +40339,7 @@ "category": "latents", "class": "invocation", "classification": "prototype", - "description": "Decode a FLUX latent with the PiD pixel-diffusion decoder.\n\nThe FLUX AutoEncoder usually denormalises the stored latent internally\nbefore its conv decoder runs (`z / scale + shift`); we apply the same\ntransform manually here so PiD sees the raw latent it was trained on.", + "description": "Decode a FLUX latent with the PiD pixel-diffusion decoder.\n\nProduces a 4x super-resolved image in a single pass. PiD takes the stored,\nFLUX-normalized latent as is.", "node_pack": "invokeai", "properties": { "board": { @@ -40497,7 +40496,7 @@ "tags": ["latents", "image", "pid", "flux", "upscale"], "title": "Latents to Image - FLUX + PiD (4x SR)", "type": "object", - "version": "1.0.0", + "version": "1.0.1", "output": { "$ref": "#/components/schemas/ImageOutput" } @@ -92223,7 +92222,7 @@ } ], "default": null, - "description": "The FLUX.1 VAE, which Z-Image also uses. PiD upscale runs the FLUX backbone and undoes FLUX VAE scaling, so a VAE of another family is not accepted.", + "description": "The FLUX.1 VAE, which Z-Image also uses. PiD's FLUX decoder was trained on latents normalized with its constants, so a VAE of another family is not accepted.", "field_kind": "input", "input": "connection", "orig_required": true @@ -92305,7 +92304,7 @@ "tags": ["upscale", "image", "pid", "super-resolution", "flux"], "title": "PiD Upscale (4x) - FLUX VAE", "type": "object", - "version": "1.0.0", + "version": "1.0.1", "output": { "$ref": "#/components/schemas/ImageOutput" } @@ -97026,7 +97025,7 @@ "category": "latents", "class": "invocation", "classification": "prototype", - "description": "Decode a Qwen-Image latent with the PiD pixel-diffusion decoder.\n\nProduces a 4x super-resolved image in a single pass. The 5D Qwen latent is\nreduced to 2D and per-channel denormalized (``z * std + mean``) before PiD.", + "description": "Decode a Qwen-Image latent with the PiD pixel-diffusion decoder.\n\nProduces a 4x super-resolved image in a single pass. The 5D Qwen latent is\nreduced to 2D and otherwise handed to PiD in its normalized form.", "node_pack": "invokeai", "properties": { "board": { @@ -97159,7 +97158,7 @@ } ], "default": null, - "description": "Qwen-Image VAE, used to read the per-channel latents_mean / latents_std. If omitted, the diffusers default Qwen-Image constants are used.", + "description": "Not used: PiD decodes the normalized latent directly. Kept so existing workflows that connect a VAE stay valid.", "field_kind": "input", "input": "connection", "orig_default": null, @@ -97200,7 +97199,7 @@ "tags": ["latents", "image", "pid", "qwen-image", "upscale"], "title": "Latents to Image - Qwen-Image + PiD (4x SR)", "type": "object", - "version": "1.0.0", + "version": "1.0.1", "output": { "$ref": "#/components/schemas/ImageOutput" } @@ -99586,7 +99585,7 @@ "tags": ["latents", "image", "pid", "sd3", "upscale"], "title": "Latents to Image - SD3 + PiD (4x SR)", "type": "object", - "version": "1.0.0", + "version": "1.0.1", "output": { "$ref": "#/components/schemas/ImageOutput" } @@ -100194,7 +100193,7 @@ "category": "latents", "class": "invocation", "classification": "prototype", - "description": "Decode an SDXL latent with the PiD pixel-diffusion decoder.\n\nProduces a 4x super-resolved image in a single pass. The SDXL latent is\n4-channel at an 8x down-factor, so it is denormalized (``z / scaling_factor``)\nand handed straight to PiD - no packing needed.", + "description": "Decode an SDXL latent with the PiD pixel-diffusion decoder.\n\nProduces a 4x super-resolved image in a single pass. The stored, scaled SDXL\nlatent goes to PiD as is.", "node_pack": "invokeai", "properties": { "board": { @@ -100327,7 +100326,7 @@ } ], "default": null, - "description": "SDXL VAE, used to read scaling_factor / shift_factor. If omitted, the SDXL fallback constants (0.13025 / 0.0) are used.", + "description": "Not used: PiD decodes the normalized latent directly. Kept so existing workflows that connect a VAE stay valid.", "field_kind": "input", "input": "connection", "orig_default": null, @@ -100368,7 +100367,7 @@ "tags": ["latents", "image", "pid", "sdxl", "upscale"], "title": "Latents to Image - SDXL + PiD (4x SR)", "type": "object", - "version": "1.0.0", + "version": "1.0.1", "output": { "$ref": "#/components/schemas/ImageOutput" } @@ -118160,7 +118159,7 @@ } ], "default": null, - "description": "Z-Image VAE used to read scaling_factor / shift_factor. If omitted, the FLUX.1 fallback constants (0.3611 / 0.1159) are used.", + "description": "Not used: PiD decodes the normalized latent directly. Kept so existing workflows that connect a VAE stay valid.", "field_kind": "input", "input": "connection", "orig_default": null, @@ -118201,7 +118200,7 @@ "tags": ["latents", "image", "pid", "z-image", "upscale"], "title": "Latents to Image - Z-Image + PiD (4x SR)", "type": "object", - "version": "1.0.0", + "version": "1.0.1", "output": { "$ref": "#/components/schemas/ImageOutput" } diff --git a/invokeai/frontend/api/schema.ts b/invokeai/frontend/api/schema.ts index a845bdf0a0f..df119e94267 100644 --- a/invokeai/frontend/api/schema.ts +++ b/invokeai/frontend/api/schema.ts @@ -14608,9 +14608,8 @@ export type components = { * @description 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. */ Flux2PiDDecodeInvocation: { /** @@ -14665,7 +14664,7 @@ export type components = { pid_decoder?: components["schemas"]["PiDDecoderField"] | null; /** * 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. * @default null */ vae?: components["schemas"]["VAEField"] | null; @@ -15857,9 +15856,8 @@ export type components = { * Latents to Image - FLUX + PiD (4x SR) * @description 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. */ FluxPiDDecodeInvocation: { /** @@ -39200,7 +39198,7 @@ export type components = { */ image?: components["schemas"]["ImageField"] | null; /** - * @description The FLUX.1 VAE, which Z-Image also uses. PiD upscale runs the FLUX backbone and undoes FLUX VAE scaling, so a VAE of another family is not accepted. + * @description The FLUX.1 VAE, which Z-Image also uses. PiD's FLUX decoder was trained on latents normalized with its constants, so a VAE of another family is not accepted. * @default null */ vae?: components["schemas"]["VAEField"] | null; @@ -41974,7 +41972,7 @@ export type components = { * @description Decode a Qwen-Image latent with the PiD pixel-diffusion decoder. * * Produces a 4x super-resolved image in a single pass. The 5D Qwen latent is - * reduced to 2D and per-channel denormalized (``z * std + mean``) before PiD. + * reduced to 2D and otherwise handed to PiD in its normalized form. */ QwenImagePiDDecodeInvocation: { /** @@ -42029,7 +42027,7 @@ export type components = { pid_decoder?: components["schemas"]["PiDDecoderField"] | null; /** * VAE - * @description Qwen-Image VAE, used to read the per-channel latents_mean / latents_std. If omitted, the diffusers default Qwen-Image constants are used. + * @description Not used: PiD decodes the normalized latent directly. Kept so existing workflows that connect a VAE stay valid. * @default null */ vae?: components["schemas"]["VAEField"] | null; @@ -43628,9 +43626,8 @@ export type components = { * Latents to Image - SDXL + PiD (4x SR) * @description Decode an SDXL latent with the PiD pixel-diffusion decoder. * - * Produces a 4x super-resolved image in a single pass. The SDXL latent is - * 4-channel at an 8x down-factor, so it is denormalized (``z / scaling_factor``) - * and handed straight to PiD - no packing needed. + * Produces a 4x super-resolved image in a single pass. The stored, scaled SDXL + * latent goes to PiD as is. */ SDXLPiDDecodeInvocation: { /** @@ -43685,7 +43682,7 @@ export type components = { pid_decoder?: components["schemas"]["PiDDecoderField"] | null; /** * VAE - * @description SDXL VAE, used to read scaling_factor / shift_factor. If omitted, the SDXL fallback constants (0.13025 / 0.0) are used. + * @description Not used: PiD decodes the normalized latent directly. Kept so existing workflows that connect a VAE stay valid. * @default null */ vae?: components["schemas"]["VAEField"] | null; @@ -53462,7 +53459,7 @@ export type components = { pid_decoder?: components["schemas"]["PiDDecoderField"] | null; /** * VAE - * @description Z-Image VAE used to read scaling_factor / shift_factor. If omitted, the FLUX.1 fallback constants (0.3611 / 0.1159) are used. + * @description Not used: PiD decodes the normalized latent directly. Kept so existing workflows that connect a VAE stay valid. * @default null */ vae?: components["schemas"]["VAEField"] | null; diff --git a/invokeai/frontend/webv2/src/features/generation/core/canvas/compileCanvasGraph.test.ts b/invokeai/frontend/webv2/src/features/generation/core/canvas/compileCanvasGraph.test.ts index 661cc95bd71..a6229a7e167 100644 --- a/invokeai/frontend/webv2/src/features/generation/core/canvas/compileCanvasGraph.test.ts +++ b/invokeai/frontend/webv2/src/features/generation/core/canvas/compileCanvasGraph.test.ts @@ -1242,6 +1242,7 @@ describe('compileCanvasGraph — PiD decode', () => { .map((edge) => [edge.destination.field, [graph.nodes[edge.source.node_id]?.type, edge.source.field]]) ); const PID_BASES = [sdxlModel, sd3Model, fluxModel, flux2Model, qwenImageModel, zImageModel]; + const sdxlVae: VaeModelConfig = { base: 'sdxl', key: 'sdxl-vae', name: 'SDXL VAE', type: 'vae' }; it('txt2img fit: the PiD chain replaces the VAE decode and a resize back to the bbox owns canvas_output', () => { const { backendGraph } = compile(sdxlModel, 'txt2img', { settings: pidSettings(sdxlModel, 'fit') }); @@ -1304,7 +1305,7 @@ describe('compileCanvasGraph — PiD decode', () => { expect(getNodeByType(backendGraph, 'core_metadata')?.width).toBe(bbox.width); }); - it('img2img fit: encodes the composite at the bbox with the VAE the PiD decode reads', () => { + it('img2img fit: encodes the composite at the bbox with the VAE the PiD decode records', () => { const { backendGraph } = compile(sdxlModel, 'img2img', { settings: pidSettings(sdxlModel, 'fit'), strength: 0.5 }); const encode = backendGraph.nodes.canvas_i2l; @@ -1352,6 +1353,26 @@ describe('compileCanvasGraph — PiD decode', () => { expect(backendGraph.nodes.canvas_processing_output).toBeUndefined(); }); + it.each(['fit', 'native'] as const)('img2img %s (SDXL): encodes and PiD decodes with the VAE override', (pidMode) => { + const { backendGraph } = compile(sdxlModel, 'img2img', { + settings: { ...pidSettings(sdxlModel, pidMode), vae: sdxlVae }, + }); + + expect(backendGraph.nodes.vae_loader?.vae_model).toEqual(sdxlVae); + expect(getEdge(backendGraph, 'canvas_i2l', 'vae')?.source).toEqual({ field: 'vae', node_id: 'vae_loader' }); + expect(getEdge(backendGraph, pidDecodeNode(backendGraph)!.id, 'vae')?.source.node_id).toBe('vae_loader'); + }); + + it.each(['fit', 'native'] as const)('img2img %s (SDXL): encodes through the seamless VAE when tiling', (pidMode) => { + const { backendGraph } = compile(sdxlModel, 'img2img', { + settings: { ...pidSettings(sdxlModel, pidMode), seamlessXAxis: true, seamlessYAxis: true }, + }); + + expect(backendGraph.nodes.seamless).toMatchObject({ seamless_x: true, seamless_y: true }); + expect(getEdge(backendGraph, 'canvas_i2l', 'vae')?.source).toEqual({ field: 'vae', node_id: 'seamless' }); + expect(getEdge(backendGraph, pidDecodeNode(backendGraph)!.id, 'vae')?.source.node_id).toBe('seamless'); + }); + it.each(PID_BASES)('resolves the encode VAE for a $base native img2img', (model) => { const { backendGraph } = compile(model, 'img2img', { settings: { ...COMPONENT_OVERRIDES[model.base], ...pidSettings(model, 'native') }, diff --git a/invokeai/frontend/webv2/src/features/generation/core/graph.test.ts b/invokeai/frontend/webv2/src/features/generation/core/graph.test.ts index a70cebf8087..d96e7f2ea18 100644 --- a/invokeai/frontend/webv2/src/features/generation/core/graph.test.ts +++ b/invokeai/frontend/webv2/src/features/generation/core/graph.test.ts @@ -1256,7 +1256,7 @@ describe('PiD decode', () => { expect(getNodeByType(graph, 'gemma2_encoder_loader')).toMatchObject({ gemma2_model: gemma2 }); }); - it('does not wire a VAE into decoders that use fixed constants', () => { + it('wires no VAE into the FLUX.1 decode, which has no VAE input', () => { const graph = compile(fluxModel, { clipEmbedModel: clipEmbed, t5EncoderModel: t5Encoder, @@ -1265,17 +1265,23 @@ describe('PiD decode', () => { }); const decode = getNodeByType(graph, 'flux_pid_decode'); - // FLUX.1 and SD3 use fixed VAE constants and expose no `vae` input. expect(getEdge(graph, decode!.id, 'vae')).toBeUndefined(); }); - it('wires a VAE into decoders that read its constants at runtime', () => { + it('wires the VAE into the Qwen-Image decode, recording the VAE the graph decodes with', () => { const graph = compile(qwenImageModel, { ...pidSettings('qwen-image') }); const decode = getNodeByType(graph, 'qwen_image_pid_decode'); expect(getEdge(graph, decode!.id, 'vae')?.source.node_id).toBe('model_loader'); }); + it('wires the VAE into the FLUX.2 decode, which normalizes with its BatchNorm statistics', () => { + const graph = compile(flux2Model, { ...pidSettings('flux2') }); + const decode = getNodeByType(graph, 'flux2_pid_decode'); + + expect(getEdge(graph, decode!.id, 'vae')?.source.node_id).toBe('model_loader'); + }); + it('uses a FLUX decoder for Z-Image, which shares its VAE', () => { const graph = compile(zImageModel, { ...pidSettings('flux') }); diff --git a/invokeai/frontend/webv2/src/features/generation/core/graph.ts b/invokeai/frontend/webv2/src/features/generation/core/graph.ts index 598c1dbbf05..6621e877fb4 100644 --- a/invokeai/frontend/webv2/src/features/generation/core/graph.ts +++ b/invokeai/frontend/webv2/src/features/generation/core/graph.ts @@ -539,10 +539,10 @@ const addDecodeOutput = ({ }; /** - * The node and field a built graph decodes with, for grafting an encode onto it. An ordinary decode reads it on - * `canvas_output.vae`; a PiD decode reads it only when its node type takes a `vae` input. The FLUX and SD-3 PiD - * decodes take none, and those builders decode with the model loader, which the fallback names. Null when no builder - * convention applies. + * The node and field a built graph decodes with, for grafting an encode onto it. An ordinary decode carries it on + * `canvas_output.vae`. A PiD decode carries it on its `vae` edge when its node type takes that input: only the FLUX.2 + * decode reads it, the others keep the edge as the record of the decode VAE. The FLUX and SD-3 PiD decodes take none, + * and those builders decode with the model loader, which the fallback names. Null when no builder convention applies. */ export const findVaeSource = ( graph: BackendGraphContract @@ -557,7 +557,8 @@ export const findVaeSource = ( const node = graph.nodes[edge.source.node_id]; return node ? { field: edge.source.field, node } : null; } - // The fallback holds because each builder wires exactly one `vae` edge into its decode node, from `model_loader`. + // A PiD decode without a `vae` edge is the FLUX or SD-3 one. Their builders decode with `model_loader`, which takes + // any VAE override as a loader input, so its `vae` output is the decode VAE. const modelLoader = graph.nodes.model_loader; const decodesWithPid = Object.values(graph.nodes).some((node) => isPidDecodeNodeType(node.type)); return modelLoader && decodesWithPid ? { field: 'vae', node: modelLoader } : null; diff --git a/invokeai/frontend/webv2/src/features/generation/core/pidGraph.ts b/invokeai/frontend/webv2/src/features/generation/core/pidGraph.ts index 03d1ed3fa44..2dcb090a656 100644 --- a/invokeai/frontend/webv2/src/features/generation/core/pidGraph.ts +++ b/invokeai/frontend/webv2/src/features/generation/core/pidGraph.ts @@ -24,7 +24,12 @@ const PID_DECODE_NODE_TYPE_SET: ReadonlySet = new Set(Object.values(PID_ export const isPidDecodeNodeType = (type: string | undefined): type is PidDecodeNodeType => type !== undefined && PID_DECODE_NODE_TYPE_SET.has(type); -/** Retain VAE edges for runtime scaling constants even though PiD replaces decode. */ +/** + * The decode node types that take a `vae` input. PiD takes the latent in the VAE's normalized space, and only + * FLUX.2 stores it denormalized, so only its decode reads the VAE (BatchNorm statistics to normalize its latent). + * The others take the stored latent as is, but the edge still records which VAE the graph decodes with: + * `findVaeSource` reads it to give a Canvas img2img encode that same VAE, override and seamless wrapping included. + */ const PID_DECODE_NODES_WITH_VAE_INPUT: ReadonlySet = new Set([ 'flux2_pid_decode', 'qwen_image_pid_decode', @@ -67,7 +72,7 @@ interface AddPidDecodeArg { positivePrompt: BackendInvocationContract; /** The seed node, reused for PiD's internal decode noise so results reproduce. */ seed: BackendInvocationContract; - /** VAE source, wired only for decode nodes that read its constants. */ + /** The VAE the graph decodes with, wired into every decode node that takes a `vae` input. */ vaeSource?: BackendInvocationContract; /** The VAE output field on `vaeSource` (model loaders expose `vae`). */ vaeField?: string; diff --git a/tests/backend/pid/test_pid_latent_normalization.py b/tests/backend/pid/test_pid_latent_normalization.py new file mode 100644 index 00000000000..659ad7d0426 --- /dev/null +++ b/tests/backend/pid/test_pid_latent_normalization.py @@ -0,0 +1,65 @@ +"""The FLUX.2 PiD node must hand PiD the latent the upstream FLUX.2 tokenizer encodes to. + +Upstream (nv-tlabs/PiD ``tokenizers/flux2_vae.py``) trains PiD on ``BatchNorm(patchify(raw))`` with the +VAE's running statistics. ``flux2_denoise`` stores the BN-*denormalized*, unpacked latent, so the node +has to undo exactly that. +""" + +import pytest +import torch +from einops import rearrange +from torch import nn + +from invokeai.backend.flux2.sampling_utils import get_flux2_bn_stats, unpack_flux2 +from invokeai.backend.pid.decode import pack_flux2_latent_for_pid + + +class _Flux2LikeVae(nn.Module): + def __init__(self) -> None: + super().__init__() + self.bn = nn.BatchNorm2d(128, eps=1e-4, affine=False, track_running_stats=True) + generator = torch.Generator().manual_seed(0) + self.bn.running_mean.copy_(torch.randn(128, generator=generator)) + self.bn.running_var.copy_(torch.rand(128, generator=generator) * 4 + 0.1) + + +def test_pack_matches_upstream_flux2_tokenizer_encode() -> None: + vae = _Flux2LikeVae() + raw = torch.randn(2, 32, 12, 8, generator=torch.Generator().manual_seed(1)) + + # The upstream tokenizer's encode, verbatim: patchify with channel order (c pi pj), then BatchNorm in eval mode. + vae.bn.eval() + expected = vae.bn(rearrange(raw, "... c (i pi) (j pj) -> ... (c pi pj) i j", pi=2, pj=2)) + + bn_stats = get_flux2_bn_stats(vae) + assert bn_stats is not None + torch.testing.assert_close(pack_flux2_latent_for_pid(raw, *bn_stats), expected) + + +def test_pack_recovers_the_latent_flux2_denoise_produced() -> None: + """Round trip through flux2_denoise's output path: denoised (normalized, packed) -> stored -> PiD input.""" + bn_mean, bn_std = get_flux2_bn_stats(_Flux2LikeVae()) + height, width = 192, 128 + packed_h, packed_w = height // 16, width // 16 + normalized = torch.randn(1, packed_h * packed_w, 128, generator=torch.Generator().manual_seed(2)) + + # What flux2_denoise stores: BN-denormalize in (B, seq, 128), then unpack to (B, 32, H/8, W/8). + stored = unpack_flux2(normalized * bn_std + bn_mean, height, width) + + expected = rearrange(normalized, "b (h w) c -> b c h w", h=packed_h, w=packed_w) + torch.testing.assert_close(pack_flux2_latent_for_pid(stored, bn_mean, bn_std), expected) + + +def test_pack_rejects_a_non_flux2_latent() -> None: + bn_mean, bn_std = torch.zeros(128), torch.ones(128) + with pytest.raises(ValueError, match="32-channel"): + pack_flux2_latent_for_pid(torch.zeros(1, 16, 8, 8), bn_mean, bn_std) + + +def test_pack_rejects_batchnorm_stats_of_another_vae() -> None: + with pytest.raises(ValueError, match="128 channels"): + pack_flux2_latent_for_pid(torch.zeros(1, 32, 8, 8), torch.zeros(64), torch.ones(64)) + + +def test_bn_stats_absent_without_a_batchnorm_layer() -> None: + assert get_flux2_bn_stats(nn.Conv2d(4, 4, 1)) is None