From fb7358f6566e86b2fd7258a8ac3ad3136a8f6807 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Sun, 2 Aug 2026 08:17:13 +0200 Subject: [PATCH] refactor(architectures): drive latent-space facts from the registry The preview chain was never really keyed on architecture. The branch comments in the dispatch chain this removes said so by hand -- "Krea-2 decodes with the Qwen-Image VAE", "Z-Image uses FLUX-compatible VAE", "Anima uses Wan 2.1 VAE". Fifteen architectures share eight latent spaces, and the five most recent introduced no new preview data at all. That one fact was written in six partial, drifting copies. This replaces four of them with LatentSpaceFacet; the remaining two are named below as follow-ups. No behaviour changes. Every projection matrix moves byte-identically, comments included -- read the move with `git diff --color-moved`. Both Wan dispatches collapse into one. step_callback.py used to pick the factors by `sample.shape[-3] == 48` and then, forty lines later, pick the spatial scale by the same test. They are two halves of one fact and could drift; now resolving one LatentSpace settles both. LatentSpaceFacet.resolve matches on channel count, which is literally what the old comment said it was doing ("the latent channel count uniquely identifies the variant"). An architecture with a single latent space short-circuits without touching the sample at all, so an unusually shaped tensor cannot start raising IndexError for the other fourteen. Ideogram 4 keeps its own preview loop -- its callback signature is step/total/packed_latents and it must unpatchify and denormalize first -- but it now reads the same declaration instead of inlining the FLUX.2 factors and a hardcoded x8. Its try/except fallback and its missing is_canceled check are left alone; the latter is a real bug, in that Ideogram generations do not stop promptly on cancel, but fixing it would change behaviour, which this PR does not. max_unet_downscale was duplicated verbatim across denoise_latents and T2IAdapterExt, comment and error string included. It is a UNet property, not VAE geometry, so it gets its own optional facet; the accessor rather than require carries the error, reproducing the old message exactly, enum repr included. The three Wan 2.1 matrices -- QWEN_IMAGE_, ANIMA_ and WAN_LATENT_RGB_FACTORS -- were byte-identical, as were their three biases. They are now one WAN21_16, asserted by object identity so the duplication cannot creep back. Two known copies stay: constants.LATENT_SCALE_FACTOR, whose blast radius is every latent node, and pid/decode.py, which is the same shape in different units (packed latents at 128ch/16x, not VAE latents at 32ch/8x). Both are follow-ups now that a correct source exists. invocation_context.py is untouched. The spec proposed collapsing flux_step_callback and flux2_step_callback into sd_step_callback, but they are already thin wrappers over the same function, so the change would be pure docs churn in docs/src/generated/invocation-context.json. Tests. test_latent_space.py carries a reference table read off the old chain, one row per architecture, plus first/last rows and column sums as exact fingerprints -- IEEE-754 addition in fixed order, so platform-independent, and the rows catch a reordering a sum is blind to. test_step_callback.py now covers diffusion_step_callback end to end for all 15 bases and both Wan cases, which is the first automated coverage the spatial-scale path has ever had. The relocated projection tests hardcode their expected pixel. The versions they replace recomputed it by summing the very matrix under test, so they would have passed against a corrupted one -- and had in fact drifted: that test's docstring claimed column sums of 0.3677/0.4577/0.9101 where the real ones are 0.3887/0.8771/1.3152. The registry fixture now empties Facet.FACET_TYPES as well as the registry. The first REQUIRED facet exposed that isolating only one of the two globals let validate() hold dummy architectures to real facets. denoise_latents.py keeps a third, divergent copy of the SDXL BGR rule: it reads the UNet's base where the other paths read each adapter's. The two disagree for an SD1 adapter on an SDXL UNet. Left as-is with a NOTE, since fixing it changes behaviour; the natural home is a bgr_input field on UNetDownscaleFacet. openapi.json, schema.ts and invocation-context.json are unchanged. Co-Authored-By: Claude Opus 5 (1M context) --- invokeai/app/invocations/denoise_latents.py | 16 +- invokeai/app/invocations/ideogram4_denoise.py | 29 +- invokeai/app/util/step_callback.py | 352 +--------------- invokeai/backend/architectures/__init__.py | 13 + invokeai/backend/architectures/defs/anima.py | 7 +- .../backend/architectures/defs/cogview4.py | 6 +- .../backend/architectures/defs/ernie_image.py | 10 +- invokeai/backend/architectures/defs/flux.py | 6 +- invokeai/backend/architectures/defs/flux2.py | 6 +- .../backend/architectures/defs/ideogram_4.py | 10 +- invokeai/backend/architectures/defs/krea_2.py | 8 +- .../backend/architectures/defs/qwen_image.py | 7 +- invokeai/backend/architectures/defs/sd_1.py | 8 +- invokeai/backend/architectures/defs/sd_2.py | 7 +- invokeai/backend/architectures/defs/sd_3.py | 6 +- invokeai/backend/architectures/defs/sdxl.py | 8 +- .../architectures/defs/sdxl_refiner.py | 7 +- invokeai/backend/architectures/defs/wan.py | 9 +- .../backend/architectures/defs/z_image.py | 7 +- .../architectures/facets/latent_space.py | 398 ++++++++++++++++++ invokeai/backend/architectures/facets/unet.py | 38 ++ .../extensions/t2i_adapter.py | 10 +- tests/app/util/test_step_callback.py | 261 +++++++----- .../architectures/test_latent_space.py | 303 +++++++++++++ tests/backend/architectures/test_registry.py | 21 +- .../architectures/test_unet_downscale.py | 34 ++ 26 files changed, 1081 insertions(+), 506 deletions(-) create mode 100644 invokeai/backend/architectures/facets/latent_space.py create mode 100644 invokeai/backend/architectures/facets/unet.py create mode 100644 tests/backend/architectures/test_latent_space.py create mode 100644 tests/backend/architectures/test_unet_downscale.py diff --git a/invokeai/app/invocations/denoise_latents.py b/invokeai/app/invocations/denoise_latents.py index 89508dfca68..d7929ab9071 100644 --- a/invokeai/app/invocations/denoise_latents.py +++ b/invokeai/app/invocations/denoise_latents.py @@ -38,6 +38,7 @@ from invokeai.app.invocations.t2i_adapter import T2IAdapterField from invokeai.app.services.shared.invocation_context import InvocationContext from invokeai.app.util.controlnet_utils import prepare_control_image +from invokeai.backend.architectures import get_max_unet_downscale from invokeai.backend.ip_adapter.ip_adapter import IPAdapter from invokeai.backend.model_manager.configs.factory import AnyModelConfig from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelVariantType @@ -637,17 +638,13 @@ def run_t2i_adapters( t2i_adapter_model_config = context.models.get_config(t2i_adapter_field.t2i_adapter_model.key) image = context.images.get_pil(t2i_adapter_field.image.image_name, mode="RGB") - # The max_unet_downscale is the maximum amount that the UNet model downscales the latent image internally. - if t2i_adapter_model_config.base == BaseModelType.StableDiffusion1: - max_unet_downscale = 8 - elif t2i_adapter_model_config.base == BaseModelType.StableDiffusionXL: - max_unet_downscale = 4 + # Raises for a base without a UNet, before the BGR swap below -- same order as before. + max_unet_downscale = get_max_unet_downscale(t2i_adapter_model_config.base) + if t2i_adapter_model_config.base == BaseModelType.StableDiffusionXL: # SDXL adapters are trained on cv2's BGR outputs r, g, b = image.split() image = Image.merge("RGB", (b, g, r)) - else: - raise ValueError(f"Unexpected T2I-Adapter base model type: '{t2i_adapter_model_config.base}'.") t2i_adapter_model: T2IAdapter with context.models.load(t2i_adapter_field.t2i_adapter_model) as t2i_adapter_model: @@ -934,6 +931,11 @@ def step_callback(state: PipelineIntermediateState) -> None: # ext = extension_field.to_extension(exit_stack, context, ext_manager) # ext_manager.add_extension(ext) self.parse_controlnet_field(exit_stack, context, self.control, ext_manager) + # NOTE: this decides the BGR swap from the *UNet's* base, while run_t2i_adapters above + # decides it from each *adapter's* base. The two disagree for an SD1 adapter on an SDXL + # UNet. Left as-is deliberately: fixing it changes behaviour, and this refactor does + # not. The natural fix is to fold `bgr_input` into UNetDownscaleFacet so both paths read + # one declaration. bgr_mode = self.unet.unet.base == BaseModelType.StableDiffusionXL self.parse_t2i_adapter_field(exit_stack, context, self.t2i_adapter, ext_manager, bgr_mode) diff --git a/invokeai/app/invocations/ideogram4_denoise.py b/invokeai/app/invocations/ideogram4_denoise.py index f9bc9ea855d..cb43b4b7f34 100644 --- a/invokeai/app/invocations/ideogram4_denoise.py +++ b/invokeai/app/invocations/ideogram4_denoise.py @@ -12,16 +12,13 @@ from invokeai.app.invocations.model import TransformerField from invokeai.app.invocations.primitives import LatentsOutput from invokeai.app.services.shared.invocation_context import InvocationContext -from invokeai.app.util.step_callback import ( - FLUX2_LATENT_RGB_BIAS, - FLUX2_LATENT_RGB_FACTORS, - sample_to_lowres_estimated_image, -) +from invokeai.backend.architectures import get_latent_space from invokeai.backend.ideogram4 import run_ideogram4_denoise from invokeai.backend.ideogram4.latent_norm import get_latent_norm from invokeai.backend.ideogram4.sampler_configs import PRESETS from invokeai.backend.ideogram4.sampling_utils import unpatchify_and_denormalize from invokeai.backend.ideogram4.transformer_pair import Ideogram4TransformerPair +from invokeai.backend.model_manager.taxonomy import BaseModelType from invokeai.backend.stable_diffusion.diffusion.conditioning_data import Ideogram4ConditioningInfo from invokeai.backend.util.devices import TorchDevice @@ -122,12 +119,13 @@ def invoke(self, context: InvocationContext) -> LatentsOutput: assert isinstance(info, Ideogram4ConditioningInfo) llm_features = info.prompt_embeds.to(device=device, dtype=torch.float32) - # Progress-preview setup: Ideogram uses a FLUX.2-style 32-channel VAE, so the FLUX.2 - # latent->RGB factors give a usable (approximate) low-res preview of the forming image at each - # step, without a full VAE decode. Denormalization params come from get_latent_norm (no VAE). + # Progress-preview setup: Ideogram's latent space gives a usable (approximate) low-res + # preview of the forming image at each step, without a full VAE decode. Denormalization + # params come from get_latent_norm (no VAE). This does not go through + # diffusion_step_callback: the callback signature here is (step, total, packed_latents) and + # the latents must be unpatchified and denormalized first. latent_shift, latent_scale = get_latent_norm() - rgb_factors = torch.tensor(FLUX2_LATENT_RGB_FACTORS, dtype=torch.float32) - rgb_bias = torch.tensor(FLUX2_LATENT_RGB_BIAS, dtype=torch.float32) + latent_space = get_latent_space(BaseModelType.Ideogram4) def step_callback(step: int, total: int, packed_latents: torch.Tensor) -> None: preview = None @@ -138,11 +136,7 @@ def step_callback(step: int, total: int, packed_latents: torch.Tensor) -> None: latent_shift.to(packed_latents.device), latent_scale.to(packed_latents.device), ) - preview = sample_to_lowres_estimated_image( - samples=vae_latent, - latent_rgb_factors=rgb_factors.to(vae_latent.device), - latent_rgb_bias=rgb_bias.to(vae_latent.device), - ) + preview = latent_space.preview(vae_latent) except Exception: # A preview must never break generation — fall back to a plain progress signal. preview = None @@ -151,7 +145,10 @@ def step_callback(step: int, total: int, packed_latents: torch.Tensor) -> None: "Running Ideogram 4 denoising", step / total, preview, - (preview.width * 8, preview.height * 8), + ( + preview.width * latent_space.spatial_compression, + preview.height * latent_space.spatial_compression, + ), ) else: context.util.signal_progress("Running Ideogram 4 denoising", step / total) diff --git a/invokeai/app/util/step_callback.py b/invokeai/app/util/step_callback.py index a6206448e9a..b3fb7f69782 100644 --- a/invokeai/app/util/step_callback.py +++ b/invokeai/app/util/step_callback.py @@ -1,287 +1,13 @@ from math import floor -from typing import Callable, Optional, TypeAlias +from typing import Callable, TypeAlias -import torch from PIL import Image from invokeai.app.services.session_processor.session_processor_common import CanceledException +from invokeai.backend.architectures import resolve_latent_space from invokeai.backend.model_manager.taxonomy import BaseModelType from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState -# See scripts/generate_vae_linear_approximation.py for generating these factors. - -# fast latents preview matrix for sdxl -# generated by @StAlKeR7779 -SDXL_LATENT_RGB_FACTORS = [ - # R G B - [0.3816, 0.4930, 0.5320], - [-0.3753, 0.1631, 0.1739], - [0.1770, 0.3588, -0.2048], - [-0.4350, -0.2644, -0.4289], -] -SDXL_SMOOTH_MATRIX = [ - [0.0358, 0.0964, 0.0358], - [0.0964, 0.4711, 0.0964], - [0.0358, 0.0964, 0.0358], -] - -# origingally adapted from code by @erucipe and @keturn here: -# https://discuss.huggingface.co/t/decoding-latents-to-rgb-without-upscaling/23204/7 -# these updated numbers for v1.5 are from @torridgristle -SD1_5_LATENT_RGB_FACTORS = [ - # R G B - [0.3444, 0.1385, 0.0670], # L1 - [0.1247, 0.4027, 0.1494], # L2 - [-0.3192, 0.2513, 0.2103], # L3 - [-0.1307, -0.1874, -0.7445], # L4 -] - -SD3_5_LATENT_RGB_FACTORS = [ - [-0.05240681, 0.03251581, 0.0749016], - [-0.0580572, 0.00759826, 0.05729818], - [0.16144888, 0.01270368, -0.03768577], - [0.14418615, 0.08460266, 0.15941818], - [0.04894035, 0.0056485, -0.06686988], - [0.05187166, 0.19222395, 0.06261094], - [0.1539433, 0.04818359, 0.07103094], - [-0.08601796, 0.09013458, 0.10893912], - [-0.12398469, -0.06766567, 0.0033688], - [-0.0439737, 0.07825329, 0.02258823], - [0.03101129, 0.06382551, 0.07753657], - [-0.01315361, 0.08554491, -0.08772475], - [0.06464487, 0.05914605, 0.13262741], - [-0.07863674, -0.02261737, -0.12761454], - [-0.09923835, -0.08010759, -0.06264447], - [-0.03392309, -0.0804029, -0.06078822], -] - -FLUX_LATENT_RGB_FACTORS = [ - [-0.0412, 0.0149, 0.0521], - [0.0056, 0.0291, 0.0768], - [0.0342, -0.0681, -0.0427], - [-0.0258, 0.0092, 0.0463], - [0.0863, 0.0784, 0.0547], - [-0.0017, 0.0402, 0.0158], - [0.0501, 0.1058, 0.1152], - [-0.0209, -0.0218, -0.0329], - [-0.0314, 0.0083, 0.0896], - [0.0851, 0.0665, -0.0472], - [-0.0534, 0.0238, -0.0024], - [0.0452, -0.0026, 0.0048], - [0.0892, 0.0831, 0.0881], - [-0.1117, -0.0304, -0.0789], - [0.0027, -0.0479, -0.0043], - [-0.1146, -0.0827, -0.0598], -] - -COGVIEW4_LATENT_RGB_FACTORS = [ - [0.00408832, -0.00082485, -0.00214816], - [0.00084172, 0.00132241, 0.00842067], - [-0.00466737, -0.00983181, -0.00699561], - [0.03698397, -0.04797235, 0.03585809], - [0.00234701, -0.00124326, 0.00080869], - [-0.00723903, -0.00388422, -0.00656606], - [-0.00970917, -0.00467356, -0.00971113], - [0.17292486, -0.03452463, -0.1457515], - [0.02330308, 0.02942557, 0.02704329], - [-0.00903131, -0.01499841, -0.01432564], - [0.01250298, 0.0019407, -0.02168986], - [0.01371188, 0.00498283, -0.01302135], - [0.42396525, 0.4280575, 0.42148206], - [0.00983825, 0.00613302, 0.00610316], - [0.00473307, -0.00889551, -0.00915924], - [-0.00955853, -0.00980067, -0.00977842], -] - -# Qwen Image uses the same VAE as Wan 2.1 (16-channel). -# Factors from ComfyUI: https://github.com/comfyanonymous/ComfyUI/blob/master/comfy/latent_formats.py -QWEN_IMAGE_LATENT_RGB_FACTORS = [ - [-0.1299, -0.1692, 0.2932], - [0.0671, 0.0406, 0.0442], - [0.3568, 0.2548, 0.1747], - [0.0372, 0.2344, 0.1420], - [0.0313, 0.0189, -0.0328], - [0.0296, -0.0956, -0.0665], - [-0.3477, -0.4059, -0.2925], - [0.0166, 0.1902, 0.1975], - [-0.0412, 0.0267, -0.1364], - [-0.1293, 0.0740, 0.1636], - [0.0680, 0.3019, 0.1128], - [0.0032, 0.0581, 0.0639], - [-0.1251, 0.0927, 0.1699], - [0.0060, -0.0633, 0.0005], - [0.3477, 0.2275, 0.2950], - [0.1984, 0.0913, 0.1861], -] - -QWEN_IMAGE_LATENT_RGB_BIAS = [-0.1835, -0.0868, -0.3360] - -# FLUX.2 uses 32 latent channels. -# Factors from ComfyUI: https://github.com/Comfy-Org/ComfyUI/blob/main/comfy/latent_formats.py -FLUX2_LATENT_RGB_FACTORS = [ - # R G B - [0.0058, 0.0113, 0.0073], - [0.0495, 0.0443, 0.0836], - [-0.0099, 0.0096, 0.0644], - [0.2144, 0.3009, 0.3652], - [0.0166, -0.0039, -0.0054], - [0.0157, 0.0103, -0.0160], - [-0.0398, 0.0902, -0.0235], - [-0.0052, 0.0095, 0.0109], - [-0.3527, -0.2712, -0.1666], - [-0.0301, -0.0356, -0.0180], - [-0.0107, 0.0078, 0.0013], - [0.0746, 0.0090, -0.0941], - [0.0156, 0.0169, 0.0070], - [-0.0034, -0.0040, -0.0114], - [0.0032, 0.0181, 0.0080], - [-0.0939, -0.0008, 0.0186], - [0.0018, 0.0043, 0.0104], - [0.0284, 0.0056, -0.0127], - [-0.0024, -0.0022, -0.0030], - [0.1207, -0.0026, 0.0065], - [0.0128, 0.0101, 0.0142], - [0.0137, -0.0072, -0.0007], - [0.0095, 0.0092, -0.0059], - [0.0000, -0.0077, -0.0049], - [-0.0465, -0.0204, -0.0312], - [0.0095, 0.0012, -0.0066], - [0.0290, -0.0034, 0.0025], - [0.0220, 0.0169, -0.0048], - [-0.0332, -0.0457, -0.0468], - [-0.0085, 0.0389, 0.0609], - [-0.0076, 0.0003, -0.0043], - [-0.0111, -0.0460, -0.0614], -] - -FLUX2_LATENT_RGB_BIAS = [-0.0329, -0.0718, -0.0851] - -# Anima uses Wan 2.1 VAE with 16 latent channels. -# Factors from ComfyUI: https://github.com/Comfy-Org/ComfyUI/blob/main/comfy/latent_formats.py -ANIMA_LATENT_RGB_FACTORS = [ - [-0.1299, -0.1692, 0.2932], - [0.0671, 0.0406, 0.0442], - [0.3568, 0.2548, 0.1747], - [0.0372, 0.2344, 0.1420], - [0.0313, 0.0189, -0.0328], - [0.0296, -0.0956, -0.0665], - [-0.3477, -0.4059, -0.2925], - [0.0166, 0.1902, 0.1975], - [-0.0412, 0.0267, -0.1364], - [-0.1293, 0.0740, 0.1636], - [0.0680, 0.3019, 0.1128], - [0.0032, 0.0581, 0.0639], - [-0.1251, 0.0927, 0.1699], - [0.0060, -0.0633, 0.0005], - [0.3477, 0.2275, 0.2950], - [0.1984, 0.0913, 0.1861], -] - -ANIMA_LATENT_RGB_BIAS = [-0.1835, -0.0868, -0.3360] - -# Wan 2.2 A14B uses the standard 16-channel Wan VAE. -# Factors come from ComfyUI's Wan21 latent_format (same VAE as A14B). -WAN_LATENT_RGB_FACTORS = [ - [-0.1299, -0.1692, 0.2932], - [0.0671, 0.0406, 0.0442], - [0.3568, 0.2548, 0.1747], - [0.0372, 0.2344, 0.1420], - [0.0313, 0.0189, -0.0328], - [0.0296, -0.0956, -0.0665], - [-0.3477, -0.4059, -0.2925], - [0.0166, 0.1902, 0.1975], - [-0.0412, 0.0267, -0.1364], - [-0.1293, 0.0740, 0.1636], - [0.0680, 0.3019, 0.1128], - [0.0032, 0.0581, 0.0639], - [-0.1251, 0.0927, 0.1699], - [0.0060, -0.0633, 0.0005], - [0.3477, 0.2275, 0.2950], - [0.1984, 0.0913, 0.1861], -] - -WAN_LATENT_RGB_BIAS = [-0.1835, -0.0868, -0.3360] - -# Wan 2.2 TI2V-5B uses Wan2.2-VAE with 48 latent channels and 16x spatial downscale. -# Factors come from ComfyUI's Wan22 latent_format. -WAN22_LATENT_RGB_FACTORS = [ - [0.0119, 0.0103, 0.0046], - [-0.1062, -0.0504, 0.0165], - [0.0140, 0.0409, 0.0491], - [-0.0813, -0.0677, 0.0607], - [0.0656, 0.0851, 0.0808], - [0.0264, 0.0463, 0.0912], - [0.0295, 0.0326, 0.0590], - [-0.0244, -0.0270, 0.0025], - [0.0443, -0.0102, 0.0288], - [-0.0465, -0.0090, -0.0205], - [0.0359, 0.0236, 0.0082], - [-0.0776, 0.0854, 0.1048], - [0.0564, 0.0264, 0.0561], - [0.0006, 0.0594, 0.0418], - [-0.0319, -0.0542, -0.0637], - [-0.0268, 0.0024, 0.0260], - [0.0539, 0.0265, 0.0358], - [-0.0359, -0.0312, -0.0287], - [-0.0285, -0.1032, -0.1237], - [0.1041, 0.0537, 0.0622], - [-0.0086, -0.0374, -0.0051], - [0.0390, 0.0670, 0.2863], - [0.0069, 0.0144, 0.0082], - [0.0006, -0.0167, 0.0079], - [0.0313, -0.0574, -0.0232], - [-0.1454, -0.0902, -0.0481], - [0.0714, 0.0827, 0.0447], - [-0.0304, -0.0574, -0.0196], - [0.0401, 0.0384, 0.0204], - [-0.0758, -0.0297, -0.0014], - [0.0568, 0.1307, 0.1372], - [-0.0055, -0.0310, -0.0380], - [0.0239, -0.0305, 0.0325], - [-0.0663, -0.0673, -0.0140], - [-0.0416, -0.0047, -0.0023], - [0.0166, 0.0112, -0.0093], - [-0.0211, 0.0011, 0.0331], - [0.1833, 0.1466, 0.2250], - [-0.0368, 0.0370, 0.0295], - [-0.3441, -0.3543, -0.2008], - [-0.0479, -0.0489, -0.0420], - [-0.0660, -0.0153, 0.0800], - [-0.0101, 0.0068, 0.0156], - [-0.0690, -0.0452, -0.0927], - [-0.0145, 0.0041, 0.0015], - [0.0421, 0.0451, 0.0373], - [0.0504, -0.0483, -0.0356], - [-0.0837, 0.0168, 0.0055], -] - -WAN22_LATENT_RGB_BIAS = [0.0317, -0.0878, -0.1388] - - -def sample_to_lowres_estimated_image( - samples: torch.Tensor, - latent_rgb_factors: torch.Tensor, - smooth_matrix: Optional[torch.Tensor] = None, - latent_rgb_bias: Optional[torch.Tensor] = None, -): - if samples.dim() == 4: - samples = samples[0] - latent_image = samples.permute(1, 2, 0) @ latent_rgb_factors - - if latent_rgb_bias is not None: - latent_image = latent_image + latent_rgb_bias - - if smooth_matrix is not None: - latent_image = latent_image.unsqueeze(0).permute(3, 0, 1, 2) - latent_image = torch.nn.functional.conv2d(latent_image, smooth_matrix.reshape((1, 1, 3, 3)), padding=1) - latent_image = latent_image.permute(1, 2, 3, 0).squeeze(0) - - latents_ubyte = ( - ((latent_image + 1) / 2).clamp(0, 1).mul(0xFF).byte() # change scale from -1..1 to 0..1 # to 0..255 - ).cpu() - - return Image.fromarray(latents_ubyte.numpy()) - def calc_percentage(intermediate_state: PipelineIntermediateState) -> float: """Calculate the percentage of completion of denoising.""" @@ -322,74 +48,14 @@ def diffusion_step_callback( else: sample = intermediate_state.latents - smooth_matrix: list[list[float]] | None = None - latent_rgb_bias: list[float] | None = None - if base_model in [BaseModelType.StableDiffusion1, BaseModelType.StableDiffusion2]: - latent_rgb_factors = SD1_5_LATENT_RGB_FACTORS - elif base_model in [BaseModelType.StableDiffusionXL, BaseModelType.StableDiffusionXLRefiner]: - latent_rgb_factors = SDXL_LATENT_RGB_FACTORS - smooth_matrix = SDXL_SMOOTH_MATRIX - elif base_model == BaseModelType.StableDiffusion3: - latent_rgb_factors = SD3_5_LATENT_RGB_FACTORS - elif base_model == BaseModelType.CogView4: - latent_rgb_factors = COGVIEW4_LATENT_RGB_FACTORS - elif base_model in [BaseModelType.QwenImage, BaseModelType.Krea2]: - # Krea-2 decodes with the Qwen-Image VAE (16 latent channels), so it shares the preview factors. - latent_rgb_factors = QWEN_IMAGE_LATENT_RGB_FACTORS - latent_rgb_bias = QWEN_IMAGE_LATENT_RGB_BIAS - elif base_model == BaseModelType.Flux: - latent_rgb_factors = FLUX_LATENT_RGB_FACTORS - elif base_model == BaseModelType.Flux2: - latent_rgb_factors = FLUX2_LATENT_RGB_FACTORS - latent_rgb_bias = FLUX2_LATENT_RGB_BIAS - elif base_model == BaseModelType.ZImage: - # Z-Image uses FLUX-compatible VAE with 16 latent channels - latent_rgb_factors = FLUX_LATENT_RGB_FACTORS - elif base_model == BaseModelType.Anima: - # Anima uses Wan 2.1 VAE with 16 latent channels - latent_rgb_factors = ANIMA_LATENT_RGB_FACTORS - latent_rgb_bias = ANIMA_LATENT_RGB_BIAS - elif base_model == BaseModelType.ErnieImage: - # ERNIE-Image uses AutoencoderKLFlux2 (same as FLUX.2) with 32 latent channels, and the - # denoise loop unpatches before previewing, so the shapes line up. The values do not: - # ERNIE denoises in BN-normalized latent space (denormalized only at VAE decode) and the - # BN stats live on the VAE, which isn't loaded here. Previews are therefore approximate - # in color/contrast. - latent_rgb_factors = FLUX2_LATENT_RGB_FACTORS - latent_rgb_bias = FLUX2_LATENT_RGB_BIAS - elif base_model == BaseModelType.Wan: - # A14B (16-ch standard Wan VAE, 8x spatial) vs TI2V-5B (48-ch Wan2.2-VAE, - # 16x spatial). The latent channel count uniquely identifies the variant. - if sample.shape[-3] == 48: - latent_rgb_factors = WAN22_LATENT_RGB_FACTORS - latent_rgb_bias = WAN22_LATENT_RGB_BIAS - else: - latent_rgb_factors = WAN_LATENT_RGB_FACTORS - latent_rgb_bias = WAN_LATENT_RGB_BIAS - else: - raise ValueError(f"Unsupported base model: {base_model}") - - latent_rgb_factors_torch = torch.tensor(latent_rgb_factors, dtype=sample.dtype, device=sample.device) - smooth_matrix_torch = ( - torch.tensor(smooth_matrix, dtype=sample.dtype, device=sample.device) if smooth_matrix else None - ) - latent_rgb_bias_torch = ( - torch.tensor(latent_rgb_bias, dtype=sample.dtype, device=sample.device) if latent_rgb_bias else None - ) - image = sample_to_lowres_estimated_image( - samples=sample, - latent_rgb_factors=latent_rgb_factors_torch, - smooth_matrix=smooth_matrix_torch, - latent_rgb_bias=latent_rgb_bias_torch, - ) + # The projection factors and the spatial scale are two halves of one fact, and used to be + # selected by two separate dispatches over base_model -- both of which had to agree about Wan's + # 16- vs 48-channel VAE. Resolving one latent space settles both. + latent_space = resolve_latent_space(base_model, sample) + image = latent_space.preview(sample) - # Spatial downscale ratio: 8x is the SD/SDXL/FLUX/Wan-A14B default; - # Wan TI2V-5B's Wan2.2-VAE uses 16x. - spatial_scale = 8 - if base_model == BaseModelType.Wan and sample.shape[-3] == 48: - spatial_scale = 16 - width = image.width * spatial_scale - height = image.height * spatial_scale + width = image.width * latent_space.spatial_compression + height = image.height * latent_space.spatial_compression percentage = calc_percentage(intermediate_state) signal_progress("Denoising", percentage, image, (width, height)) diff --git a/invokeai/backend/architectures/__init__.py b/invokeai/backend/architectures/__init__.py index 0a6bb0e7584..10d95f81758 100644 --- a/invokeai/backend/architectures/__init__.py +++ b/invokeai/backend/architectures/__init__.py @@ -32,6 +32,13 @@ z_image, # noqa: F401 ) from invokeai.backend.architectures.facet import Facet +from invokeai.backend.architectures.facets.latent_space import ( + LatentSpace, + LatentSpaceFacet, + get_latent_space, + resolve_latent_space, +) +from invokeai.backend.architectures.facets.unet import UNetDownscaleFacet, get_max_unet_downscale from invokeai.backend.architectures.registry import ( ArchitectureError, defs_module_path, @@ -46,11 +53,17 @@ __all__ = [ "ArchitectureError", "Facet", + "LatentSpace", + "LatentSpaceFacet", + "UNetDownscaleFacet", "defs_module_path", "facets_of", "generative_bases", "get", + "get_latent_space", + "get_max_unet_downscale", "register", "require", + "resolve_latent_space", "validate", ] diff --git a/invokeai/backend/architectures/defs/anima.py b/invokeai/backend/architectures/defs/anima.py index 7d99ead8510..6a062952c71 100644 --- a/invokeai/backend/architectures/defs/anima.py +++ b/invokeai/backend/architectures/defs/anima.py @@ -1,4 +1,9 @@ +from invokeai.backend.architectures.facets.latent_space import WAN21_16, LatentSpaceFacet from invokeai.backend.architectures.registry import register from invokeai.backend.model_manager.taxonomy import BaseModelType -register(BaseModelType.Anima) +register( + BaseModelType.Anima, + # Anima uses the Wan 2.1 VAE with 16 latent channels. + LatentSpaceFacet(WAN21_16), +) diff --git a/invokeai/backend/architectures/defs/cogview4.py b/invokeai/backend/architectures/defs/cogview4.py index 3f9b2a3eafc..2d176cbecfd 100644 --- a/invokeai/backend/architectures/defs/cogview4.py +++ b/invokeai/backend/architectures/defs/cogview4.py @@ -1,4 +1,8 @@ +from invokeai.backend.architectures.facets.latent_space import COGVIEW4_16, LatentSpaceFacet from invokeai.backend.architectures.registry import register from invokeai.backend.model_manager.taxonomy import BaseModelType -register(BaseModelType.CogView4) +register( + BaseModelType.CogView4, + LatentSpaceFacet(COGVIEW4_16), +) diff --git a/invokeai/backend/architectures/defs/ernie_image.py b/invokeai/backend/architectures/defs/ernie_image.py index fffddea0a07..73c384fc6e9 100644 --- a/invokeai/backend/architectures/defs/ernie_image.py +++ b/invokeai/backend/architectures/defs/ernie_image.py @@ -1,4 +1,12 @@ +from invokeai.backend.architectures.facets.latent_space import FLUX2_32, LatentSpaceFacet from invokeai.backend.architectures.registry import register from invokeai.backend.model_manager.taxonomy import BaseModelType -register(BaseModelType.ErnieImage) +register( + BaseModelType.ErnieImage, + # ERNIE-Image uses AutoencoderKLFlux2 (same as FLUX.2) with 32 latent channels, and the denoise + # loop unpatches before previewing, so the shapes line up. The values do not: ERNIE denoises in + # BN-normalized latent space (denormalized only at VAE decode) and the BN stats live on the VAE, + # which isn't loaded here. Previews are therefore approximate in color/contrast. + LatentSpaceFacet(FLUX2_32), +) diff --git a/invokeai/backend/architectures/defs/flux.py b/invokeai/backend/architectures/defs/flux.py index 3fd50410cd7..043cbf94fbe 100644 --- a/invokeai/backend/architectures/defs/flux.py +++ b/invokeai/backend/architectures/defs/flux.py @@ -1,4 +1,8 @@ +from invokeai.backend.architectures.facets.latent_space import FLUX_16, LatentSpaceFacet from invokeai.backend.architectures.registry import register from invokeai.backend.model_manager.taxonomy import BaseModelType -register(BaseModelType.Flux) +register( + BaseModelType.Flux, + LatentSpaceFacet(FLUX_16), +) diff --git a/invokeai/backend/architectures/defs/flux2.py b/invokeai/backend/architectures/defs/flux2.py index fed21e6ecf3..e55b0f37879 100644 --- a/invokeai/backend/architectures/defs/flux2.py +++ b/invokeai/backend/architectures/defs/flux2.py @@ -1,4 +1,8 @@ +from invokeai.backend.architectures.facets.latent_space import FLUX2_32, LatentSpaceFacet from invokeai.backend.architectures.registry import register from invokeai.backend.model_manager.taxonomy import BaseModelType -register(BaseModelType.Flux2) +register( + BaseModelType.Flux2, + LatentSpaceFacet(FLUX2_32), +) diff --git a/invokeai/backend/architectures/defs/ideogram_4.py b/invokeai/backend/architectures/defs/ideogram_4.py index 81b3aa53ce3..f20e66ad8ed 100644 --- a/invokeai/backend/architectures/defs/ideogram_4.py +++ b/invokeai/backend/architectures/defs/ideogram_4.py @@ -1,4 +1,12 @@ +from invokeai.backend.architectures.facets.latent_space import FLUX2_32, LatentSpaceFacet from invokeai.backend.architectures.registry import register from invokeai.backend.model_manager.taxonomy import BaseModelType -register(BaseModelType.Ideogram4) +register( + BaseModelType.Ideogram4, + # Ideogram 4 uses a FLUX.2-style 32-channel VAE. Its denoise loop drives the preview itself + # rather than going through diffusion_step_callback, because its callback signature is + # (step, total, packed_latents) and it must unpatchify and denormalize first -- but the latent + # space it ends up projecting is this one. + LatentSpaceFacet(FLUX2_32), +) diff --git a/invokeai/backend/architectures/defs/krea_2.py b/invokeai/backend/architectures/defs/krea_2.py index 20268dc5f00..c0b7524ddeb 100644 --- a/invokeai/backend/architectures/defs/krea_2.py +++ b/invokeai/backend/architectures/defs/krea_2.py @@ -1,4 +1,10 @@ +from invokeai.backend.architectures.facets.latent_space import WAN21_16, LatentSpaceFacet from invokeai.backend.architectures.registry import register from invokeai.backend.model_manager.taxonomy import BaseModelType -register(BaseModelType.Krea2) +register( + BaseModelType.Krea2, + # Krea-2 decodes with the Qwen-Image VAE, which is the Wan 2.1 VAE (16 latent channels), so it + # shares the preview factors. + LatentSpaceFacet(WAN21_16), +) diff --git a/invokeai/backend/architectures/defs/qwen_image.py b/invokeai/backend/architectures/defs/qwen_image.py index bae4874cd99..7ea0c30c4e8 100644 --- a/invokeai/backend/architectures/defs/qwen_image.py +++ b/invokeai/backend/architectures/defs/qwen_image.py @@ -1,4 +1,9 @@ +from invokeai.backend.architectures.facets.latent_space import WAN21_16, LatentSpaceFacet from invokeai.backend.architectures.registry import register from invokeai.backend.model_manager.taxonomy import BaseModelType -register(BaseModelType.QwenImage) +register( + BaseModelType.QwenImage, + # Qwen-Image decodes with the 16-channel Wan 2.1 VAE. + LatentSpaceFacet(WAN21_16), +) diff --git a/invokeai/backend/architectures/defs/sd_1.py b/invokeai/backend/architectures/defs/sd_1.py index 80a12502b73..70b90d45062 100644 --- a/invokeai/backend/architectures/defs/sd_1.py +++ b/invokeai/backend/architectures/defs/sd_1.py @@ -1,4 +1,10 @@ +from invokeai.backend.architectures.facets.latent_space import SD15_4, LatentSpaceFacet +from invokeai.backend.architectures.facets.unet import UNetDownscaleFacet from invokeai.backend.architectures.registry import register from invokeai.backend.model_manager.taxonomy import BaseModelType -register(BaseModelType.StableDiffusion1) +register( + BaseModelType.StableDiffusion1, + LatentSpaceFacet(SD15_4), + UNetDownscaleFacet(max_unet_downscale=8), +) diff --git a/invokeai/backend/architectures/defs/sd_2.py b/invokeai/backend/architectures/defs/sd_2.py index 94ef6005c53..737b45daea3 100644 --- a/invokeai/backend/architectures/defs/sd_2.py +++ b/invokeai/backend/architectures/defs/sd_2.py @@ -1,4 +1,9 @@ +from invokeai.backend.architectures.facets.latent_space import SD15_4, LatentSpaceFacet from invokeai.backend.architectures.registry import register from invokeai.backend.model_manager.taxonomy import BaseModelType -register(BaseModelType.StableDiffusion2) +register( + BaseModelType.StableDiffusion2, + # SD2 shares SD1's 4-channel latent space and preview factors. + LatentSpaceFacet(SD15_4), +) diff --git a/invokeai/backend/architectures/defs/sd_3.py b/invokeai/backend/architectures/defs/sd_3.py index 42a59b1da66..f7fa3b24988 100644 --- a/invokeai/backend/architectures/defs/sd_3.py +++ b/invokeai/backend/architectures/defs/sd_3.py @@ -1,4 +1,8 @@ +from invokeai.backend.architectures.facets.latent_space import SD3_16, LatentSpaceFacet from invokeai.backend.architectures.registry import register from invokeai.backend.model_manager.taxonomy import BaseModelType -register(BaseModelType.StableDiffusion3) +register( + BaseModelType.StableDiffusion3, + LatentSpaceFacet(SD3_16), +) diff --git a/invokeai/backend/architectures/defs/sdxl.py b/invokeai/backend/architectures/defs/sdxl.py index 8e09e8322d5..c5f9d96ff97 100644 --- a/invokeai/backend/architectures/defs/sdxl.py +++ b/invokeai/backend/architectures/defs/sdxl.py @@ -1,4 +1,10 @@ +from invokeai.backend.architectures.facets.latent_space import SDXL_4, LatentSpaceFacet +from invokeai.backend.architectures.facets.unet import UNetDownscaleFacet from invokeai.backend.architectures.registry import register from invokeai.backend.model_manager.taxonomy import BaseModelType -register(BaseModelType.StableDiffusionXL) +register( + BaseModelType.StableDiffusionXL, + LatentSpaceFacet(SDXL_4), + UNetDownscaleFacet(max_unet_downscale=4), +) diff --git a/invokeai/backend/architectures/defs/sdxl_refiner.py b/invokeai/backend/architectures/defs/sdxl_refiner.py index aca97b1b1ce..877413bf42a 100644 --- a/invokeai/backend/architectures/defs/sdxl_refiner.py +++ b/invokeai/backend/architectures/defs/sdxl_refiner.py @@ -1,4 +1,9 @@ +from invokeai.backend.architectures.facets.latent_space import SDXL_4, LatentSpaceFacet from invokeai.backend.architectures.registry import register from invokeai.backend.model_manager.taxonomy import BaseModelType -register(BaseModelType.StableDiffusionXLRefiner) +register( + BaseModelType.StableDiffusionXLRefiner, + # The refiner shares SDXL's latent space, smooth matrix included. + LatentSpaceFacet(SDXL_4), +) diff --git a/invokeai/backend/architectures/defs/wan.py b/invokeai/backend/architectures/defs/wan.py index f728e022d6e..0e9795ab4e7 100644 --- a/invokeai/backend/architectures/defs/wan.py +++ b/invokeai/backend/architectures/defs/wan.py @@ -1,4 +1,11 @@ +from invokeai.backend.architectures.facets.latent_space import WAN21_16, WAN22_48, LatentSpaceFacet from invokeai.backend.architectures.registry import register from invokeai.backend.model_manager.taxonomy import BaseModelType -register(BaseModelType.Wan) +register( + BaseModelType.Wan, + # The only architecture with more than one latent space. A14B uses the standard 16-channel Wan + # VAE at 8x spatial; TI2V-5B uses the 48-channel Wan2.2-VAE at 16x. The latent channel count + # uniquely identifies the variant, which is how `LatentSpaceFacet.resolve()` tells them apart. + LatentSpaceFacet(WAN21_16, alternates=(WAN22_48,)), +) diff --git a/invokeai/backend/architectures/defs/z_image.py b/invokeai/backend/architectures/defs/z_image.py index c2a70c61b3b..39e879a1f42 100644 --- a/invokeai/backend/architectures/defs/z_image.py +++ b/invokeai/backend/architectures/defs/z_image.py @@ -1,4 +1,9 @@ +from invokeai.backend.architectures.facets.latent_space import FLUX_16, LatentSpaceFacet from invokeai.backend.architectures.registry import register from invokeai.backend.model_manager.taxonomy import BaseModelType -register(BaseModelType.ZImage) +register( + BaseModelType.ZImage, + # Z-Image uses a FLUX-compatible VAE with 16 latent channels. + LatentSpaceFacet(FLUX_16), +) diff --git a/invokeai/backend/architectures/facets/latent_space.py b/invokeai/backend/architectures/facets/latent_space.py new file mode 100644 index 00000000000..3f7e0087466 --- /dev/null +++ b/invokeai/backend/architectures/facets/latent_space.py @@ -0,0 +1,398 @@ +"""The VAE latent geometry of an architecture, and how to project it to a preview image. + +The preview chain was never really keyed on architecture -- the branch comments in the dispatch +chain this replaces said so by hand ("Krea-2 decodes with the Qwen-Image VAE", "Z-Image uses +FLUX-compatible VAE", "Anima uses Wan 2.1 VAE"). Fifteen architectures share eight latent spaces, +and the five most recent ones introduced no new preview data at all. + +Not migrated here, deliberately: + +- `constants.LATENT_SCALE_FACTOR` (a module-level `8` with a live `HACK:` comment): its blast radius + is every latent node. A follow-up, now that a correct source exists. +- `pid/decode.py`'s `_PER_BACKBONE`: same shape, different units. It describes packed latents + (Flux2: 128 channels at 16x) where this describes VAE latents (Flux2: 32 channels at 8x). They + are not derivable from one another. +""" + +from dataclasses import dataclass + +import torch +from PIL import Image + +from invokeai.backend.architectures.facet import Facet +from invokeai.backend.architectures.registry import require +from invokeai.backend.model_manager.taxonomy import BaseModelType + +# See scripts/generate_vae_linear_approximation.py for generating these factors. + +# fast latents preview matrix for sdxl +# generated by @StAlKeR7779 +SDXL_LATENT_RGB_FACTORS = [ + # R G B + [0.3816, 0.4930, 0.5320], + [-0.3753, 0.1631, 0.1739], + [0.1770, 0.3588, -0.2048], + [-0.4350, -0.2644, -0.4289], +] +SDXL_SMOOTH_MATRIX = [ + [0.0358, 0.0964, 0.0358], + [0.0964, 0.4711, 0.0964], + [0.0358, 0.0964, 0.0358], +] + +# origingally adapted from code by @erucipe and @keturn here: +# https://discuss.huggingface.co/t/decoding-latents-to-rgb-without-upscaling/23204/7 +# these updated numbers for v1.5 are from @torridgristle +SD1_5_LATENT_RGB_FACTORS = [ + # R G B + [0.3444, 0.1385, 0.0670], # L1 + [0.1247, 0.4027, 0.1494], # L2 + [-0.3192, 0.2513, 0.2103], # L3 + [-0.1307, -0.1874, -0.7445], # L4 +] + +SD3_5_LATENT_RGB_FACTORS = [ + [-0.05240681, 0.03251581, 0.0749016], + [-0.0580572, 0.00759826, 0.05729818], + [0.16144888, 0.01270368, -0.03768577], + [0.14418615, 0.08460266, 0.15941818], + [0.04894035, 0.0056485, -0.06686988], + [0.05187166, 0.19222395, 0.06261094], + [0.1539433, 0.04818359, 0.07103094], + [-0.08601796, 0.09013458, 0.10893912], + [-0.12398469, -0.06766567, 0.0033688], + [-0.0439737, 0.07825329, 0.02258823], + [0.03101129, 0.06382551, 0.07753657], + [-0.01315361, 0.08554491, -0.08772475], + [0.06464487, 0.05914605, 0.13262741], + [-0.07863674, -0.02261737, -0.12761454], + [-0.09923835, -0.08010759, -0.06264447], + [-0.03392309, -0.0804029, -0.06078822], +] + +FLUX_LATENT_RGB_FACTORS = [ + [-0.0412, 0.0149, 0.0521], + [0.0056, 0.0291, 0.0768], + [0.0342, -0.0681, -0.0427], + [-0.0258, 0.0092, 0.0463], + [0.0863, 0.0784, 0.0547], + [-0.0017, 0.0402, 0.0158], + [0.0501, 0.1058, 0.1152], + [-0.0209, -0.0218, -0.0329], + [-0.0314, 0.0083, 0.0896], + [0.0851, 0.0665, -0.0472], + [-0.0534, 0.0238, -0.0024], + [0.0452, -0.0026, 0.0048], + [0.0892, 0.0831, 0.0881], + [-0.1117, -0.0304, -0.0789], + [0.0027, -0.0479, -0.0043], + [-0.1146, -0.0827, -0.0598], +] + +COGVIEW4_LATENT_RGB_FACTORS = [ + [0.00408832, -0.00082485, -0.00214816], + [0.00084172, 0.00132241, 0.00842067], + [-0.00466737, -0.00983181, -0.00699561], + [0.03698397, -0.04797235, 0.03585809], + [0.00234701, -0.00124326, 0.00080869], + [-0.00723903, -0.00388422, -0.00656606], + [-0.00970917, -0.00467356, -0.00971113], + [0.17292486, -0.03452463, -0.1457515], + [0.02330308, 0.02942557, 0.02704329], + [-0.00903131, -0.01499841, -0.01432564], + [0.01250298, 0.0019407, -0.02168986], + [0.01371188, 0.00498283, -0.01302135], + [0.42396525, 0.4280575, 0.42148206], + [0.00983825, 0.00613302, 0.00610316], + [0.00473307, -0.00889551, -0.00915924], + [-0.00955853, -0.00980067, -0.00977842], +] + +# The 16-channel Wan 2.1 VAE, shared by Qwen-Image, Krea-2 (which decodes with the Qwen-Image VAE), +# Anima, and Wan 2.2 A14B. Before the registry these numbers were written out three times, as +# QWEN_IMAGE_/ANIMA_/WAN_LATENT_RGB_FACTORS; they were byte-identical, as were the three biases. +# Factors from ComfyUI's Wan21 latent_format: +# https://github.com/comfyanonymous/ComfyUI/blob/master/comfy/latent_formats.py +WAN21_LATENT_RGB_FACTORS = [ + [-0.1299, -0.1692, 0.2932], + [0.0671, 0.0406, 0.0442], + [0.3568, 0.2548, 0.1747], + [0.0372, 0.2344, 0.1420], + [0.0313, 0.0189, -0.0328], + [0.0296, -0.0956, -0.0665], + [-0.3477, -0.4059, -0.2925], + [0.0166, 0.1902, 0.1975], + [-0.0412, 0.0267, -0.1364], + [-0.1293, 0.0740, 0.1636], + [0.0680, 0.3019, 0.1128], + [0.0032, 0.0581, 0.0639], + [-0.1251, 0.0927, 0.1699], + [0.0060, -0.0633, 0.0005], + [0.3477, 0.2275, 0.2950], + [0.1984, 0.0913, 0.1861], +] + +WAN21_LATENT_RGB_BIAS = [-0.1835, -0.0868, -0.3360] + +# FLUX.2 uses 32 latent channels. +# Factors from ComfyUI: https://github.com/Comfy-Org/ComfyUI/blob/main/comfy/latent_formats.py +FLUX2_LATENT_RGB_FACTORS = [ + # R G B + [0.0058, 0.0113, 0.0073], + [0.0495, 0.0443, 0.0836], + [-0.0099, 0.0096, 0.0644], + [0.2144, 0.3009, 0.3652], + [0.0166, -0.0039, -0.0054], + [0.0157, 0.0103, -0.0160], + [-0.0398, 0.0902, -0.0235], + [-0.0052, 0.0095, 0.0109], + [-0.3527, -0.2712, -0.1666], + [-0.0301, -0.0356, -0.0180], + [-0.0107, 0.0078, 0.0013], + [0.0746, 0.0090, -0.0941], + [0.0156, 0.0169, 0.0070], + [-0.0034, -0.0040, -0.0114], + [0.0032, 0.0181, 0.0080], + [-0.0939, -0.0008, 0.0186], + [0.0018, 0.0043, 0.0104], + [0.0284, 0.0056, -0.0127], + [-0.0024, -0.0022, -0.0030], + [0.1207, -0.0026, 0.0065], + [0.0128, 0.0101, 0.0142], + [0.0137, -0.0072, -0.0007], + [0.0095, 0.0092, -0.0059], + [0.0000, -0.0077, -0.0049], + [-0.0465, -0.0204, -0.0312], + [0.0095, 0.0012, -0.0066], + [0.0290, -0.0034, 0.0025], + [0.0220, 0.0169, -0.0048], + [-0.0332, -0.0457, -0.0468], + [-0.0085, 0.0389, 0.0609], + [-0.0076, 0.0003, -0.0043], + [-0.0111, -0.0460, -0.0614], +] + +FLUX2_LATENT_RGB_BIAS = [-0.0329, -0.0718, -0.0851] + +# Wan 2.2 TI2V-5B uses Wan2.2-VAE with 48 latent channels and 16x spatial downscale. +# Factors come from ComfyUI's Wan22 latent_format. +WAN22_LATENT_RGB_FACTORS = [ + [0.0119, 0.0103, 0.0046], + [-0.1062, -0.0504, 0.0165], + [0.0140, 0.0409, 0.0491], + [-0.0813, -0.0677, 0.0607], + [0.0656, 0.0851, 0.0808], + [0.0264, 0.0463, 0.0912], + [0.0295, 0.0326, 0.0590], + [-0.0244, -0.0270, 0.0025], + [0.0443, -0.0102, 0.0288], + [-0.0465, -0.0090, -0.0205], + [0.0359, 0.0236, 0.0082], + [-0.0776, 0.0854, 0.1048], + [0.0564, 0.0264, 0.0561], + [0.0006, 0.0594, 0.0418], + [-0.0319, -0.0542, -0.0637], + [-0.0268, 0.0024, 0.0260], + [0.0539, 0.0265, 0.0358], + [-0.0359, -0.0312, -0.0287], + [-0.0285, -0.1032, -0.1237], + [0.1041, 0.0537, 0.0622], + [-0.0086, -0.0374, -0.0051], + [0.0390, 0.0670, 0.2863], + [0.0069, 0.0144, 0.0082], + [0.0006, -0.0167, 0.0079], + [0.0313, -0.0574, -0.0232], + [-0.1454, -0.0902, -0.0481], + [0.0714, 0.0827, 0.0447], + [-0.0304, -0.0574, -0.0196], + [0.0401, 0.0384, 0.0204], + [-0.0758, -0.0297, -0.0014], + [0.0568, 0.1307, 0.1372], + [-0.0055, -0.0310, -0.0380], + [0.0239, -0.0305, 0.0325], + [-0.0663, -0.0673, -0.0140], + [-0.0416, -0.0047, -0.0023], + [0.0166, 0.0112, -0.0093], + [-0.0211, 0.0011, 0.0331], + [0.1833, 0.1466, 0.2250], + [-0.0368, 0.0370, 0.0295], + [-0.3441, -0.3543, -0.2008], + [-0.0479, -0.0489, -0.0420], + [-0.0660, -0.0153, 0.0800], + [-0.0101, 0.0068, 0.0156], + [-0.0690, -0.0452, -0.0927], + [-0.0145, 0.0041, 0.0015], + [0.0421, 0.0451, 0.0373], + [0.0504, -0.0483, -0.0356], + [-0.0837, 0.0168, 0.0055], +] + +WAN22_LATENT_RGB_BIAS = [0.0317, -0.0878, -0.1388] + + +def sample_to_lowres_estimated_image( + samples: torch.Tensor, + latent_rgb_factors: torch.Tensor, + smooth_matrix: torch.Tensor | None = None, + latent_rgb_bias: torch.Tensor | None = None, +) -> Image.Image: + if samples.dim() == 4: + samples = samples[0] + latent_image = samples.permute(1, 2, 0) @ latent_rgb_factors + + if latent_rgb_bias is not None: + latent_image = latent_image + latent_rgb_bias + + if smooth_matrix is not None: + latent_image = latent_image.unsqueeze(0).permute(3, 0, 1, 2) + latent_image = torch.nn.functional.conv2d(latent_image, smooth_matrix.reshape((1, 1, 3, 3)), padding=1) + latent_image = latent_image.permute(1, 2, 3, 0).squeeze(0) + + latents_ubyte = ( + ((latent_image + 1) / 2).clamp(0, 1).mul(0xFF).byte() # change scale from -1..1 to 0..1 # to 0..255 + ).cpu() + + return Image.fromarray(latents_ubyte.numpy()) + + +@dataclass(frozen=True) +class LatentSpace: + """One VAE latent geometry, and the linear projection that turns it into a preview image.""" + + name: str + """Stable identifier. Used in diagnostics and to keep the closed set closed in tests.""" + + channels: int + """Latent channel count. Also how `LatentSpaceFacet` tells its latent spaces apart at runtime.""" + + spatial_compression: int + """Image pixels per latent pixel, per axis. 8 everywhere except Wan's Wan2.2-VAE.""" + + rgb_factors: list[list[float]] + """A `channels` x 3 linear projection from latent space to RGB.""" + + rgb_bias: list[float] | None = None + smooth_matrix: list[list[float]] | None = None + + def __post_init__(self) -> None: + if len(self.rgb_factors) != self.channels: + raise ValueError( + f"{self.name}: declares {self.channels} channels but carries {len(self.rgb_factors)} rgb_factors rows." + ) + + def preview(self, sample: torch.Tensor) -> Image.Image: + """Project `sample` to a preview image at *latent* resolution. + + The caller scales the result by `spatial_compression` to get the image-space size. Keeping + this a method rather than a bare factor matrix means an architecture that eventually needs a + real tiny-VAE decode is not forced through a linear projection. + """ + return sample_to_lowres_estimated_image( + samples=sample, + latent_rgb_factors=torch.tensor(self.rgb_factors, dtype=sample.dtype, device=sample.device), + smooth_matrix=( + torch.tensor(self.smooth_matrix, dtype=sample.dtype, device=sample.device) + if self.smooth_matrix + else None + ), + latent_rgb_bias=( + torch.tensor(self.rgb_bias, dtype=sample.dtype, device=sample.device) if self.rgb_bias else None + ), + ) + + +SD15_4 = LatentSpace(name="SD15_4", channels=4, spatial_compression=8, rgb_factors=SD1_5_LATENT_RGB_FACTORS) + +SDXL_4 = LatentSpace( + name="SDXL_4", + channels=4, + spatial_compression=8, + rgb_factors=SDXL_LATENT_RGB_FACTORS, + smooth_matrix=SDXL_SMOOTH_MATRIX, +) + +SD3_16 = LatentSpace(name="SD3_16", channels=16, spatial_compression=8, rgb_factors=SD3_5_LATENT_RGB_FACTORS) + +COGVIEW4_16 = LatentSpace( + name="COGVIEW4_16", channels=16, spatial_compression=8, rgb_factors=COGVIEW4_LATENT_RGB_FACTORS +) + +FLUX_16 = LatentSpace(name="FLUX_16", channels=16, spatial_compression=8, rgb_factors=FLUX_LATENT_RGB_FACTORS) + +WAN21_16 = LatentSpace( + name="WAN21_16", + channels=16, + spatial_compression=8, + rgb_factors=WAN21_LATENT_RGB_FACTORS, + rgb_bias=WAN21_LATENT_RGB_BIAS, +) + +# The FLUX.2 VAE is 16x on packed latents, but every denoise loop that reaches a preview unpacks +# first, so in the units this class uses it is 8x. Do not "correct" this to 16. +FLUX2_32 = LatentSpace( + name="FLUX2_32", + channels=32, + spatial_compression=8, + rgb_factors=FLUX2_LATENT_RGB_FACTORS, + rgb_bias=FLUX2_LATENT_RGB_BIAS, +) + +WAN22_48 = LatentSpace( + name="WAN22_48", + channels=48, + spatial_compression=16, + rgb_factors=WAN22_LATENT_RGB_FACTORS, + rgb_bias=WAN22_LATENT_RGB_BIAS, +) + + +@dataclass(frozen=True) +class LatentSpaceFacet(Facet): + """The latent space, or spaces, an architecture denoises in.""" + + REQUIRED = True + + default: LatentSpace + + alternates: tuple[LatentSpace, ...] = () + """Further latent spaces this architecture may denoise in, told apart at runtime by the sample's + channel count. Only Wan has any: A14B uses the 16-channel Wan 2.1 VAE, TI2V-5B the 48-channel + Wan2.2-VAE, and before the registry the code picked between them by hand -- in two places, once + for the factors and once for the spatial scale, which is exactly how those two could drift.""" + + def __post_init__(self) -> None: + channels = [latent_space.channels for latent_space in (self.default, *self.alternates)] + if len(set(channels)) != len(channels): + raise ValueError( + f"{self.default.name}: latent spaces are selected by channel count, so no two may share one " + f"(got {channels})." + ) + + def resolve(self, sample: torch.Tensor) -> LatentSpace: + """The latent space matching `sample`'s channel count. + + An architecture with a single latent space returns it without inspecting the sample at all. + That preserves the pre-registry behaviour: only the Wan branch ever read `sample.shape[-3]`, + so an unusually shaped sample must not start raising IndexError for the other fourteen. + """ + if not self.alternates: + return self.default + channels = sample.shape[-3] + for latent_space in self.alternates: + if latent_space.channels == channels: + return latent_space + return self.default + + +def get_latent_space(base: BaseModelType) -> LatentSpace: + """The latent space `base` denoises in; for the one architecture with several, its default. + + Use this where no sample is available yet -- to hoist the projection out of a per-step closure, + for instance. Use `resolve_latent_space()` where there is a sample. + """ + return require(base, LatentSpaceFacet).default + + +def resolve_latent_space(base: BaseModelType, sample: torch.Tensor) -> LatentSpace: + """The latent space `base` denoises in, disambiguated by `sample` where it has more than one.""" + return require(base, LatentSpaceFacet).resolve(sample) diff --git a/invokeai/backend/architectures/facets/unet.py b/invokeai/backend/architectures/facets/unet.py new file mode 100644 index 00000000000..7af269f18de --- /dev/null +++ b/invokeai/backend/architectures/facets/unet.py @@ -0,0 +1,38 @@ +"""How far a UNet-based architecture downscales the latent image internally. + +Kept apart from `LatentSpaceFacet` on purpose: this is a property of the *UNet* (SD1 downscales 8x +internally, SDXL 4x), unrelated to the VAE latent geometry next door. The two happen to both be +small integers about downscaling, which is exactly why they should not share a field. +""" + +from dataclasses import dataclass + +from invokeai.backend.architectures.facet import Facet +from invokeai.backend.architectures.registry import get +from invokeai.backend.model_manager.taxonomy import BaseModelType + + +@dataclass(frozen=True) +class UNetDownscaleFacet(Facet): + """Optional: only SD1 and SDXL have a UNet in the sense T2I-Adapter conditioning needs. + + Every other architecture legitimately does not declare it, so this facet is not `REQUIRED` and + the accessor -- not `require()` -- carries the error, preserving the message the two duplicated + dispatches it replaces raised. + """ + + max_unet_downscale: int + + +def get_max_unet_downscale(base: BaseModelType) -> int: + """The maximum amount the UNet downscales the latent image internally. + + Raises for architectures without a UNet, which is what the T2I-Adapter call sites did before + this facet existed. + """ + facet = get(base, UNetDownscaleFacet) + if facet is None: + # The message is reproduced verbatim, including how the enum renders: BaseModelType is a + # `str, Enum` mixin rather than a StrEnum, so this interpolates as "BaseModelType.Flux". + raise ValueError(f"Unexpected T2I-Adapter base model type: '{base}'.") + return facet.max_unet_downscale diff --git a/invokeai/backend/stable_diffusion/extensions/t2i_adapter.py b/invokeai/backend/stable_diffusion/extensions/t2i_adapter.py index 67fede93664..d2a8a07339c 100644 --- a/invokeai/backend/stable_diffusion/extensions/t2i_adapter.py +++ b/invokeai/backend/stable_diffusion/extensions/t2i_adapter.py @@ -8,7 +8,7 @@ from PIL.Image import Image from invokeai.app.util.controlnet_utils import prepare_control_image -from invokeai.backend.model_manager.taxonomy import BaseModelType +from invokeai.backend.architectures import get_max_unet_downscale from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ConditioningMode from invokeai.backend.stable_diffusion.extension_callback_type import ExtensionCallbackType from invokeai.backend.stable_diffusion.extensions.base import ExtensionBase, callback @@ -43,14 +43,8 @@ def __init__( self._adapter_state: Optional[List[torch.Tensor]] = None - # The max_unet_downscale is the maximum amount that the UNet model downscales the latent image internally. model_config = self._node_context.models.get_config(self._model_id.key) - if model_config.base == BaseModelType.StableDiffusion1: - self._max_unet_downscale = 8 - elif model_config.base == BaseModelType.StableDiffusionXL: - self._max_unet_downscale = 4 - else: - raise ValueError(f"Unexpected T2I-Adapter base model type: '{model_config.base}'.") + self._max_unet_downscale = get_max_unet_downscale(model_config.base) @callback(ExtensionCallbackType.SETUP) def setup(self, ctx: DenoiseContext): diff --git a/tests/app/util/test_step_callback.py b/tests/app/util/test_step_callback.py index 3235d4d4699..31391f6a9b0 100644 --- a/tests/app/util/test_step_callback.py +++ b/tests/app/util/test_step_callback.py @@ -1,119 +1,168 @@ -"""Tests for diffusion step callback preview image generation.""" +"""End-to-end tests for `diffusion_step_callback`. -import torch -from PIL import Image - -from invokeai.app.util.step_callback import ( - QWEN_IMAGE_LATENT_RGB_BIAS, - QWEN_IMAGE_LATENT_RGB_FACTORS, - sample_to_lowres_estimated_image, -) - - -class TestSampleToLowresEstimatedImage: - """Test the latent-to-preview-image conversion used during denoising.""" +The signal this pins is the reported image size: `image.size * spatial_compression`. Before the +registry that factor came from a second dispatch over `base_model`, separate from the one choosing +the projection factors, and the two had to agree about Wan's 16- vs 48-channel VAE by hand. Nothing +covered it. - def test_qwen_image_preview_produces_valid_image(self): - """A synthetic Qwen latent tensor produces a valid RGB preview image.""" - # Create a small 1x16x4x4 latent tensor (batch=1, channels=16, 4x4 spatial) - torch.manual_seed(42) - sample = torch.randn(1, 16, 4, 4) +The projection itself is tested next to the data it uses, in +tests/backend/architectures/test_latent_space.py. +""" - factors = torch.tensor(QWEN_IMAGE_LATENT_RGB_FACTORS, dtype=sample.dtype) - bias = torch.tensor(QWEN_IMAGE_LATENT_RGB_BIAS, dtype=sample.dtype) +import pytest +import torch +from PIL import Image - image = sample_to_lowres_estimated_image( - samples=sample, - latent_rgb_factors=factors, - latent_rgb_bias=bias, +from invokeai.app.services.session_processor.session_processor_common import CanceledException +from invokeai.app.util.step_callback import calc_percentage, diffusion_step_callback +from invokeai.backend.architectures import ArchitectureError, generative_bases +from invokeai.backend.model_manager.taxonomy import BaseModelType +from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState + +# base -> (latent channels, expected spatial compression). Read off the pre-registry dispatch +# chain; every base used 8 except Wan's 48-channel TI2V-5B. +CHANNELS_AND_SCALE = { + BaseModelType.StableDiffusion1: (4, 8), + BaseModelType.StableDiffusion2: (4, 8), + BaseModelType.StableDiffusionXL: (4, 8), + BaseModelType.StableDiffusionXLRefiner: (4, 8), + BaseModelType.StableDiffusion3: (16, 8), + BaseModelType.CogView4: (16, 8), + BaseModelType.Flux: (16, 8), + BaseModelType.ZImage: (16, 8), + BaseModelType.QwenImage: (16, 8), + BaseModelType.Krea2: (16, 8), + BaseModelType.Anima: (16, 8), + BaseModelType.Flux2: (32, 8), + BaseModelType.ErnieImage: (32, 8), + BaseModelType.Ideogram4: (32, 8), + BaseModelType.Wan: (16, 8), +} + +LATENT_HEIGHT = 4 +LATENT_WIDTH = 6 + + +class _Spy: + def __init__(self) -> None: + self.calls: list[tuple[str, float | None, Image.Image | None, tuple[int, int] | None]] = [] + + def __call__( + self, + message: str, + percentage: float | None = None, + image: Image.Image | None = None, + image_size: tuple[int, int] | None = None, + ) -> None: + self.calls.append((message, percentage, image, image_size)) + + +def _state(channels: int) -> PipelineIntermediateState: + return PipelineIntermediateState( + step=1, + order=1, + total_steps=4, + timestep=0, + latents=torch.zeros(1, channels, LATENT_HEIGHT, LATENT_WIDTH), + ) + + +def _run(base: BaseModelType, channels: int) -> _Spy: + spy = _Spy() + diffusion_step_callback( + signal_progress=spy, + intermediate_state=_state(channels), + base_model=base, + is_canceled=lambda: False, + ) + return spy + + +def test_the_table_covers_every_architecture() -> None: + assert set(CHANNELS_AND_SCALE) == set(generative_bases()) + + +@pytest.mark.parametrize("base", sorted(CHANNELS_AND_SCALE, key=lambda b: b.value)) +def test_reports_the_image_at_latent_resolution_scaled_by_the_compression(base: BaseModelType) -> None: + channels, scale = CHANNELS_AND_SCALE[base] + + spy = _run(base, channels) + + assert len(spy.calls) == 1 + message, percentage, image, image_size = spy.calls[0] + assert message == "Denoising" + assert percentage == 0.25 + assert image is not None + assert image.size == (LATENT_WIDTH, LATENT_HEIGHT) + assert image_size == (LATENT_WIDTH * scale, LATENT_HEIGHT * scale) + + +def test_wan_switches_to_16x_for_the_48_channel_vae() -> None: + """The one runtime-resolved case: TI2V-5B's Wan2.2-VAE is 48 channels at 16x, A14B 16 at 8x.""" + _, _, _, a14b_size = _run(BaseModelType.Wan, 16).calls[0] + _, _, _, ti2v_size = _run(BaseModelType.Wan, 48).calls[0] + + assert a14b_size == (LATENT_WIDTH * 8, LATENT_HEIGHT * 8) + assert ti2v_size == (LATENT_WIDTH * 16, LATENT_HEIGHT * 16) + + +def test_predicted_original_is_preferred_over_the_noisy_latents() -> None: + state = _state(16) + state.predicted_original = torch.ones(1, 16, 2, 2) + + spy = _Spy() + diffusion_step_callback( + signal_progress=spy, + intermediate_state=state, + base_model=BaseModelType.Flux, + is_canceled=lambda: False, + ) + + _, _, image, _ = spy.calls[0] + assert image is not None + assert image.size == (2, 2) + + +def test_cancellation_is_checked_before_anything_else() -> None: + spy = _Spy() + + with pytest.raises(CanceledException): + diffusion_step_callback( + signal_progress=spy, + intermediate_state=_state(16), + base_model=BaseModelType.Flux, + is_canceled=lambda: True, ) - assert isinstance(image, Image.Image) - assert image.size == (4, 4) - assert image.mode == "RGB" - - def test_qwen_image_preview_deterministic(self): - """The same input tensor always produces the same preview image.""" - sample = torch.ones(1, 16, 2, 2) - - factors = torch.tensor(QWEN_IMAGE_LATENT_RGB_FACTORS, dtype=sample.dtype) - bias = torch.tensor(QWEN_IMAGE_LATENT_RGB_BIAS, dtype=sample.dtype) - - image1 = sample_to_lowres_estimated_image(samples=sample, latent_rgb_factors=factors, latent_rgb_bias=bias) - image2 = sample_to_lowres_estimated_image(samples=sample, latent_rgb_factors=factors, latent_rgb_bias=bias) - - assert list(image1.getdata()) == list(image2.getdata()) + assert spy.calls == [] - def test_qwen_image_preview_known_value(self): - """Verify the preview computation against a hand-calculated expected value. - With a 1x16x1x1 tensor of all ones: - - latent_image = [1,1,...,1] @ factors = sum of each column of factors - - R = sum(col 0) = 0.3677, G = sum(col 1) = 0.4577, B = sum(col 2) = 0.9101 - - After bias: R = 0.1842, G = 0.3709, B = 0.5741 - - After scale ((x+1)/2): R = 0.5921, G = 0.6855, B = 0.7871 - - After quantize (*255): R = 151, G = 175, B = 201 - """ - sample = torch.ones(1, 16, 1, 1) +@pytest.mark.parametrize("base", [BaseModelType.Any, BaseModelType.External, BaseModelType.Unknown]) +def test_a_non_architecture_base_raises_with_a_fix_it_message(base: BaseModelType) -> None: + """This replaces `raise ValueError(f"Unsupported base model: {base}")`. - factors = torch.tensor(QWEN_IMAGE_LATENT_RGB_FACTORS, dtype=sample.dtype) - bias = torch.tensor(QWEN_IMAGE_LATENT_RGB_BIAS, dtype=sample.dtype) + ArchitectureError is a ValueError, so any caller catching ValueError is unaffected -- but the + message now names the file to create rather than only the base that failed. + """ + with pytest.raises(ArchitectureError, match="is not registered") as exc_info: + _run(base, 16) - image = sample_to_lowres_estimated_image(samples=sample, latent_rgb_factors=factors, latent_rgb_bias=bias) + assert "invokeai/backend/architectures/defs/" in str(exc_info.value) + assert isinstance(exc_info.value, ValueError) - assert image.size == (1, 1) - pixel = image.getpixel((0, 0)) - # Compute expected values - col_sums = [sum(row[c] for row in QWEN_IMAGE_LATENT_RGB_FACTORS) for c in range(3)] - expected = [] - for c in range(3): - val = col_sums[c] + QWEN_IMAGE_LATENT_RGB_BIAS[c] - val = (val + 1) / 2 # scale from [-1,1] to [0,1] - val = max(0.0, min(1.0, val)) # clamp - expected.append(int(val * 255)) - - assert pixel == tuple(expected), f"Expected {tuple(expected)}, got {pixel}" - - def test_qwen_image_preview_zeros_tensor(self): - """A zero tensor with bias produces a valid image reflecting just the bias.""" - sample = torch.zeros(1, 16, 2, 2) - - factors = torch.tensor(QWEN_IMAGE_LATENT_RGB_FACTORS, dtype=sample.dtype) - bias = torch.tensor(QWEN_IMAGE_LATENT_RGB_BIAS, dtype=sample.dtype) - - image = sample_to_lowres_estimated_image(samples=sample, latent_rgb_factors=factors, latent_rgb_bias=bias) - - assert isinstance(image, Image.Image) - assert image.size == (2, 2) - - # All pixels should be identical (uniform zero input) - pixels = [image.getpixel((x, y)) for y in range(image.height) for x in range(image.width)] - assert all(p == pixels[0] for p in pixels) - - # With zero input, result = bias, scaled: ((bias + 1) / 2) * 255 - expected = [] - for c in range(3): - val = (QWEN_IMAGE_LATENT_RGB_BIAS[c] + 1) / 2 - val = max(0.0, min(1.0, val)) - expected.append(int(val * 255)) - assert pixels[0] == tuple(expected) - - def test_qwen_image_factors_have_correct_shape(self): - """Qwen Image uses 16 latent channels, so factors should be 16x3.""" - assert len(QWEN_IMAGE_LATENT_RGB_FACTORS) == 16 - for row in QWEN_IMAGE_LATENT_RGB_FACTORS: - assert len(row) == 3 - assert len(QWEN_IMAGE_LATENT_RGB_BIAS) == 3 - - def test_3d_input_accepted(self): - """sample_to_lowres_estimated_image accepts 3D input (no batch dim).""" - sample = torch.randn(16, 4, 4) # no batch dimension - - factors = torch.tensor(QWEN_IMAGE_LATENT_RGB_FACTORS, dtype=sample.dtype) - bias = torch.tensor(QWEN_IMAGE_LATENT_RGB_BIAS, dtype=sample.dtype) - - image = sample_to_lowres_estimated_image(samples=sample, latent_rgb_factors=factors, latent_rgb_bias=bias) +@pytest.mark.parametrize( + ("step", "order", "total_steps", "expected"), + [ + (1, 1, 4, 0.25), + (0, 1, 0, 0.0), + (2, 2, 4, 0.5), + (1, 2, 1, 0.0), + ], +) +def test_calc_percentage(step: int, order: int, total_steps: int, expected: float) -> None: + state = PipelineIntermediateState( + step=step, order=order, total_steps=total_steps, timestep=0, latents=torch.zeros(1, 4, 1, 1) + ) - assert isinstance(image, Image.Image) - assert image.size == (4, 4) + assert calc_percentage(state) == expected diff --git a/tests/backend/architectures/test_latent_space.py b/tests/backend/architectures/test_latent_space.py new file mode 100644 index 00000000000..e947f0bd3df --- /dev/null +++ b/tests/backend/architectures/test_latent_space.py @@ -0,0 +1,303 @@ +"""The latent-space data, pinned against what the dispatch chain in step_callback.py used to select. + +The reference table below is the point of this file. Each row was read off the pre-registry +`elif` chain, and together they are what makes moving that data provably behaviour-preserving. + +Fingerprints rather than 600 lines of copied matrices: the per-column sums are plain IEEE-754 +addition in a fixed order, so they are exact and platform-independent, and the first and last rows +catch a reordering that a column sum is blind to. +""" + +from dataclasses import dataclass + +import pytest +import torch +from PIL import Image + +from invokeai.backend.architectures import generative_bases, get_latent_space, require, resolve_latent_space +from invokeai.backend.architectures.facets.latent_space import ( + LatentSpace, + LatentSpaceFacet, + sample_to_lowres_estimated_image, +) +from invokeai.backend.model_manager.taxonomy import BaseModelType + +# The closed set. A new name here means a genuinely new VAE geometry, not a new architecture. +LATENT_SPACE_NAMES = {"SD15_4", "SDXL_4", "SD3_16", "COGVIEW4_16", "FLUX_16", "WAN21_16", "FLUX2_32", "WAN22_48"} + + +@dataclass(frozen=True) +class Expected: + space: str + channels: int + spatial_compression: int + has_bias: bool + has_smooth: bool + first_row: list[float] + last_row: list[float] + + +# base -> what step_callback.py's chain selected for it before the registry existed. +REFERENCE = { + BaseModelType.StableDiffusion1: Expected( + "SD15_4", 4, 8, False, False, [0.3444, 0.1385, 0.0670], [-0.1307, -0.1874, -0.7445] + ), + BaseModelType.StableDiffusion2: Expected( + "SD15_4", 4, 8, False, False, [0.3444, 0.1385, 0.0670], [-0.1307, -0.1874, -0.7445] + ), + BaseModelType.StableDiffusionXL: Expected( + "SDXL_4", 4, 8, False, True, [0.3816, 0.4930, 0.5320], [-0.4350, -0.2644, -0.4289] + ), + BaseModelType.StableDiffusionXLRefiner: Expected( + "SDXL_4", 4, 8, False, True, [0.3816, 0.4930, 0.5320], [-0.4350, -0.2644, -0.4289] + ), + BaseModelType.StableDiffusion3: Expected( + "SD3_16", 16, 8, False, False, [-0.05240681, 0.03251581, 0.0749016], [-0.03392309, -0.0804029, -0.06078822] + ), + BaseModelType.CogView4: Expected( + "COGVIEW4_16", + 16, + 8, + False, + False, + [0.00408832, -0.00082485, -0.00214816], + [-0.00955853, -0.00980067, -0.00977842], + ), + BaseModelType.Flux: Expected( + "FLUX_16", 16, 8, False, False, [-0.0412, 0.0149, 0.0521], [-0.1146, -0.0827, -0.0598] + ), + BaseModelType.ZImage: Expected( + "FLUX_16", 16, 8, False, False, [-0.0412, 0.0149, 0.0521], [-0.1146, -0.0827, -0.0598] + ), + BaseModelType.QwenImage: Expected( + "WAN21_16", 16, 8, True, False, [-0.1299, -0.1692, 0.2932], [0.1984, 0.0913, 0.1861] + ), + BaseModelType.Krea2: Expected("WAN21_16", 16, 8, True, False, [-0.1299, -0.1692, 0.2932], [0.1984, 0.0913, 0.1861]), + BaseModelType.Anima: Expected("WAN21_16", 16, 8, True, False, [-0.1299, -0.1692, 0.2932], [0.1984, 0.0913, 0.1861]), + BaseModelType.Wan: Expected("WAN21_16", 16, 8, True, False, [-0.1299, -0.1692, 0.2932], [0.1984, 0.0913, 0.1861]), + BaseModelType.Flux2: Expected( + "FLUX2_32", 32, 8, True, False, [0.0058, 0.0113, 0.0073], [-0.0111, -0.0460, -0.0614] + ), + BaseModelType.ErnieImage: Expected( + "FLUX2_32", 32, 8, True, False, [0.0058, 0.0113, 0.0073], [-0.0111, -0.0460, -0.0614] + ), + # Ideogram 4 had no branch in the chain -- it drives its own preview loop -- but it inlined the + # FLUX.2 factors and a hardcoded x8, which is this row. + BaseModelType.Ideogram4: Expected( + "FLUX2_32", 32, 8, True, False, [0.0058, 0.0113, 0.0073], [-0.0111, -0.0460, -0.0614] + ), +} + +WAN21_BIAS = [-0.1835, -0.0868, -0.3360] +FLUX2_BIAS = [-0.0329, -0.0718, -0.0851] +WAN22_BIAS = [0.0317, -0.0878, -0.1388] + + +def _column_sums(factors: list[list[float]]) -> list[float]: + return [round(sum(row[c] for row in factors), 9) for c in range(3)] + + +def test_the_reference_table_covers_every_architecture() -> None: + assert set(REFERENCE) == set(generative_bases()) + + +@pytest.mark.parametrize("base", sorted(REFERENCE, key=lambda b: b.value)) +def test_latent_space_matches_the_pre_registry_dispatch(base: BaseModelType) -> None: + expected = REFERENCE[base] + latent_space = get_latent_space(base) + + assert latent_space.name == expected.space + assert latent_space.channels == expected.channels + assert latent_space.spatial_compression == expected.spatial_compression + assert (latent_space.rgb_bias is not None) is expected.has_bias + assert (latent_space.smooth_matrix is not None) is expected.has_smooth + assert len(latent_space.rgb_factors) == expected.channels + assert latent_space.rgb_factors[0] == expected.first_row + assert latent_space.rgb_factors[-1] == expected.last_row + + +def test_wan22_matches_the_pre_registry_dispatch() -> None: + # The only latent space no base uses by default, so it is not in the table above. + wan22 = resolve_latent_space(BaseModelType.Wan, torch.zeros(1, 48, 4, 4)) + + assert wan22.name == "WAN22_48" + assert wan22.channels == 48 + assert wan22.spatial_compression == 16 + assert wan22.rgb_factors[0] == [0.0119, 0.0103, 0.0046] + assert wan22.rgb_factors[-1] == [-0.0837, 0.0168, 0.0055] + assert wan22.rgb_bias == WAN22_BIAS + + +@pytest.mark.parametrize( + ("base", "expected_bias"), + [ + (BaseModelType.QwenImage, WAN21_BIAS), + (BaseModelType.Krea2, WAN21_BIAS), + (BaseModelType.Anima, WAN21_BIAS), + (BaseModelType.Wan, WAN21_BIAS), + (BaseModelType.Flux2, FLUX2_BIAS), + (BaseModelType.ErnieImage, FLUX2_BIAS), + (BaseModelType.Ideogram4, FLUX2_BIAS), + ], +) +def test_biases_match_the_pre_registry_dispatch(base: BaseModelType, expected_bias: list[float]) -> None: + assert get_latent_space(base).rgb_bias == expected_bias + + +@pytest.mark.parametrize( + ("base", "expected_sums"), + [ + (BaseModelType.StableDiffusion1, [0.0192, 0.6051, -0.3178]), + (BaseModelType.StableDiffusionXL, [-0.2517, 0.7505, 0.0722]), + (BaseModelType.QwenImage, [0.3887, 0.8771, 1.3152]), + ], +) +def test_column_sums_are_an_exact_fingerprint(base: BaseModelType, expected_sums: list[float]) -> None: + """Hand-checked totals for three spaces, as a guard against a single edited number. + + Note that the docstring of the test this replaces claimed 0.3677/0.4577/0.9101 for the Wan 2.1 + space. Those numbers were stale: that test recomputed its expectation from the same constants + it was checking, so it would have passed against a corrupted matrix and nothing ever caught the + drift. These totals are the real ones. + """ + assert [round(s, 4) for s in _column_sums(get_latent_space(base).rgb_factors)] == expected_sums + + +def test_qwen_krea_anima_and_wan_share_one_latent_space_object() -> None: + """The three byte-identical Wan 2.1 matrices were merged into one; this is the proof. + + Identity, not equality: they must be the same object, or the duplication has crept back. + """ + spaces = [ + get_latent_space(base) + for base in (BaseModelType.QwenImage, BaseModelType.Krea2, BaseModelType.Anima, BaseModelType.Wan) + ] + assert all(space is spaces[0] for space in spaces) + + +@pytest.mark.parametrize( + "bases", + [ + (BaseModelType.StableDiffusion1, BaseModelType.StableDiffusion2), + (BaseModelType.StableDiffusionXL, BaseModelType.StableDiffusionXLRefiner), + (BaseModelType.Flux, BaseModelType.ZImage), + (BaseModelType.Flux2, BaseModelType.ErnieImage, BaseModelType.Ideogram4), + ], +) +def test_architectures_sharing_a_latent_space_share_the_object(bases: tuple[BaseModelType, ...]) -> None: + spaces = [get_latent_space(base) for base in bases] + assert all(space is spaces[0] for space in spaces) + + +def test_the_latent_space_set_is_closed() -> None: + """Every declared latent space, defaults and alternates alike, is one of the eight.""" + names: set[str] = set() + for base in generative_bases(): + facet = require(base, LatentSpaceFacet) + names.add(facet.default.name) + names.update(alternate.name for alternate in facet.alternates) + + assert names == LATENT_SPACE_NAMES + + +# --- runtime resolution --------------------------------------------------------------------------- + + +def test_wan_resolves_by_channel_count() -> None: + assert resolve_latent_space(BaseModelType.Wan, torch.zeros(1, 48, 4, 4)).name == "WAN22_48" + assert resolve_latent_space(BaseModelType.Wan, torch.zeros(1, 16, 4, 4)).name == "WAN21_16" + + +def test_wan_falls_back_to_the_default_for_an_unexpected_channel_count() -> None: + # The pre-registry code was an `if shape[-3] == 48 / else`, so anything not 48 took the A14B + # branch. That must not become an error. + assert resolve_latent_space(BaseModelType.Wan, torch.zeros(1, 32, 4, 4)).name == "WAN21_16" + + +def test_a_single_space_architecture_never_inspects_the_sample() -> None: + # Only the Wan branch ever read sample.shape[-3]. A rank-2 tensor would raise IndexError if the + # resolver looked at it unconditionally. + assert resolve_latent_space(BaseModelType.StableDiffusion1, torch.zeros(2, 2)).name == "SD15_4" + + +def test_latent_spaces_within_a_facet_must_have_distinct_channel_counts() -> None: + space = get_latent_space(BaseModelType.Flux) + + with pytest.raises(ValueError, match="selected by channel count"): + LatentSpaceFacet(space, alternates=(space,)) + + +def test_channel_count_must_match_the_factor_rows() -> None: + with pytest.raises(ValueError, match="declares 4 channels but carries 1 rgb_factors rows"): + LatentSpace(name="BROKEN", channels=4, spatial_compression=8, rgb_factors=[[0.0, 0.0, 0.0]]) + + +# --- the projection itself, relocated from tests/app/util/test_step_callback.py ------------------- + + +def test_preview_produces_a_valid_rgb_image() -> None: + torch.manual_seed(42) + sample = torch.randn(1, 16, 4, 4) + + image = get_latent_space(BaseModelType.QwenImage).preview(sample) + + assert isinstance(image, Image.Image) + assert image.size == (4, 4) + assert image.mode == "RGB" + + +def test_preview_is_deterministic() -> None: + sample = torch.ones(1, 16, 2, 2) + latent_space = get_latent_space(BaseModelType.QwenImage) + + assert latent_space.preview(sample).tobytes() == latent_space.preview(sample).tobytes() + + +def test_preview_known_value() -> None: + """Hand-calculated pixel for a 1x16x1x1 tensor of ones, on the Wan 2.1 space. + + latent_image = [1,...,1] @ factors = the column sums, 0.3887 / 0.8771 / 1.3152. + After bias: 0.2052 / 0.7903 / 0.9792 + After ((x+1)/2): 0.6026 / 0.8952 / 0.9896 + After clamp, *255: 153 / 228 / 252 + + Hardcoded on purpose. The test this replaces derived its expectation by summing the very + matrix under test, which made it a tautology. + """ + image = get_latent_space(BaseModelType.QwenImage).preview(torch.ones(1, 16, 1, 1)) + + assert image.size == (1, 1) + assert image.getpixel((0, 0)) == (153, 228, 252) + + +def test_preview_of_zeros_reflects_only_the_bias() -> None: + latent_space = get_latent_space(BaseModelType.QwenImage) + assert latent_space.rgb_bias is not None + + image = latent_space.preview(torch.zeros(1, 16, 2, 2)) + + pixels = [image.getpixel((x, y)) for y in range(image.height) for x in range(image.width)] + assert all(pixel == pixels[0] for pixel in pixels) + assert pixels[0] == tuple(int(max(0.0, min(1.0, (b + 1) / 2)) * 255) for b in latent_space.rgb_bias) + + +def test_preview_accepts_input_without_a_batch_dimension() -> None: + image = get_latent_space(BaseModelType.QwenImage).preview(torch.randn(16, 4, 4)) + + assert image.size == (4, 4) + + +def test_preview_applies_the_smooth_matrix_only_where_declared() -> None: + # SDXL is the only latent space with one, and it must actually change the result. + sample = torch.zeros(1, 4, 3, 3) + sample[0, 0, 1, 1] = 4.0 + sdxl = get_latent_space(BaseModelType.StableDiffusionXL) + assert sdxl.smooth_matrix is not None + + smoothed = sdxl.preview(sample) + unsmoothed = sample_to_lowres_estimated_image( + samples=sample, + latent_rgb_factors=torch.tensor(sdxl.rgb_factors, dtype=sample.dtype), + ) + + assert smoothed.tobytes() != unsmoothed.tobytes() diff --git a/tests/backend/architectures/test_registry.py b/tests/backend/architectures/test_registry.py index eb57161daa3..f9854a7e900 100644 --- a/tests/backend/architectures/test_registry.py +++ b/tests/backend/architectures/test_registry.py @@ -4,7 +4,6 @@ architectures -- so they keep passing as architectures and facets are added. """ -from collections.abc import Iterator from dataclasses import dataclass import pytest @@ -26,19 +25,19 @@ class _BetaFacet(Facet): @pytest.fixture -def isolated_registry(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: - """Run a test against an empty registry, and restore the facet-type table afterwards. +def isolated_registry(monkeypatch: pytest.MonkeyPatch) -> None: + """Run a test against an empty registry *and* an empty facet-type table. - Both are process-global and mutated at class-creation/import time. Defining a facet subclass - anywhere -- even inside a test function -- adds it to `Facet.FACET_TYPES` permanently. Without - the restore, a dummy facet with `REQUIRED = True` would make `validate()` fail for every real - architecture in every test that runs afterwards. + Both are process-global and both matter. Emptying only the registry would leave the real facets + in `Facet.FACET_TYPES`, so `validate()` would hold the dummy architectures below to required + facets they were never meant to declare -- exactly the coupling these tests exist to avoid. + Emptying only the facet table would leave the 15 real architectures registered. + + Note that defining a facet subclass anywhere, including inside a test function, adds it to + whichever table is installed at class-creation time; monkeypatch puts the real one back. """ monkeypatch.setattr(registry, "_ARCHITECTURES", {}) - facet_types = dict(Facet.FACET_TYPES) - yield - Facet.FACET_TYPES.clear() - Facet.FACET_TYPES.update(facet_types) + monkeypatch.setattr(Facet, "FACET_TYPES", {}) def test_register_and_get_roundtrip(isolated_registry: None) -> None: diff --git a/tests/backend/architectures/test_unet_downscale.py b/tests/backend/architectures/test_unet_downscale.py new file mode 100644 index 00000000000..7642f93159c --- /dev/null +++ b/tests/backend/architectures/test_unet_downscale.py @@ -0,0 +1,34 @@ +"""`get_max_unet_downscale` replaces two verbatim-duplicated dispatches. + +They lived in `denoise_latents.run_t2i_adapters` and `T2IAdapterExt.__init__`, identical down to +the comment and the error string. The message is reproduced exactly, because it is user-facing. +""" + +import pytest + +from invokeai.backend.architectures import generative_bases, get_max_unet_downscale +from invokeai.backend.model_manager.taxonomy import BaseModelType + +HAS_UNET = {BaseModelType.StableDiffusion1: 8, BaseModelType.StableDiffusionXL: 4} + + +@pytest.mark.parametrize(("base", "expected"), sorted(HAS_UNET.items(), key=lambda item: item[0].value)) +def test_returns_the_declared_downscale(base: BaseModelType, expected: int) -> None: + assert get_max_unet_downscale(base) == expected + + +@pytest.mark.parametrize("base", sorted(set(generative_bases()) - set(HAS_UNET), key=lambda b: b.value)) +def test_raises_for_architectures_without_a_unet(base: BaseModelType) -> None: + # Verbatim, including the quoting and how the enum renders. BaseModelType is a `str, Enum` + # mixin rather than a StrEnum, so it interpolates as "BaseModelType.Flux", not "flux". + with pytest.raises(ValueError) as exc_info: + get_max_unet_downscale(base) + + assert str(exc_info.value) == f"Unexpected T2I-Adapter base model type: '{base}'." + + +def test_the_sd1_and_sdxl_values_are_the_pre_registry_ones() -> None: + # SD1's UNet downscales 8x internally, SDXL's 4x. Pinned separately from the parametrized test + # so that an edit to HAS_UNET cannot silently redefine what is being asserted. + assert get_max_unet_downscale(BaseModelType.StableDiffusion1) == 8 + assert get_max_unet_downscale(BaseModelType.StableDiffusionXL) == 4