diff --git a/invokeai/app/invocations/denoise_latents.py b/invokeai/app/invocations/denoise_latents.py index 68e8dfc0dcd..95d79f01c33 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 @@ -676,17 +677,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: @@ -985,6 +982,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