diff --git a/docs/src/content/docs/Users Guide/03.Models/03.model-format-support.mdx b/docs/src/content/docs/Users Guide/03.Models/03.model-format-support.mdx index 5b8574b1859..2cad94c150c 100644 --- a/docs/src/content/docs/Users Guide/03.Models/03.model-format-support.mdx +++ b/docs/src/content/docs/Users Guide/03.Models/03.model-format-support.mdx @@ -1,7 +1,7 @@ --- title: Model Format Support description: Which file formats and quantizations InvokeAI loads for every model family, text encoder, VAE and adapter. -lastUpdated: 2026-10-18 +lastUpdated: 2026-10-19 sidebar: order: 3 --- @@ -102,7 +102,7 @@ not check, so a quantized file there fails with an unclear error or decodes inco | SDXL | ✓ (LyCORIS, Diffusers, OMI) | ✓ | IP-Adapter, T2I-Adapter, textual inversion | | FLUX.1 | ✓ (LyCORIS, Diffusers, OMI) | ✓ | Control LoRA, IP-Adapter, Redux | | FLUX.2, Z-Image | ✓ (LyCORIS, Diffusers) | FLUX.2 [dev] (single-file checkpoint), Z-Image | | -| Qwen-Image, Krea-2, Wan 2.2, Anima | ✓ (LyCORIS/kohya, including LoKR where noted on the family page) | Qwen-Image (single-file checkpoint), Anima (ControlNet-LLLite) | | +| Qwen-Image, Krea-2, Wan 2.2, Anima | ✓ (LyCORIS/kohya, including LoKR where noted on the family page) | Qwen-Image (single-file checkpoint: Fun and InstantX ControlNet Unions; InstantX Inpainting is refused), Anima (ControlNet-LLLite) | | | LTX-2, MiniMax H3 | ✓ (plain low-rank; LoKR, LoHA and DoRA are refused) | ✗ | | | SD 3, CogView4, Ideogram 4, ERNIE-Image, Qwen-Image-2.1 | ✗ | ✗ | | diff --git a/docs/src/content/docs/Users Guide/03.Models/Local Models/qwen-image.mdx b/docs/src/content/docs/Users Guide/03.Models/Local Models/qwen-image.mdx index 093193b3b50..db8db23ad31 100644 --- a/docs/src/content/docs/Users Guide/03.Models/Local Models/qwen-image.mdx +++ b/docs/src/content/docs/Users Guide/03.Models/Local Models/qwen-image.mdx @@ -1,7 +1,7 @@ --- title: Qwen Image -description: Generate and edit images with the Qwen Image 2512 text-to-image model and the Qwen Image Edit 2511 reference-image editing model, including GGUF installs, the Lightning LoRAs and the Fun-ControlNet-Union. -lastUpdated: 2026-10-18 +description: Generate and edit images with the Qwen Image 2512 text-to-image model and the Qwen Image Edit 2511 reference-image editing model, including GGUF installs, the Lightning LoRAs and the Fun and InstantX ControlNet Unions. +lastUpdated: 2026-10-19 sidebar: order: 8 --- @@ -22,8 +22,8 @@ and only the text-to-image variant takes a [ControlNet](#controlnet). ## License -Qwen Image 2512, Qwen Image Edit 2511, the Lightning LoRAs, the Qwen2.5-VL text encoder and the Qwen Image -ControlNet Union are released under **Apache 2.0**, which allows commercial use. None of the downloads is gated. +Qwen Image 2512, Qwen Image Edit 2511, the Lightning LoRAs, the Qwen2.5-VL text encoder and both Qwen Image +ControlNet Unions are released under **Apache 2.0**, which allows commercial use. None of the downloads is gated. This is a summary, not legal advice. Read the full license on the model's HuggingFace page before using a model commercially. @@ -139,26 +139,42 @@ In the workflow editor use **Apply LoRA - Qwen Image** or **Apply LoRA Collectio ## ControlNet -The **Qwen Image ControlNet Union** starter model is alibaba-pai's *Fun-ControlNet-Union* for Qwen Image 2512 -(~3.5 GB, Apache 2.0). One model handles every condition: **Canny**, **HED**, **Depth**, **Pose**, **MLSD**, -**Scribble** and **Gray** control images, plus **inpainting**. It is trained for the text-to-image transformer: -with a Qwen Image Edit model selected, a control layer's settings and the **Invoke** button say that the -ControlNet is not compatible, and nothing is generated until the layer is disabled or the model changed. +Two ControlNets are starter models, both a single model for several conditions (~3.5 GB each, Apache 2.0): + +| Starter model | Design | Conditions | Recommended weight | +| ------------- | ------ | ---------- | ------------------ | +| **Qwen Image ControlNet Union** | alibaba-pai's *Fun-ControlNet-Union* for Qwen Image 2512 | **Canny**, **HED**, **Depth**, **Pose**, **MLSD**, **Scribble** and **Gray** control images, plus **inpainting** | `0.70` to `0.95` | +| **Qwen Image InstantX ControlNet Union** | InstantX's *ControlNet-Union*, trained on the original Qwen Image | **Canny**, **Soft Edge** (HED), **Depth** and **Pose** control images; control only | about `1.5` to `2.0` on Qwen Image 2512 (`0.80` to `1.00` on the original Qwen Image) | + +The two work differently. The Fun union is a control branch inside the transformer that reads the control image +and, when inpainting, the region being regenerated. The InstantX union is a five-block copy of the transformer that +runs beside it once per step and adds its output to the transformer's blocks; it takes the control image and +nothing else; InstantX recommends 30 steps at CFG 4 for it. InstantX trained it on the original Qwen Image, and on +Qwen Image 2512 it guides weakly at its card's `0.80` to `1.00`: a pose skeleton barely moves the figure at `1.0` +and is followed at `1.5` to `2.0`. The Fun union is trained for Qwen Image 2512 directly. InvokeAI tells the two +apart from the file on install; a copy of the InstantX union installed by an earlier version as a Qwen Image main model is fixed by +**Re-identify** in the Model Manager. InstantX's separate *Inpainting* ControlNet is not supported yet and is +refused on install. + +Both are trained for the text-to-image transformer: with a Qwen Image Edit model selected, a control layer's +settings and the **Invoke** button say that the ControlNet is not compatible, and nothing is generated until the +layer is disabled or the model changed. On the Canvas, add a **Control Layer** while a Qwen Image model is selected. Its adapter type is **Qwen Image -ControlNet** and its pixels are the control image; the layer's **Filter** turns a raster into one of the -conditions (Canny, HED, Depth Anything, DW Openpose, MLSD, Scribble or Grayscale). The layer's **Weight** is -the model's control scale: it ranges from `0` to `2`, starts at `0.85`, and `0.70` to `0.95` is the -recommended range. **Only one** Qwen Image control layer can be used per generation. +ControlNet**, its model either union, and its pixels are the control image; the layer's **Filter** turns a raster +into one of the conditions (Canny, HED, Depth Anything, DW Openpose, MLSD, Scribble or Grayscale). The layer's +**Weight** is the model's control scale: it ranges from `0` to `2` and starts at `0.85`. **Only one** Qwen Image +control layer can be used per generation. -When the Canvas inpaints or outpaints with a control layer present, the ControlNet also receives the image being -regenerated and the mask of the region as its inpainting condition. Without a control layer, inpainting and -outpainting work as before. +When the Canvas inpaints or outpaints with a Fun union layer present, the ControlNet also receives the image being +regenerated and the mask of the region as its inpainting condition. An InstantX union layer only guides the +regenerated region with its control image. Without a control layer, inpainting and outpainting work as before. In the workflow editor, the same thing is the **Qwen Image ControlNet** node feeding the **Control** input of **Denoise - Qwen Image**. Give it a control **Image**, an **Inpaint Image** together with its **Inpaint Mask** -(white is regenerated), or both. **Denoise - Qwen Image** then also needs its **VAE** input connected, which it -uses to encode these images. +(white is regenerated), or both; with the InstantX union, a control **Image** only, and the denoise refuses inpaint +inputs. **Denoise - Qwen Image** then also needs its **VAE** input connected, which it uses to encode these +images. ## PiD super-resolution decode @@ -173,7 +189,7 @@ Decoder Qwen-Image (2K to 4K)* decoders (including an int8 build of the latter) | **Main Model - Qwen Image** | Loads the transformer, VAE and Qwen VL encoder (or a component source) | | **Prompt - Qwen Image** | Encodes the prompt and any reference images | | **Denoise - Qwen Image** | Runs sampling; accepts reference latents, img2img latents, masks and a ControlNet | -| **Qwen Image ControlNet** | Configures the ControlNet Union: control image and/or inpaint image with mask, scale, step range | +| **Qwen Image ControlNet** | Configures either ControlNet Union: control image and/or inpaint image with mask (Fun only), scale, step range | | **Image to Latents - Qwen Image** | VAE encode | | **Latents to Image - Qwen Image** | VAE decode | | **Latents to Image - Qwen-Image + PiD (4x SR)** | PiD super-resolution decode | diff --git a/invokeai/app/invocations/qwen_image/qwen_image_control.py b/invokeai/app/invocations/qwen_image/qwen_image_control.py index 86404e17962..386a4b27fbb 100644 --- a/invokeai/app/invocations/qwen_image/qwen_image_control.py +++ b/invokeai/app/invocations/qwen_image/qwen_image_control.py @@ -1,4 +1,5 @@ -"""Qwen Image ControlNet: the alibaba-pai Fun-ControlNet-Union for spatial control and inpainting.""" +"""Qwen Image ControlNet: alibaba-pai's Fun-ControlNet-Union for spatial control and inpainting, or InstantX's +ControlNet-Union for spatial control.""" from typing import Optional @@ -18,7 +19,10 @@ from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelType from invokeai.backend.util.qwen_image_vae import QWEN_IMAGE_VAE_MIN_TILE_SIZE -_CONTROL_SCALE_DESCRIPTION = "Strength of the control signal. Recommended range: 0.70-0.95." +_CONTROL_SCALE_DESCRIPTION = ( + "Strength of the control signal. Recommended range: 0.70-0.95 for the Fun-ControlNet-Union; for the InstantX " + "ControlNet-Union 0.80-1.00 on the original Qwen Image, about 1.5-2.0 on Qwen Image 2512." +) _BEGIN_DESCRIPTION = "When the control is first applied (% of total steps)" _END_DESCRIPTION = "When the control is last applied (% of total steps)" _TILE_SIZE_DESCRIPTION = ( @@ -28,7 +32,8 @@ class QwenImageControlField(BaseModel): - """Qwen Image Fun-ControlNet conditioning: a control image, an inpaint image with its mask, or both.""" + """Qwen Image ControlNet conditioning: a control image, an inpaint image with its mask, or both (the InstantX + ControlNet takes the control image alone).""" image_name: str | None = Field(default=None, description="The preprocessed control image (Canny, depth, pose, ...)") inpaint_image_name: str | None = Field(default=None, description="The image to inpaint") @@ -61,10 +66,12 @@ class QwenImageControlOutput(BaseInvocationOutput): classification=Classification.Prototype, ) class QwenImageControlInvocation(BaseInvocation): - """Configure the Qwen Image Fun-ControlNet-Union for spatial control, inpainting, or both. + """Configure a Qwen Image ControlNet for spatial control, inpainting, or both. Takes a preprocessed control image (Canny, HED, depth, pose, MLSD, ...) and/or an image to inpaint with its mask - (white is regenerated), plus the Qwen Image ControlNet model. Recommended control scale: 0.70-0.95. + (white is regenerated), plus the Qwen Image ControlNet model. The Fun-ControlNet-Union takes either or both + (recommended scale 0.70-0.95); the InstantX ControlNet-Union takes the control image only (0.80-1.00 on the + original Qwen Image, about 1.5-2.0 on Qwen Image 2512). """ image: Optional[ImageField] = InputField( diff --git a/invokeai/app/invocations/qwen_image/qwen_image_denoise.py b/invokeai/app/invocations/qwen_image/qwen_image_denoise.py index 167c2661701..b4517a70c3e 100644 --- a/invokeai/app/invocations/qwen_image/qwen_image_denoise.py +++ b/invokeai/app/invocations/qwen_image/qwen_image_denoise.py @@ -1,10 +1,11 @@ import math from contextlib import ExitStack from pathlib import Path -from typing import TYPE_CHECKING, Callable, ClassVar, Iterator, Optional +from typing import TYPE_CHECKING, Callable, ClassVar, Iterator, Optional, TypeAlias import torch import torchvision.transforms as tv_transforms +from diffusers.models.controlnets.controlnet_qwenimage import QwenImageControlNetModel from diffusers.models.transformers.transformer_qwenimage import QwenImageTransformer2DModel from PIL import Image from torchvision.transforms.functional import resize as tv_resize @@ -68,6 +69,10 @@ "time_shift_type": "exponential", } +QwenImageControlModel: TypeAlias = QwenImageFunControlBranch | QwenImageControlNetModel +"""Either Qwen-Image ControlNet design: the Fun branch hooked into the transformer's blocks, or the InstantX +ControlNet run beside it.""" + @invocation( "qwen_image_denoise", @@ -313,7 +318,7 @@ def _run_diffusion(self, context: InvocationContext): getattr(transformer_info.model, "zero_cond_t", False) or getattr(transformer_info.model.config, "zero_cond_t", False) ) - control_branch_info = self._load_control_branch(context, has_zero_cond_t) + control_info, control_model = self._load_control_model(context, has_zero_cond_t) # Load conditioning pos_prompt_embeds, pos_prompt_mask = self._load_text_conditioning( @@ -485,12 +490,6 @@ def _run_diffusion(self, context: InvocationContext): model_input_seq_len = noisy_seq_len + (ref_latents_packed.shape[1] if ref_latents_packed is not None else 0) transformer_config = context.models.get_config(self.transformer.transformer) - num_control_blocks = 0 - control_hidden_size = 0 - if control_branch_info is not None: - assert isinstance(control_branch_info.model, QwenImageFunControlBranch) - num_control_blocks = len(control_branch_info.model.control_blocks) - control_hidden_size = control_branch_info.model.hidden_size # The forward's activations, which the cache keeps free by offloading as much of the transformer as that # takes, plus the per-forward dequantization transient of an nvfp4 or GGUF build, alive alongside them and # invisible to the estimate. All read from the unlocked model, before the lock; the transient is zero for a @@ -502,8 +501,7 @@ def _run_diffusion(self, context: InvocationContext): image_seq_len=model_input_seq_len, text_seq_len=text_seq_len, element_size=inference_dtype.itemsize, - num_control_blocks=num_control_blocks, - control_hidden_size=control_hidden_size, + control_model=control_model, ) # Where SDPA has no fused kernel for the joint [text ; image] sequence (ROCm on Windows, MPS) it materializes # the score matrices, which the linear term above was not priced for: 24 heads x 4608^2 scores is ~1 GB at @@ -519,8 +517,8 @@ def _run_diffusion(self, context: InvocationContext): # Encoded ahead of the transformer's lock, so the VAE does not have to fit beside it. control_context: torch.Tensor | None = None - if control_branch_info is not None: - control_context = self._build_control_context(context, latent_height, latent_width) + if control_model is not None: + control_context = self._build_control_context(context, control_model, latent_height, latent_width) control_context = control_context.to(device=device, dtype=inference_dtype) with ExitStack() as exit_stack: @@ -529,24 +527,30 @@ def _run_diffusion(self, context: InvocationContext): ) assert isinstance(transformer, QwenImageTransformer2DModel) control_hooks: FunControlHooks | None = None - if control_branch_info is not None: + controlnet: QwenImageControlNetModel | None = None + if control_info is not None: assert self.control is not None and control_context is not None - # Locked with the same working memory: the transformer can no longer be offloaded, so a branch placed - # with only the default reservation would fill the room kept for the forward. - (_, control_branch) = exit_stack.enter_context( - control_branch_info.model_on_device(working_mem_bytes=working_mem_bytes) + # Locked with the same working memory: the transformer can no longer be offloaded, so a control model + # placed with only the default reservation would fill the room kept for the forward. + (_, locked_control) = exit_stack.enter_context( + control_info.model_on_device(working_mem_bytes=working_mem_bytes) ) - assert isinstance(control_branch, QwenImageFunControlBranch) - # Dtype invariant: the loader cast the branch to the device's bfloat16-safe dtype, the policy the - # transformer loaders follow, which is `inference_dtype` wherever this node runs. Should the safe - # dtype ever be another one, the branch still runs its chain in its own dtype and casts every tensor - # on the way in and the hint on the way out (`QwenImageFunControlBranch.block_runner` / `add_hint`), - # so no cache-managed module is cast here. - control_hooks = exit_stack.enter_context( - install_fun_control( - transformer, control_branch, control_context, scale=self.control.control_context_scale + if isinstance(locked_control, QwenImageControlNetModel): + # Run beside the transformer and fed what it is fed: loaded by the same dtype policy, it takes the + # same activations and hands its residuals back in the dtype the transformer adds them in. + controlnet = locked_control + else: + assert isinstance(locked_control, QwenImageFunControlBranch) + # Dtype invariant: the loader cast the branch to the device's bfloat16-safe dtype, the policy the + # transformer loaders follow, which is `inference_dtype` wherever this node runs. Should the safe + # dtype ever be another one, the branch still runs its chain in its own dtype and casts every + # tensor on the way in and the hint on the way out (`QwenImageFunControlBranch.block_runner` / + # `add_hint`), so no cache-managed module is cast here. + control_hooks = exit_stack.enter_context( + install_fun_control( + transformer, locked_control, control_context, scale=self.control.control_context_scale + ) ) - ) # Quantized weights (packed nvfp4 Linears, GGUF, SDNQ) cannot take a direct patch, so their LoRAs ride as # sidecars. Asked of the loaded module tree: a `checkpoint` build may hold packed Linears. model_is_quantized = requires_sidecar_patching(transformer, transformer_config.format) @@ -566,12 +570,12 @@ def _run_diffusion(self, context: InvocationContext): exit_stack.enter_context(sage_attention_scope()) for step_idx, t in enumerate(tqdm(timesteps_sched)): + # One gate per step: both guidance branches run with the control, or neither does. + control_active = self.control is not None and should_apply( + step_idx, total_steps, self.control.begin_step_percent, self.control.end_step_percent + ) if control_hooks is not None: - assert self.control is not None - # One gate per step: both guidance branches run with the control, or neither does. - control_hooks.active = should_apply( - step_idx, total_steps, self.control.begin_step_percent, self.control.end_step_percent - ) + control_hooks.active = control_active # The pipeline passes timestep / 1000 to the transformer timestep = t.expand(latents.shape[0]).to(inference_dtype) @@ -583,12 +587,29 @@ def _run_diffusion(self, context: InvocationContext): else: model_input = latents + # The InstantX ControlNet runs once per step, on the positive prompt, and both guidance branches add + # its residuals, already scaled: diffusers' QwenImageControlNetPipeline. Skipped outside the step range. + controlnet_block_samples: list[torch.Tensor] | None = None + if controlnet is not None and control_active: + assert self.control is not None + controlnet_block_samples = controlnet( + hidden_states=model_input, + controlnet_cond=control_context, + conditioning_scale=self.control.control_context_scale, + encoder_hidden_states=pos_prompt_embeds, + encoder_hidden_states_mask=pos_prompt_mask, + timestep=timestep / 1000, + img_shapes=img_shapes, + return_dict=False, + ) + noise_pred_cond = transformer( hidden_states=model_input, encoder_hidden_states=pos_prompt_embeds, encoder_hidden_states_mask=pos_prompt_mask, timestep=timestep / 1000, img_shapes=img_shapes, + controlnet_block_samples=controlnet_block_samples, return_dict=False, )[0] # Only keep the noisy-latent portion of the output @@ -601,6 +622,7 @@ def _run_diffusion(self, context: InvocationContext): encoder_hidden_states_mask=neg_prompt_mask, timestep=timestep / 1000, img_shapes=img_shapes, + controlnet_block_samples=controlnet_block_samples, return_dict=False, )[0] noise_pred_uncond = noise_pred_uncond[:, :noisy_seq_len] @@ -637,29 +659,41 @@ def _run_diffusion(self, context: InvocationContext): latents = latents.unsqueeze(2) return latents - def _load_control_branch(self, context: InvocationContext, has_zero_cond_t: bool) -> LoadedModel | None: - """The control branch in RAM, ahead of the transformer's lock so the reservation can size its hints.""" + def _load_control_model( + self, context: InvocationContext, has_zero_cond_t: bool + ) -> tuple[LoadedModel, QwenImageControlModel] | tuple[None, None]: + """The ControlNet in RAM, and the model it holds, ahead of the transformer's lock so the reservation can size + what the ControlNet keeps alive.""" if self.control is None: - return None + return None, None if self.vae is None: raise ValueError( "Qwen Image ControlNet needs the VAE that encodes its images: connect one to the 'vae' input." ) if self.reference_latents is not None: raise ValueError("Qwen Image ControlNet cannot be combined with reference latents.") - # The Fun branch was trained beside the generation transformer. An Edit transformer takes a - # [noisy ; reference] sequence and modulates the two apart (zero_cond_t), a layout the branch never saw. + # Both ControlNets were trained beside the generation transformer. An Edit transformer takes a + # [noisy ; reference] sequence and modulates the two apart (zero_cond_t), a layout neither ever saw. if has_zero_cond_t: raise ValueError( "Qwen Image ControlNet applies to the Qwen Image generation transformer, not to an Edit transformer." ) - control_branch_info = context.models.load(self.control.control_model) - if not isinstance(control_branch_info.model, QwenImageFunControlBranch): + control_info = context.models.load(self.control.control_model) + control_model = control_info.model + if isinstance(control_model, QwenImageControlNetModel): + # Its control input is the control image's latent and nothing else: no mask or masked-image channels + # an inpainting condition could go into. + if self.control.inpaint_image_name is not None or self.control.inpaint_mask_name is not None: + raise ValueError( + f"'{self.control.control_model.name}' is an InstantX Qwen Image ControlNet, which takes a control " + "image only and has no inpainting condition: disconnect the inpaint image and mask." + ) + elif not isinstance(control_model, QwenImageFunControlBranch): raise TypeError( f"Expected a Qwen Image ControlNet for '{self.control.control_model.name}', " - f"got {type(control_branch_info.model).__name__}." + f"got {type(control_model).__name__}." ) - return control_branch_info + return control_info, control_model @staticmethod def _estimate_working_memory( @@ -667,8 +701,7 @@ def _estimate_working_memory( text_seq_len: int, element_size: int, *, - num_control_blocks: int = 0, - control_hidden_size: int = 0, + control_model: QwenImageControlModel | None = None, ) -> int: """Peak transformer activations, in bytes, for the model cache to keep free. @@ -677,15 +710,31 @@ def _estimate_working_memory( for the allocator. The per-token term is a bf16 figure and scales with `element_size`, the compute dtype's width. The guidance branches run one after the other, so CFG adds nothing to the peak. - A control branch runs its blocks over the image sequence and keeps one hint per block alive until the last - listed base layer has consumed it. `num_control_blocks + 4` hint-sized tensors is what Z-Image measured for - its adapter over the plain forward (`z_image_denoise.py`), pending the same measurement here. + A ControlNet adds image-sized tensors (one batch item, `image_seq_len` tokens at its width): + + - The Fun branch runs its blocks over the image sequence and keeps one hint per block alive until the last + listed base layer has consumed it. `blocks + 4` is what Z-Image measured for its adapter over the plain + forward (`z_image_denoise.py`), pending the same measurement here. + - The InstantX ControlNet adds no fitted coefficient, only what provably stays alive. It runs its forward + before the transformer's: its blocks are the transformer's at the same joint sequence and never run beside + them, so their activations are inside the base term, and so is the transient at the end of its forward + (block outputs, projections, scaled copies: three image-sized tensors per block, under a tenth of the + per-token base at Qwen-Image's width). What stays alive through both guidance forwards is the list of + scaled residuals, one per block: exactly `blocks` tensors. """ mib, gib = 1024**2, 1024**3 width = element_size / 2 estimated = int((image_seq_len + text_seq_len) * 0.25 * mib * width) + gib - if num_control_blocks: - estimated += (num_control_blocks + 4) * image_seq_len * control_hidden_size * element_size + if control_model is None: + return estimated + if isinstance(control_model, QwenImageFunControlBranch): + estimated += ( + (len(control_model.control_blocks) + 4) * image_seq_len * control_model.hidden_size * element_size + ) + else: + # The alias's other member. Not `elif` + `assert_never`: mypy cannot narrow a diffusers model class out of + # a union, so the remainder would never type as `Never`. + estimated += len(control_model.controlnet_blocks) * image_seq_len * control_model.inner_dim * element_size return estimated @staticmethod @@ -712,8 +761,12 @@ def _attention_score_bytes( via_diffusers_dispatch=True, ) - def _build_control_context(self, context: InvocationContext, latent_height: int, latent_width: int) -> torch.Tensor: - """The packed ``(1, S_img, 132)`` control context at the generation size, from the control field's images.""" + def _build_control_context( + self, context: InvocationContext, control_model: QwenImageControlModel, latent_height: int, latent_width: int + ) -> torch.Tensor: + """The packed control context at the generation size, from the control field's images, as ``control_model`` + takes it: ``(1, S_img, 132)`` control, keep-mask and masked-image latents for the Fun branch, ``(1, S_img, 64)`` + the control image's latents alone for the InstantX ControlNet (diffusers' QwenImageControlNetPipeline).""" assert self.control is not None and self.vae is not None control = self.control vae_info = context.models.load(self.vae.vae) @@ -733,14 +786,19 @@ def pixels(image_name: str, mode: str) -> torch.Tensor: # (1, C, H, W): pixels in [-1, 1], a mask in [0, 1]. return image_resized_to_grid_as_tensor(image, normalize=mode == "RGB", multiple_of=16).unsqueeze(0) - context_4d = build_fun_control_context( - encode=encode, - control_image=pixels(control.image_name, "RGB") if control.image_name is not None else None, - inpaint_image=( - pixels(control.inpaint_image_name, "RGB") if control.inpaint_image_name is not None else None - ), - inpaint_mask=pixels(control.inpaint_mask_name, "L") if control.inpaint_mask_name is not None else None, - ) + if isinstance(control_model, QwenImageFunControlBranch): + context_4d = build_fun_control_context( + encode=encode, + control_image=pixels(control.image_name, "RGB") if control.image_name is not None else None, + inpaint_image=( + pixels(control.inpaint_image_name, "RGB") if control.inpaint_image_name is not None else None + ), + inpaint_mask=pixels(control.inpaint_mask_name, "L") if control.inpaint_mask_name is not None else None, + ) + else: + # The InstantX ControlNet, the alias's other member (see `_estimate_working_memory`). + assert control.image_name is not None + context_4d = encode(pixels(control.image_name, "RGB")) return self._pack_latents(context_4d, 1, context_4d.shape[1], latent_height, latent_width) def _build_step_callback(self, context: InvocationContext) -> Callable[[PipelineIntermediateState], None]: diff --git a/invokeai/backend/flux/controlnet/state_dict_utils.py b/invokeai/backend/flux/controlnet/state_dict_utils.py index 87eae5a96bc..6f9f4ee6487 100644 --- a/invokeai/backend/flux/controlnet/state_dict_utils.py +++ b/invokeai/backend/flux/controlnet/state_dict_utils.py @@ -3,6 +3,7 @@ import torch from invokeai.backend.flux.model import FluxParams +from invokeai.backend.model_manager.qwen_image_controlnet_keys import is_qwen_image_instantx_controlnet def is_state_dict_xlabs_controlnet(sd: dict[str | int, Any]) -> bool: @@ -26,7 +27,7 @@ def is_state_dict_xlabs_controlnet(sd: dict[str | int, Any]) -> bool: def is_state_dict_instantx_controlnet(sd: dict[str | int, Any]) -> bool: - """Is the state dict for an InstantX ControlNet model? + """Is the state dict for an InstantX FLUX ControlNet model? This is intended to be a reasonably high-precision detector, but it is not guaranteed to have perfect precision. """ @@ -38,9 +39,11 @@ def is_state_dict_instantx_controlnet(sd: dict[str | int, Any]) -> bool: "controlnet_x_embedder.weight", } - if expected_keys.issubset(sd.keys()): - return True - return False + if not expected_keys.issubset(sd.keys()): + return False + # InstantX's Qwen-Image ControlNets use the same names over Qwen-Image's own blocks. ComfyUI tells the two apart + # the same way, before its FLUX fallback. + return not is_qwen_image_instantx_controlnet(sd) def _fuse_weights(*t: torch.Tensor) -> torch.Tensor: diff --git a/invokeai/backend/model_manager/configs/controlnet.py b/invokeai/backend/model_manager/configs/controlnet.py index 24faa05a88f..fa4db7630c9 100644 --- a/invokeai/backend/model_manager/configs/controlnet.py +++ b/invokeai/backend/model_manager/configs/controlnet.py @@ -12,6 +12,7 @@ ) from invokeai.backend.model_manager.configs.base import Checkpoint_Config_Base, Config_Base, Diffusers_Config_Base from invokeai.backend.model_manager.configs.identification_utils import ( + InvalidMatchError, NotAMatchError, common_config_paths, get_config_dict_or_raise, @@ -22,6 +23,7 @@ state_dict_has_any_keys_starting_with, ) from invokeai.backend.model_manager.model_on_disk import ModelOnDisk +from invokeai.backend.model_manager.qwen_image_controlnet_keys import is_qwen_image_instantx_controlnet from invokeai.backend.model_manager.taxonomy import ( BaseModelType, ModelFormat, @@ -300,21 +302,56 @@ def _has_qwen_image_fun_control_keys(state_dict: dict) -> bool: return "control_blocks.0.img_mod.1.weight" in state_dict +QwenImageControlArchitecture = Literal["fun_union", "instantx_union"] +"""The two Qwen-Image ControlNet designs, which share a config and a loader but not a recipe. + +`fun_union` is alibaba-pai's branch hooked into the transformer's blocks, conditioned on control and inpaint latents; +`instantx_union` is InstantX's five-block copy of the transformer run beside it, conditioned on control latents only. +""" + +# InstantX's `controlnet_x_embedder` projects the 64 packed features of a 16-channel latent into Qwen-Image's +# 3072-wide stream; its inpainting sibling adds the 4 packed features of a mask (68). +_QWEN_IMAGE_INSTANTX_WIDTH = 3072 +_QWEN_IMAGE_INSTANTX_CONTROL_FEATURES = 64 + + +def _qwen_image_instantx_control_features(state_dict: dict) -> int | None: + """The features per token an InstantX-layout Qwen-Image ControlNet takes, or None when it is not one. + + The key signature is shared with the probes that have to refuse the file (`qwen_image_controlnet_keys.py`); the + width is what makes it a ControlNet for this base. + """ + if not is_qwen_image_instantx_controlnet(state_dict): + return None + shape = getattr(state_dict["controlnet_x_embedder.weight"], "shape", None) + if shape is None or len(shape) != 2 or shape[0] != _QWEN_IMAGE_INSTANTX_WIDTH: + return None + return int(shape[1]) + + class ControlNet_Checkpoint_QwenImage_Config(Checkpoint_Config_Base, Config_Base): - """Model config for Qwen-Image Fun-ControlNet adapter models (Safetensors checkpoint). + """Model config for Qwen-Image ControlNet models (Safetensors checkpoint), in either of two designs. - Qwen-Image Fun-ControlNets are standalone adapters containing only the control branch - (control_img_in, control_blocks) that extends the base Qwen-Image transformer with spatial - conditioning capabilities. + alibaba-pai's Fun-ControlNet-Union is a control branch (control_img_in, control_blocks) hooked into the base + Qwen-Image transformer. Supports Canny, HED, Depth, Pose, MLSD, Scribble, Gray and inpainting; recommended + control_context_scale 0.70-0.95. - Supports: Canny, HED, Depth, Pose, MLSD, Scribble, Gray, inpainting. - Recommended control_context_scale: 0.70-0.95. + InstantX's ControlNet-Union is a five-block copy of the transformer (controlnet_x_embedder, transformer_blocks, + controlnet_blocks) whose block outputs are added to the base image stream. Supports Canny, Soft Edge, Depth and + Pose; recommended control_context_scale 0.80-1.00 on the original Qwen Image it was trained on, about 1.5-2.0 on + Qwen Image 2512. It takes no inpainting condition. """ type: Literal[ModelType.ControlNet] = Field(default=ModelType.ControlNet) format: Literal[ModelFormat.Checkpoint] = Field(default=ModelFormat.Checkpoint) base: Literal[BaseModelType.QwenImage] = Field(default=BaseModelType.QwenImage) default_settings: ControlAdapterDefaultSettings | None = Field(None) + control_architecture: QwenImageControlArchitecture = Field( + default="fun_union", + description="Which Qwen-Image ControlNet design the file holds: 'fun_union' (alibaba-pai Fun-ControlNet-Union, " + "control and inpainting) or 'instantx_union' (InstantX ControlNet-Union, control only). Records written " + "before the field existed are Fun-ControlNets.", + ) @classmethod def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self: @@ -322,10 +359,23 @@ def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) - raise_for_override_fields(cls, override_fields) - if not _has_qwen_image_fun_control_keys(mod.load_state_dict()): - raise NotAMatchError("state dict does not look like a Qwen-Image Fun-ControlNet") + architecture = cls._get_control_architecture_or_raise(mod.load_state_dict()) + return cls(**{**override_fields, "control_architecture": architecture}) - return cls(**override_fields) + @classmethod + def _get_control_architecture_or_raise(cls, state_dict: dict) -> QwenImageControlArchitecture: + if _has_qwen_image_fun_control_keys(state_dict): + return "fun_union" + features = _qwen_image_instantx_control_features(state_dict) + if features is None: + raise NotAMatchError("state dict does not look like a Qwen-Image Fun or InstantX ControlNet") + if features != _QWEN_IMAGE_INSTANTX_CONTROL_FEATURES: + # Recognised and refused, so the installer shows why instead of registering a model nothing can run. + raise InvalidMatchError( + "This is the InstantX Qwen-Image inpainting ControlNet, which Invoke does not support yet. Install the " + "InstantX Qwen-Image ControlNet Union instead." + ) + return "instantx_union" # `control_img_in` takes 260 packed channels into FLUX.2 [dev]'s 6144-wide stream. Klein is 3072 or 4096 wide, so diff --git a/invokeai/backend/model_manager/configs/main.py b/invokeai/backend/model_manager/configs/main.py index bd297396714..0960e896444 100644 --- a/invokeai/backend/model_manager/configs/main.py +++ b/invokeai/backend/model_manager/configs/main.py @@ -46,6 +46,7 @@ ) from invokeai.backend.model_manager.configs.qwen3_encoder import _SDNQ_LOADABLE_QWEN_ARCHITECTURES from invokeai.backend.model_manager.model_on_disk import ModelOnDisk +from invokeai.backend.model_manager.qwen_image_controlnet_keys import is_qwen_image_instantx_controlnet from invokeai.backend.model_manager.taxonomy import ( AnimaVariantType, BaseModelType, @@ -2217,6 +2218,10 @@ def _has_qwen_image_keys(state_dict: dict[str | int, Any]) -> bool: and has_txt_norm and has_img_in and not has_context_embedder + # InstantX's Qwen-Image ControlNet is a trimmed copy of this transformer, `img_in`/`txt_in`/`txt_norm` + # included; the ControlNet config claims it. Asked of the stripped keys, as the loader asks: a prefixed copy + # identified as a main model would be refused at load, and identified as one again by Re-identify. + and not is_qwen_image_instantx_controlnet(set(keys)) and not _has_qwen_image21_keys(state_dict) ) diff --git a/invokeai/backend/model_manager/load/model_loaders/fun_controlnet.py b/invokeai/backend/model_manager/load/model_loaders/fun_controlnet.py index fb0a97926e6..09d956dbe98 100644 --- a/invokeai/backend/model_manager/load/model_loaders/fun_controlnet.py +++ b/invokeai/backend/model_manager/load/model_loaders/fun_controlnet.py @@ -1,4 +1,5 @@ -"""Loaders for Fun-ControlNet branches (alibaba-pai VideoX-Fun): Qwen-Image and FLUX.2 [dev].""" +"""The FLUX.2 [dev] Fun-ControlNet loader (alibaba-pai VideoX-Fun), and the dense-checkpoint read every ControlNet +branch loader shares. The Qwen-Image registration, which also loads InstantX's design, is `qwen_image_controlnet.py`.""" from pathlib import Path from typing import Optional @@ -7,7 +8,6 @@ from safetensors.torch import load_file from invokeai.backend.fun_controlnet.flux2 import Flux2FunControlBranch -from invokeai.backend.fun_controlnet.qwen_image import QwenImageFunControlBranch from invokeai.backend.model_manager.configs.base import Checkpoint_Config_Base from invokeai.backend.model_manager.configs.factory import AnyModelConfig from invokeai.backend.model_manager.load.fp8_capability import Unimplemented @@ -15,14 +15,14 @@ from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry from invokeai.backend.model_manager.taxonomy import AnyModel, BaseModelType, ModelFormat, ModelType, SubModelType from invokeai.backend.quantization.fp8_scaled import reject_quantized_side_channel -from invokeai.backend.util.devices import TorchDevice -def _read_dense_checkpoint(config: AnyModelConfig, label: str) -> dict[str, torch.Tensor]: - """The branch's tensors, refused if any carries a quantization side channel. +def read_dense_checkpoint(config: AnyModelConfig, label: str) -> dict[str, torch.Tensor]: + """The ControlNet's tensors, refused if any carries a quantization side channel. - Refused before the geometry probe: the shapes it reads mean nothing on a packed weight, and the branches assign the - file's tensors as they are, so no scale would be folded in. alibaba-pai publishes dense bf16 files only. + Refused before the geometry probe: the shapes it reads mean nothing on a packed weight, and the loaders assign the + file's tensors as they are, so no scale would be folded in. alibaba-pai and Comfy-Org publish dense bf16 files + only. """ if not isinstance(config, Checkpoint_Config_Base): raise ValueError("Only checkpoint configs are supported here.") @@ -32,29 +32,6 @@ def _read_dense_checkpoint(config: AnyModelConfig, label: str) -> dict[str, torc return sd -@ModelLoaderRegistry.register( - base=BaseModelType.QwenImage, - type=ModelType.ControlNet, - format=ModelFormat.Checkpoint, - fp8_storage=Unimplemented( - "the branch is cast to the compute dtype at load and runs as its own module tree, so the layerwise cast would " - "fit; at 3.5 GB in bf16 it is worth closing" - ), -) -class QwenImageFunControlNetLoader(ModelLoader): - """Loads the Qwen-Image Fun-ControlNet-Union branch from its safetensors file. - - The file carries no config: the branch's geometry is read off its tensors and the base-layer map is the family's - constant (`QWEN_IMAGE_FUN_CONTROL_LAYERS`). - """ - - def _load_model(self, config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> AnyModel: - return QwenImageFunControlBranch.from_state_dict( - _read_dense_checkpoint(config, "Qwen Image"), - dtype=TorchDevice.choose_bfloat16_safe_dtype(self._torch_device), - ) - - @ModelLoaderRegistry.register( base=BaseModelType.Flux2, type=ModelType.ControlNet, @@ -73,4 +50,4 @@ class Flux2FunControlNetLoader(ModelLoader): """ def _load_model(self, config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> AnyModel: - return Flux2FunControlBranch.from_state_dict(_read_dense_checkpoint(config, "FLUX.2"), dtype=torch.bfloat16) + return Flux2FunControlBranch.from_state_dict(read_dense_checkpoint(config, "FLUX.2"), dtype=torch.bfloat16) diff --git a/invokeai/backend/model_manager/load/model_loaders/qwen_image.py b/invokeai/backend/model_manager/load/model_loaders/qwen_image.py index 1c7e6bb7ff7..275082d91f4 100644 --- a/invokeai/backend/model_manager/load/model_loaders/qwen_image.py +++ b/invokeai/backend/model_manager/load/model_loaders/qwen_image.py @@ -22,6 +22,7 @@ _strip_quantization_metadata, ) from invokeai.backend.model_manager.load.model_loaders.generic_diffusers import GenericDiffusersLoader +from invokeai.backend.model_manager.qwen_image_controlnet_keys import is_qwen_image_instantx_controlnet from invokeai.backend.model_manager.taxonomy import ( AnyModel, BaseModelType, @@ -283,6 +284,14 @@ def _load_from_singlefile(self, config: AnyModelConfig) -> AnyModel: sd = load_file(str(model_path)) sd = CheckpointPrefix.detect(sd).strip(sd) + # Earlier versions identified InstantX's Qwen-Image ControlNet, a trimmed copy of this transformer, as a main + # model. Such a record still exists; the file only needs identifying again. + if is_qwen_image_instantx_controlnet(sd): + raise ValueError( + f"'{config.name}' is a Qwen Image ControlNet, not a main model: an earlier version of Invoke " + "installed it as one. Run Re-identify model on it in the Model Manager to fix the record, then pick " + "it in a control layer." + ) # Comfy's nvfp4 build keeps the image stream's attention and MLP in nvfp4, beside scaled fp8. Take those layers # out before the fold: `_dequantize_comfyui_fp8` multiplies every `.weight_scale` into its weight, nvfp4's diff --git a/invokeai/backend/model_manager/load/model_loaders/qwen_image_controlnet.py b/invokeai/backend/model_manager/load/model_loaders/qwen_image_controlnet.py new file mode 100644 index 00000000000..c5d073781a9 --- /dev/null +++ b/invokeai/backend/model_manager/load/model_loaders/qwen_image_controlnet.py @@ -0,0 +1,109 @@ +"""The loader for Qwen-Image ControlNets, in either design the config records: alibaba-pai's Fun-ControlNet-Union +branch or InstantX's ControlNet-Union.""" + +from collections.abc import Mapping +from typing import Optional, assert_never + +import accelerate +import torch +from diffusers.models.controlnets.controlnet_qwenimage import QwenImageControlNetModel + +from invokeai.backend.fun_controlnet.qwen_image import QwenImageFunControlBranch +from invokeai.backend.model_manager.configs.controlnet import ControlNet_Checkpoint_QwenImage_Config +from invokeai.backend.model_manager.configs.factory import AnyModelConfig +from invokeai.backend.model_manager.load.fp8_capability import Unimplemented +from invokeai.backend.model_manager.load.load_default import ModelLoader +from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry +from invokeai.backend.model_manager.load.model_loaders.fun_controlnet import read_dense_checkpoint +from invokeai.backend.model_manager.taxonomy import AnyModel, BaseModelType, ModelFormat, ModelType, SubModelType +from invokeai.backend.util.devices import TorchDevice + +# Qwen-Image's RoPE splits a 128-wide head into frame, height and width axes. The ControlNet runs its own RoPE over the +# same positions as the transformer and the file does not record the split, so it is a constant of the family. +QWEN_IMAGE_AXES_DIMS_ROPE = (16, 56, 56) + + +def load_qwen_image_instantx_controlnet( + sd: Mapping[str, torch.Tensor], *, dtype: torch.dtype +) -> QwenImageControlNetModel: + """Build the `QwenImageControlNetModel` the checkpoint describes and load it strictly, cast to ``dtype``. + + The geometry is read off the tensors: the width and latent features from `img_in`, the text width from `txt_in`, + the head width from the query norm, the block count from the highest `controlnet_blocks.N` index. A control input + wider than the latent is the inpainting layout, whose extra mask features nothing here produces, so it is refused. + """ + required = ( + "img_in.weight", + "txt_in.weight", + "controlnet_x_embedder.weight", + "transformer_blocks.0.attn.norm_q.weight", + ) + absent = [key for key in required if key not in sd] + if absent: + raise ValueError(f"Not an InstantX Qwen Image ControlNet checkpoint: {', '.join(absent)} missing.") + inner_dim, in_channels = sd["img_in.weight"].shape + control_features = sd["controlnet_x_embedder.weight"].shape[1] + if control_features != in_channels: + raise ValueError( + f"This InstantX Qwen Image ControlNet takes {control_features} control features per token, {in_channels} " + "for the latent plus a mask: the inpainting layout, which is not supported yet." + ) + head_dim = sd["transformer_blocks.0.attn.norm_q.weight"].shape[0] + if head_dim != sum(QWEN_IMAGE_AXES_DIMS_ROPE) or inner_dim % head_dim: + raise ValueError( + f"A {inner_dim}-wide ControlNet with {head_dim}-wide attention heads is not Qwen-Image's geometry " + f"({sum(QWEN_IMAGE_AXES_DIMS_ROPE)}-wide heads)." + ) + block_indices = { + int(index) for key in sd if key.startswith("controlnet_blocks.") and (index := key.split(".")[1]).isdigit() + } + num_layers = max(block_indices, default=-1) + 1 + if num_layers == 0: + raise ValueError("Not an InstantX Qwen Image ControlNet checkpoint: it has no `controlnet_blocks`.") + + with accelerate.init_empty_weights(): + model = QwenImageControlNetModel( + in_channels=in_channels, + num_layers=num_layers, + attention_head_dim=head_dim, + num_attention_heads=inner_dim // head_dim, + joint_attention_dim=sd["txt_in.weight"].shape[1], + axes_dims_rope=QWEN_IMAGE_AXES_DIMS_ROPE, + ) + missing, unexpected = model.load_state_dict(dict(sd), strict=False, assign=True) + if missing or unexpected: + raise ValueError( + "This checkpoint does not match the InstantX Qwen Image ControlNet layout: " + f"missing {sorted(missing)[:5]}, unexpected {sorted(unexpected)[:5]}." + ) + return model.to(dtype=dtype).eval() + + +@ModelLoaderRegistry.register( + base=BaseModelType.QwenImage, + type=ModelType.ControlNet, + format=ModelFormat.Checkpoint, + fp8_storage=Unimplemented( + "both designs are cast to the compute dtype at load and run as their own module trees, so the layerwise cast " + "would fit; at 3.5 GB in bf16 each it is worth closing" + ), +) +class QwenImageControlNetLoader(ModelLoader): + """Loads a Qwen-Image ControlNet from its safetensors file, in the design its record names. + + Neither file carries a config: the geometry is read off the tensors, and the Fun branch's base-layer map is the + family's constant (`QWEN_IMAGE_FUN_CONTROL_LAYERS`). + """ + + def _load_model(self, config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> AnyModel: + if not isinstance(config, ControlNet_Checkpoint_QwenImage_Config): + raise ValueError("Only Qwen Image ControlNet checkpoint configs are supported here.") + sd = read_dense_checkpoint(config, "Qwen Image") + dtype = TorchDevice.choose_bfloat16_safe_dtype(self._torch_device) + match config.control_architecture: + case "fun_union": + return QwenImageFunControlBranch.from_state_dict(sd, dtype=dtype) + case "instantx_union": + return load_qwen_image_instantx_controlnet(sd, dtype=dtype) + case _: + assert_never(config.control_architecture) diff --git a/invokeai/backend/model_manager/qwen_image_controlnet_keys.py b/invokeai/backend/model_manager/qwen_image_controlnet_keys.py new file mode 100644 index 00000000000..7f11ea0c167 --- /dev/null +++ b/invokeai/backend/model_manager/qwen_image_controlnet_keys.py @@ -0,0 +1,27 @@ +"""The key signature of InstantX's Qwen-Image ControlNets, as one predicate. + +The layout is diffusers' `QwenImageControlNetModel`: a trimmed copy of the Qwen-Image transformer (`img_in`, `txt_in`, +`txt_norm`, `transformer_blocks.N`) plus `controlnet_x_embedder` and one `controlnet_blocks.N` projection per block. +That puts it inside two other probes' nets: the FLUX InstantX ControlNet probe reads the same `controlnet_*` names, and +the Qwen-Image main model probe the same `img_in`/`txt_in`/`txt_norm`. Both have to refuse exactly what the Qwen-Image +ControlNet config claims, and the main checkpoint loader has to recognise a file an earlier version recorded as a main +model, so all four read the signature here rather than restating it. + +Imports nothing, like `checkpoint_prefix.py`: identification and the FLUX state-dict utilities both reach it. +""" + +from collections.abc import Container + +QWEN_IMAGE_INSTANTX_CONTROLNET_KEYS = ( + "controlnet_x_embedder.weight", + "controlnet_blocks.0.weight", + # Qwen-Image's dual-stream block and text input norm: a FLUX ControlNet has `norm1`/`ff` blocks and a + # `context_embedder` instead. + "transformer_blocks.0.img_mod.1.weight", + "txt_norm.weight", +) + + +def is_qwen_image_instantx_controlnet(state_dict: Container[object]) -> bool: + """Whether the state dict has InstantX's Qwen-Image ControlNet layout, its inpainting variant included.""" + return all(key in state_dict for key in QWEN_IMAGE_INSTANTX_CONTROLNET_KEYS) diff --git a/invokeai/backend/model_manager/starter_models/__init__.py b/invokeai/backend/model_manager/starter_models/__init__.py index 75c721d2f02..fd905b1cafe 100644 --- a/invokeai/backend/model_manager/starter_models/__init__.py +++ b/invokeai/backend/model_manager/starter_models/__init__.py @@ -214,6 +214,7 @@ qwen_image_gguf_q4_k_m, qwen_image_gguf_q6_k, qwen_image_gguf_q8_0, + qwen_image_instantx_controlnet_union, qwen_image_lightning_4step, qwen_image_lightning_8step, qwen_image_vae, @@ -465,6 +466,7 @@ qwen_image_lightning_4step, qwen_image_lightning_8step, qwen_image_fun_controlnet_union, + qwen_image_instantx_controlnet_union, flux_krea, flux_krea_quantized, z_image_turbo, diff --git a/invokeai/backend/model_manager/starter_models/qwen_image.py b/invokeai/backend/model_manager/starter_models/qwen_image.py index 028be0b828c..41c089ca817 100644 --- a/invokeai/backend/model_manager/starter_models/qwen_image.py +++ b/invokeai/backend/model_manager/starter_models/qwen_image.py @@ -206,3 +206,12 @@ "control images, plus inpainting. Works best at weights 0.70-0.95. Apache-2.0. ~3.5GB", type=ModelType.ControlNet, ) +qwen_image_instantx_controlnet_union = StarterModel( + name="Qwen Image InstantX ControlNet Union", + base=BaseModelType.QwenImage, + source="https://huggingface.co/Comfy-Org/Qwen-Image-InstantX-ControlNets/resolve/main/split_files/controlnet/Qwen-Image-InstantX-ControlNet-Union.safetensors", + description="InstantX's ControlNet for Qwen Image supporting Canny, Soft Edge, Depth and Pose control images; " + "no inpainting. Trained on the original Qwen Image: on Qwen Image 2512 use weights of about 1.5-2.0. Apache-2.0. " + "~3.5GB", + type=ModelType.ControlNet, +) diff --git a/invokeai/frontend/api/openapi.json b/invokeai/frontend/api/openapi.json index cb42543e7ee..9fcd05f072e 100644 --- a/invokeai/frontend/api/openapi.json +++ b/invokeai/frontend/api/openapi.json @@ -28391,6 +28391,13 @@ "type": "null" } ] + }, + "control_architecture": { + "type": "string", + "enum": ["fun_union", "instantx_union"], + "title": "Control Architecture", + "description": "Which Qwen-Image ControlNet design the file holds: 'fun_union' (alibaba-pai Fun-ControlNet-Union, control and inpainting) or 'instantx_union' (InstantX ControlNet-Union, control only). Records written before the field existed are Fun-ControlNets.", + "default": "fun_union" } }, "type": "object", @@ -28410,10 +28417,11 @@ "type", "format", "base", - "default_settings" + "default_settings", + "control_architecture" ], "title": "ControlNet_Checkpoint_QwenImage_Config", - "description": "Model config for Qwen-Image Fun-ControlNet adapter models (Safetensors checkpoint).\n\nQwen-Image Fun-ControlNets are standalone adapters containing only the control branch\n(control_img_in, control_blocks) that extends the base Qwen-Image transformer with spatial\nconditioning capabilities.\n\nSupports: Canny, HED, Depth, Pose, MLSD, Scribble, Gray, inpainting.\nRecommended control_context_scale: 0.70-0.95." + "description": "Model config for Qwen-Image ControlNet models (Safetensors checkpoint), in either of two designs.\n\nalibaba-pai's Fun-ControlNet-Union is a control branch (control_img_in, control_blocks) hooked into the base\nQwen-Image transformer. Supports Canny, HED, Depth, Pose, MLSD, Scribble, Gray and inpainting; recommended\ncontrol_context_scale 0.70-0.95.\n\nInstantX's ControlNet-Union is a five-block copy of the transformer (controlnet_x_embedder, transformer_blocks,\ncontrolnet_blocks) whose block outputs are added to the base image stream. Supports Canny, Soft Edge, Depth and\nPose; recommended control_context_scale 0.80-1.00 on the original Qwen Image it was trained on, about 1.5-2.0 on\nQwen Image 2512. It takes no inpainting condition." }, "ControlNet_Checkpoint_SD1_Config": { "properties": { @@ -96937,7 +96945,7 @@ "type": "object" }, "QwenImageControlField": { - "description": "Qwen Image Fun-ControlNet conditioning: a control image, an inpaint image with its mask, or both.", + "description": "Qwen Image ControlNet conditioning: a control image, an inpaint image with its mask, or both (the InstantX\nControlNet takes the control image alone).", "properties": { "image_name": { "anyOf": [ @@ -96984,7 +96992,7 @@ }, "control_context_scale": { "default": 0.85, - "description": "Strength of the control signal. Recommended range: 0.70-0.95.", + "description": "Strength of the control signal. Recommended range: 0.70-0.95 for the Fun-ControlNet-Union; for the InstantX ControlNet-Union 0.80-1.00 on the original Qwen Image, about 1.5-2.0 on Qwen Image 2512.", "maximum": 2.0, "minimum": 0.0, "title": "Control Context Scale", @@ -97029,7 +97037,7 @@ "category": "conditioning", "class": "invocation", "classification": "prototype", - "description": "Configure the Qwen Image Fun-ControlNet-Union for spatial control, inpainting, or both.\n\nTakes a preprocessed control image (Canny, HED, depth, pose, MLSD, ...) and/or an image to inpaint with its mask\n(white is regenerated), plus the Qwen Image ControlNet model. Recommended control scale: 0.70-0.95.", + "description": "Configure a Qwen Image ControlNet for spatial control, inpainting, or both.\n\nTakes a preprocessed control image (Canny, HED, depth, pose, MLSD, ...) and/or an image to inpaint with its mask\n(white is regenerated), plus the Qwen Image ControlNet model. The Fun-ControlNet-Union takes either or both\n(recommended scale 0.70-0.95); the InstantX ControlNet-Union takes the control image only (0.80-1.00 on the\noriginal Qwen Image, about 1.5-2.0 on Qwen Image 2512).", "node_pack": "invokeai", "properties": { "id": { @@ -97124,7 +97132,7 @@ }, "control_context_scale": { "default": 0.85, - "description": "Strength of the control signal. Recommended range: 0.70-0.95.", + "description": "Strength of the control signal. Recommended range: 0.70-0.95 for the Fun-ControlNet-Union; for the InstantX ControlNet-Union 0.80-1.00 on the original Qwen Image, about 1.5-2.0 on Qwen Image 2512.", "field_kind": "input", "input": "any", "maximum": 2.0, diff --git a/invokeai/frontend/api/schema.ts b/invokeai/frontend/api/schema.ts index ce3e700a895..4867daf363d 100644 --- a/invokeai/frontend/api/schema.ts +++ b/invokeai/frontend/api/schema.ts @@ -9921,14 +9921,16 @@ export type components = { }; /** * ControlNet_Checkpoint_QwenImage_Config - * @description Model config for Qwen-Image Fun-ControlNet adapter models (Safetensors checkpoint). + * @description Model config for Qwen-Image ControlNet models (Safetensors checkpoint), in either of two designs. * - * Qwen-Image Fun-ControlNets are standalone adapters containing only the control branch - * (control_img_in, control_blocks) that extends the base Qwen-Image transformer with spatial - * conditioning capabilities. + * alibaba-pai's Fun-ControlNet-Union is a control branch (control_img_in, control_blocks) hooked into the base + * Qwen-Image transformer. Supports Canny, HED, Depth, Pose, MLSD, Scribble, Gray and inpainting; recommended + * control_context_scale 0.70-0.95. * - * Supports: Canny, HED, Depth, Pose, MLSD, Scribble, Gray, inpainting. - * Recommended control_context_scale: 0.70-0.95. + * InstantX's ControlNet-Union is a five-block copy of the transformer (controlnet_x_embedder, transformer_blocks, + * controlnet_blocks) whose block outputs are added to the base image stream. Supports Canny, Soft Edge, Depth and + * Pose; recommended control_context_scale 0.80-1.00 on the original Qwen Image it was trained on, about 1.5-2.0 on + * Qwen Image 2512. It takes no inpainting condition. */ ControlNet_Checkpoint_QwenImage_Config: { /** @@ -10007,6 +10009,13 @@ export type components = { */ base: "qwen-image"; default_settings: components["schemas"]["ControlAdapterDefaultSettings"] | null; + /** + * Control Architecture + * @description Which Qwen-Image ControlNet design the file holds: 'fun_union' (alibaba-pai Fun-ControlNet-Union, control and inpainting) or 'instantx_union' (InstantX ControlNet-Union, control only). Records written before the field existed are Fun-ControlNets. + * @default fun_union + * @enum {string} + */ + control_architecture: "fun_union" | "instantx_union"; }; /** ControlNet_Checkpoint_SD1_Config */ ControlNet_Checkpoint_SD1_Config: { @@ -41928,7 +41937,8 @@ export type components = { }; /** * QwenImageControlField - * @description Qwen Image Fun-ControlNet conditioning: a control image, an inpaint image with its mask, or both. + * @description Qwen Image ControlNet conditioning: a control image, an inpaint image with its mask, or both (the InstantX + * ControlNet takes the control image alone). */ QwenImageControlField: { /** @@ -41953,7 +41963,7 @@ export type components = { control_model: components["schemas"]["ModelIdentifierField"]; /** * Control Context Scale - * @description Strength of the control signal. Recommended range: 0.70-0.95. + * @description Strength of the control signal. Recommended range: 0.70-0.95 for the Fun-ControlNet-Union; for the InstantX ControlNet-Union 0.80-1.00 on the original Qwen Image, about 1.5-2.0 on Qwen Image 2512. * @default 0.85 */ control_context_scale?: number; @@ -41984,10 +41994,12 @@ export type components = { }; /** * Qwen Image ControlNet - * @description Configure the Qwen Image Fun-ControlNet-Union for spatial control, inpainting, or both. + * @description Configure a Qwen Image ControlNet for spatial control, inpainting, or both. * * Takes a preprocessed control image (Canny, HED, depth, pose, MLSD, ...) and/or an image to inpaint with its mask - * (white is regenerated), plus the Qwen Image ControlNet model. Recommended control scale: 0.70-0.95. + * (white is regenerated), plus the Qwen Image ControlNet model. The Fun-ControlNet-Union takes either or both + * (recommended scale 0.70-0.95); the InstantX ControlNet-Union takes the control image only (0.80-1.00 on the + * original Qwen Image, about 1.5-2.0 on Qwen Image 2512). */ QwenImageControlInvocation: { /** @@ -42030,7 +42042,7 @@ export type components = { control_model?: components["schemas"]["ModelIdentifierField"] | null; /** * Control Scale - * @description Strength of the control signal. Recommended range: 0.70-0.95. + * @description Strength of the control signal. Recommended range: 0.70-0.95 for the Fun-ControlNet-Union; for the InstantX ControlNet-Union 0.80-1.00 on the original Qwen Image, about 1.5-2.0 on Qwen Image 2512. * @default 0.85 */ control_context_scale?: number; diff --git a/invokeai/frontend/webv2/src/features/generation/core/__snapshots__/qwenImageControlCanvasGraphs.json b/invokeai/frontend/webv2/src/features/generation/core/__snapshots__/qwenImageControlCanvasGraphs.json index 6aef96e01ff..25de26524ec 100644 --- a/invokeai/frontend/webv2/src/features/generation/core/__snapshots__/qwenImageControlCanvasGraphs.json +++ b/invokeai/frontend/webv2/src/features/generation/core/__snapshots__/qwenImageControlCanvasGraphs.json @@ -410,6 +410,392 @@ } } }, + "inpaint_instantx_control_layer": { + "edges": [ + { + "destination": { + "field": "vae", + "node_id": "canvas_l2i" + }, + "source": { + "field": "vae", + "node_id": "model_loader" + } + }, + { + "destination": { + "field": "latents", + "node_id": "canvas_l2i" + }, + "source": { + "field": "latents", + "node_id": "denoise_latents" + } + }, + { + "destination": { + "field": "transformer", + "node_id": "denoise_latents" + }, + "source": { + "field": "transformer", + "node_id": "model_loader" + } + }, + { + "destination": { + "field": "vae", + "node_id": "denoise_latents" + }, + "source": { + "field": "vae", + "node_id": "model_loader" + } + }, + { + "destination": { + "field": "qwen_vl_encoder", + "node_id": "pos_cond" + }, + "source": { + "field": "qwen_vl_encoder", + "node_id": "model_loader" + } + }, + { + "destination": { + "field": "prompt", + "node_id": "pos_cond" + }, + "source": { + "field": "value", + "node_id": "positive_prompt" + } + }, + { + "destination": { + "field": "positive_conditioning", + "node_id": "denoise_latents" + }, + "source": { + "field": "conditioning", + "node_id": "pos_cond" + } + }, + { + "destination": { + "field": "qwen_vl_encoder", + "node_id": "neg_cond" + }, + "source": { + "field": "qwen_vl_encoder", + "node_id": "model_loader" + } + }, + { + "destination": { + "field": "prompt", + "node_id": "neg_cond" + }, + "source": { + "field": "value", + "node_id": "negative_prompt" + } + }, + { + "destination": { + "field": "negative_conditioning", + "node_id": "denoise_latents" + }, + "source": { + "field": "conditioning", + "node_id": "neg_cond" + } + }, + { + "destination": { + "field": "seed", + "node_id": "denoise_latents" + }, + "source": { + "field": "value", + "node_id": "seed" + } + }, + { + "destination": { + "field": "seed", + "node_id": "core_metadata_0_i" + }, + "source": { + "field": "value", + "node_id": "seed" + } + }, + { + "destination": { + "field": "positive_prompt", + "node_id": "core_metadata_0_i" + }, + "source": { + "field": "value", + "node_id": "positive_prompt" + } + }, + { + "destination": { + "field": "negative_prompt", + "node_id": "core_metadata_0_i" + }, + "source": { + "field": "value", + "node_id": "negative_prompt" + } + }, + { + "destination": { + "field": "metadata", + "node_id": "canvas_output" + }, + "source": { + "field": "metadata", + "node_id": "core_metadata_0_i" + } + }, + { + "destination": { + "field": "vae", + "node_id": "canvas_i2l" + }, + "source": { + "field": "vae", + "node_id": "model_loader" + } + }, + { + "destination": { + "field": "latents", + "node_id": "denoise_latents" + }, + "source": { + "field": "latents", + "node_id": "canvas_i2l" + } + }, + { + "destination": { + "field": "vae", + "node_id": "create_gradient_mask" + }, + "source": { + "field": "vae", + "node_id": "model_loader" + } + }, + { + "destination": { + "field": "denoise_mask", + "node_id": "denoise_latents" + }, + "source": { + "field": "denoise_mask", + "node_id": "create_gradient_mask" + } + }, + { + "destination": { + "field": "mask", + "node_id": "expand_mask" + }, + "source": { + "field": "expanded_mask_area", + "node_id": "create_gradient_mask" + } + }, + { + "destination": { + "field": "layer_upper", + "node_id": "canvas_output" + }, + "source": { + "field": "image", + "node_id": "canvas_l2i" + } + }, + { + "destination": { + "field": "mask", + "node_id": "canvas_output" + }, + "source": { + "field": "image", + "node_id": "expand_mask" + } + }, + { + "destination": { + "field": "control", + "node_id": "denoise_latents" + }, + "source": { + "field": "control", + "node_id": "qwen_image_control_q-control" + } + } + ], + "id": "qwen_image_control_canvas_inpaint_instantx", + "nodes": { + "negative_prompt": { + "is_intermediate": true, + "use_cache": true, + "id": "negative_prompt", + "type": "string" + }, + "positive_prompt": { + "is_intermediate": true, + "use_cache": true, + "id": "positive_prompt", + "type": "string" + }, + "seed": { + "is_intermediate": true, + "use_cache": true, + "id": "seed", + "type": "integer" + }, + "model_loader": { + "is_intermediate": true, + "use_cache": true, + "id": "model_loader", + "model": { + "base": "qwen-image", + "format": "diffusers", + "key": "qwen-model", + "name": "Qwen Image", + "type": "main", + "hash": "blake3:qwen" + }, + "type": "qwen_image_model_loader" + }, + "pos_cond": { + "is_intermediate": true, + "use_cache": true, + "id": "pos_cond", + "type": "qwen_image_text_encoder" + }, + "neg_cond": { + "is_intermediate": true, + "use_cache": true, + "id": "neg_cond", + "type": "qwen_image_text_encoder" + }, + "denoise_latents": { + "is_intermediate": true, + "use_cache": true, + "cfg_scale": 4, + "height": 1024, + "id": "denoise_latents", + "steps": 40, + "type": "qwen_image_denoise", + "width": 768, + "denoising_end": 1, + "denoising_start": 0.4 + }, + "core_metadata_0_i": { + "is_intermediate": true, + "use_cache": true, + "cfg_scale": 4, + "cfg_rescale_multiplier": 0, + "generation_mode": "qwen_image_inpaint", + "height": 1024, + "id": "core_metadata_0_i", + "model": { + "base": "qwen-image", + "format": "diffusers", + "key": "qwen-model", + "name": "Qwen Image", + "type": "main", + "hash": "blake3:qwen" + }, + "rand_device": "cpu", + "seamless_x": false, + "seamless_y": false, + "steps": 40, + "type": "core_metadata", + "width": 768, + "qwen_image_quantization": "none", + "qwen_image_shift": "auto", + "clip_skip": 0, + "strength": 0.6 + }, + "canvas_i2l": { + "is_intermediate": true, + "use_cache": true, + "id": "canvas_i2l", + "type": "qwen_image_i2l", + "image": { + "image_name": "canvas-composite.png" + } + }, + "create_gradient_mask": { + "is_intermediate": true, + "use_cache": true, + "coherence_mode": "Gaussian Blur", + "edge_radius": 16, + "fp32": false, + "id": "create_gradient_mask", + "image": { + "image_name": "canvas-composite.png" + }, + "minimum_denoise": 0, + "type": "create_gradient_mask", + "mask": { + "image_name": "canvas-mask.png" + } + }, + "expand_mask": { + "is_intermediate": true, + "use_cache": true, + "fade_size_px": 16, + "id": "expand_mask", + "type": "expand_mask_with_fade" + }, + "canvas_l2i": { + "is_intermediate": true, + "use_cache": false, + "id": "canvas_l2i", + "type": "qwen_image_l2i" + }, + "canvas_output": { + "is_intermediate": true, + "use_cache": false, + "layer_base": { + "image_name": "canvas-composite.png" + }, + "type": "invokeai_img_blend", + "id": "canvas_output" + }, + "qwen_image_control_q-control": { + "is_intermediate": true, + "use_cache": true, + "begin_step_percent": 0, + "control_context_scale": 0.85, + "control_model": { + "base": "qwen-image", + "key": "qwen-instantx", + "name": "Qwen Image InstantX ControlNet Union", + "type": "controlnet", + "hash": "blake3:instantx" + }, + "end_step_percent": 1, + "id": "qwen_image_control_q-control", + "image": { + "image_name": "q-control.png" + }, + "type": "qwen_image_control" + } + } + }, "outpaint_control_layer": { "edges": [ { diff --git a/invokeai/frontend/webv2/src/features/generation/core/canvas/addControlLayers.test.ts b/invokeai/frontend/webv2/src/features/generation/core/canvas/addControlLayers.test.ts index 7312b1a8147..a46f2f180c4 100644 --- a/invokeai/frontend/webv2/src/features/generation/core/canvas/addControlLayers.test.ts +++ b/invokeai/frontend/webv2/src/features/generation/core/canvas/addControlLayers.test.ts @@ -449,6 +449,38 @@ describe('addControlLayers — Qwen-Image control', () => { ]); }); + it('hands an InstantX union the control image alone, as the canvas inpaints', () => { + const graph = run({ + base: 'qwen-image', + inpaint, + layers: [qwenLayer('Q1', { modelControlArchitecture: 'instantx_union' })], + }); + + expect(graph.nodes.qwen_image_control_Q1).toEqual({ + begin_step_percent: 0, + control_context_scale: 0.85, + control_model: union, + end_step_percent: 1, + id: 'qwen_image_control_Q1', + image: { image_name: 'Q1.png' }, + is_intermediate: true, + type: 'qwen_image_control', + use_cache: true, + }); + expect(Object.keys(graph.nodes)).toEqual(['denoise_latents', 'qwen_image_control_Q1']); + expect(graph.edges.map((edge) => edge.destination)).toEqual([{ field: 'control', node_id: 'denoise_latents' }]); + }); + + it.each(['fun_union', null, undefined] as const)( + 'hands a Fun union (record says %s) the inpainted region as well', + (modelControlArchitecture) => { + const graph = run({ base: 'qwen-image', inpaint, layers: [qwenLayer('Q1', { modelControlArchitecture })] }); + + expect(graph.nodes.qwen_image_control_Q1).toMatchObject({ inpaint_image: { image_name: 'composite.png' } }); + expect(edgesTo(graph, 'qwen_image_control_Q1', 'inpaint_mask')).toHaveLength(1); + } + ); + it('leaves the inpaint inputs to the kinds that condition on the region', () => { const graph = run({ base: 'z-image', diff --git a/invokeai/frontend/webv2/src/features/generation/core/canvas/addControlLayers.ts b/invokeai/frontend/webv2/src/features/generation/core/canvas/addControlLayers.ts index fe080f9c8ac..2b0ed6e8c17 100644 --- a/invokeai/frontend/webv2/src/features/generation/core/canvas/addControlLayers.ts +++ b/invokeai/frontend/webv2/src/features/generation/core/canvas/addControlLayers.ts @@ -2,6 +2,7 @@ import type { SupportedGenerateBase } from '@features/generation/core/baseGenerationPolicies'; import type { BackendGraphContract, BackendInvocationContract } from '@features/generation/core/contracts'; +import type { ModelConfig } from '@features/models'; import { addEdge, addNode } from '@features/generation/core/graphBuilder'; @@ -40,6 +41,11 @@ export interface ControlLayerGraphInput { controlMode: 'balanced' | 'more_prompt' | 'more_control' | 'unbalanced' | null; /** The resolved model's `cond_in_channels`; required to validate an Anima ControlNet-LLLite layer. */ modelCondInChannels?: number | null; + /** + * The resolved model's `control_architecture` (Qwen-Image ControlNets). `instantx_union` takes the control image + * only; a Fun union, or a record written before the field existed, also takes the inpainted region. + */ + modelControlArchitecture?: ModelConfig['control_architecture'] | null; } /** Kind support is separate from control-mode/model compatibility. */ @@ -84,14 +90,16 @@ export interface AddControlLayersOptions { layers: readonly ControlLayerGraphInput[]; /** * Present when the canvas compiles in inpaint or outpaint mode. The Qwen-Image and FLUX.2 control nodes also - * condition on the inpainted region and take them; the other kinds ignore them. + * condition on the inpainted region and take them, except an InstantX Qwen-Image union; the other kinds ignore + * them. */ inpaint?: CanvasInpaintSources; } /** * The Fun-ControlNet-Union nodes (`qwen_image_control`, `flux2_control`): the control image, plus the region being - * regenerated when the canvas inpaints or outpaints, wired straight to the denoise's single `control` field. + * regenerated when the canvas inpaints or outpaints (not for an InstantX Qwen-Image union, which conditions on the + * control image alone), wired straight to the denoise's single `control` field. * Validation admits one per generation, so there is no collector. */ const addFunUnionControl = ( @@ -110,9 +118,10 @@ const addFunUnionControl = ( image: { image_name: layer.imageName }, type, }); - if (inpaint) { - // The union's inpainting condition. The canvas mask is white = keep and the node wants white = regenerate, - // so `img_lerp` 255 → 0 inverts it first, as the Anima inpaint adapter does. + if (inpaint && layer.modelControlArchitecture !== 'instantx_union') { + // The Fun union's inpainting condition; InstantX's Qwen-Image union has no input for one. The canvas mask is + // white = keep and the node wants white = regenerate, so `img_lerp` 255 → 0 inverts it first, as the Anima + // inpaint adapter does. const invertMask = addNode(graph, { id: `${type}_mask_${layer.id}`, max: 0, diff --git a/invokeai/frontend/webv2/src/features/generation/core/canvas/compileCanvasGraph.test.ts b/invokeai/frontend/webv2/src/features/generation/core/canvas/compileCanvasGraph.test.ts index ea3a660c815..58d9154ce15 100644 --- a/invokeai/frontend/webv2/src/features/generation/core/canvas/compileCanvasGraph.test.ts +++ b/invokeai/frontend/webv2/src/features/generation/core/canvas/compileCanvasGraph.test.ts @@ -940,6 +940,16 @@ describe('compileCanvasGraph — Qwen-Image ControlNet (integration)', () => { model: { base: 'qwen-image', key: 'qwen-union', name: 'Qwen Image ControlNet Union', type: 'controlnet' }, weight: 0.85, }; + const instantxLayer: ControlLayerGraphInput = { + ...layer, + model: { + base: 'qwen-image', + key: 'qwen-instantx', + name: 'Qwen Image InstantX ControlNet Union', + type: 'controlnet', + }, + modelControlArchitecture: 'instantx_union', + }; const CONTROL = 'qwen_image_control_q-control'; const MASK = 'qwen_image_control_mask_q-control'; const edgesInto = (graph: BackendGraphContract, nodeId: string) => @@ -1034,6 +1044,21 @@ describe('compileCanvasGraph — Qwen-Image ControlNet (integration)', () => { expect(getEdge(scaled, CONTROL, 'inpaint_image')?.source.node_id).toBe('infill'); }); + it.each(['inpaint', 'outpaint'] as const)('%s: an InstantX union reads the control image alone', (mode) => { + const { backendGraph } = compile(qwenImageModel, mode, { controlLayers: [instantxLayer] }); + + expect(backendGraph.nodes[CONTROL]).toMatchObject({ + control_model: instantxLayer.model, + image: { image_name: 'q-control.png' }, + }); + expect(backendGraph.nodes[CONTROL]).not.toHaveProperty('inpaint_image'); + expect(edgesInto(backendGraph, CONTROL)).toEqual([]); + expect(backendGraph.nodes[MASK]).toBeUndefined(); + // The region is still inpainted: the denoise keeps its mask. + expect(getEdge(backendGraph, 'denoise_latents', 'denoise_mask')?.source.node_id).toBe('create_gradient_mask'); + expect(getEdge(backendGraph, 'denoise_latents', 'control')?.source).toEqual({ field: 'control', node_id: CONTROL }); + }); + it('rejects the control on a Qwen Image Edit model, as the denoise does', () => { expect(() => compile({ ...qwenImageModel, variant: 'edit' }, 'inpaint', { controlLayers: [layer] })).toThrow( /qwen_image_edit_control/ @@ -1055,15 +1080,19 @@ describe('compileCanvasGraph — Qwen-Image ControlNet (integration)', () => { const hashed = (config: T, hash: string): T => ({ ...config, hash }); const model = hashed(qwenImageModel, 'blake3:qwen'); const controlLayers = [{ ...layer, model: hashed(layer.model, 'blake3:union') }]; + const instantxLayers = [{ ...instantxLayer, model: hashed(instantxLayer.model, 'blake3:instantx') }]; let txt2img: BackendGraphContract; let inpaint: BackendGraphContract; let outpaint: BackendGraphContract; + let instantxInpaint: BackendGraphContract; try { txt2img = compile(model, 'txt2img', { controlLayers }).backendGraph; // Without Scale Before Processing the control reads the canvas uploads as fixed image fields. inpaint = compile(model, 'inpaint', { controlLayers }).backendGraph; // The infilled image and the combined mask reach the control through edges. outpaint = compile(model, 'outpaint', { controlLayers, maskImageName: 'canvas-mask.png' }).backendGraph; + // An InstantX union inpaints with the control image alone. + instantxInpaint = compile(model, 'inpaint', { controlLayers: instantxLayers }).backendGraph; } finally { clock.mockRestore(); random.mockRestore(); @@ -1072,6 +1101,7 @@ describe('compileCanvasGraph — Qwen-Image ControlNet (integration)', () => { _comment: 'Generated by compileCanvasGraph.test.ts; regenerate with `vitest -u`. Consumed by tests/app/invocations/test_frontend_fun_control_graphs.py, which validates the graphs against the backend invocations.', inpaint_control_layer: { ...inpaint, id: 'qwen_image_control_canvas_inpaint' }, + inpaint_instantx_control_layer: { ...instantxInpaint, id: 'qwen_image_control_canvas_inpaint_instantx' }, outpaint_control_layer: { ...outpaint, id: 'qwen_image_control_canvas_outpaint' }, txt2img_control_layer: { ...txt2img, id: 'qwen_image_control_canvas_txt2img' }, }; diff --git a/invokeai/frontend/webv2/src/features/generation/core/canvas/controlValidation.ts b/invokeai/frontend/webv2/src/features/generation/core/canvas/controlValidation.ts index 677036fa7dc..71b8231063d 100644 --- a/invokeai/frontend/webv2/src/features/generation/core/canvas/controlValidation.ts +++ b/invokeai/frontend/webv2/src/features/generation/core/canvas/controlValidation.ts @@ -179,7 +179,7 @@ export const getControlValidationReason = (params: { return 'flux_fill_control_lora'; } if (kind === 'qwen_image_control' && mainBase === 'qwen-image' && mainVariant === 'edit') { - // The Fun-ControlNet branch was trained beside the generation transformer; `qwen_image_denoise` refuses it on + // Both ControlNet designs were trained beside the generation transformer; `qwen_image_denoise` refuses either on // an Edit transformer. return 'qwen_image_edit_control'; } diff --git a/invokeai/frontend/webv2/src/features/models/core/types.ts b/invokeai/frontend/webv2/src/features/models/core/types.ts index 1bdb5d300ac..41e304687db 100644 --- a/invokeai/frontend/webv2/src/features/models/core/types.ts +++ b/invokeai/frontend/webv2/src/features/models/core/types.ts @@ -152,6 +152,8 @@ export interface ModelConfig { image_encoder_model_id?: string; /** Anima ControlNet-LLLite: 3 = control image, 4 = inpainting; null when installed before it was recorded. */ cond_in_channels?: number | null; + /** Qwen-Image ControlNets: `fun_union` also conditions on an inpainted region, `instantx_union` does not. */ + control_architecture?: 'fun_union' | 'instantx_union'; /** Spandrel upscalers: native scale (1 = same size); null when installed before it was recorded. */ scale?: number | null; /** External API models. */ diff --git a/invokeai/frontend/webv2/src/workbench/controlAdapters.ts b/invokeai/frontend/webv2/src/workbench/controlAdapters.ts index 8e146c3a127..b407e75a688 100644 --- a/invokeai/frontend/webv2/src/workbench/controlAdapters.ts +++ b/invokeai/frontend/webv2/src/workbench/controlAdapters.ts @@ -34,7 +34,8 @@ export const CONTROL_ADAPTER_DEFAULTS: Readonly< model: null, weight: 0.75, }, - // The Fun-ControlNet-Union's recommended 0.70–0.95 context scale, centred. + // The Fun-ControlNet-Union's recommended 0.70–0.95 context scale, centred. An InstantX union on Qwen Image 2512 + // wants about 1.5–2.0, which the user sets. qwen_image_control: { beginEndStepPct: [0, 1], controlMode: null, diff --git a/invokeai/frontend/webv2/src/workbench/widgets/canvas/invoke/prepareCanvasInvocation.test.ts b/invokeai/frontend/webv2/src/workbench/widgets/canvas/invoke/prepareCanvasInvocation.test.ts index a343e591c06..7e580ff1fc1 100644 --- a/invokeai/frontend/webv2/src/workbench/widgets/canvas/invoke/prepareCanvasInvocation.test.ts +++ b/invokeai/frontend/webv2/src/workbench/widgets/canvas/invoke/prepareCanvasInvocation.test.ts @@ -1072,6 +1072,34 @@ describe('runCanvasInvocation', () => { }); }); + it('hands an InstantX union, as its record says it is, the control image alone when the canvas inpaints', async () => { + const instantx = { + ...union, + control_architecture: 'instantx_union' as const, + hash: 'qwen-instantx-hash', + key: 'qwen-instantx', + name: 'Qwen Image InstantX ControlNet Union', + }; + const harness = makeHarness({ + document: docWithLayers([ + rasterLayer('base'), + qwenLayer('q-layer', { model: instantx.key }), + inpaintMaskLayer('mask'), + ]), + model: qwenImage, + models: [union, instantx], + }); + + await runCanvasInvocation(harness.deps); + + expect(harness.notices()).toHaveLength(0); + const nodes = harness.submittedGraphs()[0]!.graph.backendGraph!.nodes; + expect(nodes['qwen_image_control_q-layer']).toMatchObject({ control_model: { key: 'qwen-instantx' } }); + expect(nodes['qwen_image_control_q-layer']).not.toHaveProperty('inpaint_image'); + expect(nodes['qwen_image_control_mask_q-layer']).toBeUndefined(); + expect(nodes.create_gradient_mask).toBeDefined(); + }); + it('blocks a second enabled nonempty control', async () => { const harness = makeHarness({ document: docWithLayers([qwenLayer('first'), qwenLayer('second')]), diff --git a/invokeai/frontend/webv2/src/workbench/widgets/canvas/invoke/prepareCanvasInvocation.ts b/invokeai/frontend/webv2/src/workbench/widgets/canvas/invoke/prepareCanvasInvocation.ts index ae525b4b8e3..e5ce0cb74d5 100644 --- a/invokeai/frontend/webv2/src/workbench/widgets/canvas/invoke/prepareCanvasInvocation.ts +++ b/invokeai/frontend/webv2/src/workbench/widgets/canvas/invoke/prepareCanvasInvocation.ts @@ -167,6 +167,7 @@ const createControlLayerCollector = ( ...(typeof resolved.hash === 'string' ? { hash: resolved.hash } : {}), }, modelCondInChannels: resolved.cond_in_channels, + modelControlArchitecture: resolved.control_architecture, weight: adapter.weight, }); return true; diff --git a/invokeai/frontend/webv2/src/workbench/widgets/layers/ControlLayerSettings.browser.test.tsx b/invokeai/frontend/webv2/src/workbench/widgets/layers/ControlLayerSettings.browser.test.tsx index fc5047aa9e9..e5e9286c2fb 100644 --- a/invokeai/frontend/webv2/src/workbench/widgets/layers/ControlLayerSettings.browser.test.tsx +++ b/invokeai/frontend/webv2/src/workbench/widgets/layers/ControlLayerSettings.browser.test.tsx @@ -490,18 +490,25 @@ describe('ControlLayerSettings on a Qwen Image main model', () => { const QWEN_EDIT_MAIN = { ...QWEN_MAIN, key: 'qwen-edit', name: 'Qwen Image Edit 2511', variant: 'edit' }; const QWEN_MODELS = [ { base: 'qwen-image', key: 'qwen-union', name: 'Qwen Image ControlNet Union', type: 'controlnet' }, + { + base: 'qwen-image', + control_architecture: 'instantx_union', + key: 'qwen-instantx', + name: 'Qwen Image InstantX ControlNet Union', + type: 'controlnet', + }, { base: 'z-image', key: 'z-union', name: 'Z-Image ControlNet Union', type: 'controlnet' }, ]; /** A layer as the Layers panel creates it under a Qwen Image main model. */ const qwenLayer = (id: string) => createControlLayer(id, id, 'qwen-image', 'qwen-union'); - it('offers the Qwen Image ControlNet kind and its models, and accepts one without a warning', async () => { + it('offers the Qwen Image ControlNet kind and both of its designs, and accepts one without a warning', async () => { await renderWith(QWEN_MAIN, QWEN_MODELS, qwenLayer('control-1')); expect(offersKind('qwen_image_control')).toBe(true); expect(offersKind('controlnet')).toBe(false); expect(offersKind('z_image_control')).toBe(false); - expect(offeredModels()).toEqual(['qwen-union']); + expect(offeredModels()).toEqual(['qwen-union', 'qwen-instantx']); expect(host!.textContent).toContain('Qwen Image ControlNet Union'); expect(textOf('alert')).toBe(''); }); diff --git a/tests/app/invocations/test_frontend_fun_control_graphs.py b/tests/app/invocations/test_frontend_fun_control_graphs.py index a973c3fef13..f04184b13b0 100644 --- a/tests/app/invocations/test_frontend_fun_control_graphs.py +++ b/tests/app/invocations/test_frontend_fun_control_graphs.py @@ -4,7 +4,8 @@ `compileCanvasGraph.test.ts` (regenerate with `vitest -u`) from the real canvas compiler, one per family: a text-to-image submission with one control layer, an inpainting submission where that layer also hands the control the canvas composite and the inverted canvas mask, and an outpainting submission where it hands over the infilled image -and the combined mask through edges. This module checks more than names: pydantic validates every node against its +and the combined mask through edges. Qwen-Image adds an inpainting submission whose layer holds an InstantX union, +which gets the control image alone. This module checks more than names: pydantic validates every node against its invocation (bounds such as `control_context_scale` in [0, 2] included), `Graph.validate_self` checks every edge type into the denoise's `control` and `vae`, the control node's own `invoke` runs on the fields it would receive, and `core_metadata` is serialised the way the gallery persists it. No model is loaded. @@ -53,6 +54,8 @@ class Family: control_image: str scale: float metadata_prefix: str + # Graphs beyond the three every family has, by name. + extra_graphs: tuple[str, ...] = () FAMILIES = { @@ -65,6 +68,7 @@ class Family: control_image="q-control.png", scale=0.85, metadata_prefix="qwen_image", + extra_graphs=("inpaint_instantx_control_layer",), ), "flux2": Family( fixture="flux2ControlCanvasGraphs.json", @@ -122,7 +126,9 @@ def _control_field(graph: Graph, node: Any) -> Any: @pytest.mark.parametrize("family", sorted(FAMILIES)) def test_the_fixtures_cover_every_canvas_mode_with_one_control_layer(family: str) -> None: spec = FAMILIES[family] - assert sorted(GRAPHS[family]) == ["inpaint_control_layer", "outpaint_control_layer", "txt2img_control_layer"] + assert sorted(GRAPHS[family]) == sorted( + ["inpaint_control_layer", "outpaint_control_layer", "txt2img_control_layer", *spec.extra_graphs] + ) for name, raw in GRAPHS[family].items(): types = [node["type"] for node in raw["nodes"].values()] assert types.count(spec.control.get_type()) == 1, name @@ -230,6 +236,26 @@ def test_outpainting_adds_the_infilled_image_and_the_inverted_combined_mask(fami assert field.inpaint_image_name is not None and field.inpaint_mask_name is not None +def test_a_qwen_image_instantx_inpaint_sends_the_control_image_alone_and_keeps_the_denoise_mask() -> None: + """The InstantX union has no inpainting condition and the denoise refuses one, so the canvas leaves the inpaint + inputs off its node; the region is still regenerated through the denoise mask.""" + spec = FAMILIES["qwen_image"] + graph = _graph("qwen_image", "inpaint_instantx_control_layer") + control = _only(graph, spec.control) + + assert control.control_model.key == "qwen-instantx" + assert control.image is not None and control.image.image_name == spec.control_image + assert control.inpaint_image is None and control.inpaint_mask is None + assert _inputs(graph, control.id) == {} + assert not any(isinstance(node, ImageLerpInvocation) for node in graph.nodes.values()) + gradient_mask = graph.nodes[_inputs(graph, _only(graph, spec.denoise).id)["denoise_mask"].node_id] + assert isinstance(gradient_mask, CreateGradientMaskInvocation) + assert gradient_mask.mask.image_name == "canvas-mask.png" + field = _control_field(graph, control) + assert field.image_name == spec.control_image + assert field.inpaint_image_name is None and field.inpaint_mask_name is None + + @pytest.mark.parametrize("family", sorted(FAMILIES)) def test_the_control_node_refuses_the_half_pair_the_frontend_never_sends(family: str) -> None: """The frontend sends the inpaint image and mask together or not at all; pin the node constraint it mirrors.""" @@ -247,5 +273,5 @@ def test_persisted_metadata_records_the_canvas_mode(family: str, name: str) -> N persisted = metadata.invoke(None).metadata.root # type: ignore[arg-type] - mode = name.removesuffix("_control_layer") + mode = name.split("_", 1)[0] assert persisted["generation_mode"] == f"{FAMILIES[family].metadata_prefix}_{mode}" diff --git a/tests/app/invocations/test_qwen_image_control.py b/tests/app/invocations/test_qwen_image_control.py index d711745d87c..3b745994ea5 100644 --- a/tests/app/invocations/test_qwen_image_control.py +++ b/tests/app/invocations/test_qwen_image_control.py @@ -3,9 +3,10 @@ reservation inside the exit stack, gates it per step, and builds the packed control context from the field's images the way it packs its latents. -The denoise runs its transformer in bfloat16, which the CPU test fleet cannot matmul, so every run here stops at the -transformer's lock or at the LoRA patcher, or runs the loop with the transformer's forward stood in. The forward with -the hooks is covered in `tests/backend/fun_controlnet`. +The denoise runs its transformer in bfloat16, which the CPU test fleet cannot matmul, so the Fun runs here stop at the +transformer's lock or at the LoRA patcher, or run the loop with the transformer's forward stood in; the forward with +the hooks is covered in `tests/backend/fun_controlnet`. The InstantX runs go through the real tiny transformer and +ControlNet, each widening what the node feeds it to float32 (`_Float32Compute`). """ from contextlib import contextmanager, nullcontext @@ -15,6 +16,7 @@ import pytest import torch from diffusers.models.autoencoders.autoencoder_kl_qwenimage import AutoencoderKLQwenImage +from diffusers.models.controlnets.controlnet_qwenimage import QwenImageControlNetModel from diffusers.models.transformers.transformer_qwenimage import QwenImageTransformer2DModel from PIL import Image from pydantic import ValidationError @@ -314,12 +316,13 @@ def test_with_control_the_lock_reserves_the_hints_too_and_the_images_are_encoded with pytest.raises(_Stop): _denoise()._run_diffusion(context) - expected = QwenImageDenoiseInvocation._estimate_working_memory( - image_seq_len=16, - text_seq_len=TEXT_TOKENS, - element_size=torch.bfloat16.itemsize, - num_control_blocks=len(CONTROL_LAYERS), - control_hidden_size=HIDDEN, + expected = ( + QwenImageDenoiseInvocation._estimate_working_memory( + image_seq_len=16, + text_seq_len=TEXT_TOKENS, + element_size=torch.bfloat16.itemsize, + ) + + (len(CONTROL_LAYERS) + 4) * 16 * HIDDEN * torch.bfloat16.itemsize ) assert transformer_info.model_on_device.call_args.kwargs["working_mem_bytes"] == expected # The control image and the masked inpaint image went through the VAE before the transformer was locked. @@ -463,7 +466,7 @@ def recording(**kwargs): monkeypatch.setattr(QwenImageImageToLatentsInvocation, "vae_encode", staticmethod(recording)) _denoise(control=_control_field(tiled=tiled, tile_size=tile_size))._build_control_context( - context, LATENT_SIDE, LATENT_SIDE + context, _tiny_branch(), LATENT_SIDE, LATENT_SIDE ) # Once for the control image, once for the masked inpaint image, both with the field's tiling. @@ -482,7 +485,7 @@ def test_the_control_context_is_the_encoded_images_packed_like_the_latents(tmp_p vae_info = _loaded(vae, _yielding(vae)) context = _context(tmp_path, vae_info=vae_info) - packed = _denoise()._build_control_context(context, LATENT_SIDE, LATENT_SIDE) + packed = _denoise()._build_control_context(context, _tiny_branch(), LATENT_SIDE, LATENT_SIDE) assert packed.shape == (1, (LATENT_SIDE // 2) ** 2, 33 * 4) # Channel-major within a token: feature c * 4 + dy * 2 + dx of token (i, j) is channel c at (2i + dy, 2j + dx). @@ -502,3 +505,258 @@ def encode(pixels: torch.Tensor) -> torch.Tensor: masked = _as_denoise_input(_inpaint_image(), "RGB") masked[..., : SIDE // 2, : SIDE // 2] = 0.0 assert torch.equal(context_4d[:, 17:], encode(masked)) + + +# --- the InstantX ControlNet --- + +INSTANTX_BLOCKS = 2 +INSTANTX_SCALE = 0.9 + + +def _widened(value): + if isinstance(value, torch.Tensor): + return value.float() if value.is_floating_point() else value + if isinstance(value, (list, tuple)): + return type(value)(_widened(item) for item in value) + return value + + +class _Float32Compute: + """Runs the real forward in float32 on what the node feeds it in bfloat16, and hands the result back in the + dtype it was fed: the bfloat16 GEMM the CPU fleet lacks, done wider. The direct restatement calls the same + models, so both sides see the same arithmetic.""" + + def forward(self, **kwargs): + dtype = kwargs["hidden_states"].dtype + output = super().forward(**{name: _widened(value) for name, value in kwargs.items()}) # type: ignore[misc] + return [tensor.to(dtype) for tensor in output] + + +class _InstantXTransformer(_Float32Compute, QwenImageTransformer2DModel): + """The tiny transformer, recording every forward the node makes: its arguments and its output.""" + + def forward(self, **kwargs): + output = super().forward(**kwargs) + self.calls.append(SimpleNamespace(kwargs=kwargs, output=output[0])) + return (output[0],) + + +class _InstantXControlNet(_Float32Compute, QwenImageControlNetModel): + """The tiny ControlNet, counting its forwards.""" + + def forward(self, **kwargs): + self.calls += 1 + return super().forward(**kwargs) + + +def _randomized(model: torch.nn.Module, seed: int) -> torch.nn.Module: + """Every parameter from N(0, 0.1): diffusers zero-initialises `controlnet_blocks` and `controlnet_x_embedder`, + which would make every residual zero and every comparison here vacuous, and leaves the modulation gates near + zero.""" + torch.manual_seed(seed) + for parameter in model.parameters(): + torch.nn.init.normal_(parameter, std=0.1) + return model.eval() + + +def _instantx_transformer() -> _InstantXTransformer: + model = _InstantXTransformer( + patch_size=2, + in_channels=64, + out_channels=16, + num_layers=4, + attention_head_dim=8, + num_attention_heads=2, + joint_attention_dim=8, + axes_dims_rope=(2, 4, 2), + ) + model.calls = [] + _randomized(model, seed=10) + return model + + +def _instantx_controlnet() -> _InstantXControlNet: + model = _InstantXControlNet( + patch_size=2, + in_channels=64, + out_channels=16, + num_layers=INSTANTX_BLOCKS, + attention_head_dim=8, + num_attention_heads=2, + joint_attention_dim=8, + axes_dims_rope=(2, 4, 2), + ) + model.calls = 0 + _randomized(model, seed=11) + return model + + +# Distinct prompts, the positive one padded, so a ControlNet fed the negative prompt or its mask would show. +_PROMPTS = { + "pos": (torch.randn(1, TEXT_TOKENS, 8, generator=torch.Generator().manual_seed(20)), [True, True, True, False]), + "neg": (torch.randn(1, TEXT_TOKENS, 8, generator=torch.Generator().manual_seed(21)), [True, True, True, True]), +} + + +def _prompt(name: str) -> tuple[torch.Tensor, torch.Tensor]: + """A prompt as the denoise holds it: embeddings in bfloat16, the mask as is.""" + embeds, mask = _PROMPTS[name] + return embeds.to(torch.bfloat16), torch.tensor([mask]) + + +def _instantx_context(tmp_path, transformer, controlnet, vae) -> MagicMock: + context = _context( + tmp_path, + _loaded(transformer, _yielding(transformer)), + _loaded(controlnet, _yielding(controlnet)), + _loaded(vae, _yielding(vae)), + ) + context.conditioning.load.side_effect = lambda name: ConditioningFieldData( + conditionings=[ + QwenImageConditioningInfo( + prompt_embeds=_PROMPTS[name][0], prompt_embeds_mask=torch.tensor([_PROMPTS[name][1]]) + ) + ] + ) + return context + + +def _instantx_field(**overrides) -> QwenImageControlField: + fields = {"inpaint_image_name": None, "inpaint_mask_name": None, "control_context_scale": INSTANTX_SCALE} + fields.update(overrides) + return _control_field(**fields) + + +def _expected_control_latents(vae) -> torch.Tensor: + """diffusers' pipeline: the control image at generation size in [-1, 1], encoded and normalised, packed 2x2 into + `(1, tokens, 64)` and cast to the prompt's dtype.""" + pixels = _as_denoise_input(_control_image(), "RGB") + latents = QwenImageImageToLatentsInvocation.vae_encode(_loaded(vae, _yielding(vae)), pixels).squeeze(2) + half = LATENT_SIDE // 2 + packed = latents.view(1, 16, half, 2, half, 2).permute(0, 2, 4, 1, 3, 5).reshape(1, half * half, 64) + return packed.to(torch.bfloat16) + + +@torch.no_grad() +def test_an_instantx_step_is_the_diffusers_controlnet_pipeline_step(tmp_path) -> None: + """Every step's two transformer forwards are what QwenImageControlNetPipeline computes from that step's + latents: the ControlNet once, on the positive prompt and its mask at `timestep / 1000`, and its scaled residuals + added in both guidance branches.""" + transformer, controlnet, vae = _instantx_transformer(), _instantx_controlnet(), _tiny_vae() + context = _instantx_context(tmp_path, transformer, controlnet, vae) + + _denoise(control=_instantx_field())._run_diffusion(context) + + calls = list(transformer.calls) + assert len(calls) == 4 and controlnet.calls == 2 + control_latents = _expected_control_latents(vae) + img_shapes = [[(1, LATENT_SIDE // 2, LATENT_SIDE // 2)]] + pos_embeds, pos_mask = _prompt("pos") + neg_embeds, neg_mask = _prompt("neg") + for positive, negative in (calls[0:2], calls[2:4]): + latents, timestep = positive.kwargs["hidden_states"], positive.kwargs["timestep"] + samples = controlnet( + hidden_states=latents, + controlnet_cond=control_latents, + conditioning_scale=INSTANTX_SCALE, + encoder_hidden_states=pos_embeds, + encoder_hidden_states_mask=pos_mask, + timestep=timestep, + img_shapes=img_shapes, + return_dict=False, + ) + shared = {"hidden_states": latents, "timestep": timestep, "img_shapes": img_shapes, "return_dict": False} + expected_positive = transformer( + encoder_hidden_states=pos_embeds, + encoder_hidden_states_mask=pos_mask, + controlnet_block_samples=samples, + **shared, + )[0] + expected_negative = transformer( + encoder_hidden_states=neg_embeds, + encoder_hidden_states_mask=neg_mask, + controlnet_block_samples=samples, + **shared, + )[0] + uncontrolled = transformer(encoder_hidden_states=pos_embeds, encoder_hidden_states_mask=pos_mask, **shared)[0] + + assert torch.equal(negative.kwargs["hidden_states"], latents) + assert torch.equal(positive.output, expected_positive) + assert torch.equal(negative.output, expected_negative) + # The residuals move the prediction, so the comparisons above can fail. + assert not torch.equal(positive.output, uncontrolled) + + +@torch.no_grad() +def test_an_instantx_control_at_scale_zero_generates_what_no_control_does(tmp_path) -> None: + vae = _tiny_vae() + + def generate(**overrides) -> torch.Tensor: + transformer, controlnet = _instantx_transformer(), _instantx_controlnet() + context = _instantx_context(tmp_path, transformer, controlnet, vae) + return _denoise(**overrides)._run_diffusion(context) + + uncontrolled = generate(control=None, vae=None) + + assert torch.equal(generate(control=_instantx_field(control_context_scale=0.0)), uncontrolled) + assert not torch.equal(generate(control=_instantx_field()), uncontrolled) + + +@torch.no_grad() +def test_an_instantx_control_outside_its_step_range_runs_no_controlnet(tmp_path) -> None: + """Four steps, the second half controlled: the ControlNet runs on the two steps inside the range only, and both + guidance forwards of a step take the same residuals or none.""" + transformer, controlnet, vae = _instantx_transformer(), _instantx_controlnet(), _tiny_vae() + context = _instantx_context(tmp_path, transformer, controlnet, vae) + + _denoise(steps=4, control=_instantx_field(begin_step_percent=0.5, end_step_percent=1.0))._run_diffusion(context) + + assert controlnet.calls == 2 + samples = [call.kwargs["controlnet_block_samples"] for call in transformer.calls] + assert samples[:4] == [None] * 4 + assert all(isinstance(residuals, list) and len(residuals) == INSTANTX_BLOCKS for residuals in samples[4:]) + assert samples[4] is samples[5] and samples[6] is samples[7] + + +@pytest.mark.parametrize("overrides", [{}, {"image_name": None}], ids=["control-and-inpaint", "inpaint-alone"]) +def test_an_instantx_control_refuses_inpaint_inputs(tmp_path, overrides) -> None: + transformer, controlnet, vae = _instantx_transformer(), _instantx_controlnet(), _tiny_vae() + context = _instantx_context(tmp_path, transformer, controlnet, vae) + vae_info = context.models.load(_model(ModelType.VAE)) + + with pytest.raises(ValueError, match="no inpainting condition"): + _denoise(control=_control_field(**overrides))._run_diffusion(context) + + # Refused before anything was encoded or run. + assert vae_info.model_on_device.call_count == 0 + assert transformer.calls == [] and controlnet.calls == 0 + + +def test_an_instantx_control_reserves_its_residuals_and_is_locked_with_the_same_reservation( + tmp_path, monkeypatch: pytest.MonkeyPatch +) -> None: + transformer, controlnet, vae = _instantx_transformer(), _instantx_controlnet(), _tiny_vae() + context = _instantx_context(tmp_path, transformer, controlnet, vae) + transformer_info = context.models.load(_model(ModelType.Main)) + control_info = context.models.load(_model(ModelType.ControlNet)) + hooks_at_patching: list[int] = [] + + def stop(**_kwargs): + hooks_at_patching.append(_hook_count(transformer)) + raise _Stop + + monkeypatch.setattr(LayerPatcher, "apply_smart_model_patches", stop) + + with pytest.raises(_Stop): + _denoise(control=_instantx_field())._run_diffusion(context) + + plain = QwenImageDenoiseInvocation._estimate_working_memory( + image_seq_len=16, text_seq_len=TEXT_TOKENS, element_size=torch.bfloat16.itemsize + ) + # The scaled residuals, one per block, alive through both guidance forwards: 16 tokens at the ControlNet's width. + residuals = INSTANTX_BLOCKS * 16 * HIDDEN * torch.bfloat16.itemsize + reserved = transformer_info.model_on_device.call_args.kwargs["working_mem_bytes"] + assert reserved == plain + residuals + assert control_info.model_on_device.call_args.kwargs["working_mem_bytes"] == reserved + # Nothing is hooked into the transformer: the ControlNet runs beside it. + assert hooks_at_patching == [0] diff --git a/tests/backend/flux/controlnet/test_state_dict_utils.py b/tests/backend/flux/controlnet/test_state_dict_utils.py index 2688dd3bebd..a47f328466c 100644 --- a/tests/backend/flux/controlnet/test_state_dict_utils.py +++ b/tests/backend/flux/controlnet/test_state_dict_utils.py @@ -13,6 +13,9 @@ ) from tests.backend.flux.controlnet.instantx_flux_controlnet_state_dict import instantx_config, instantx_sd_shapes from tests.backend.flux.controlnet.xlabs_flux_controlnet_state_dict import xlabs_sd_shapes +from tests.backend.model_manager.load.state_dicts.qwen_image_instantx_controlnet_keys import ( + state_dict_keys as qwen_image_instantx_keys, +) @pytest.mark.parametrize( @@ -33,6 +36,8 @@ def test_is_state_dict_xlabs_controlnet(sd_shapes: dict[str, list[int]], expecte [ (instantx_sd_shapes, True), (xlabs_sd_shapes, False), + # InstantX's Qwen-Image ControlNet: the same `controlnet_*` keys over Qwen-Image's own blocks. + (qwen_image_instantx_keys, False), (["foo"], False), ], ) diff --git a/tests/backend/model_manager/configs/test_qwen_image_controlnet_config.py b/tests/backend/model_manager/configs/test_qwen_image_controlnet_config.py new file mode 100644 index 00000000000..8e50b0ca056 --- /dev/null +++ b/tests/backend/model_manager/configs/test_qwen_image_controlnet_config.py @@ -0,0 +1,265 @@ +"""Identification of the two Qwen-Image ControlNet designs, through the config factory. + +alibaba-pai's Fun-ControlNet-Union is a VACE-style control branch: `control_img_in` projects 132 packed channels into +Qwen-Image's 3072-wide stream, and five `control_blocks` copy the dual-stream `QwenImageTransformerBlock` with +`before_proj`/`after_proj` taps. Its siblings reuse the naming -- the Qwen-Image-2.1 branch is single-stream at 4096 +wide, FLUX.2's keeps `control_transformer_blocks` at 6144. + +InstantX's ControlNet-Union is a five-block copy of the Qwen-Image transformer with `controlnet_x_embedder` and +`controlnet_blocks`: the key family of InstantX's FLUX ControlNets, and with `img_in`/`txt_in`/`txt_norm` the keys the +Qwen-Image main model probe reads. Both used to claim it, and the factory ranks main models first. + +Both are one config, the design recorded as `control_architecture`. Every case is a header written back without its +data: identification reads nothing else, and each config sees the real key names, shapes and dtypes. +""" + +from pathlib import Path + +import pytest +from pydantic import TypeAdapter + +from invokeai.backend.model_manager.configs.controlnet import ( + ControlNet_Checkpoint_FLUX_Config, + ControlNet_Checkpoint_QwenImage_Config, +) +from invokeai.backend.model_manager.configs.factory import AnyModelConfig, ModelConfigFactory +from invokeai.backend.model_manager.configs.identification_utils import InvalidMatchError, NotAMatchError +from invokeai.backend.model_manager.configs.main import Main_Checkpoint_QwenImage_Config +from invokeai.backend.model_manager.model_on_disk import ModelOnDisk +from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelType +from tests.backend.flux.controlnet.instantx_flux_controlnet_state_dict import instantx_sd_shapes +from tests.backend.model_manager.load.state_dicts import ( + flux2_fun_controlnet_keys, + qwen_image_fun_controlnet_keys, + qwen_image_instantx_controlnet_keys, +) +from tests.backend.model_manager.load.state_dicts.utils import write_header_only_safetensors + +Keys = dict[str, tuple[list[int], str]] + +FUN_FILE_NAME = "Qwen-Image-2512-Fun-Controlnet-Union-2602.safetensors" +INSTANTX_FILE_NAME = "Qwen-Image-InstantX-ControlNet-Union.safetensors" + +_OVERRIDE_FIELDS: dict[str, object] = { + "hash": "blake3:fakehash", + "path": "/fake/models/controlnet.safetensors", + "file_size": 1000, + "name": "controlnet", + "description": "test", + "source": "test", + "source_type": "path", + "key": "test-key", +} + + +def _fun_keys() -> Keys: + return dict(qwen_image_fun_controlnet_keys.state_dict_keys) + + +def _instantx_keys() -> Keys: + return dict(qwen_image_instantx_controlnet_keys.state_dict_keys) + + +def _bf16(shapes: dict[str, list[int]]) -> Keys: + return {key: (shape, "BF16") for key, shape in shapes.items()} + + +# Sibling layouts with no captured header, reduced to the keys the probes read. +QWEN_IMAGE_2_1_FUN_KEYS = _bf16( + { + "control_img_in.weight": [4096, 129], + "control_img_in.bias": [4096], + "control_blocks.0.before_proj.weight": [4096, 4096], + "control_blocks.0.after_proj.weight": [4096, 4096], + "control_blocks.0.attn.to_q.weight": [4096, 4096], + "control_blocks.0.img_mlp.gate_layer.weight": [12288, 4096], + "control_blocks.0.img_mlp.proj.weight": [12288, 4096], + "control_blocks.0.img_mlp.out.weight": [4096, 12288], + "control_blocks.15.after_proj.weight": [4096, 4096], + } +) +Z_IMAGE_FUN_KEYS = _bf16( + { + "control_layers.0.attention.qkv.weight": [11520, 3840], + "control_layers.0.attention.out.weight": [3840, 3840], + "control_all_x_embedder.2-1.weight": [3840, 64], + "control_noise_refiner.0.attention.qkv.weight": [11520, 3840], + } +) +DIFFSYNTH_KEYS = _bf16( + { + "img_in.weight": [3072, 64], + "controlnet_blocks.0.input_proj.weight": [3072, 3072], + "controlnet_blocks.0.output_proj.weight": [3072, 3072], + "controlnet_blocks.0.x_rms.weight": [3072], + "controlnet_blocks.0.y_rms.weight": [3072], + } +) +# InstantX's FLUX.1 ControlNet-Union, diffusers layout: the same `controlnet_*` names over FLUX blocks. +FLUX_INSTANTX_KEYS = _bf16(instantx_sd_shapes) + + +def _write(tmp_path: Path, name: str, keys: Keys) -> Path: + return write_header_only_safetensors(tmp_path / name, keys) + + +def _instantx_inpainting_keys() -> Keys: + """InstantX's inpainting sibling: the control input also takes the 4 packed features of a mask.""" + return {**_instantx_keys(), "controlnet_x_embedder.weight": ([3072, 68], "BF16")} + + +@pytest.mark.parametrize( + ("name", "keys", "architecture"), + [(FUN_FILE_NAME, _fun_keys(), "fun_union"), (INSTANTX_FILE_NAME, _instantx_keys(), "instantx_union")], + ids=["fun", "instantx"], +) +def test_each_released_file_is_a_qwen_image_controlnet_of_its_design_and_nothing_else( + tmp_path: Path, name: str, keys: Keys, architecture: str +) -> None: + result = ModelConfigFactory.from_model_on_disk(_write(tmp_path, name, keys), allow_unknown=False) + + assert type(result.config) is ControlNet_Checkpoint_QwenImage_Config, result.details + assert result.match_count == 1, result.details + assert result.config.control_architecture == architecture + assert result.config.base is BaseModelType.QwenImage + assert result.config.type is ModelType.ControlNet + assert result.config.format is ModelFormat.Checkpoint + + +def test_the_flux_instantx_controlnet_is_still_a_flux_controlnet_and_not_a_qwen_image_one(tmp_path: Path) -> None: + path = _write(tmp_path, "FLUX.1-dev-Controlnet-Union.safetensors", FLUX_INSTANTX_KEYS) + + result = ModelConfigFactory.from_model_on_disk(path, allow_unknown=False) + + assert type(result.config) is ControlNet_Checkpoint_FLUX_Config, result.details + assert result.match_count == 1, result.details + with pytest.raises(NotAMatchError): + ControlNet_Checkpoint_QwenImage_Config.from_model_on_disk(ModelOnDisk(path), dict(_OVERRIDE_FIELDS)) + + +def test_the_instantx_inpainting_layout_is_refused_with_a_reason_rather_than_registered(tmp_path: Path) -> None: + """Recognised and unusable, so nothing is registered -- not even as an unknown model -- and the reason is the + installer's message.""" + result = ModelConfigFactory.from_model_on_disk( + _write(tmp_path, "Qwen-Image-InstantX-ControlNet-Inpainting.safetensors", _instantx_inpainting_keys()) + ) + + assert result.config is None + refusal = result.details[ControlNet_Checkpoint_QwenImage_Config.__name__] + assert isinstance(refusal, InvalidMatchError) + assert "does not support yet" in str(refusal) and "ControlNet Union instead" in str(refusal) + + +@pytest.mark.parametrize("prefix", ["", "model.diffusion_model."], ids=["bare", "comfyui-prefixed"]) +def test_the_qwen_image_main_probe_refuses_the_instantx_union(tmp_path: Path, prefix: str) -> None: + """The probe strips a ComfyUI prefix before it looks for the transformer's keys, so it has to before it looks + for the ControlNet's too, as the main loader does: otherwise a prefixed copy would be identified as a main + model, refused at load, and identified as one again by Re-identify.""" + keys = {f"{prefix}{name}": spec for name, spec in _instantx_keys().items()} + mod = ModelOnDisk(_write(tmp_path, INSTANTX_FILE_NAME, keys)) + + with pytest.raises(NotAMatchError): + Main_Checkpoint_QwenImage_Config.from_model_on_disk(mod, dict(_OVERRIDE_FIELDS)) + + +@pytest.mark.parametrize( + ("name", "preprocessor"), + [(FUN_FILE_NAME, None), ("Qwen-Image-Fun-Controlnet-Depth.safetensors", "depth_anything_image_processor")], + ids=["union", "depth"], +) +def test_the_preprocessor_is_chosen_from_the_name_as_for_every_controlnet( + tmp_path: Path, name: str, preprocessor: str | None +) -> None: + """A union adapter takes any control image, so none is pre-selected; a single-condition build picks its own.""" + config = ModelConfigFactory.from_model_on_disk(_write(tmp_path, name, _fun_keys()), allow_unknown=False).config + + assert isinstance(config, ControlNet_Checkpoint_QwenImage_Config) + assert config.default_settings is not None + assert config.default_settings.preprocessor == preprocessor + + +@pytest.mark.parametrize( + ("name", "keys"), [(FUN_FILE_NAME, _fun_keys()), (INSTANTX_FILE_NAME, _instantx_keys())], ids=["fun", "instantx"] +) +def test_a_stored_record_reads_back_as_the_same_qwen_image_controlnet(tmp_path: Path, name: str, keys: Keys) -> None: + """The record crosses the database and the API as a plain dict, so the config has to be a member of + `AnyModelConfig` under a discriminator of its own, and the design has to survive the round trip.""" + identified = ModelConfigFactory.from_model_on_disk(_write(tmp_path, name, keys), allow_unknown=False).config + assert identified is not None + + restored = TypeAdapter(AnyModelConfig).validate_python(identified.model_dump(mode="json")) + + assert type(restored) is ControlNet_Checkpoint_QwenImage_Config + assert restored.base is BaseModelType.QwenImage + assert restored.control_architecture == identified.control_architecture + + +def test_a_record_stored_before_the_design_was_recorded_reads_back_as_the_fun_union(tmp_path: Path) -> None: + """Every Qwen-Image ControlNet installed until now is a Fun-ControlNet: no other design was identified.""" + identified = ModelConfigFactory.from_model_on_disk( + _write(tmp_path, FUN_FILE_NAME, _fun_keys()), allow_unknown=False + ).config + assert identified is not None + stored = identified.model_dump(mode="json") + del stored["control_architecture"] + + restored = TypeAdapter(AnyModelConfig).validate_python(stored) + + assert type(restored) is ControlNet_Checkpoint_QwenImage_Config + assert restored.control_architecture == "fun_union" + + +@pytest.mark.parametrize( + ("name", "keys"), + [ + ("Qwen-Image-2.1-Fun-Controlnet-Union.safetensors", QWEN_IMAGE_2_1_FUN_KEYS), + ("FLUX.2-dev-Fun-Controlnet-Union-2602.safetensors", dict(flux2_fun_controlnet_keys.state_dict_keys)), + ("Z-Image-Turbo-Fun-Controlnet-Union.safetensors", Z_IMAGE_FUN_KEYS), + ("Qwen-Image-DiffSynth-Controlnet.safetensors", DIFFSYNTH_KEYS), + ], + ids=["qwen_image_2_1_fun", "flux2_fun", "z_image_fun", "diffsynth"], +) +def test_the_other_control_designs_are_refused(tmp_path: Path, name: str, keys: Keys) -> None: + mod = ModelOnDisk(_write(tmp_path, name, keys)) + + with pytest.raises(NotAMatchError): + ControlNet_Checkpoint_QwenImage_Config.from_model_on_disk(mod, dict(_OVERRIDE_FIELDS)) + + +def _without_img_mod(keys: Keys) -> Keys: + return {key: value for key, value in keys.items() if ".img_mod." not in key} + + +def _with_control_img_in(keys: Keys, shape: list[int]) -> Keys: + return {**keys, "control_img_in.weight": (shape, "BF16")} + + +def _without(keys: Keys, key: str) -> Keys: + return {name: value for name, value in keys.items() if name != key} + + +@pytest.mark.parametrize( + "keys", + [ + _without_img_mod(_fun_keys()), + _with_control_img_in(_fun_keys(), [4096, 129]), + _without_img_mod(_instantx_keys()), + _without(_instantx_keys(), "txt_norm.weight"), + {**_instantx_keys(), "controlnet_x_embedder.weight": ([4096, 64], "BF16")}, + ], + ids=[ + "fun_single_stream_blocks", + "fun_wider_projection", + "instantx_single_stream_blocks", + "instantx_without_text_norm", + "instantx_wider_embedder", + ], +) +def test_a_controlnet_with_the_names_but_not_the_fingerprint_is_refused(tmp_path: Path, keys: Keys) -> None: + """Every part of a fingerprint carries weight: `control_blocks` at Qwen-Image's width without `img_mod` would be + a single-stream branch, `img_mod` behind a wider projection a branch for a wider base, and `controlnet_blocks` + without Qwen-Image's dual-stream block and text norm a FLUX ControlNet.""" + mod = ModelOnDisk(_write(tmp_path, FUN_FILE_NAME, keys)) + + with pytest.raises(NotAMatchError): + ControlNet_Checkpoint_QwenImage_Config.from_model_on_disk(mod, dict(_OVERRIDE_FIELDS)) diff --git a/tests/backend/model_manager/configs/test_qwen_image_fun_controlnet_config.py b/tests/backend/model_manager/configs/test_qwen_image_fun_controlnet_config.py deleted file mode 100644 index 701a0302967..00000000000 --- a/tests/backend/model_manager/configs/test_qwen_image_fun_controlnet_config.py +++ /dev/null @@ -1,165 +0,0 @@ -"""Identification of alibaba-pai's Qwen-Image Fun-ControlNet-Union, through the config factory. - -The adapter is a VACE-style control branch: `control_img_in` projects 132 packed channels into Qwen-Image's -3072-wide stream, and five `control_blocks` copy the dual-stream `QwenImageTransformerBlock` with -`before_proj`/`after_proj` taps. Its siblings reuse the naming -- the Qwen-Image-2.1 branch is single-stream at -4096 wide, FLUX.2's keeps `control_transformer_blocks` at 6144 -- and the other Qwen-Image ControlNet designs -(InstantX, DiffSynth) share key families with FLUX ControlNets. So every case is a header written back without its -data: identification reads nothing else, and each config sees the real key names, shapes and dtypes. -""" - -from pathlib import Path - -import pytest -from pydantic import TypeAdapter - -from invokeai.backend.model_manager.configs.controlnet import ControlNet_Checkpoint_QwenImage_Config -from invokeai.backend.model_manager.configs.factory import AnyModelConfig, ModelConfigFactory -from invokeai.backend.model_manager.configs.identification_utils import NotAMatchError -from invokeai.backend.model_manager.model_on_disk import ModelOnDisk -from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelType -from tests.backend.model_manager.load.state_dicts import ( - flux2_fun_controlnet_keys, - qwen_image_fun_controlnet_keys, - qwen_image_instantx_controlnet_keys, -) -from tests.backend.model_manager.load.state_dicts.utils import write_header_only_safetensors - -Keys = dict[str, tuple[list[int], str]] - -FUN_FILE_NAME = "Qwen-Image-2512-Fun-Controlnet-Union-2602.safetensors" - -_OVERRIDE_FIELDS: dict[str, object] = { - "hash": "blake3:fakehash", - "path": "/fake/models/controlnet.safetensors", - "file_size": 1000, - "name": "controlnet", - "description": "test", - "source": "test", - "source_type": "path", - "key": "test-key", -} - - -def _fun_keys() -> Keys: - return dict(qwen_image_fun_controlnet_keys.state_dict_keys) - - -def _bf16(shapes: dict[str, list[int]]) -> Keys: - return {key: (shape, "BF16") for key, shape in shapes.items()} - - -# Sibling layouts with no captured header, reduced to the keys the probes read. -QWEN_IMAGE_2_1_FUN_KEYS = _bf16( - { - "control_img_in.weight": [4096, 129], - "control_img_in.bias": [4096], - "control_blocks.0.before_proj.weight": [4096, 4096], - "control_blocks.0.after_proj.weight": [4096, 4096], - "control_blocks.0.attn.to_q.weight": [4096, 4096], - "control_blocks.0.img_mlp.gate_layer.weight": [12288, 4096], - "control_blocks.0.img_mlp.proj.weight": [12288, 4096], - "control_blocks.0.img_mlp.out.weight": [4096, 12288], - "control_blocks.15.after_proj.weight": [4096, 4096], - } -) -Z_IMAGE_FUN_KEYS = _bf16( - { - "control_layers.0.attention.qkv.weight": [11520, 3840], - "control_layers.0.attention.out.weight": [3840, 3840], - "control_all_x_embedder.2-1.weight": [3840, 64], - "control_noise_refiner.0.attention.qkv.weight": [11520, 3840], - } -) -DIFFSYNTH_KEYS = _bf16( - { - "img_in.weight": [3072, 64], - "controlnet_blocks.0.input_proj.weight": [3072, 3072], - "controlnet_blocks.0.output_proj.weight": [3072, 3072], - "controlnet_blocks.0.x_rms.weight": [3072], - "controlnet_blocks.0.y_rms.weight": [3072], - } -) - - -def _write(tmp_path: Path, name: str, keys: Keys) -> Path: - return write_header_only_safetensors(tmp_path / name, keys) - - -def test_the_released_file_is_a_qwen_image_controlnet_and_nothing_else(tmp_path: Path) -> None: - result = ModelConfigFactory.from_model_on_disk(_write(tmp_path, FUN_FILE_NAME, _fun_keys()), allow_unknown=False) - - assert type(result.config) is ControlNet_Checkpoint_QwenImage_Config, result.details - assert result.match_count == 1, result.details - assert result.config.base is BaseModelType.QwenImage - assert result.config.type is ModelType.ControlNet - assert result.config.format is ModelFormat.Checkpoint - - -@pytest.mark.parametrize( - ("name", "preprocessor"), - [(FUN_FILE_NAME, None), ("Qwen-Image-Fun-Controlnet-Depth.safetensors", "depth_anything_image_processor")], - ids=["union", "depth"], -) -def test_the_preprocessor_is_chosen_from_the_name_as_for_every_controlnet( - tmp_path: Path, name: str, preprocessor: str | None -) -> None: - """A union adapter takes any control image, so none is pre-selected; a single-condition build picks its own.""" - config = ModelConfigFactory.from_model_on_disk(_write(tmp_path, name, _fun_keys()), allow_unknown=False).config - - assert isinstance(config, ControlNet_Checkpoint_QwenImage_Config) - assert config.default_settings is not None - assert config.default_settings.preprocessor == preprocessor - - -def test_a_stored_record_reads_back_as_a_qwen_image_controlnet(tmp_path: Path) -> None: - """The record crosses the database and the API as a plain dict, so the config has to be a member of - `AnyModelConfig` under a discriminator of its own.""" - stored = ModelConfigFactory.from_model_on_disk( - _write(tmp_path, FUN_FILE_NAME, _fun_keys()), allow_unknown=False - ).config.model_dump(mode="json") - - restored = TypeAdapter(AnyModelConfig).validate_python(stored) - - assert type(restored) is ControlNet_Checkpoint_QwenImage_Config - assert restored.base is BaseModelType.QwenImage - - -@pytest.mark.parametrize( - ("name", "keys"), - [ - ("Qwen-Image-2.1-Fun-Controlnet-Union.safetensors", QWEN_IMAGE_2_1_FUN_KEYS), - ("FLUX.2-dev-Fun-Controlnet-Union-2602.safetensors", dict(flux2_fun_controlnet_keys.state_dict_keys)), - ("Z-Image-Turbo-Fun-Controlnet-Union.safetensors", Z_IMAGE_FUN_KEYS), - ("Qwen-Image-InstantX-ControlNet-Union.safetensors", dict(qwen_image_instantx_controlnet_keys.state_dict_keys)), - ("Qwen-Image-DiffSynth-Controlnet.safetensors", DIFFSYNTH_KEYS), - ], - ids=["qwen_image_2_1_fun", "flux2_fun", "z_image_fun", "instantx", "diffsynth"], -) -def test_the_other_control_designs_are_refused(tmp_path: Path, name: str, keys: Keys) -> None: - mod = ModelOnDisk(_write(tmp_path, name, keys)) - - with pytest.raises(NotAMatchError): - ControlNet_Checkpoint_QwenImage_Config.from_model_on_disk(mod, dict(_OVERRIDE_FIELDS)) - - -def _without_img_mod(keys: Keys) -> Keys: - return {key: value for key, value in keys.items() if ".img_mod." not in key} - - -def _with_control_img_in(keys: Keys, shape: list[int]) -> Keys: - return {**keys, "control_img_in.weight": (shape, "BF16")} - - -@pytest.mark.parametrize( - "keys", - [_without_img_mod(_fun_keys()), _with_control_img_in(_fun_keys(), [4096, 129])], - ids=["single_stream_blocks", "wider_projection"], -) -def test_a_branch_with_the_names_but_not_the_fingerprint_is_refused(tmp_path: Path, keys: Keys) -> None: - """Both halves of the fingerprint carry weight: `control_blocks` at Qwen-Image's width without `img_mod` - would be a single-stream branch, and `img_mod` behind a wider projection a branch for a wider base.""" - mod = ModelOnDisk(_write(tmp_path, FUN_FILE_NAME, keys)) - - with pytest.raises(NotAMatchError): - ControlNet_Checkpoint_QwenImage_Config.from_model_on_disk(mod, dict(_OVERRIDE_FIELDS)) diff --git a/tests/backend/model_manager/load/state_dicts/qwen_image_instantx_controlnet_keys.py b/tests/backend/model_manager/load/state_dicts/qwen_image_instantx_controlnet_keys.py index ac6d5eb68b5..3f3ff708834 100644 --- a/tests/backend/model_manager/load/state_dicts/qwen_image_instantx_controlnet_keys.py +++ b/tests/backend/model_manager/load/state_dicts/qwen_image_instantx_controlnet_keys.py @@ -10,8 +10,8 @@ A different design from the Fun branch: a trimmed copy of the Qwen-Image transformer itself (`img_in`, `txt_in` at Qwen2.5-VL's 3584 width, `txt_norm`, `time_text_embed`, `transformer_blocks.N`) with `controlnet_x_embedder` for the control latent and one `controlnet_blocks.N` output projection per block -- -the key family InstantX's FLUX ControlNets use. It has no `control_img_in`, which is what keeps it out of -the Qwen-Image Fun config. +the key family InstantX's FLUX ControlNets use, which is why the FLUX probe has to refuse it. The full file +has five of each stack. """ state_dict_keys: dict[str, tuple[list[int], str]] = { diff --git a/tests/backend/model_manager/load/test_flux2_fun_controlnet_loader.py b/tests/backend/model_manager/load/test_flux2_fun_controlnet_loader.py new file mode 100644 index 00000000000..cc41bfc6ee5 --- /dev/null +++ b/tests/backend/model_manager/load/test_flux2_fun_controlnet_loader.py @@ -0,0 +1,92 @@ +"""What the FLUX.2 [dev] Fun-ControlNet loader makes of a checkpoint: the geometry it reads off the tensors, the dtype +it casts to, and the quantized files it refuses rather than loads unscaled. + +Like the Z-Image ControlNet loader it builds under `init_empty_weights` and assigns the file's tensors, so an fp8 +weight with its scale dropped would become the parameter as it is (`test_z_image_control_adapter_loader.py`). No +quantized build of this adapter is published (alibaba-pai ships dense bf16 files), so refusing costs nothing. The +Qwen-Image loader is covered in `test_qwen_image_controlnet_loader.py`. +""" + +from pathlib import Path + +import pytest +import torch + +from invokeai.backend.fun_controlnet.flux2 import FLUX2_FUN_CONTROL_LAYERS, Flux2FunControlBranch +from invokeai.backend.fun_controlnet.qwen_image import QwenImageFunControlBranch +from invokeai.backend.model_manager.configs.controlnet import ControlNet_Checkpoint_Flux2_Config +from invokeai.backend.model_manager.load.model_loaders import fun_controlnet +from invokeai.backend.model_manager.load.model_loaders.fun_controlnet import Flux2FunControlNetLoader +from tests.fixtures.loader_seams import Seam, prepare +from tests.fixtures.quantized_payloads import quantize_scaled_fp8 + +# One 128-wide head: the narrowest branch the loader's head rule accepts, with the real four-block layout. +HIDDEN = 128 + + +def _checkpoint() -> dict[str, torch.Tensor]: + """A Qwen-Image Fun branch's file, which the FLUX.2 loader has to refuse.""" + torch.manual_seed(0) + branch = QwenImageFunControlBranch(hidden_size=HIDDEN, num_attention_heads=1, attention_head_dim=128) + return {name: tensor.detach().clone() for name, tensor in branch.state_dict().items()} + + +FLUX2_SEAM = Seam( + loader=Flux2FunControlNetLoader, + module=fun_controlnet, + entry="_load_model", + compute_dtype=torch.bfloat16, + # The loader casts to bfloat16, the dtype the FLUX.2 denoise runs in, whatever the device would pick. + patches_device=False, + sets_torch_dtype=False, + casts_fp8_storage=False, +) +FLUX2_QUANTIZED_LAYER = "control_transformer_blocks.0.attn.to_q" + + +def _flux2_checkpoint() -> dict[str, torch.Tensor]: + """A float32 file with the real four-block layout, under the names alibaba-pai's file uses.""" + torch.manual_seed(0) + branch = Flux2FunControlBranch(hidden_size=HIDDEN, num_attention_heads=1, attention_head_dim=128) + return { + name.replace("control_blocks.", "control_transformer_blocks.", 1): tensor.detach().clone() + for name, tensor in branch.state_dict().items() + } + + +def _flux2_driver(monkeypatch: pytest.MonkeyPatch, tmp_path: Path, state_dict: dict[str, torch.Tensor]): + checkpoint = tmp_path / "flux2_fun_controlnet.safetensors" + checkpoint.touch() + config = ControlNet_Checkpoint_Flux2_Config.model_construct(path=str(checkpoint), name="flux2-control") + return prepare(FLUX2_SEAM, monkeypatch, state_dict=state_dict), config + + +def test_a_dense_flux2_checkpoint_loads_with_its_geometry_in_bfloat16(monkeypatch, tmp_path) -> None: + run, config = _flux2_driver(monkeypatch, tmp_path, _flux2_checkpoint()) + + model = run.load(config) + + assert isinstance(model, Flux2FunControlBranch) + assert model.control_layers == FLUX2_FUN_CONTROL_LAYERS + first = model.control_blocks[0] + assert (model.hidden_size, model.control_in_dim, first.attn.heads, first.mlp_hidden_dim) == (HIDDEN, 260, 1, 384) + assert {parameter.dtype for parameter in model.parameters()} == {torch.bfloat16} + assert not model.training + + +def test_a_scaled_fp8_flux2_checkpoint_is_refused_rather_than_loaded_unscaled(monkeypatch, tmp_path) -> None: + payload = quantize_scaled_fp8(torch.randn(HIDDEN, HIDDEN)) + state_dict = _flux2_checkpoint() + state_dict[f"{FLUX2_QUANTIZED_LAYER}.weight"] = payload.codes + state_dict[f"{FLUX2_QUANTIZED_LAYER}.weight_scale"] = payload.scale + run, config = _flux2_driver(monkeypatch, tmp_path, state_dict) + + with pytest.raises(ValueError, match="quantization side channel"): + run.load(config) + + +def test_the_qwen_image_recipe_is_refused_by_the_flux2_loader(monkeypatch, tmp_path) -> None: + run, config = _flux2_driver(monkeypatch, tmp_path, _checkpoint()) + + with pytest.raises(ValueError, match="takes 132"): + run.load(config) diff --git a/tests/backend/model_manager/load/test_fun_controlnet_loader.py b/tests/backend/model_manager/load/test_fun_controlnet_loader.py deleted file mode 100644 index 4734546f2e9..00000000000 --- a/tests/backend/model_manager/load/test_fun_controlnet_loader.py +++ /dev/null @@ -1,168 +0,0 @@ -"""What the Fun-ControlNet loaders (Qwen-Image, FLUX.2 [dev]) make of a checkpoint: the geometry they read off the -tensors, the dtype they cast to, and the quantized files they refuse rather than load unscaled. - -Like the Z-Image ControlNet loader it builds under `init_empty_weights` and assigns the file's tensors, so an fp8 -weight with its scale dropped would become the parameter as it is (`test_z_image_control_adapter_loader.py`). No -quantized build of this adapter is published (alibaba-pai ships two dense bf16 files), so refusing costs nothing. -""" - -from pathlib import Path - -import pytest -import torch - -from invokeai.backend.fun_controlnet.flux2 import FLUX2_FUN_CONTROL_LAYERS, Flux2FunControlBranch -from invokeai.backend.fun_controlnet.qwen_image import QWEN_IMAGE_FUN_CONTROL_LAYERS, QwenImageFunControlBranch -from invokeai.backend.model_manager.configs.controlnet import ( - ControlNet_Checkpoint_Flux2_Config, - ControlNet_Checkpoint_QwenImage_Config, -) -from invokeai.backend.model_manager.load.model_loaders import fun_controlnet -from invokeai.backend.model_manager.load.model_loaders.fun_controlnet import ( - Flux2FunControlNetLoader, - QwenImageFunControlNetLoader, -) -from tests.fixtures.loader_seams import Seam, prepare -from tests.fixtures.quantized_payloads import comfy_quant_marker, quantize_convrot, quantize_scaled_fp8 - -# One 128-wide head: the narrowest branch the loader's head rule accepts, with the real five-block layout. -HIDDEN = 128 -QUANTIZED_LAYER = "control_blocks.0.attn.to_q" - -SEAM = Seam( - loader=QwenImageFunControlNetLoader, - module=fun_controlnet, - entry="_load_model", - compute_dtype=torch.bfloat16, - # The loader takes its dtype from the device and never reads `self._torch_dtype`. - patches_device=True, - sets_torch_dtype=False, - casts_fp8_storage=False, -) - - -def _checkpoint() -> dict[str, torch.Tensor]: - """A float32 file with the real layout; the loader has to cast it.""" - torch.manual_seed(0) - branch = QwenImageFunControlBranch(hidden_size=HIDDEN, num_attention_heads=1, attention_head_dim=128) - return {name: tensor.detach().clone() for name, tensor in branch.state_dict().items()} - - -def _driver(monkeypatch: pytest.MonkeyPatch, tmp_path: Path, state_dict: dict[str, torch.Tensor]): - checkpoint = tmp_path / "qwen_image_fun_controlnet.safetensors" - checkpoint.touch() - config = ControlNet_Checkpoint_QwenImage_Config.model_construct(path=str(checkpoint), name="qwen-image-control") - return prepare(SEAM, monkeypatch, state_dict=state_dict), config - - -def test_a_dense_checkpoint_loads_with_its_geometry_in_the_compute_dtype(monkeypatch, tmp_path) -> None: - run, config = _driver(monkeypatch, tmp_path, _checkpoint()) - - model = run.load(config) - - assert isinstance(model, QwenImageFunControlBranch) - assert model.control_layers == QWEN_IMAGE_FUN_CONTROL_LAYERS - assert (model.hidden_size, model.control_in_dim, model.control_blocks[0].attn.heads) == (HIDDEN, 132, 1) - assert {parameter.dtype for parameter in model.parameters()} == {torch.bfloat16} - assert not model.training - - -def test_a_scaled_fp8_checkpoint_is_refused_rather_than_loaded_unscaled(monkeypatch, tmp_path) -> None: - payload = quantize_scaled_fp8(torch.randn(HIDDEN, HIDDEN)) - state_dict = _checkpoint() - state_dict[f"{QUANTIZED_LAYER}.weight"] = payload.codes - state_dict[f"{QUANTIZED_LAYER}.weight_scale"] = payload.scale - run, config = _driver(monkeypatch, tmp_path, state_dict) - - with pytest.raises(ValueError, match="quantization side channel"): - run.load(config) - - -def test_an_int8_checkpoint_is_refused_by_its_marker(monkeypatch, tmp_path) -> None: - codes, scale, _ = quantize_convrot(torch.randn(HIDDEN, 256)) - state_dict = _checkpoint() - state_dict[f"{QUANTIZED_LAYER}.weight"] = codes - state_dict[f"{QUANTIZED_LAYER}.weight_scale"] = scale - state_dict[f"{QUANTIZED_LAYER}.comfy_quant"] = comfy_quant_marker({"format": "int8_tensorwise", "convrot": True}) - run, config = _driver(monkeypatch, tmp_path, state_dict) - - with pytest.raises(ValueError, match="quantization side channel"): - run.load(config) - - -@pytest.mark.parametrize( - ("rewrite", "message"), - [ - (lambda sd: {k: v for k, v in sd.items() if not k.startswith("control_blocks.4.")}, "4 control blocks"), - (lambda sd: {**sd, "control_img_in.weight": torch.zeros(HIDDEN, 64)}, "takes 64"), - ], - ids=["four-blocks", "sixteen-channel-context"], -) -def test_another_recipe_is_refused_by_name(monkeypatch, tmp_path, rewrite, message) -> None: - run, config = _driver(monkeypatch, tmp_path, rewrite(_checkpoint())) - - with pytest.raises(ValueError, match=message): - run.load(config) - - -# --- FLUX.2 [dev] --- - -FLUX2_SEAM = Seam( - loader=Flux2FunControlNetLoader, - module=fun_controlnet, - entry="_load_model", - compute_dtype=torch.bfloat16, - # The loader casts to bfloat16, the dtype the FLUX.2 denoise runs in, whatever the device would pick. - patches_device=False, - sets_torch_dtype=False, - casts_fp8_storage=False, -) -FLUX2_QUANTIZED_LAYER = "control_transformer_blocks.0.attn.to_q" - - -def _flux2_checkpoint() -> dict[str, torch.Tensor]: - """A float32 file with the real four-block layout, under the names alibaba-pai's file uses.""" - torch.manual_seed(0) - branch = Flux2FunControlBranch(hidden_size=HIDDEN, num_attention_heads=1, attention_head_dim=128) - return { - name.replace("control_blocks.", "control_transformer_blocks.", 1): tensor.detach().clone() - for name, tensor in branch.state_dict().items() - } - - -def _flux2_driver(monkeypatch: pytest.MonkeyPatch, tmp_path: Path, state_dict: dict[str, torch.Tensor]): - checkpoint = tmp_path / "flux2_fun_controlnet.safetensors" - checkpoint.touch() - config = ControlNet_Checkpoint_Flux2_Config.model_construct(path=str(checkpoint), name="flux2-control") - return prepare(FLUX2_SEAM, monkeypatch, state_dict=state_dict), config - - -def test_a_dense_flux2_checkpoint_loads_with_its_geometry_in_bfloat16(monkeypatch, tmp_path) -> None: - run, config = _flux2_driver(monkeypatch, tmp_path, _flux2_checkpoint()) - - model = run.load(config) - - assert isinstance(model, Flux2FunControlBranch) - assert model.control_layers == FLUX2_FUN_CONTROL_LAYERS - first = model.control_blocks[0] - assert (model.hidden_size, model.control_in_dim, first.attn.heads, first.mlp_hidden_dim) == (HIDDEN, 260, 1, 384) - assert {parameter.dtype for parameter in model.parameters()} == {torch.bfloat16} - assert not model.training - - -def test_a_scaled_fp8_flux2_checkpoint_is_refused_rather_than_loaded_unscaled(monkeypatch, tmp_path) -> None: - payload = quantize_scaled_fp8(torch.randn(HIDDEN, HIDDEN)) - state_dict = _flux2_checkpoint() - state_dict[f"{FLUX2_QUANTIZED_LAYER}.weight"] = payload.codes - state_dict[f"{FLUX2_QUANTIZED_LAYER}.weight_scale"] = payload.scale - run, config = _flux2_driver(monkeypatch, tmp_path, state_dict) - - with pytest.raises(ValueError, match="quantization side channel"): - run.load(config) - - -def test_the_qwen_image_recipe_is_refused_by_the_flux2_loader(monkeypatch, tmp_path) -> None: - run, config = _flux2_driver(monkeypatch, tmp_path, _checkpoint()) - - with pytest.raises(ValueError, match="takes 132"): - run.load(config) diff --git a/tests/backend/model_manager/load/test_qwen_image_controlnet_loader.py b/tests/backend/model_manager/load/test_qwen_image_controlnet_loader.py new file mode 100644 index 00000000000..f6f59ff4e76 --- /dev/null +++ b/tests/backend/model_manager/load/test_qwen_image_controlnet_loader.py @@ -0,0 +1,250 @@ +"""What the Qwen-Image ControlNet loader makes of a checkpoint: which design it builds (the one the record names), the +geometry it reads off the tensors, the dtype it casts to, and the files it refuses rather than loads wrongly. + +Both designs build under `init_empty_weights` and assign the file's tensors, so an fp8 weight with its scale dropped +would become the parameter as it is (`test_z_image_control_adapter_loader.py`). No quantized build of either is +published (alibaba-pai and Comfy-Org ship dense bf16 files), so refusing costs nothing. +""" + +from pathlib import Path + +import pytest +import safetensors.torch +import torch +from diffusers.models.controlnets.controlnet_qwenimage import QwenImageControlNetModel + +from invokeai.backend.fun_controlnet.qwen_image import QWEN_IMAGE_FUN_CONTROL_LAYERS, QwenImageFunControlBranch +from invokeai.backend.model_manager.configs.controlnet import ControlNet_Checkpoint_QwenImage_Config +from invokeai.backend.model_manager.configs.main import Main_Checkpoint_QwenImage_Config +from invokeai.backend.model_manager.load.model_loaders import fun_controlnet, qwen_image, qwen_image_controlnet +from invokeai.backend.model_manager.load.model_loaders.qwen_image import QwenImageCheckpointModel +from invokeai.backend.model_manager.load.model_loaders.qwen_image_controlnet import ( + QWEN_IMAGE_AXES_DIMS_ROPE, + QwenImageControlNetLoader, +) +from tests.backend.model_manager.load.state_dicts.qwen_image_instantx_controlnet_keys import state_dict_keys +from tests.fixtures.loader_seams import Seam, prepare +from tests.fixtures.quantized_payloads import comfy_quant_marker, quantize_convrot, quantize_scaled_fp8 + +# One 128-wide head: the narrowest ControlNet the loaders' head rule accepts, with the real five-block layouts. +HIDDEN = 128 +QUANTIZED_LAYER = "control_blocks.0.attn.to_q" +INSTANTX_BLOCKS = 5 +INSTANTX_TEXT_DIM = 24 + +SEAM = Seam( + loader=QwenImageControlNetLoader, + module=qwen_image_controlnet, + entry="_load_model", + # The file is read by the dense-checkpoint reader the ControlNet loaders share. + load_file_host=fun_controlnet, + compute_dtype=torch.bfloat16, + # The loader takes its dtype from the device and never reads `self._torch_dtype`. + patches_device=True, + sets_torch_dtype=False, + casts_fp8_storage=False, +) + + +def _checkpoint() -> dict[str, torch.Tensor]: + """A float32 Fun-ControlNet file with the real layout; the loader has to cast it.""" + torch.manual_seed(0) + branch = QwenImageFunControlBranch(hidden_size=HIDDEN, num_attention_heads=1, attention_head_dim=128) + return {name: tensor.detach().clone() for name, tensor in branch.state_dict().items()} + + +def _instantx_checkpoint() -> dict[str, torch.Tensor]: + """A float32 InstantX file: diffusers' model at one 128-wide head, with every tensor random so a load that + dropped or swapped one shows.""" + torch.manual_seed(1) + model = QwenImageControlNetModel( + num_layers=INSTANTX_BLOCKS, + attention_head_dim=128, + num_attention_heads=1, + joint_attention_dim=INSTANTX_TEXT_DIM, + axes_dims_rope=QWEN_IMAGE_AXES_DIMS_ROPE, + ) + return {name: torch.randn_like(tensor) for name, tensor in model.state_dict().items()} + + +def _driver(monkeypatch: pytest.MonkeyPatch, tmp_path: Path, state_dict: dict[str, torch.Tensor], **record): + checkpoint = tmp_path / "qwen_image_controlnet.safetensors" + checkpoint.touch() + config = ControlNet_Checkpoint_QwenImage_Config.model_construct( + path=str(checkpoint), name="qwen-image-control", **record + ) + return prepare(SEAM, monkeypatch, state_dict=state_dict), config + + +def test_a_record_without_a_design_loads_the_fun_branch_with_its_geometry_in_the_compute_dtype( + monkeypatch, tmp_path +) -> None: + run, config = _driver(monkeypatch, tmp_path, _checkpoint()) + + model = run.load(config) + + assert isinstance(model, QwenImageFunControlBranch) + assert model.control_layers == QWEN_IMAGE_FUN_CONTROL_LAYERS + assert (model.hidden_size, model.control_in_dim, model.control_blocks[0].attn.heads) == (HIDDEN, 132, 1) + assert {parameter.dtype for parameter in model.parameters()} == {torch.bfloat16} + assert not model.training + + +def test_an_instantx_record_loads_every_tensor_into_the_controlnet_its_keys_describe(monkeypatch, tmp_path) -> None: + checkpoint = _instantx_checkpoint() + run, config = _driver(monkeypatch, tmp_path, checkpoint, control_architecture="instantx_union") + + model = run.load(config) + + assert isinstance(model, QwenImageControlNetModel) + assert (len(model.transformer_blocks), len(model.controlnet_blocks)) == (INSTANTX_BLOCKS, INSTANTX_BLOCKS) + assert (model.inner_dim, model.transformer_blocks[0].attn.heads) == (HIDDEN, 1) + assert (model.config.joint_attention_dim, model.config.in_channels) == (INSTANTX_TEXT_DIM, 64) + assert model.config.extra_condition_channels == 0 + assert not model.training + loaded = model.state_dict() + assert loaded.keys() == checkpoint.keys() + for name, tensor in checkpoint.items(): + assert torch.equal(loaded[name], tensor.to(torch.bfloat16)), name + + +def test_the_released_instantx_header_loads_strictly_into_the_controlnet_built_from_it(monkeypatch, tmp_path) -> None: + """The captured header widened to the file's five blocks (its 181 tensors), on the meta device so nothing of the + 3.5 GB is allocated: every name and shape lands, and the geometry is Qwen-Image's.""" + full = { + name.replace(".0.", f".{index}.", 1): spec + for name, spec in state_dict_keys.items() + for index in ( + range(INSTANTX_BLOCKS) if name.startswith(("transformer_blocks.0.", "controlnet_blocks.0.")) else [0] + ) + } + assert len(full) == 181 + with torch.device("meta"): + checkpoint = {name: torch.empty(shape, dtype=torch.bfloat16) for name, (shape, _dtype) in full.items()} + run, config = _driver(monkeypatch, tmp_path, checkpoint, control_architecture="instantx_union") + + model = run.load(config) + + assert isinstance(model, QwenImageControlNetModel) + assert {name: list(tensor.shape) for name, tensor in model.state_dict().items()} == { + name: shape for name, (shape, _dtype) in full.items() + } + assert (model.config.num_layers, model.config.num_attention_heads, model.config.joint_attention_dim) == ( + 5, + 24, + 3584, + ) + + +def test_the_record_decides_the_design_and_a_wrong_one_fails_by_name(monkeypatch, tmp_path) -> None: + run, config = _driver(monkeypatch, tmp_path, _instantx_checkpoint(), control_architecture="fun_union") + with pytest.raises(ValueError, match="control_img_in.weight"): + run.load(config) + + run, config = _driver(monkeypatch, tmp_path, _checkpoint(), control_architecture="instantx_union") + with pytest.raises(ValueError, match="Not an InstantX Qwen Image ControlNet"): + run.load(config) + + +def test_the_instantx_inpainting_layout_is_refused_by_name(monkeypatch, tmp_path) -> None: + checkpoint = _instantx_checkpoint() + checkpoint["controlnet_x_embedder.weight"] = torch.zeros(HIDDEN, 68) + run, config = _driver(monkeypatch, tmp_path, checkpoint, control_architecture="instantx_union") + + with pytest.raises(ValueError, match="inpainting layout"): + run.load(config) + + +@pytest.mark.parametrize( + ("rewrite", "message"), + [ + (lambda sd: {k: v for k, v in sd.items() if not k.startswith("transformer_blocks.4.")}, "missing"), + (lambda sd: {**sd, "controlnet_blocks.5.weight": torch.zeros(HIDDEN, HIDDEN)}, "missing"), + (lambda sd: {**sd, "pos_embed_input.weight": torch.zeros(HIDDEN, 64)}, "unexpected"), + (lambda sd: {**sd, "transformer_blocks.0.attn.norm_q.weight": torch.zeros(64)}, "attention heads"), + ], + ids=["fewer-blocks-than-projections", "more-projections-than-blocks", "stray-key", "other-head-width"], +) +def test_an_instantx_checkpoint_that_is_not_the_layout_is_refused(monkeypatch, tmp_path, rewrite, message) -> None: + run, config = _driver(monkeypatch, tmp_path, rewrite(_instantx_checkpoint()), control_architecture="instantx_union") + + with pytest.raises(ValueError, match=message): + run.load(config) + + +@pytest.mark.parametrize( + ("checkpoint", "layer", "architecture"), + [ + (_checkpoint, QUANTIZED_LAYER, "fun_union"), + (_instantx_checkpoint, "transformer_blocks.0.attn.to_q", "instantx_union"), + ], + ids=["fun", "instantx"], +) +def test_a_scaled_fp8_checkpoint_is_refused_rather_than_loaded_unscaled( + monkeypatch, tmp_path, checkpoint, layer, architecture +) -> None: + payload = quantize_scaled_fp8(torch.randn(HIDDEN, HIDDEN)) + state_dict = checkpoint() + state_dict[f"{layer}.weight"] = payload.codes + state_dict[f"{layer}.weight_scale"] = payload.scale + run, config = _driver(monkeypatch, tmp_path, state_dict, control_architecture=architecture) + + with pytest.raises(ValueError, match="quantization side channel"): + run.load(config) + + +def test_an_int8_checkpoint_is_refused_by_its_marker(monkeypatch, tmp_path) -> None: + codes, scale, _ = quantize_convrot(torch.randn(HIDDEN, 256)) + state_dict = _checkpoint() + state_dict[f"{QUANTIZED_LAYER}.weight"] = codes + state_dict[f"{QUANTIZED_LAYER}.weight_scale"] = scale + state_dict[f"{QUANTIZED_LAYER}.comfy_quant"] = comfy_quant_marker({"format": "int8_tensorwise", "convrot": True}) + run, config = _driver(monkeypatch, tmp_path, state_dict) + + with pytest.raises(ValueError, match="quantization side channel"): + run.load(config) + + +@pytest.mark.parametrize( + ("rewrite", "message"), + [ + (lambda sd: {k: v for k, v in sd.items() if not k.startswith("control_blocks.4.")}, "4 control blocks"), + (lambda sd: {**sd, "control_img_in.weight": torch.zeros(HIDDEN, 64)}, "takes 64"), + ], + ids=["four-blocks", "sixteen-channel-context"], +) +def test_another_fun_recipe_is_refused_by_name(monkeypatch, tmp_path, rewrite, message) -> None: + run, config = _driver(monkeypatch, tmp_path, rewrite(_checkpoint())) + + with pytest.raises(ValueError, match=message): + run.load(config) + + +def test_an_instantx_file_an_earlier_version_recorded_as_a_main_model_is_refused_with_the_fix( + monkeypatch, tmp_path +) -> None: + """Earlier versions identified the InstantX union as a Qwen Image main model, and such records outlive the fix + to identification: loading one as a transformer says what it is and how to repair the record, before any room + is reserved for a transformer that cannot be built.""" + with torch.device("meta"): + checkpoint = { + name: torch.empty(shape, dtype=torch.bfloat16) for name, (shape, _dtype) in state_dict_keys.items() + } + path = tmp_path / "Qwen-Image-InstantX-ControlNet-Union.safetensors" + path.touch() + seam = Seam( + loader=QwenImageCheckpointModel, + module=qwen_image, + entry="_load_from_singlefile", + load_file_host=safetensors.torch, + compute_dtype=torch.bfloat16, + patches_device=True, + sets_torch_dtype=False, + casts_fp8_storage=False, + ) + run = prepare(seam, monkeypatch, state_dict=checkpoint, metadata=None) + + with pytest.raises(ValueError, match="is a Qwen Image ControlNet, not a main model.*Re-identify"): + run.load(Main_Checkpoint_QwenImage_Config.model_construct(path=str(path), name="qwen-instantx")) + + assert run.reserved == [] diff --git a/tests/model_identification/stripped_models/10359709-71d3-4f27-9761-c975cd47cfb7/Qwen-Image-InstantX-ControlNet-Union.safetensors b/tests/model_identification/stripped_models/10359709-71d3-4f27-9761-c975cd47cfb7/Qwen-Image-InstantX-ControlNet-Union.safetensors new file mode 100644 index 00000000000..057eb1a2f79 --- /dev/null +++ b/tests/model_identification/stripped_models/10359709-71d3-4f27-9761-c975cd47cfb7/Qwen-Image-InstantX-ControlNet-Union.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:073abb19a85ac86b267e6865d6bd51ea11359cfc8d3788b066b0761e257ed28d +size 27457 diff --git a/tests/model_identification/stripped_models/10359709-71d3-4f27-9761-c975cd47cfb7/__test_metadata__.json b/tests/model_identification/stripped_models/10359709-71d3-4f27-9761-c975cd47cfb7/__test_metadata__.json new file mode 100644 index 00000000000..beac4aadea6 --- /dev/null +++ b/tests/model_identification/stripped_models/10359709-71d3-4f27-9761-c975cd47cfb7/__test_metadata__.json @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:cf049fe66e4cd5367fdaaa8950c7a2310759d6b999d006cff56ce98ec64bc958 +size 685