From 3250c61adb6dc227e07978da0ae99acb032fa7c6 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Mon, 19 Oct 2026 06:33:03 +0200 Subject: [PATCH 1/2] feat(qwen-image-2-1): add edits from reference images The prompt encoder reads up to four references through the vision tower; their latents fill the image slots. The encoder reserves memory for the references' attention; the denoise checks slot grids before loading. Reference images are offered only on Diffusers models, which keep the vision tower. --- .../03.Models/02.model-families.mdx | 2 +- .../03.Models/Local Models/qwen-image-2-1.mdx | 11 +- .../04.Image Generation/generate-panel.mdx | 2 +- .../contributing/new-model-integration.mdx | 4 +- .../docs/development/Guides/recall-api.mdx | 5 +- .../qwen_image_2_1/qwen_image_2_1_denoise.py | 194 ++++++++--- .../qwen_image_2_1_text_encoder.py | 89 ++++- .../vae/qwen_image_2_1_image_to_latents.py | 53 ++- .../backend/architectures/capabilities.py | 8 +- .../architectures/defs/qwen_image_2_1.py | 3 + .../backend/architectures/facets/features.py | 5 +- .../load/model_loaders/qwen_image_2_1.py | 20 +- .../backend/qwen_image_2_1/text_encoding.py | 110 +++++- invokeai/backend/qwen_image_2_1/vae.py | 41 ++- .../diffusion/conditioning_data.py | 10 + invokeai/frontend/api/openapi.json | 68 +++- invokeai/frontend/api/schema.ts | 33 +- .../architectureCapabilities.json | 33 +- .../core/architectureCapabilities.ts | 1 + .../core/baseGenerationPolicies.test.ts | 35 ++ .../generation/core/baseGenerationPolicies.ts | 33 +- .../features/generation/core/graph.test.ts | 60 ++++ .../src/features/generation/core/graph.ts | 73 ++++ .../src/features/generation/core/settings.ts | 4 + .../src/features/generation/core/types.ts | 1 + .../webv2/src/workbench/canvasMigration.ts | 1 + .../invocations/test_qwen_image_2_1_nodes.py | 314 +++++++++++++++++- tests/backend/architectures/test_features.py | 15 +- .../load/test_qwen3_vl_visual_tower.py | 21 ++ .../qwen_image_2_1/test_text_encoding.py | 97 +++++- tests/backend/qwen_image_2_1/test_vae.py | 36 +- 31 files changed, 1252 insertions(+), 130 deletions(-) diff --git a/docs/src/content/docs/Users Guide/03.Models/02.model-families.mdx b/docs/src/content/docs/Users Guide/03.Models/02.model-families.mdx index 757cd7a3545..aad924586b4 100644 --- a/docs/src/content/docs/Users Guide/03.Models/02.model-families.mdx +++ b/docs/src/content/docs/Users Guide/03.Models/02.model-families.mdx @@ -123,7 +123,7 @@ faster generation. | **Z-Image** | Images | 6B model in a fast, few-step Turbo variant and a Base variant. Understands English and Chinese prompts, with ControlNets. | Apache 2.0 | [Z-Image](/users-guide/models/local-models/z-image/) | | **ERNIE-Image** | Images | Baidu's 8B text-to-image model, in standard and distilled Turbo variants. | Apache 2.0 | [ERNIE-Image](/users-guide/models/local-models/ernie-image/) | | **Qwen Image** | Images | Text-to-image (2512) and instruction-based image editing (Edit 2511), with fast Lightning variants. | Apache 2.0 | [Qwen Image](/users-guide/models/local-models/qwen-image/) | -| **Qwen-Image-2.1** | Images | 7B text-to-image model with strong text rendering and transparent (RGBA) output, in a base and an 8-step Turbo variant. | Non-commercial | [Qwen-Image-2.1](/users-guide/models/local-models/qwen-image-2-1/) | +| **Qwen-Image-2.1** | Images | 7B text-to-image model with strong text rendering and transparent (RGBA) output, in a base and an 8-step Turbo variant. Edits from reference images. | Non-commercial | [Qwen-Image-2.1](/users-guide/models/local-models/qwen-image-2-1/) | | **Krea-2** | Images | 12B model in fast Turbo and undistilled Raw variants, with training-free style reference. | Community, free under $1M revenue | [Krea-2](/users-guide/models/local-models/krea-2/) | | **Ideogram 4** | Images | Structured prompts that place described elements in regions of the image. | Non-commercial | [Ideogram 4](/users-guide/models/local-models/ideogram4/) | | **Anima** | Images | Compact 2B anime-focused model, with ControlNet-LLLite adapters for inpainting and sketch guidance. | Non-commercial | [Anima](/users-guide/models/local-models/anima/) | diff --git a/docs/src/content/docs/Users Guide/03.Models/Local Models/qwen-image-2-1.mdx b/docs/src/content/docs/Users Guide/03.Models/Local Models/qwen-image-2-1.mdx index f7d9298c745..ecbb1190186 100644 --- a/docs/src/content/docs/Users Guide/03.Models/Local Models/qwen-image-2-1.mdx +++ b/docs/src/content/docs/Users Guide/03.Models/Local Models/qwen-image-2-1.mdx @@ -14,7 +14,7 @@ Two variants are available, and the variant is detected on install: - **Qwen-Image-2.1** — the base model. 40 steps by default. - **Qwen-Image-2.1 Turbo** — distilled for 8 steps on a fixed schedule of its own. -Text-to-image, image-to-image, inpainting and outpainting on the Canvas are supported. Reference images (image editing), LoRAs and regional guidance are not supported yet. +Text-to-image, image-to-image, inpainting and outpainting on the Canvas are supported, and with the Diffusers models also edits from [reference images](#reference-images). LoRAs and regional guidance are not supported yet. ### License @@ -57,6 +57,15 @@ The text encoder runs first and is unloaded from the GPU when the transformer ne The VAE needs about 8 GB of working memory to decode a 1024×1024 image, and is priced the same to encode one for image-to-image, inpainting or outpainting (twice that with `precision: float32`). When an image is too large for the card to handle in one piece, Invoke decodes and encodes it in tiles; see [Large decodes](/configuration/optimization/low-vram-mode/). Invoke picks 1024-pixel tiles, and 768 or 512 only on cards with less memory: measured at 2048×2048, 1024-pixel tiles came close to an untiled decode (47.7 dB PSNR), while the 256-pixel tiles other VAEs default to left visible seams (35.6 dB). +### Reference images + +With a Qwen-Image-2.1 **Diffusers** model, add up to four images with **Add reference image** in the Guidance section (see [Adding reference images](/users-guide/image-generation/generate-panel/#adding-reference-images)) and say in the prompt what to do with them, for example *put the jacket from the first image on the person in the second* or *the same cat, as a watercolor*. The model reads each reference twice: its Qwen3-VL encoder looks at it alongside the prompt, and the transformer sees it as clean latents next to the image it generates. + +- **Diffusers models only.** Only the encoder that comes with the Diffusers model keeps the Qwen3-VL vision tower, which reads the images. A single-file or GGUF transformer with a standalone encoder does not offer reference images, and on a Diffusers model the references refuse to run while a standalone **Qwen3-VL Encoder** is selected under **Components**; clear it to edit. +- **Size.** Each reference is read at about one megapixel at its own aspect ratio, whatever its size. The generated image takes the size set in the panel; **Use size** on a reference's thumbnail matches it. +- **Transparency.** The encoder sees a transparent reference over white, as the model was trained; its latents keep the transparency. +- **Memory.** The vision tower adds about 1.1 GB to the Diffusers model's encoder, on every generation with it. Each reference adds about 4,100 tokens to the start of the sequence, which the model keeps for the whole generation: about 2 GB per reference, twice that with **CFG** above 1. Where that would take more than half of the GPU, Invoke reprocesses the references at every step instead of keeping them, which uses less memory but takes noticeably longer. Reading the references also needs encoder memory that grows quickly with their number: about 7 GB beyond its weights for four, so on a 24 GB card part of the encoder then runs from system memory. + ### Transparent images The Qwen-Image-2.1 VAE decodes an alpha channel along with the colors, so a prompt such as *a sticker of a red fox on a transparent background* produces a cut-out image. An image is saved with transparency only when it has some: when more than 0.1% of its pixels are noticeably transparent. Opaque images are saved as ordinary RGB, so the faint edges tiled decoding leaves in the alpha channel never turn an opaque image transparent. diff --git a/docs/src/content/docs/Users Guide/04.Image Generation/generate-panel.mdx b/docs/src/content/docs/Users Guide/04.Image Generation/generate-panel.mdx index 1ec9cbe5742..b5f3de5a4ce 100644 --- a/docs/src/content/docs/Users Guide/04.Image Generation/generate-panel.mdx +++ b/docs/src/content/docs/Users Guide/04.Image Generation/generate-panel.mdx @@ -120,7 +120,7 @@ What a reference does, and what options it has, depends on the model family: | **SD 1.5, SDXL** | An IP-Adapter model carries the reference's look into the image. | **Mode**: **Style**, **Composition** or **Both**. **Weight**. Under **Advanced**: the IP-Adapter model, a **Style variant** and the **Active steps** range. | | **FLUX.1** | An IP-Adapter, or a **FLUX Redux** model, which transfers the reference's content. | IP-Adapter options as above, or **Image influence** (Lowest to Highest) for Redux. | | **FLUX.1 Kontext** | Kontext edits or recombines the reference images according to your prompt. | None beyond the model. | -| **FLUX.2, Qwen Image Edit** | The model reads the references directly and uses them as your prompt instructs, for example *put the jacket from the first image on the person in the second*. | None. | +| **FLUX.2, Qwen Image Edit, Qwen-Image-2.1** | The model reads the references directly and uses them as your prompt instructs, for example *put the jacket from the first image on the person in the second*. | None. | | **Krea-2** | One reference's look — palette, texture, rendering — is transferred into the image while the prompt keeps driving the content; see [Krea-2](/users-guide/models/local-models/krea-2/#style-reference). | **Style strength** (0–2; 0 turns the reference off). | Reference images require an extra model for SD 1.5, SDXL and FLUX.1 (an IP-Adapter or FLUX Redux model). These are available as starter models in the Model Manager, and an installed IP-Adapter is selected automatically. diff --git a/docs/src/content/docs/contributing/new-model-integration.mdx b/docs/src/content/docs/contributing/new-model-integration.mdx index 04e4530b2d8..aea40112611 100644 --- a/docs/src/content/docs/contributing/new-model-integration.mdx +++ b/docs/src/content/docs/contributing/new-model-integration.mdx @@ -884,8 +884,8 @@ Paths below are relative to `invokeai/frontend/webv2/src/features/generation/cor ### Reference images -- Backend: set `FeaturesFacet.max_reference_images` (and `reference_images_require_variant` if only one variant supports them). Wire the images into the text encoder and/or the denoise node, as the reference pipeline does. -- webv2: extend the reference-image config union (`types.ts`), `getDefaultReferenceImageConfig` and `getReferenceImageConfigSupported` (`baseGenerationPolicies.ts`), and wire the images in the builder. +- Backend: set `FeaturesFacet.max_reference_images`, plus `reference_images_require_variant` if only one variant supports them or `reference_images_require_format` if only one format does. Wire the images into the text encoder and/or the denoise node, as the reference pipeline does. +- webv2: extend the reference-image config union (`types.ts`); the config-type set and the normalize, clone and sync switches in `settings.ts`; the persisted union in `src/workbench/canvasMigration.ts`; and `getDefaultReferenceImageConfig` and `getReferenceImageConfigSupported` (`baseGenerationPolicies.ts`). Wire the images in the builder, and report any setting that keeps them from running in `getReferenceImageValidationReasons`, not only from the builder. ### Control and regional guidance diff --git a/docs/src/content/docs/development/Guides/recall-api.mdx b/docs/src/content/docs/development/Guides/recall-api.mdx index 8478a3fe4a9..771f1d917de 100644 --- a/docs/src/content/docs/development/Guides/recall-api.mdx +++ b/docs/src/content/docs/development/Guides/recall-api.mdx @@ -250,12 +250,13 @@ separate adapter model: - **FLUX.2 Klein** — built-in reference image support. - **FLUX Kontext** — reference image associated with the main model. - **Qwen Image Edit** — reference image associated with the main model. +- **Qwen-Image-2.1** — reference images read by its own encoder, on the Diffusers model only. Because there is no adapter model to resolve, these entries carry only `image_name`. When the frontend receives them, it picks the appropriate config flavor (`flux2_reference_image`, `flux_kontext_reference_image`, -or `qwen_image_reference_image`) based on the currently-selected main -model, matching the behavior of a manual drag-and-drop. +`qwen_image_reference_image` or `qwen_image_2_1_reference_image`) based on +the currently-selected main model, matching the behavior of a manual drag-and-drop. ## Usage Examples diff --git a/invokeai/app/invocations/qwen_image_2_1/qwen_image_2_1_denoise.py b/invokeai/app/invocations/qwen_image_2_1/qwen_image_2_1_denoise.py index 9c428c8a729..7f69766c476 100644 --- a/invokeai/app/invocations/qwen_image_2_1/qwen_image_2_1_denoise.py +++ b/invokeai/app/invocations/qwen_image_2_1/qwen_image_2_1_denoise.py @@ -1,3 +1,4 @@ +import itertools import math from typing import Optional @@ -63,6 +64,61 @@ def clip_sigmas( return list(sigmas[first_step : last_step + 1]), first_step +def slot_runs(image_pad_mask: torch.Tensor | None) -> list[int]: + """The lengths of the runs of image slots in a prompt's mask: one run per reference image, in order.""" + if image_pad_mask is None: + return [] + return [len(list(run)) for is_slot, run in itertools.groupby(image_pad_mask[0].tolist()) if is_slot] + + +def _describe_grids(grids: list[tuple[int, int]]) -> str: + # A slot is 32x32 pixels of the reference as the encoder read it. + return ", ".join(f"{columns * 32}x{rows * 32}" for rows, columns in grids) or "none" + + +def check_reference_slots( + info: QwenImage21ConditioningInfo, reference_shapes: list[tuple[int, int, int]], prompt: str +) -> None: + """Refuse references whose latents do not match the prompt's image slots, reference by reference. + + Each reference's latents fill one run of the prompt's slots, 2x2 latents per slot, at the slot grid the + encoder read it at. A mismatch means the prompt was encoded with other images, in another order or at another + size than the latents: the transformer would place one reference's latents in another's slots. + """ + expected = [(h // 2, w // 2) for _, h, w in reference_shapes] + found = list(info.reference_grids) + if found != expected or slot_runs(info.image_pad_mask) != [rows * columns for rows, columns in found]: + raise ValueError( + f"The {prompt} prompt was encoded with {len(found)} reference image(s) ({_describe_grids(found)}), but " + f"the denoise node received {len(expected)} reference latent(s) ({_describe_grids(expected)}). Encode " + "each reference with Image to Latents in reference mode, from the same images and in the same order " + "as the prompts." + ) + + +def prefix_length(info: QwenImage21ConditioningInfo) -> int: + """Tokens before the target in the joint sequence: the prompt, each reference slot widened to 2x2 latents.""" + slots = 0 if info.image_pad_mask is None else int(info.image_pad_mask.sum()) + return info.prompt_embeds.shape[1] + (TOKENS_PER_SLOT - 1) * slots + + +# K and V of one prefix token in all 32 blocks: 2 x 4096 x 32 x 2 bytes in bf16. +KV_BYTES_PER_PREFIX_TOKEN = 2 * 4096 * 32 * 2 + + +def prefix_cache_fits(cache_bytes: int, device: torch.device) -> bool: + """Whether the prefix caches can stay on the GPU for the whole run. + + They cannot be offloaded the way the transformer's weights can, so caches larger than half the card would + leave too little for the model however it streams: four references at CFG above 1 hold ~16 GiB. Past that the + prefix is recomputed every step instead -- slower, but within the memory of an ordinary run. Memory that is + not a dedicated card's (CPU, MPS) is not limited here. + """ + if device.type != "cuda": + return True + return cache_bytes <= torch.cuda.get_device_properties(device).total_memory // 2 + + def pack_latents(latents: torch.Tensor) -> torch.Tensor: """(B, C, h, w) -> (B, h*w, C): one token per latent pixel, in raster order (patch size 1).""" batch, channels, height, width = latents.shape @@ -80,7 +136,7 @@ def unpack_latents(latents: torch.Tensor, height: int, width: int) -> torch.Tens title="Denoise - Qwen-Image-2.1", tags=["image", "qwen_image_2_1", "qwen-image-2.1"], category="image", - version="1.0.1", + version="1.1.0", classification=Classification.Prototype, ) class QwenImage21DenoiseInvocation(BaseInvocation, WithMetadata, WithBoard): @@ -92,6 +148,13 @@ class QwenImage21DenoiseInvocation(BaseInvocation, WithMetadata, WithBoard): denoise_mask: Optional[DenoiseMaskField] = InputField( default=None, description=FieldDescriptions.denoise_mask, input=Input.Connection ) + reference_latents: LatentsField | list[LatentsField] | None = InputField( + default=None, + description="Latents of the reference images an edit reads, from Image to Latents in reference mode, in the " + "order the prompts were encoded with them.", + input=Input.Connection, + title="Reference Latents", + ) denoising_start: float = InputField(default=0.0, ge=0, le=1, description=FieldDescriptions.denoising_start) denoising_end: float = InputField(default=1.0, ge=0, le=1, description=FieldDescriptions.denoising_end) transformer: TransformerField = InputField( @@ -135,11 +198,30 @@ def invoke(self, context: InvocationContext) -> LatentsOutput: def _load_conditioning( self, context: InvocationContext, field: QwenImage21ConditioningField, dtype: torch.dtype, device: torch.device - ) -> torch.Tensor: + ) -> QwenImage21ConditioningInfo: data = context.conditioning.load(field.conditioning_name) if len(data.conditionings) != 1 or not isinstance(data.conditionings[0], QwenImage21ConditioningInfo): raise ValueError("Expected exactly one Qwen-Image-2.1 conditioning.") - return data.conditionings[0].to(dtype=dtype, device=device).prompt_embeds + return data.conditionings[0].to(dtype=dtype, device=device) + + def _load_references(self, context: InvocationContext, device: torch.device) -> list[torch.Tensor]: + """The reference latents, `(1, C, h, w)` each, in the order the prompts were encoded with them.""" + if self.reference_latents is None: + return [] + fields = self.reference_latents if isinstance(self.reference_latents, list) else [self.reference_latents] + references = [] + for field in fields: + reference = context.tensors.load(field.latents_name).to(device=device, dtype=torch.float32) + if reference.dim() == 5: + reference = reference.squeeze(2) + if reference.shape[1] != LATENT_CHANNELS or reference.shape[-2] % 2 or reference.shape[-1] % 2: + raise ValueError( + f"Reference latents are {tuple(reference.shape[1:])}; Qwen-Image-2.1 reads {LATENT_CHANNELS} " + "channels with even sides. Encode the reference with Image to Latents - Qwen-Image-2.1 in " + "reference mode." + ) + references.append(reference) + return references def _prepare_cfg_scale(self, num_steps: int) -> list[float]: if isinstance(self.cfg_scale, list): @@ -174,7 +256,6 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: device = TorchDevice.choose_torch_device() dtype = TorchDevice.choose_bfloat16_safe_dtype(device) transformer_config = context.models.get_config(self.transformer.transformer) - transformer_info = context.models.load(self.transformer.transformer) latent_height, latent_width = self.height // VAE_SCALE_FACTOR, self.width // VAE_SCALE_FACTOR image_seq_len = latent_height * latent_width @@ -190,13 +271,26 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: sigmas = torch.tensor(sigmas, dtype=torch.float32, device=device) cfg_scale = self._prepare_cfg_scale(len(full_sigmas) - 1)[first_step : first_step + len(sigmas) - 1] - pos_embeds = self._load_conditioning(context, self.positive_conditioning, dtype, device) + pos = self._load_conditioning(context, self.positive_conditioning, dtype, device) do_cfg = self.negative_conditioning is not None and any(value > 1.0 for value in cfg_scale) - neg_embeds = ( + neg = ( self._load_conditioning(context, self.negative_conditioning, dtype, device) if do_cfg and self.negative_conditioning is not None else None ) + pos_embeds = pos.prompt_embeds + neg_embeds = neg.prompt_embeds if neg is not None else None + + references = self._load_references(context, device) + reference_shapes = [(1, r.shape[-2], r.shape[-1]) for r in references] + check_reference_slots(pos, reference_shapes, "positive") + if neg is not None: + check_reference_slots(neg, reference_shapes, "negative") + # Clean latents, condition images first: the transformer drops them into the prompt's image slots. + reference_tokens = torch.cat([pack_latents(r) for r in references], dim=1).to(dtype) if references else None + + # Only now, with every input checked: loading reads the whole transformer into RAM. + transformer_info = context.models.load(self.transformer.transformer) # Drawn as the pipeline draws it, (B, frame, C, h, w), on the CPU so a seed means the same on any device. noise = torch.randn( @@ -234,29 +328,41 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: # upcast); matching that is what makes this node reproduce its outputs. latents = pack_latents(latents).to(dtype) - def image_slot_mask(text_len: int) -> torch.Tensor: - # The joint sequence's slots: the text, then one slot per 2x2 group of target latents. Condition - # images would sit between the two. - text = torch.zeros(1, text_len, dtype=torch.bool, device=device) + def image_slot_mask(info: QwenImage21ConditioningInfo) -> torch.Tensor: + # The joint sequence's slots: the prompt with its reference slots, then one slot per 2x2 group of + # target latents. + text_len = info.prompt_embeds.shape[1] + prompt = ( + info.image_pad_mask.bool() + if info.image_pad_mask is not None + else torch.zeros(1, text_len, dtype=torch.bool, device=device) + ) target = torch.ones(1, image_seq_len // TOKENS_PER_SLOT, dtype=torch.bool, device=device) - return torch.cat([text, target], dim=1) - - img_shapes = [[(1, latent_height, latent_width)]] - pos_img_mask = image_slot_mask(pos_embeds.shape[1]) - neg_img_mask = image_slot_mask(neg_embeds.shape[1]) if neg_embeds is not None else None - + return torch.cat([prompt, target], dim=1) + + img_shapes = [[*reference_shapes, (1, latent_height, latent_width)]] + pos_img_mask = image_slot_mask(pos) + neg_img_mask = image_slot_mask(neg) if neg is not None else None + + # Each reference slot widens to 2x2 latent tokens in the joint sequence; with the text they are the + # prefix the cache holds. + pos_prefix = prefix_length(pos) + neg_prefix = prefix_length(neg) if neg is not None else None + # Text and reference keys and values do not depend on the step under `causal_condition`: the first step + # prefills them, the rest compute the image tokens only -- when the caches fit. + cache_bytes = KV_BYTES_PER_PREFIX_TOKEN * (pos_prefix + (neg_prefix or 0)) + use_cache = bool(getattr(transformer_info.model.config, "causal_condition", False)) + if use_cache and not prefix_cache_fits(cache_bytes, device): + use_cache = False + logger.info( + f"Qwen-Image-2.1: the prefix caches would hold {cache_bytes / 1024**3:.1f} GiB, more than half of " + "this GPU; recomputing the prefix every step instead, which is slower." + ) working_memory = self._estimate_working_memory( - image_seq_len, - pos_embeds.shape[1], - None if neg_embeds is None else neg_embeds.shape[1], - do_cfg, - dtype.itemsize, + image_seq_len, pos_prefix, neg_prefix, do_cfg, use_cache, dtype.itemsize ) working_memory += self._attention_score_bytes( - transformer_info.model, - image_seq_len + max(pos_embeds.shape[1], 0 if neg_embeds is None else neg_embeds.shape[1]), - device, - dtype, + transformer_info.model, image_seq_len, max(pos_prefix, neg_prefix or 0), device, dtype ) working_memory += peak_dequant_transient_bytes(transformer_info.model, dtype) log_attention_backends(logger, device) @@ -274,16 +380,15 @@ def image_slot_mask(text_len: int) -> torch.Tensor: ) with transformer_info.model_on_device(working_mem_bytes=working_memory) as (_, transformer): - # Text keys and values do not depend on the step under `causal_condition`: the first step prefills - # them, the rest compute the image tokens only. One cache per guidance branch, for this run only. + # One cache per guidance branch, for this run only. num_blocks = len(transformer.transformer_blocks) - use_cache = bool(transformer.config.causal_condition) pos_cache = QwenImage21KVCache(num_blocks) if use_cache else None neg_cache = QwenImage21KVCache(num_blocks) if use_cache and neg_embeds is not None else None def predict(embeds: torch.Tensor, img_mask: torch.Tensor, cache, mode, timestep) -> torch.Tensor: + hidden_states = latents if reference_tokens is None else torch.cat([reference_tokens, latents], dim=1) out = transformer( - hidden_states=latents, + hidden_states=hidden_states, encoder_hidden_states=embeds, encoder_hidden_states_mask=None, timestep=timestep, @@ -335,26 +440,36 @@ def predict(embeds: torch.Tensor, img_mask: torch.Tensor, cache, mode, timestep) @staticmethod def _estimate_working_memory( - image_seq_len: int, pos_text_len: int, neg_text_len: int | None, do_cfg: bool, element_size: int + image_seq_len: int, + pos_prefix: int, + neg_prefix: int | None, + do_cfg: bool, + use_cache: bool, + element_size: int, ) -> int: """Peak transformer activations, in bytes, for the model cache to keep free. Measured on an RTX 4090 in bf16 with the transformer fully resident: 0.55 GiB of activations at 1024x1024 (4096 image tokens) and 2.2 GiB at 2048x2048 (16384), about 0.14 MiB per token. 0.25 MiB per token plus a 1 GiB base leaves room for the allocator and the prefix cache, which holds K and V of every - text token in all 32 blocks (~0.5 MiB per text token, per guidance branch). Both per-token terms are bf16 - figures and scale with `element_size`, the compute dtype's width: a float32 fallback needs twice as much. + prefix token -- text and reference latents -- in all 32 blocks (~0.5 MiB per token, per guidance branch). + Both per-token terms are bf16 figures and scale with `element_size`, the compute dtype's width: a + float32 fallback needs twice as much. """ mib, gib = 1024**2, 1024**3 width = element_size / 2 - text_len = max(pos_text_len, neg_text_len or 0) - estimated = int((image_seq_len + text_len) * 0.25 * mib * width) + gib - branches = 2 if do_cfg else 1 - estimated += int(branches * (pos_text_len if neg_text_len is None else text_len) * 0.5 * mib * width) + prefix = max(pos_prefix, neg_prefix or 0) + estimated = int((image_seq_len + prefix) * 0.25 * mib * width) + gib + if use_cache: + branches = 2 if do_cfg else 1 + prefix_tokens = branches * (pos_prefix if neg_prefix is None else prefix) + estimated += int(prefix_tokens * KV_BYTES_PER_PREFIX_TOKEN * width) return estimated @staticmethod - def _attention_score_bytes(transformer: object, seq_len: int, device: torch.device, dtype: torch.dtype) -> int: + def _attention_score_bytes( + transformer: object, image_seq_len: int, prefix: int, device: torch.device, dtype: torch.dtype + ) -> int: """The score matrix attention materializes on a build without a fused kernel for it (0 where it has one).""" config = getattr(transformer, "config", None) num_heads = getattr(config, "num_attention_heads", None) @@ -363,13 +478,14 @@ def _attention_score_bytes(transformer: object, seq_len: int, device: torch.devi return 0 # The largest call, the image queries over every key, goes unmasked both in the prefill and on decode steps: # only the text segments carry a mask, and they are as short as the prompt. All of them go through diffusers' - # attention dispatch. + # attention dispatch. Its score matrix is image x (prefix + image), which the square helper prices at the + # side with the same area. return sdpa_score_matrix_bytes( device=device, dtype=dtype, num_heads=num_heads, head_dim=head_dim, - seq_len=seq_len, + seq_len=math.isqrt(image_seq_len * (image_seq_len + prefix)), via_diffusers_dispatch=True, ) diff --git a/invokeai/app/invocations/text_encoder/qwen_image_2_1_text_encoder.py b/invokeai/app/invocations/text_encoder/qwen_image_2_1_text_encoder.py index 0bea270aa49..7b59a30b17e 100644 --- a/invokeai/app/invocations/text_encoder/qwen_image_2_1_text_encoder.py +++ b/invokeai/app/invocations/text_encoder/qwen_image_2_1_text_encoder.py @@ -1,56 +1,115 @@ +from contextlib import nullcontext + import torch from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, invocation -from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField, UIComponent -from invokeai.app.invocations.model import Qwen3VLEncoderField +from invokeai.app.invocations.fields import FieldDescriptions, ImageField, Input, InputField, UIComponent +from invokeai.app.invocations.model import ModelIdentifierField, Qwen3VLEncoderField from invokeai.app.invocations.primitives import QwenImage21ConditioningOutput from invokeai.app.services.shared.invocation_context import InvocationContext from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device +from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelType, SubModelType from invokeai.backend.quantization.dequantizing_linear import peak_dequant_transient_bytes -from invokeai.backend.qwen_image_2_1.text_encoding import encode_prompt +from invokeai.backend.qwen_image_2_1.text_encoding import ( + encode_prompt, + has_vision_tower, + sequence_length, + working_memory_bytes, +) +from invokeai.backend.qwen_image_2_1.vae import to_reference from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ( ConditioningFieldData, QwenImage21ConditioningInfo, ) from invokeai.backend.util.devices import TorchDevice +_NO_VISION_TOWER = ( + "Reference images are read by the Qwen3-VL vision tower, which only the Qwen-Image-2.1 Diffusers model's own " + "encoder brings. Clear the standalone encoder, or use a Qwen-Image-2.1 Diffusers model." +) + @invocation( "qwen_image_2_1_text_encoder", title="Prompt - Qwen-Image-2.1", tags=["prompt", "conditioning", "qwen_image_2_1", "qwen-image-2.1"], category="conditioning", - version="1.0.0", + version="1.1.0", classification=Classification.Prototype, idle_gpu_offloadable=True, ) class QwenImage21TextEncoderInvocation(BaseInvocation): - """Encodes a text prompt for Qwen-Image-2.1 with its Qwen3-VL-8B encoder.""" + """Encodes a prompt for Qwen-Image-2.1 with its Qwen3-VL-8B encoder, with the reference images an edit reads. + + Each reference also goes to the denoise node, encoded by Image to Latents in its reference mode, in the same + order. + """ prompt: str = InputField(description="Text prompt describing the desired image.", ui_component=UIComponent.Textarea) qwen3_vl_encoder: Qwen3VLEncoderField = InputField( title="Qwen3-VL Encoder", description=FieldDescriptions.qwen3_vl_encoder, input=Input.Connection ) + reference_images: list[ImageField] = InputField( + default=[], + description="Reference images to edit or draw from, in the order the denoise node receives their latents. " + "Needs a Qwen-Image-2.1 Diffusers model's own encoder.", + ) @torch.no_grad() def invoke(self, context: InvocationContext) -> QwenImage21ConditioningOutput: - tokenizer_info = context.models.load(self.qwen3_vl_encoder.tokenizer) + images = [to_reference(context.images.get_pil(field.image_name)) for field in self.reference_images] + processor_id = self._processor_identifier(context) if images else None + + # The encoder first: loading it evicts unlocked small entries from the RAM cache, which would otherwise drop + # the tokenizer and processor loaded before it. text_encoder_info = context.models.load(self.qwen3_vl_encoder.text_encoder) + tokenizer_info = context.models.load(self.qwen3_vl_encoder.tokenizer) + processor_info = context.models.load(processor_id) if processor_id is not None else None # A GGUF or int8 encoder dequantizes one layer at a time; that transient is the working memory. dequant_bytes = peak_dequant_transient_bytes( text_encoder_info.model, TorchDevice.choose_bfloat16_safe_dtype(text_encoder_info.compute_device) ) + # Before the encoder moves to the GPU: a graph can pair this pipeline's tokenizer with another encoder. + if images and not has_vision_tower(text_encoder_info.model): + raise ValueError(_NO_VISION_TOWER) context.util.signal_progress("Running Qwen3-VL text encoder") - with ( - tokenizer_info as tokenizer, - text_encoder_info.model_on_device(working_mem_bytes=dequant_bytes) as ( - _, - text_encoder, - ), - ): - embeds = encode_prompt(text_encoder, tokenizer, self.prompt, get_effective_device(text_encoder)) + with tokenizer_info as tokenizer: + # Reference images make the sequence long, and its attention grows with the square of it: room the + # model cache has to make, or the driver spills the encode to system memory and it crawls. + working_memory = dequant_bytes + working_memory_bytes(sequence_length(tokenizer, self.prompt, images)) + with ( + processor_info.model_on_device() if processor_info is not None else nullcontext((None, None)) as ( + _, + processor, + ), + text_encoder_info.model_on_device(working_mem_bytes=working_memory) as (_, text_encoder), + ): + encoding = encode_prompt( + text_encoder, + tokenizer, + self.prompt, + get_effective_device(text_encoder), + processor=processor, + images=images, + ) - conditioning = QwenImage21ConditioningInfo(prompt_embeds=embeds.detach().cpu()) + conditioning = QwenImage21ConditioningInfo( + prompt_embeds=encoding.embeds.detach().cpu(), + image_pad_mask=None if encoding.image_pad_mask is None else encoding.image_pad_mask.cpu(), + reference_grids=encoding.reference_grids, + ) name = context.conditioning.save(ConditioningFieldData(conditionings=[conditioning])) return QwenImage21ConditioningOutput.build(name) + + def _processor_identifier(self, context: InvocationContext) -> ModelIdentifierField: + """The image processor of the model the tokenizer came from, which has to be a Qwen-Image-2.1 pipeline.""" + tokenizer = self.qwen3_vl_encoder.tokenizer + config = context.models.get_config(tokenizer) + if ( + config.base is not BaseModelType.QwenImage21 + or config.type is not ModelType.Main + or config.format is not ModelFormat.Diffusers + ): + raise ValueError(_NO_VISION_TOWER) + return tokenizer.model_copy(update={"submodel_type": SubModelType.Processor}) diff --git a/invokeai/app/invocations/vae/qwen_image_2_1_image_to_latents.py b/invokeai/app/invocations/vae/qwen_image_2_1_image_to_latents.py index 991e744184d..e8aa7d57e9b 100644 --- a/invokeai/app/invocations/vae/qwen_image_2_1_image_to_latents.py +++ b/invokeai/app/invocations/vae/qwen_image_2_1_image_to_latents.py @@ -14,7 +14,13 @@ from invokeai.app.invocations.model import VAEField from invokeai.app.invocations.primitives import LatentsOutput from invokeai.app.services.shared.invocation_context import InvocationContext -from invokeai.backend.qwen_image_2_1.vae import SPATIAL_SCALE, choose_tile_size, to_vae_input, working_memory_bytes +from invokeai.backend.qwen_image_2_1.vae import ( + SPATIAL_SCALE, + choose_tile_size, + to_reference, + to_vae_input, + working_memory_bytes, +) from invokeai.backend.util.devices import TorchDevice from invokeai.backend.util.qwen_image_vae import patch_qwen_image_vae_tiling @@ -27,26 +33,41 @@ title="Image to Latents - Qwen-Image-2.1", tags=["image", "latents", "vae", "i2l", "qwen_image_2_1"], category="image", - version="1.0.0", + version="1.1.0", classification=Classification.Prototype, ) class QwenImage21ImageToLatentsInvocation(BaseInvocation, WithMetadata, WithBoard): - """Encodes an image into Qwen-Image-2.1 latents. The image is encoded opaque; its alpha is not read.""" + """Encodes an image into Qwen-Image-2.1 latents. + + For image-to-image and the canvas, the image is encoded opaque at its own size; its alpha is not read. As a + reference for an edit, it is resized to the area the text encoder reads it at and keeps its alpha. + """ image: ImageField = InputField(description="The image to encode.") vae: VAEField = InputField(description=FieldDescriptions.vae, input=Input.Connection) tiled: bool = InputField(default=False, description=FieldDescriptions.tiled) tile_size: int = InputField(default=0, ge=0, multiple_of=16, description=FieldDescriptions.vae_tile_size) + reference: bool = InputField( + default=False, + description="Encode as a reference image for the denoise node's reference latents: resized to ~1 megapixel, " + "as the text encoder reads it, with its alpha kept.", + ) @torch.no_grad() def invoke(self, context: InvocationContext) -> LatentsOutput: image = context.images.get_pil(self.image.image_name) - # Crop-free resize onto the 32px grid the transformer needs, as the canvas sizes its own images. - width, height = image.width // _GRID * _GRID, image.height // _GRID * _GRID - if width == 0 or height == 0: - raise ValueError(f"Qwen-Image-2.1 encodes images of at least {_GRID}x{_GRID}; this one is {image.size}.") - if (width, height) != image.size: - image = image.resize((width, height), resample=PILImage.LANCZOS) + if self.reference: + image = to_reference(image) + width, height = image.size + else: + # Crop-free resize onto the 32px grid the transformer needs, as the canvas sizes its own images. + width, height = image.width // _GRID * _GRID, image.height // _GRID * _GRID + if width == 0 or height == 0: + raise ValueError( + f"Qwen-Image-2.1 encodes images of at least {_GRID}x{_GRID}; this one is {image.size}." + ) + if (width, height) != image.size: + image = image.resize((width, height), resample=PILImage.LANCZOS) vae_info = context.models.load(self.vae.vae) device = vae_info.compute_device @@ -66,12 +87,16 @@ def invoke(self, context: InvocationContext) -> LatentsOutput: working_memory = working_memory_bytes(height, width, tile_size, element_size) with vae_info.model_on_device(working_mem_bytes=working_memory) as (_, vae): - pixels = to_vae_input(image).to(device=device, dtype=vae.dtype).unsqueeze(2) + pixels = to_vae_input(image, keep_alpha=self.reference).to(device=device, dtype=vae.dtype).unsqueeze(2) with torch.inference_mode(), patch_qwen_image_vae_tiling(vae, tile_size, SPATIAL_SCALE): - z = vae.encode(pixels).latent_dist.mode().float()[:, :, 0] - mean = torch.tensor(vae.config.latents_mean, device=z.device).view(1, -1, 1, 1) - std = torch.tensor(vae.config.latents_std, device=z.device).view(1, -1, 1, 1) - latents = ((z - mean) / std).cpu() + z = vae.encode(pixels).latent_dist.mode()[:, :, 0] + # A reference is normalized in the VAE's dtype, as the pipeline normalizes it, so the transformer reads + # the same values; the canvas's latents only seed a sampler, and keep fp32. + z = z if self.reference else z.float() + # float32 first, then the latents' dtype: straight from Python floats to bf16 can round differently. + mean = torch.tensor(vae.config.latents_mean).to(z.device, z.dtype).view(1, -1, 1, 1) + std = torch.tensor(vae.config.latents_std).to(z.device, z.dtype).view(1, -1, 1, 1) + latents = ((z - mean) / std).float().cpu() TorchDevice.empty_cache() name = context.tensors.save(tensor=latents) diff --git a/invokeai/backend/architectures/capabilities.py b/invokeai/backend/architectures/capabilities.py index cbfccd56192..e72665da0d4 100644 --- a/invokeai/backend/architectures/capabilities.py +++ b/invokeai/backend/architectures/capabilities.py @@ -72,6 +72,10 @@ class ArchitectureFeatures(BaseModel): default=None, description="If set, reference images are only accepted for models of this variant.", ) + reference_images_require_format: str | None = Field( + default=None, + description="If set, reference images are only accepted for models of this format.", + ) supports_regional_guidance: bool = False regional_negative: bool = Field( default=False, description="Whether a region's negative prompt is masked, rather than applied globally." @@ -137,6 +141,7 @@ def _features_of(facet: FeaturesFacet, variant: AnyVariant | None = None) -> Arc control_kinds=sorted(facet.control_kinds), max_reference_images=facet.max_reference_images, reference_images_require_variant=facet.reference_images_require_variant, + reference_images_require_format=facet.reference_images_require_format, supports_regional_guidance=facet.supports_regional_guidance, regional_negative=facet.regional_negative, clip_skip_max=facet.clip_skip_max, @@ -171,7 +176,8 @@ def architecture_capabilities() -> list[ArchitectureCapabilities]: is rendered in full, so a client never has to know which fields a variant row is allowed to omit. Differences too small to have earned a mapping are expressed on the base row instead, by - `features.reference_images_require_variant`; Qwen-Image is the only one. + `features.reference_images_require_variant` (Qwen-Image) and + `features.reference_images_require_format` (Qwen-Image-2.1). Sorted by base value, then variant, so the response is stable and diffable. """ diff --git a/invokeai/backend/architectures/defs/qwen_image_2_1.py b/invokeai/backend/architectures/defs/qwen_image_2_1.py index 0a523835091..8dea696ac0b 100644 --- a/invokeai/backend/architectures/defs/qwen_image_2_1.py +++ b/invokeai/backend/architectures/defs/qwen_image_2_1.py @@ -34,6 +34,9 @@ dimension_grid=32, guidance_label="CFG", scheduler_set="flow", + # Each reference holds ~1 megapixel of latents in the prefix the cache keeps: ~2 GiB per guidance branch. + max_reference_images=4, + reference_images_require_format="diffusers", ), # Values are qwen_image_21_* because variant strings resolve without the base and must be unique. VariantFacet({ModelType.Main: QwenImage21VariantType}), diff --git a/invokeai/backend/architectures/facets/features.py b/invokeai/backend/architectures/facets/features.py index 0cae3669ef3..95887564fad 100644 --- a/invokeai/backend/architectures/facets/features.py +++ b/invokeai/backend/architectures/facets/features.py @@ -111,8 +111,11 @@ class FeaturesFacet(Facet): control_kinds: frozenset[ControlKind] = frozenset() max_reference_images: int = 0 reference_images_require_variant: str | None = None - """Qwen-Image accepts reference images only as the `edit` variant — the one feature in this + """Qwen-Image accepts reference images only as the `edit` variant — one of the two features in this table that a base alone cannot answer.""" + reference_images_require_format: str | None = None + """The other: Qwen-Image-2.1 reads reference images only as a Diffusers pipeline, whose encoder is + the only one that keeps the vision tower.""" supports_regional_guidance: bool = False regional_negative: bool = False diff --git a/invokeai/backend/model_manager/load/model_loaders/qwen_image_2_1.py b/invokeai/backend/model_manager/load/model_loaders/qwen_image_2_1.py index f94637da230..08ff7d29996 100644 --- a/invokeai/backend/model_manager/load/model_loaders/qwen_image_2_1.py +++ b/invokeai/backend/model_manager/load/model_loaders/qwen_image_2_1.py @@ -18,7 +18,6 @@ from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry from invokeai.backend.model_manager.load.model_loaders.generic_diffusers import GenericDiffusersLoader from invokeai.backend.model_manager.taxonomy import AnyModel, BaseModelType, ModelFormat, ModelType, SubModelType -from invokeai.backend.model_manager.util.qwen3_vl import drop_qwen3vl_visual_tower from invokeai.backend.quantization.fp8_scaled import ( attach_fp8_scales, cast_state_dict, @@ -62,10 +61,13 @@ def _build_transformer(sd: dict[str, Any]) -> torch.nn.Module: class QwenImage21DiffusersModel(GenericDiffusersLoader): """Loads the submodels of a Qwen-Image-2.1 diffusers pipeline. - The tokenizer lives in `processor/` (the pipeline declares a `Qwen3VLProcessor`, of which text prompts - need only the tokenizer). The text encoder loads as `Qwen3VLModel` rather than the declared - `Qwen3VLForConditionalGeneration`: the LM head is ~1.2 GiB the encoder never runs, and `Qwen3VLModel` is - what the standalone Qwen3-VL loaders return, so the text encoder node sees one class. + The tokenizer lives in `processor/` (the pipeline declares a `Qwen3VLProcessor`); the processor itself, which + also cuts reference images into patches, loads only for an edit. The text encoder loads as `Qwen3VLModel` + rather than the declared `Qwen3VLForConditionalGeneration`: the LM head is ~1.2 GiB the encoder never runs, + and `Qwen3VLModel` is what the standalone Qwen3-VL loaders return, so the text encoder node sees one class. + + Unlike those loaders, this one keeps the vision tower (~1.1 GiB): reference images reach the encoder only + through it, and only this pipeline's encoder carries one into Qwen-Image-2.1. """ def _load_model(self, config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> AnyModel: @@ -79,6 +81,10 @@ def _load_model(self, config: AnyModelConfig, submodel_type: Optional[SubModelTy model_path = Path(config.path) if submodel_type is SubModelType.Tokenizer: return AutoTokenizer.from_pretrained(model_path / "processor", local_files_only=True) + if submodel_type is SubModelType.Processor: + from transformers import AutoProcessor + + return AutoProcessor.from_pretrained(model_path / "processor", local_files_only=True) if submodel_type is SubModelType.VAE: # The configured precision, except float16, which is unstable on this autoencoder: bfloat16 where the @@ -108,10 +114,6 @@ def _load_model(self, config: AnyModelConfig, submodel_type: Optional[SubModelTy else: raise - if submodel_type is SubModelType.TextEncoder: - # Text prompts never reach the vision tower; reference images (which would) are not wired up. - drop_qwen3vl_visual_tower(result) - return self._apply_fp8_layerwise_casting(result, config, submodel_type) diff --git a/invokeai/backend/qwen_image_2_1/text_encoding.py b/invokeai/backend/qwen_image_2_1/text_encoding.py index 45f95084271..a9f145ef78b 100644 --- a/invokeai/backend/qwen_image_2_1/text_encoding.py +++ b/invokeai/backend/qwen_image_2_1/text_encoding.py @@ -1,17 +1,48 @@ -"""Qwen-Image-2.1's prompt encoding, as `QwenImage21Pipeline._get_qwen_prompt_embeds` does it for text prompts.""" +"""Qwen-Image-2.1's prompt encoding, as `QwenImage21Pipeline._get_qwen_prompt_embeds` does it.""" + +from collections.abc import Sequence +from typing import NamedTuple import torch +from PIL import Image SYSTEM_PROMPT = "Comprehend and analyze the provided prompt." SYSTEM_BLOCK = f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n" # A raw string handed to the tokenizer, not `apply_chat_template`: the two tokenize differently, and the # checkpoint was trained on this one. T2I_TEMPLATE = SYSTEM_BLOCK + "<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n" +IMAGE_PAD = "<|image_pad|>" +# Pixels per side of one vision token: 16px patches, merged 2x2. +_PIXELS_PER_TOKEN = 32 + +# What one encode needs above the weights for a sequence of T tokens, measured on an RTX 4090 (torch 2.13, +# transformers 5.10): ~286 B * T^2 + 224 KiB * T at peak. The quadratic term is the language model's attention, +# which SDPA runs on its math kernel for grouped-query heads; four 1-megapixel references (~4.1k tokens) need +# ~5.4 GiB. The allocator reserves about a fifth more than it hands out. +_BYTES_PER_TOKEN_PAIR = 286 +_BYTES_PER_TOKEN = 224 * 1024 +_RESERVED_OVER_ALLOCATED = 1.25 -def format_prompt(prompt: str) -> str: +def format_prompt(prompt: str, num_images: int = 0) -> str: + """The prompt as the checkpoint was trained on it, with one vision placeholder per reference image first.""" + placeholders = " ".join(f"<|vision_start|>{IMAGE_PAD}<|vision_end|>" for i in range(1, num_images + 1)) # Qwen has no BOS token, so an empty prompt would leave the encoder nothing to read. - return T2I_TEMPLATE.format(prompt or " ") + return T2I_TEMPLATE.format(placeholders + (prompt or " ")) + + +def sequence_length(tokenizer, prompt: str, images: Sequence[Image.Image] = ()) -> int: + """Tokens the encoder reads: the formatted prompt, each reference's placeholder one token per 32x32 pixels. + + `images` are the references at `vae.reference_size`, which the image processor leaves at their size. + """ + text = len(tokenizer(format_prompt(prompt, len(images))).input_ids) - len(images) + return text + sum((w // _PIXELS_PER_TOKEN) * (h // _PIXELS_PER_TOKEN) for w, h in (i.size for i in images)) + + +def working_memory_bytes(num_tokens: int) -> int: + """Device memory an encode of `num_tokens` tokens needs above the encoder's weights.""" + return int((_BYTES_PER_TOKEN_PAIR * num_tokens**2 + _BYTES_PER_TOKEN * num_tokens) * _RESERVED_OVER_ALLOCATED) def system_prefix_length(tokenizer) -> int: @@ -29,20 +60,83 @@ def _base_model(text_encoder: torch.nn.Module) -> torch.nn.Module: return text_encoder if hasattr(text_encoder, "language_model") else text_encoder.model +def has_vision_tower(text_encoder: torch.nn.Module) -> bool: + """Whether this encoder can read reference images: only the Qwen-Image-2.1 pipeline's keeps its vision tower. + + The standalone Qwen3-VL loaders replace the tower with an `Identity` rather than deleting it. + """ + visual = getattr(_base_model(text_encoder), "visual", None) + return isinstance(visual, torch.nn.Module) and not isinstance(visual, torch.nn.Identity) + + +def on_white(image: Image.Image) -> Image.Image: + """A reference as the vision encoder reads it: alpha composited over white, as the checkpoint was trained.""" + if image.mode != "RGBA": + return image.convert("RGB") + white = Image.new("RGB", image.size, (255, 255, 255)) + white.paste(image, mask=image.getchannel("A")) + return white + + +class PromptEncoding(NamedTuple): + embeds: torch.Tensor + """`(1, L, hidden)`, the system turn dropped.""" + image_pad_mask: torch.Tensor | None + """`(1, L)`, True at each `<|image_pad|>` slot, where the transformer places 2x2 of a reference's latents.""" + reference_grids: tuple[tuple[int, int], ...] + """Per reference, in order: its slots as (rows, columns), one slot per 32x32 pixels.""" + + @torch.no_grad() -def encode_prompt(text_encoder: torch.nn.Module, tokenizer, prompt: str, device: torch.device) -> torch.Tensor: - """Encode a prompt into `(1, L, hidden)` embeddings, as the pipeline encodes a single prompt. +def encode_prompt( + text_encoder: torch.nn.Module, + tokenizer, + prompt: str, + device: torch.device, + *, + processor=None, + images: Sequence[Image.Image] = (), +) -> PromptEncoding: + """Encode a prompt, and the reference images it edits. The transformer reads the last decoder layer's output BEFORE the final RMSNorm. transformers 5 ties the last hidden state to the normalized output, so a forward hook returning the norm's input neutralizes the norm for this call -- the pipeline's approach. The system turn every prompt shares is dropped from the result. + + `images` are the references, already at `vae.reference_size`; reading them needs the vision tower and the + `processor` that cuts them into patches. """ - inputs = tokenizer([format_prompt(prompt)], return_tensors="pt").to(device) base = _base_model(text_encoder) + if images: + if processor is None or not has_vision_tower(text_encoder): + raise ValueError("Reading reference images needs the Qwen3-VL vision tower and its image processor.") + inputs = processor( + text=[format_prompt(prompt, len(images))], images=[on_white(i) for i in images], return_tensors="pt" + ).to(device) + forward_kwargs = { + "input_ids": inputs.input_ids, + "attention_mask": inputs.attention_mask, + "pixel_values": inputs.pixel_values, + "image_grid_thw": inputs.image_grid_thw, + } + if "mm_token_type_ids" in inputs: + forward_kwargs["mm_token_type_ids"] = inputs.mm_token_type_ids + else: + inputs = tokenizer([format_prompt(prompt)], return_tensors="pt").to(device) + forward_kwargs = {"input_ids": inputs.input_ids, "attention_mask": inputs.attention_mask} + handle = base.language_model.norm.register_forward_hook(lambda module, args, output: args[0]) try: - hidden = base(input_ids=inputs.input_ids, attention_mask=inputs.attention_mask).last_hidden_state + # One forward, so the KV cache the language model would build over the whole sequence is waste. + hidden = base(**forward_kwargs, use_cache=False).last_hidden_state finally: handle.remove() - return hidden[:, system_prefix_length(tokenizer) :] + + drop = system_prefix_length(tokenizer) + if not images: + return PromptEncoding(hidden[:, drop:], None, ()) + image_pad_id = tokenizer.convert_tokens_to_ids(IMAGE_PAD) + merge = processor.image_processor.merge_size + grids = tuple((int(h) // merge, int(w) // merge) for _, h, w in inputs.image_grid_thw.tolist()) + return PromptEncoding(hidden[:, drop:], (inputs.input_ids == image_pad_id)[:, drop:], grids) diff --git a/invokeai/backend/qwen_image_2_1/vae.py b/invokeai/backend/qwen_image_2_1/vae.py index f787f2bb445..4fe48a69c90 100644 --- a/invokeai/backend/qwen_image_2_1/vae.py +++ b/invokeai/backend/qwen_image_2_1/vae.py @@ -4,6 +4,8 @@ and Wan VAEs (same tiling attributes, 16x instead of 8x). """ +import math + import numpy as np import torch from PIL import Image @@ -96,11 +98,38 @@ def to_image(decoded: torch.Tensor) -> Image.Image: return Image.fromarray(pixels, mode="RGBA") -def to_vae_input(image: Image.Image) -> torch.Tensor: - """An image as the VAE's `(1, 4, H, W)` input in [-1, 1], encoded opaque. +def to_vae_input(image: Image.Image, keep_alpha: bool = False) -> torch.Tensor: + """An image as the VAE's `(1, 4, H, W)` input in [-1, 1]. - Alpha is set to 1 rather than read from the image: the canvas hands its masks over in alpha, and a - transparent region there means "repaint", not "transparent pixels". + For image-to-image and the canvas, alpha is set to 1 rather than read from the image: the canvas hands its masks + over in alpha, and a transparent region there means "repaint", not "transparent pixels". A reference image keeps + its alpha (`keep_alpha`), as the pipeline encodes it. """ - rgb = torch.from_numpy(np.asarray(image.convert("RGB"), dtype=np.float32)).permute(2, 0, 1) / 127.5 - 1 - return torch.cat([rgb, torch.ones_like(rgb[:1])], dim=0)[None] + rgba = image.convert("RGBA") + # In the pipeline's `VaeImageProcessor` order, so the result matches it to the bit. + pixels = torch.from_numpy(np.asarray(rgba).astype(np.float32) / 255.0).permute(2, 0, 1) * 2.0 - 1.0 + if not keep_alpha: + pixels[3] = 1.0 + return pixels[None] + + +# The area every reference image is resized to, as the pipeline's default `output_resolution` sets it. +REFERENCE_AREA = 1024 * 1024 + + +def reference_size(width: int, height: int) -> tuple[int, int]: + """The size a `width` x `height` reference is read at: `REFERENCE_AREA` at its aspect ratio, on the 32px grid. + + The text encoder and the VAE must both read it at this size: each `<|image_pad|>` slot the encoder emits + stands for 2x2 of the reference's latents, and the transformer drops the latents into those slots. + """ + ratio = width / height + scaled_width = math.sqrt(REFERENCE_AREA * ratio) + # Python's round (half to even) is the pipeline's `calculate_dimensions`. + return round(scaled_width / 32) * 32, round(scaled_width / ratio / 32) * 32 + + +def to_reference(image: Image.Image) -> Image.Image: + """A reference image as both the text encoder and the VAE read it: RGBA, resized to `reference_size`.""" + rgba = image.convert("RGBA") + return rgba.resize(reference_size(*rgba.size), resample=Image.LANCZOS) diff --git a/invokeai/backend/stable_diffusion/diffusion/conditioning_data.py b/invokeai/backend/stable_diffusion/diffusion/conditioning_data.py index 04591b0f238..43c24a6d11a 100644 --- a/invokeai/backend/stable_diffusion/diffusion/conditioning_data.py +++ b/invokeai/backend/stable_diffusion/diffusion/conditioning_data.py @@ -166,8 +166,18 @@ class QwenImage21ConditioningInfo: prompt_embeds: torch.Tensor """Shape: (1, seq_len, hidden=4096). One prompt, so no padding and no mask.""" + image_pad_mask: torch.Tensor | None = None + """Shape: (1, seq_len), True at the reference images' `<|image_pad|>` slots; None without references. Each + slot stands for 2x2 latent tokens of a reference, in the order the references were encoded.""" + + reference_grids: tuple[tuple[int, int], ...] = () + """Per reference, in order: its slots as (rows, columns), so the denoise node can match each reference's + latents by shape, not only by count.""" + def to(self, device: torch.device | None = None, dtype: torch.dtype | None = None): self.prompt_embeds = self.prompt_embeds.to(device=device, dtype=dtype) + if self.image_pad_mask is not None: + self.image_pad_mask = self.image_pad_mask.to(device=device) return self diff --git a/invokeai/frontend/api/openapi.json b/invokeai/frontend/api/openapi.json index 537d0836d51..3bd8fae5101 100644 --- a/invokeai/frontend/api/openapi.json +++ b/invokeai/frontend/api/openapi.json @@ -20698,6 +20698,18 @@ "title": "Reference Images Require Variant", "description": "If set, reference images are only accepted for models of this variant." }, + "reference_images_require_format": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Reference Images Require Format", + "description": "If set, reference images are only accepted for models of this format." + }, "supports_regional_guidance": { "type": "boolean", "title": "Supports Regional Guidance", @@ -95409,6 +95421,29 @@ "orig_default": null, "orig_required": false }, + "reference_latents": { + "anyOf": [ + { + "$ref": "#/components/schemas/LatentsField" + }, + { + "items": { + "$ref": "#/components/schemas/LatentsField" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Latents of the reference images an edit reads, from Image to Latents in reference mode, in the order the prompts were encoded with them.", + "field_kind": "input", + "input": "connection", + "orig_default": null, + "orig_required": false, + "title": "Reference Latents" + }, "denoising_start": { "default": 0.0, "description": "When to start denoising, expressed a percentage of total steps", @@ -95557,7 +95592,7 @@ "tags": ["image", "qwen_image_2_1", "qwen-image-2.1"], "title": "Denoise - Qwen-Image-2.1", "type": "object", - "version": "1.0.1", + "version": "1.1.0", "output": { "$ref": "#/components/schemas/LatentsOutput" } @@ -95566,7 +95601,7 @@ "category": "image", "class": "invocation", "classification": "prototype", - "description": "Encodes an image into Qwen-Image-2.1 latents. The image is encoded opaque; its alpha is not read.", + "description": "Encodes an image into Qwen-Image-2.1 latents.\n\nFor image-to-image and the canvas, the image is encoded opaque at its own size; its alpha is not read. As a\nreference for an edit, it is resized to the area the text encoder reads it at and keeps its alpha.", "node_pack": "invokeai", "properties": { "board": { @@ -95677,6 +95712,16 @@ "title": "Tile Size", "type": "integer" }, + "reference": { + "default": false, + "description": "Encode as a reference image for the denoise node's reference latents: resized to ~1 megapixel, as the text encoder reads it, with its alpha kept.", + "field_kind": "input", + "input": "any", + "orig_default": false, + "orig_required": false, + "title": "Reference", + "type": "boolean" + }, "type": { "const": "qwen_image_2_1_i2l", "default": "qwen_image_2_1_i2l", @@ -95689,7 +95734,7 @@ "tags": ["image", "latents", "vae", "i2l", "qwen_image_2_1"], "title": "Image to Latents - Qwen-Image-2.1", "type": "object", - "version": "1.0.0", + "version": "1.1.0", "output": { "$ref": "#/components/schemas/LatentsOutput" } @@ -95963,7 +96008,7 @@ "category": "conditioning", "class": "invocation", "classification": "prototype", - "description": "Encodes a text prompt for Qwen-Image-2.1 with its Qwen3-VL-8B encoder.", + "description": "Encodes a prompt for Qwen-Image-2.1 with its Qwen3-VL-8B encoder, with the reference images an edit reads.\n\nEach reference also goes to the denoise node, encoded by Image to Latents in its reference mode, in the same\norder.", "node_pack": "invokeai", "properties": { "id": { @@ -96023,6 +96068,19 @@ "orig_required": true, "title": "Qwen3-VL Encoder" }, + "reference_images": { + "default": [], + "description": "Reference images to edit or draw from, in the order the denoise node receives their latents. Needs a Qwen-Image-2.1 Diffusers model's own encoder.", + "field_kind": "input", + "input": "any", + "items": { + "$ref": "#/components/schemas/ImageField" + }, + "orig_default": [], + "orig_required": false, + "title": "Reference Images", + "type": "array" + }, "type": { "const": "qwen_image_2_1_text_encoder", "default": "qwen_image_2_1_text_encoder", @@ -96035,7 +96093,7 @@ "tags": ["prompt", "conditioning", "qwen_image_2_1", "qwen-image-2.1"], "title": "Prompt - Qwen-Image-2.1", "type": "object", - "version": "1.0.0", + "version": "1.1.0", "output": { "$ref": "#/components/schemas/QwenImage21ConditioningOutput" } diff --git a/invokeai/frontend/api/schema.ts b/invokeai/frontend/api/schema.ts index a845bdf0a0f..76c5e0e9de5 100644 --- a/invokeai/frontend/api/schema.ts +++ b/invokeai/frontend/api/schema.ts @@ -5755,6 +5755,11 @@ export type components = { * @description If set, reference images are only accepted for models of this variant. */ reference_images_require_variant?: string | null; + /** + * Reference Images Require Format + * @description If set, reference images are only accepted for models of this format. + */ + reference_images_require_format?: string | null; /** * Supports Regional Guidance * @default false @@ -41186,6 +41191,12 @@ export type components = { * @default null */ denoise_mask?: components["schemas"]["DenoiseMaskField"] | null; + /** + * Reference Latents + * @description Latents of the reference images an edit reads, from Image to Latents in reference mode, in the order the prompts were encoded with them. + * @default null + */ + reference_latents?: components["schemas"]["LatentsField"] | components["schemas"]["LatentsField"][] | null; /** * Denoising Start * @description When to start denoising, expressed a percentage of total steps @@ -41253,7 +41264,10 @@ export type components = { }; /** * Image to Latents - Qwen-Image-2.1 - * @description Encodes an image into Qwen-Image-2.1 latents. The image is encoded opaque; its alpha is not read. + * @description Encodes an image into Qwen-Image-2.1 latents. + * + * For image-to-image and the canvas, the image is encoded opaque at its own size; its alpha is not read. As a + * reference for an edit, it is resized to the area the text encoder reads it at and keeps its alpha. */ QwenImage21ImageToLatentsInvocation: { /** @@ -41305,6 +41319,12 @@ export type components = { * @default 0 */ tile_size?: number; + /** + * Reference + * @description Encode as a reference image for the denoise node's reference latents: resized to ~1 megapixel, as the text encoder reads it, with its alpha kept. + * @default false + */ + reference?: boolean; /** * type * @default qwen_image_2_1_i2l @@ -41451,7 +41471,10 @@ export type components = { }; /** * Prompt - Qwen-Image-2.1 - * @description Encodes a text prompt for Qwen-Image-2.1 with its Qwen3-VL-8B encoder. + * @description Encodes a prompt for Qwen-Image-2.1 with its Qwen3-VL-8B encoder, with the reference images an edit reads. + * + * Each reference also goes to the denoise node, encoded by Image to Latents in its reference mode, in the same + * order. */ QwenImage21TextEncoderInvocation: { /** @@ -41483,6 +41506,12 @@ export type components = { * @default null */ qwen3_vl_encoder?: components["schemas"]["Qwen3VLEncoderField"] | null; + /** + * Reference Images + * @description Reference images to edit or draw from, in the order the denoise node receives their latents. Needs a Qwen-Image-2.1 Diffusers model's own encoder. + * @default [] + */ + reference_images?: components["schemas"]["ImageField"][]; /** * type * @default qwen_image_2_1_text_encoder diff --git a/invokeai/frontend/webv2/src/features/generation/core/__fixtures__/architectureCapabilities.json b/invokeai/frontend/webv2/src/features/generation/core/__fixtures__/architectureCapabilities.json index 005e93e4968..7be2f86da66 100644 --- a/invokeai/frontend/webv2/src/features/generation/core/__fixtures__/architectureCapabilities.json +++ b/invokeai/frontend/webv2/src/features/generation/core/__fixtures__/architectureCapabilities.json @@ -18,6 +18,7 @@ ], "max_reference_images": 0, "reference_images_require_variant": null, + "reference_images_require_format": null, "supports_regional_guidance": true, "regional_negative": false, "clip_skip_max": null, @@ -76,6 +77,7 @@ ], "max_reference_images": 0, "reference_images_require_variant": null, + "reference_images_require_format": null, "supports_regional_guidance": true, "regional_negative": false, "clip_skip_max": null, @@ -132,6 +134,7 @@ "control_kinds": [], "max_reference_images": 0, "reference_images_require_variant": null, + "reference_images_require_format": null, "supports_regional_guidance": false, "regional_negative": false, "clip_skip_max": null, @@ -173,6 +176,7 @@ "control_kinds": [], "max_reference_images": 0, "reference_images_require_variant": null, + "reference_images_require_format": null, "supports_regional_guidance": false, "regional_negative": false, "clip_skip_max": null, @@ -224,6 +228,7 @@ ], "max_reference_images": 5, "reference_images_require_variant": null, + "reference_images_require_format": null, "supports_regional_guidance": true, "regional_negative": false, "clip_skip_max": null, @@ -268,6 +273,7 @@ ], "max_reference_images": 5, "reference_images_require_variant": null, + "reference_images_require_format": null, "supports_regional_guidance": true, "regional_negative": false, "clip_skip_max": null, @@ -312,6 +318,7 @@ ], "max_reference_images": 5, "reference_images_require_variant": null, + "reference_images_require_format": null, "supports_regional_guidance": true, "regional_negative": false, "clip_skip_max": null, @@ -353,6 +360,7 @@ "control_kinds": [], "max_reference_images": 5, "reference_images_require_variant": null, + "reference_images_require_format": null, "supports_regional_guidance": true, "regional_negative": false, "clip_skip_max": null, @@ -394,6 +402,7 @@ "control_kinds": [], "max_reference_images": 5, "reference_images_require_variant": null, + "reference_images_require_format": null, "supports_regional_guidance": true, "regional_negative": false, "clip_skip_max": null, @@ -435,6 +444,7 @@ "control_kinds": [], "max_reference_images": 5, "reference_images_require_variant": null, + "reference_images_require_format": null, "supports_regional_guidance": true, "regional_negative": false, "clip_skip_max": null, @@ -476,6 +486,7 @@ "control_kinds": [], "max_reference_images": 5, "reference_images_require_variant": null, + "reference_images_require_format": null, "supports_regional_guidance": true, "regional_negative": false, "clip_skip_max": null, @@ -517,6 +528,7 @@ "control_kinds": [], "max_reference_images": 0, "reference_images_require_variant": null, + "reference_images_require_format": null, "supports_regional_guidance": false, "regional_negative": false, "clip_skip_max": null, @@ -565,6 +577,7 @@ "control_kinds": [], "max_reference_images": 1, "reference_images_require_variant": null, + "reference_images_require_format": null, "supports_regional_guidance": true, "regional_negative": false, "clip_skip_max": null, @@ -617,6 +630,7 @@ "control_kinds": [], "max_reference_images": 1, "reference_images_require_variant": null, + "reference_images_require_format": null, "supports_regional_guidance": true, "regional_negative": false, "clip_skip_max": null, @@ -669,6 +683,7 @@ "control_kinds": [], "max_reference_images": 0, "reference_images_require_variant": null, + "reference_images_require_format": null, "supports_regional_guidance": false, "regional_negative": false, "clip_skip_max": null, @@ -710,6 +725,7 @@ "control_kinds": [], "max_reference_images": 0, "reference_images_require_variant": null, + "reference_images_require_format": null, "supports_regional_guidance": false, "regional_negative": false, "clip_skip_max": null, @@ -751,6 +767,7 @@ "control_kinds": [], "max_reference_images": 0, "reference_images_require_variant": null, + "reference_images_require_format": null, "supports_regional_guidance": false, "regional_negative": false, "clip_skip_max": null, @@ -792,6 +809,7 @@ "control_kinds": [], "max_reference_images": 5, "reference_images_require_variant": "edit", + "reference_images_require_format": null, "supports_regional_guidance": false, "regional_negative": false, "clip_skip_max": null, @@ -842,8 +860,9 @@ "scheduler_set": "flow", "scheduler_applies_to_graph": false, "control_kinds": [], - "max_reference_images": 0, + "max_reference_images": 4, "reference_images_require_variant": null, + "reference_images_require_format": "diffusers", "supports_regional_guidance": false, "regional_negative": false, "clip_skip_max": null, @@ -883,8 +902,9 @@ "scheduler_set": "flow", "scheduler_applies_to_graph": false, "control_kinds": [], - "max_reference_images": 0, + "max_reference_images": 4, "reference_images_require_variant": null, + "reference_images_require_format": "diffusers", "supports_regional_guidance": false, "regional_negative": false, "clip_skip_max": null, @@ -929,6 +949,7 @@ ], "max_reference_images": 5, "reference_images_require_variant": null, + "reference_images_require_format": null, "supports_regional_guidance": true, "regional_negative": true, "clip_skip_max": 12, @@ -970,6 +991,7 @@ "control_kinds": [], "max_reference_images": 0, "reference_images_require_variant": null, + "reference_images_require_format": null, "supports_regional_guidance": true, "regional_negative": true, "clip_skip_max": 24, @@ -1011,6 +1033,7 @@ "control_kinds": [], "max_reference_images": 0, "reference_images_require_variant": null, + "reference_images_require_format": null, "supports_regional_guidance": false, "regional_negative": false, "clip_skip_max": null, @@ -1055,6 +1078,7 @@ ], "max_reference_images": 5, "reference_images_require_variant": null, + "reference_images_require_format": null, "supports_regional_guidance": true, "regional_negative": true, "clip_skip_max": null, @@ -1096,6 +1120,7 @@ "control_kinds": [], "max_reference_images": 0, "reference_images_require_variant": null, + "reference_images_require_format": null, "supports_regional_guidance": false, "regional_negative": false, "clip_skip_max": null, @@ -1137,6 +1162,7 @@ "control_kinds": [], "max_reference_images": 0, "reference_images_require_variant": null, + "reference_images_require_format": null, "supports_regional_guidance": false, "regional_negative": false, "clip_skip_max": null, @@ -1185,6 +1211,7 @@ "control_kinds": [], "max_reference_images": 0, "reference_images_require_variant": null, + "reference_images_require_format": null, "supports_regional_guidance": false, "regional_negative": false, "clip_skip_max": null, @@ -1235,6 +1262,7 @@ ], "max_reference_images": 0, "reference_images_require_variant": null, + "reference_images_require_format": null, "supports_regional_guidance": true, "regional_negative": false, "clip_skip_max": null, @@ -1285,6 +1313,7 @@ ], "max_reference_images": 0, "reference_images_require_variant": null, + "reference_images_require_format": null, "supports_regional_guidance": true, "regional_negative": false, "clip_skip_max": null, diff --git a/invokeai/frontend/webv2/src/features/generation/core/architectureCapabilities.ts b/invokeai/frontend/webv2/src/features/generation/core/architectureCapabilities.ts index fdae256b518..794f1767002 100644 --- a/invokeai/frontend/webv2/src/features/generation/core/architectureCapabilities.ts +++ b/invokeai/frontend/webv2/src/features/generation/core/architectureCapabilities.ts @@ -26,6 +26,7 @@ export interface ArchitectureCapabilitiesRow { control_kinds: string[]; max_reference_images: number; reference_images_require_variant: string | null; + reference_images_require_format: string | null; supports_regional_guidance: boolean; regional_negative: boolean; clip_skip_max: number | null; diff --git a/invokeai/frontend/webv2/src/features/generation/core/baseGenerationPolicies.test.ts b/invokeai/frontend/webv2/src/features/generation/core/baseGenerationPolicies.test.ts index d4b75ad47ac..ff4436edbf0 100644 --- a/invokeai/frontend/webv2/src/features/generation/core/baseGenerationPolicies.test.ts +++ b/invokeai/frontend/webv2/src/features/generation/core/baseGenerationPolicies.test.ts @@ -35,6 +35,7 @@ import { isReferenceImageSupported, isGenerateModelSelectable, isSupportedGenerateModel, + QWEN_IMAGE_21_STANDALONE_ENCODER_REASON, } from './baseGenerationPolicies'; import { SUPPORTED_GENERATE_BASES } from './supportedBases'; @@ -1106,6 +1107,15 @@ describe('component policies', () => { expect(isReferenceImageSupported(undefined)).toBe(false); }); + it('offers Qwen-Image-2.1 reference images only on its Diffusers pipeline, whose encoder reads images', () => { + const diffusers = createModel('qwen-image-2-1', { format: 'diffusers' }); + + expect(isReferenceImageSupported(diffusers)).toBe(true); + expect(getMaxReferenceImages(diffusers)).toBe(4); + expect(isReferenceImageSupported(createModel('qwen-image-2-1', { format: 'gguf_quantized' }))).toBe(false); + expect(isReferenceImageSupported(createModel('qwen-image-2-1', { format: 'checkpoint' }))).toBe(false); + }); + it('derives the reference image limit from external provider capabilities', () => { expect(getMaxReferenceImages(createModel('sdxl'))).toBe(5); expect(getMaxReferenceImages(createModel('cogview4'))).toBe(0); @@ -1171,6 +1181,31 @@ describe('component policies', () => { ).toContain('Reference Image #1 is not supported by sd-3 model.'); }); + it('blocks Qwen-Image-2.1 references while a standalone encoder replaces the one that reads them', () => { + const model = createModel('qwen-image-2-1', { format: 'diffusers' }); + const encoder = { + base: 'any', + key: 'qwen3-vl-8b', + name: 'Qwen3-VL 8B', + type: 'qwen3_vl_encoder', + variant: 'qwen3_vl_8b', + } as const; + const referenceImages = [ + { + config: { image: referenceImage, type: 'qwen_image_2_1_reference_image' as const }, + id: 'ref-1', + isEnabled: true, + }, + ]; + + expect( + getGenerationValidationReasons(model, createSettings(model, { qwen3VLEncoderModel: encoder, referenceImages })) + ).toContain(QWEN_IMAGE_21_STANDALONE_ENCODER_REASON); + expect(getGenerationValidationReasons(model, createSettings(model, { referenceImages }))).not.toContain( + QWEN_IMAGE_21_STANDALONE_ENCODER_REASON + ); + }); + it('reports missing reference image model dependencies', () => { const model = createModel('sd-1'); diff --git a/invokeai/frontend/webv2/src/features/generation/core/baseGenerationPolicies.ts b/invokeai/frontend/webv2/src/features/generation/core/baseGenerationPolicies.ts index c70f7ca260a..b3aa6651369 100644 --- a/invokeai/frontend/webv2/src/features/generation/core/baseGenerationPolicies.ts +++ b/invokeai/frontend/webv2/src/features/generation/core/baseGenerationPolicies.ts @@ -1048,7 +1048,8 @@ const getBaseComponentSectionPolicy = ( { ...qwen3VlEncoderSlot( 'Qwen-Image-2.1 conditions on the Qwen3-VL 8B encoder, not the 4B one Krea-2 uses. Required for ' + - 'non-Diffusers Qwen-Image-2.1 models.', + "non-Diffusers Qwen-Image-2.1 models. On a Diffusers model it replaces the model's own encoder, the " + + 'only one that reads reference images.', isQwen3Vl8bEncoder ), required: (ctx) => ctx.model.format !== 'diffusers', @@ -1308,9 +1309,11 @@ export const isReferenceImageSupported = (model: GenerateModelConfig | undefined return false; } - // The backend declares which reference-image features require a specific variant. + // The backend declares which reference-image features require a specific variant or format. return ( - features.reference_images_require_variant === null || model.variant === features.reference_images_require_variant + (features.reference_images_require_variant === null || + model.variant === features.reference_images_require_variant) && + (features.reference_images_require_format === null || model.format === features.reference_images_require_format) ); }; @@ -1352,6 +1355,10 @@ export const getDefaultReferenceImageConfig = ( return { image, styleStrength: DEFAULT_KREA2_STYLE_STRENGTH, type: 'krea2_reference_image' }; } + if (modelBase === 'qwen-image-2-1') { + return { image, type: 'qwen_image_2_1_reference_image' }; + } + if (isFluxKontextModel(model)) { return { image, model, type: 'flux_kontext_reference_image' }; } @@ -1383,6 +1390,10 @@ const getReferenceImageConfigSupported = ( return model.type !== 'external_image_generator' && model.base === 'qwen-image' && model.variant === 'edit'; case 'krea2_reference_image': return model.type !== 'external_image_generator' && model.base === 'krea-2'; + case 'qwen_image_2_1_reference_image': + return ( + model.type !== 'external_image_generator' && model.base === 'qwen-image-2-1' && isReferenceImageSupported(model) + ); case 'flux_kontext_reference_image': return isFluxKontextModel(model); case 'flux_redux': @@ -1677,6 +1688,11 @@ const getPidValidationReasons = (model: GenerateModelConfig, settings: GenerateS return reasons; }; +/** Only the encoder of a Qwen-Image-2.1 Diffusers model keeps the vision tower that reads references. */ +export const QWEN_IMAGE_21_STANDALONE_ENCODER_REASON = + "Reference images need the Qwen-Image-2.1 model's own encoder, the only one with a vision tower. " + + 'Clear the Qwen3-VL Encoder component.'; + const getReferenceImageValidationReasons = (model: GenerateModelConfig, settings: GenerateSettings): string[] => { const reasons: string[] = []; const enabled = settings.referenceImages.filter((referenceImage) => referenceImage.isEnabled); @@ -1721,6 +1737,17 @@ const getReferenceImageValidationReasons = (model: GenerateModelConfig, settings } }); + // The encoder slot stays filled across models (Ideogram 4 and the GGUF builds share it), so a Diffusers model can + // inherit one that cannot read the references. + if ( + model.base === 'qwen-image-2-1' && + settings.qwen3VLEncoderModel && + isQwen3Vl8bEncoder(settings.qwen3VLEncoderModel) && + enabled.some((referenceImage) => referenceImage.config.type === 'qwen_image_2_1_reference_image') + ) { + reasons.push(QWEN_IMAGE_21_STANDALONE_ENCODER_REASON); + } + return reasons; }; diff --git a/invokeai/frontend/webv2/src/features/generation/core/graph.test.ts b/invokeai/frontend/webv2/src/features/generation/core/graph.test.ts index a70cebf8087..5a9454a5958 100644 --- a/invokeai/frontend/webv2/src/features/generation/core/graph.test.ts +++ b/invokeai/frontend/webv2/src/features/generation/core/graph.test.ts @@ -4,6 +4,7 @@ import { afterEach, describe, expect, it, vi } from 'vitest'; import type { ComponentModelConfig, GenerateModelConfig, + GenerateReferenceImage, GenerateSettings, LoraModelConfig, MainModelConfig, @@ -744,6 +745,65 @@ describe('compileGenerateGraph', () => { expect(getEdge(graph, 'pos_cond', 'reference_images')).toBeDefined(); expect(getEdge(graph, 'denoise_latents', 'reference_latents')).toBeDefined(); }); + + describe('Qwen-Image-2.1 reference images', () => { + const qwenImage21Model: MainModelConfig = { + base: 'qwen-image-2-1', + format: 'diffusers', + key: 'qwen21', + name: 'Qwen-Image-2.1', + type: 'main', + }; + const references = (...names: string[]): GenerateReferenceImage[] => + names.map((name) => ({ + config: { + image: { original: { image: { height: 768, image_name: name, width: 512 } } }, + type: 'qwen_image_2_1_reference_image', + }, + id: name, + isEnabled: true, + })); + + it('reads each reference in both prompts and, in the same order, as latents in denoise', () => { + const graph = compile(qwenImage21Model, { cfgScale: 4, referenceImages: references('a.png', 'b.png') }); + + const nameOf = (imageNode: string) => + (graph.nodes[imageNode]?.image as { image_name: string } | undefined)?.image_name; + // Each chain is followed back from its end: the last collect holds b, the one before it a. + const chain = (lastCollect: string) => [ + getEdge(graph, lastCollect, 'collection')?.source.node_id ?? '', + lastCollect, + ]; + + const imagesInto = getEdge(graph, 'pos_cond', 'reference_images')?.source.node_id ?? ''; + expect(getEdge(graph, 'neg_cond', 'reference_images')?.source.node_id).toBe(imagesInto); + expect(chain(imagesInto).map((collect) => nameOf(getEdge(graph, collect, 'item')?.source.node_id ?? ''))).toEqual( + ['a.png', 'b.png'] + ); + + const latentsInto = getEdge(graph, 'denoise_latents', 'reference_latents')?.source.node_id ?? ''; + const latentImages = chain(latentsInto).map((collect) => { + const i2l = getEdge(graph, collect, 'item')?.source.node_id ?? ''; + expect(graph.nodes[i2l]).toMatchObject({ reference: true, type: 'qwen_image_2_1_i2l' }); + return nameOf(getEdge(graph, i2l, 'image')?.source.node_id ?? ''); + }); + expect(latentImages).toEqual(['a.png', 'b.png']); + }); + + it('refuses references with a standalone encoder, which has no vision tower', () => { + const encoder: ComponentModelConfig = { + base: 'any', + key: 'qwen3-vl-8b', + name: 'Qwen3-VL 8B', + type: 'qwen3_vl_encoder', + variant: 'qwen3_vl_8b', + }; + + expect(() => + compile(qwenImage21Model, { qwen3VLEncoderModel: encoder, referenceImages: references('a.png') }) + ).toThrow(/vision tower/); + }); + }); }); describe('generate seeds', () => { diff --git a/invokeai/frontend/webv2/src/features/generation/core/graph.ts b/invokeai/frontend/webv2/src/features/generation/core/graph.ts index 598c1dbbf05..55a600b3a9f 100644 --- a/invokeai/frontend/webv2/src/features/generation/core/graph.ts +++ b/invokeai/frontend/webv2/src/features/generation/core/graph.ts @@ -25,6 +25,7 @@ import { getFlux2DiffusersComponentSource, getGenerationDimensions, getGenerationValidationReasons, + QWEN_IMAGE_21_STANDALONE_ENCODER_REASON, type SupportedGenerateBase, } from './baseGenerationPolicies'; import { @@ -1438,14 +1439,86 @@ const buildQwenImage21Graph = ( } addEdge(graph, seed, 'value', denoise, 'seed'); addEdge(graph, denoise, 'latents', output, 'latents'); + addQwenImage21ReferenceImages(graph, settings, { + conditionings: negCond ? [posCond, negCond] : [posCond], + denoise, + hasStandaloneEncoder: Boolean(qwen3VlEncoderModel), + modelLoader, + }); addMetadata(graph, output, settings, model, 'qwen_image_2_1_txt2img', projectSettings, { qwen3_vl_encoder: qwen3VlEncoderModel ?? undefined, vae: vaeModel ?? undefined, }); + addReferenceImageMetadata(graph, output, settings); return graph; }; +/** + * Qwen-Image-2.1 reads each reference twice, at one size: the prompt encoder sees it through the vision tower, and + * the denoise node places its clean latents in the prompt's image slots. Both prompts are encoded with the + * references, and the latents follow in the same order. + */ +const addQwenImage21ReferenceImages = ( + graph: BackendGraphContract, + settings: GenerateSettings, + nodes: { + conditionings: BackendInvocationContract[]; + denoise: BackendInvocationContract; + hasStandaloneEncoder: boolean; + modelLoader: BackendInvocationContract; + } +): void => { + const referenceImages = getEnabledReferenceImages(settings, 'qwen_image_2_1_reference_image'); + if (referenceImages.length === 0) { + return; + } + // A backstop: validation reports this before Invoke. + if (nodes.hasStandaloneEncoder) { + throw new Error(QWEN_IMAGE_21_STANDALONE_ENCODER_REASON); + } + + let imageCollect: BackendInvocationContract | null = null; + let latentsCollect: BackendInvocationContract | null = null; + + for (const { config } of referenceImages) { + if (!config.image) { + continue; + } + + const imageNode = addNode(graph, { + id: createId('qwen21_ref_img'), + image: toImageField(config.image), + type: 'image', + }); + const i2l = addNode(graph, { id: createId('qwen21_ref_i2l'), reference: true, type: 'qwen_image_2_1_i2l' }); + addEdge(graph, imageNode, 'image', i2l, 'image'); + addEdge(graph, nodes.modelLoader, 'vae', i2l, 'vae'); + + const nextImageCollect = addNode(graph, { id: createId('qwen21_ref_img_collect'), type: 'collect' }); + addEdge(graph, imageNode, 'image', nextImageCollect, 'item'); + if (imageCollect) { + addEdge(graph, imageCollect, 'collection', nextImageCollect, 'collection'); + } + imageCollect = nextImageCollect; + + const nextLatentsCollect = addNode(graph, { id: createId('qwen21_ref_latents_collect'), type: 'collect' }); + addEdge(graph, i2l, 'latents', nextLatentsCollect, 'item'); + if (latentsCollect) { + addEdge(graph, latentsCollect, 'collection', nextLatentsCollect, 'collection'); + } + latentsCollect = nextLatentsCollect; + } + + if (!imageCollect || !latentsCollect) { + return; + } + for (const conditioning of nodes.conditionings) { + addEdge(graph, imageCollect, 'collection', conditioning, 'reference_images'); + } + addEdge(graph, latentsCollect, 'collection', nodes.denoise, 'reference_latents'); +}; + const buildIdeogram4Graph = ( settings: GenerateSettings, model: MainModelConfig, diff --git a/invokeai/frontend/webv2/src/features/generation/core/settings.ts b/invokeai/frontend/webv2/src/features/generation/core/settings.ts index dea964ea145..178d589f339 100644 --- a/invokeai/frontend/webv2/src/features/generation/core/settings.ts +++ b/invokeai/frontend/webv2/src/features/generation/core/settings.ts @@ -312,6 +312,7 @@ const REFERENCE_IMAGE_CONFIG_TYPES = new Set([ 'ip_adapter', 'krea2_reference_image', 'qwen_image_reference_image', + 'qwen_image_2_1_reference_image', ]); const isClipVisionModel = (value: unknown): value is 'ViT-H' | 'ViT-G' | 'ViT-L' => @@ -347,6 +348,7 @@ const normalizeReferenceImageConfig = (value: unknown): GenerateReferenceImageCo case 'external_reference_image': case 'flux2_reference_image': case 'qwen_image_reference_image': + case 'qwen_image_2_1_reference_image': return { image, type: value.type }; case 'flux_kontext_reference_image': return { image, model: getMainModelOrNull(value.model), type: value.type }; @@ -419,6 +421,7 @@ const cloneReferenceImageConfig = (config: GenerateReferenceImageConfig): Genera case 'flux2_reference_image': case 'krea2_reference_image': case 'qwen_image_reference_image': + case 'qwen_image_2_1_reference_image': return { ...config, image: config.image ? cloneCroppableImage(config.image) : null }; case 'flux_kontext_reference_image': return { @@ -589,6 +592,7 @@ const syncReferenceImageConfigWithModels = ( case 'flux2_reference_image': case 'krea2_reference_image': case 'qwen_image_reference_image': + case 'qwen_image_2_1_reference_image': return config; } }; diff --git a/invokeai/frontend/webv2/src/features/generation/core/types.ts b/invokeai/frontend/webv2/src/features/generation/core/types.ts index a2ba70b4727..ca7597e15ff 100644 --- a/invokeai/frontend/webv2/src/features/generation/core/types.ts +++ b/invokeai/frontend/webv2/src/features/generation/core/types.ts @@ -142,6 +142,7 @@ export type GenerateReferenceImageConfig = | { type: 'qwen_image_reference_image'; image: GenerateReferenceImageAsset | null } /** Krea-2 training-free style transfer: no adapter model, a single reference, `styleStrength` 0–2 (0 bypasses). */ | { type: 'krea2_reference_image'; image: GenerateReferenceImageAsset | null; styleStrength: number } + | { type: 'qwen_image_2_1_reference_image'; image: GenerateReferenceImageAsset | null } | { type: 'external_reference_image'; image: GenerateReferenceImageAsset | null }; export interface GenerateReferenceImage { diff --git a/invokeai/frontend/webv2/src/workbench/canvasMigration.ts b/invokeai/frontend/webv2/src/workbench/canvasMigration.ts index 9d91441f0c5..710dd33fa88 100644 --- a/invokeai/frontend/webv2/src/workbench/canvasMigration.ts +++ b/invokeai/frontend/webv2/src/workbench/canvasMigration.ts @@ -183,6 +183,7 @@ const zReferenceImage = z.object({ }), z.object({ image: zGeneratedImage.nullable(), type: z.literal('flux2_reference_image') }), z.object({ image: zGeneratedImage.nullable(), type: z.literal('qwen_image_reference_image') }), + z.object({ image: zGeneratedImage.nullable(), type: z.literal('qwen_image_2_1_reference_image') }), z.object({ image: zGeneratedImage.nullable(), type: z.literal('external_reference_image') }), ]), id: z.string(), diff --git a/tests/app/invocations/test_qwen_image_2_1_nodes.py b/tests/app/invocations/test_qwen_image_2_1_nodes.py index 20e7fa0ad0d..cb6e29a747e 100644 --- a/tests/app/invocations/test_qwen_image_2_1_nodes.py +++ b/tests/app/invocations/test_qwen_image_2_1_nodes.py @@ -4,29 +4,36 @@ an image-to-image run starts on each variant's schedule, and which components the loader refuses. """ -from contextlib import contextmanager +from contextlib import contextmanager, nullcontext from types import SimpleNamespace import numpy as np import pytest import torch +from PIL import Image -from invokeai.app.invocations.fields import LatentsField, QwenImage21ConditioningField -from invokeai.app.invocations.model import ModelIdentifierField, TransformerField +from invokeai.app.invocations.fields import ImageField, LatentsField, QwenImage21ConditioningField +from invokeai.app.invocations.model import ModelIdentifierField, Qwen3VLEncoderField, TransformerField, VAEField from invokeai.app.invocations.qwen_image_2_1.qwen_image_2_1_denoise import ( + KV_BYTES_PER_PREFIX_TOKEN, LATENT_CHANNELS, QwenImage21DenoiseInvocation, clip_sigmas, pack_latents, + prefix_cache_fits, + prefix_length, unpack_latents, ) from invokeai.app.invocations.qwen_image_2_1.qwen_image_2_1_model_loader import QwenImage21ModelLoaderInvocation +from invokeai.app.invocations.text_encoder.qwen_image_2_1_text_encoder import QwenImage21TextEncoderInvocation +from invokeai.app.invocations.vae.qwen_image_2_1_image_to_latents import QwenImage21ImageToLatentsInvocation from invokeai.backend.model_manager.taxonomy import ( BaseModelType, ModelFormat, ModelType, Qwen3VLVariantType, QwenImage21VariantType, + SubModelType, ) from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ( ConditioningFieldData, @@ -52,13 +59,28 @@ class _Transformer: def __init__(self) -> None: self.calls: list[tuple[str, int, str | None, float]] = [] + self.inputs: list[SimpleNamespace] = [] def modules(self): return iter(()) - def __call__(self, *, hidden_states, encoder_hidden_states, timestep, kv_cache, kv_cache_mode, **_kwargs): + def __call__( + self, + *, + hidden_states, + encoder_hidden_states, + timestep, + kv_cache, + kv_cache_mode, + img_shapes, + img_mask, + **_kwargs, + ): prompt = "pos" if float(encoder_hidden_states.mean()) > 0 else "neg" self.calls.append((prompt, id(kv_cache), kv_cache_mode, round(float(timestep[0]), 3))) + self.inputs.append( + SimpleNamespace(hidden_states=hidden_states.clone(), img_shapes=img_shapes, img_mask=img_mask) + ) # The transformer returns the joint sequence; the node keeps the image tokens at its end. return (torch.zeros(1, encoder_hidden_states.shape[1] + hidden_states.shape[1], hidden_states.shape[2]),) @@ -72,18 +94,27 @@ def model_on_device(self, **_kwargs): yield ({}, self.model) -def _context(transformer: _Transformer, variant: QwenImage21VariantType, init_latents: torch.Tensor | None = None): - conditionings = { - "pos": ConditioningFieldData(conditionings=[QwenImage21ConditioningInfo(prompt_embeds=torch.ones(1, 3, 8))]), - "neg": ConditioningFieldData(conditionings=[QwenImage21ConditioningInfo(prompt_embeds=-torch.ones(1, 2, 8))]), +def _context( + transformer: _Transformer, + variant: QwenImage21VariantType, + init_latents: torch.Tensor | None = None, + *, + infos: dict[str, QwenImage21ConditioningInfo] | None = None, + tensors: dict[str, torch.Tensor] | None = None, +): + infos = infos or { + "pos": QwenImage21ConditioningInfo(prompt_embeds=torch.ones(1, 3, 8)), + "neg": QwenImage21ConditioningInfo(prompt_embeds=-torch.ones(1, 2, 8)), } + conditionings = {name: ConditioningFieldData(conditionings=[info]) for name, info in infos.items()} + tensors = tensors or {} return SimpleNamespace( models=SimpleNamespace( load=lambda _identifier: _TransformerInfo(transformer), get_config=lambda _identifier: SimpleNamespace(variant=variant), ), conditioning=SimpleNamespace(load=lambda name: conditionings[name]), - tensors=SimpleNamespace(load=lambda _name: init_latents), + tensors=SimpleNamespace(load=lambda name: tensors.get(name, init_latents)), util=SimpleNamespace(sd_step_callback=lambda *_args: None), ) @@ -96,6 +127,7 @@ def _denoise(**fields) -> QwenImage21DenoiseInvocation: "negative_conditioning": QwenImage21ConditioningField(conditioning_name="neg"), "latents": None, "denoise_mask": None, + "reference_latents": None, "denoising_start": 0.0, "denoising_end": 1.0, "cfg_scale": 1.0, @@ -201,8 +233,8 @@ def test_a_strength_on_base_names_whole_steps_and_always_runs_one( def test_a_float32_fallback_reserves_twice_the_bfloat16_activations() -> None: estimate = QwenImage21DenoiseInvocation._estimate_working_memory base = 1024**3 - bf16 = estimate(4096, 300, 300, True, 2) - base - assert estimate(4096, 300, 300, True, 4) - base == 2 * bf16 + bf16 = estimate(4096, 300, 300, True, True, 2) - base + assert estimate(4096, 300, 300, True, True, 4) - base == 2 * bf16 @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) @@ -349,3 +381,263 @@ def test_the_4b_encoder_is_refused() -> None: loader = _loader(qwen3_vl_encoder_model=("encoder", BaseModelType.Any, ModelType.Qwen3VLEncoder)) with pytest.raises(ValueError, match="8B"): loader.invoke(_loader_context(ModelFormat.Diffusers, encoder_variant=Qwen3VLVariantType.Qwen3VL_4B)) + + +def _edit_prompt(sign: float, grids: list[tuple[int, int]]) -> QwenImage21ConditioningInfo: + """A prompt encoded with references: a token, then one run of image slots per reference, then a token.""" + mask = [False] + for rows, columns in grids: + mask += [True] * (rows * columns) + [False] + return QwenImage21ConditioningInfo( + prompt_embeds=sign * torch.ones(1, len(mask), 8), + image_pad_mask=torch.tensor([mask]), + reference_grids=tuple(grids), + ) + + +def _edit_context(transformer, prompt_grids, references: dict[str, torch.Tensor], negative=None): + infos = {"pos": _edit_prompt(1.0, prompt_grids), "neg": negative or _edit_prompt(-1.0, prompt_grids)} + return _context(transformer, QwenImage21VariantType.Base, infos=infos, tensors=references) + + +def test_references_lead_the_image_tokens_and_fill_the_prompts_slots() -> None: + transformer = _Transformer() + # A 4x8 reference is 32 latent tokens in 2x4 slots; the 64x64 target is 4x4. + reference = torch.randn(1, LATENT_CHANNELS, 4, 8) + _denoise(reference_latents=[LatentsField(latents_name="ref")], steps=1)._run_diffusion( + _edit_context(transformer, [(2, 4)], {"ref": reference}) + ) + + seen = transformer.inputs[0] + assert seen.img_shapes == [[(1, 4, 8), (1, 4, 4)]] + assert seen.hidden_states.shape[1] == 32 + 16 + assert torch.equal(seen.hidden_states[:, :32], pack_latents(reference)) + # The prompt's own mask, then one slot per 2x2 group of the target's 4x4 latents. + assert seen.img_mask.tolist() == [[False] + [True] * 8 + [False] + [True] * 4] + + +def test_references_that_do_not_fill_the_slots_are_refused() -> None: + context = _edit_context(_Transformer(), [(2, 2)], {"ref": torch.zeros(1, LATENT_CHANNELS, 4, 8)}) + with pytest.raises(ValueError, match=r"positive prompt was encoded with 1 reference image\(s\) \(64x64\)"): + _denoise(reference_latents=LatentsField(latents_name="ref"))._run_diffusion(context) + + +def test_two_references_of_one_size_in_swapped_order_are_refused() -> None: + # Landscape then portrait in the prompt; portrait then landscape as latents. The slot counts agree. + context = _edit_context( + _Transformer(), + [(2, 4), (4, 2)], + {"portrait": torch.zeros(1, LATENT_CHANNELS, 8, 4), "landscape": torch.zeros(1, LATENT_CHANNELS, 4, 8)}, + ) + latents = [LatentsField(latents_name="portrait"), LatentsField(latents_name="landscape")] + with pytest.raises(ValueError, match="in reference mode, from the same images and in the same order"): + _denoise(reference_latents=latents)._run_diffusion(context) + + +def test_latents_from_another_vae_are_refused() -> None: + context = _edit_context(_Transformer(), [(2, 4)], {"ref": torch.zeros(1, 16, 4, 8)}) + with pytest.raises(ValueError, match="64 channels"): + _denoise(reference_latents=LatentsField(latents_name="ref"))._run_diffusion(context) + + +def test_a_negative_prompt_encoded_without_the_references_is_refused() -> None: + context = _edit_context( + _Transformer(), + [(2, 4)], + {"ref": torch.zeros(1, LATENT_CHANNELS, 4, 8)}, + negative=QwenImage21ConditioningInfo(prompt_embeds=-torch.ones(1, 2, 8)), + ) + with pytest.raises(ValueError, match="negative prompt was encoded with 0"): + _denoise(reference_latents=LatentsField(latents_name="ref"), cfg_scale=4.0)._run_diffusion(context) + + +def test_the_prefix_counts_each_reference_slot_as_four_latent_tokens() -> None: + # A 1 MP reference is 32x32 slots: 1024 slots in the prompt, 4096 tokens in the joint sequence. + info = _edit_prompt(1.0, [(32, 32)]) + assert prefix_length(info) == info.prompt_embeds.shape[1] + 3 * 1024 + + +def test_the_prefix_caches_are_reserved_only_when_they_are_kept() -> None: + estimate = QwenImage21DenoiseInvocation._estimate_working_memory + prefix = 4200 + cached = estimate(4096, prefix, prefix, True, True, 2) + recomputed = estimate(4096, prefix, prefix, True, False, 2) + # K and V of every prefix token in all 32 blocks, per guidance branch: ~2 GiB per 1 MP reference. + assert cached - recomputed == 2 * prefix * KV_BYTES_PER_PREFIX_TOKEN + + +def test_caches_larger_than_half_the_card_are_not_kept(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(torch.cuda, "get_device_properties", lambda _device: SimpleNamespace(total_memory=24 * 2**30)) + cuda = torch.device("cuda") + # Two references at CFG fit a 24 GB card; four do not. + assert prefix_cache_fits(2 * 2 * 4200 * KV_BYTES_PER_PREFIX_TOKEN, cuda) + assert not prefix_cache_fits(2 * 4 * 4200 * KV_BYTES_PER_PREFIX_TOKEN, cuda) + + +def test_without_room_for_the_caches_every_step_runs_the_whole_prefix(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(f"{DENOISE}.prefix_cache_fits", lambda _bytes, _device: False) + transformer = _Transformer() + _denoise(reference_latents=LatentsField(latents_name="ref"), cfg_scale=4.0)._run_diffusion( + _edit_context(transformer, [(2, 4)], {"ref": torch.zeros(1, LATENT_CHANNELS, 4, 8)}) + ) + assert {(cache, mode) for _prompt, cache, mode, _t in transformer.calls} == {(id(None), None)} + + +def test_a_standalone_encoder_cannot_read_references() -> None: + encoder = ModelIdentifierField( + key="enc", hash="h", name="qwen3vl_8b", base=BaseModelType.Any, type=ModelType.Qwen3VLEncoder + ) + node = QwenImage21TextEncoderInvocation.model_construct( + prompt="swap them", + qwen3_vl_encoder=Qwen3VLEncoderField(tokenizer=encoder, text_encoder=encoder, loras=[]), + reference_images=[ImageField(image_name="ref.png")], + ) + context = SimpleNamespace( + images=SimpleNamespace(get_pil=lambda _name: Image.new("RGB", (64, 64))), + models=SimpleNamespace( + get_config=lambda _identifier: SimpleNamespace( + base=BaseModelType.Any, type=ModelType.Qwen3VLEncoder, format=ModelFormat.Checkpoint + ) + ), + ) + with pytest.raises(ValueError, match="vision tower"): + node.invoke(context) + + +def test_the_processor_comes_from_the_qwen_image_2_1_pipeline_the_tokenizer_is_from() -> None: + main = ModelIdentifierField( + key="qwen21", hash="h", name="Qwen-Image-2.1", base=BaseModelType.QwenImage21, type=ModelType.Main + ) + tokenizer = main.model_copy(update={"submodel_type": SubModelType.Tokenizer}) + node = QwenImage21TextEncoderInvocation.model_construct( + prompt="swap them", + qwen3_vl_encoder=Qwen3VLEncoderField(tokenizer=tokenizer, text_encoder=tokenizer, loras=[]), + reference_images=[ImageField(image_name="ref.png")], + ) + context = SimpleNamespace( + models=SimpleNamespace( + get_config=lambda _identifier: SimpleNamespace( + base=BaseModelType.QwenImage21, type=ModelType.Main, format=ModelFormat.Diffusers + ) + ) + ) + processor = node._processor_identifier(context) + assert (processor.key, processor.submodel_type) == ("qwen21", SubModelType.Processor) + + +class _ReservationRecorded(Exception): + pass + + +@pytest.mark.parametrize("num_references", [0, 4]) +def test_the_encoder_reserves_room_for_the_attention_its_references_need(num_references: int) -> None: + from invokeai.backend.qwen3_vl.qwen3_vl_assets import load_bundled_qwen3_vl_tokenizer + + main = ModelIdentifierField( + key="qwen21", hash="h", name="Qwen-Image-2.1", base=BaseModelType.QwenImage21, type=ModelType.Main + ) + tokenizer = main.model_copy(update={"submodel_type": SubModelType.Tokenizer}) + encoder = main.model_copy(update={"submodel_type": SubModelType.TextEncoder}) + reserved: list[int] = [] + + @contextmanager + def on_device(*, working_mem_bytes: int): + reserved.append(working_mem_bytes) + raise _ReservationRecorded + yield # pragma: no cover + + text_encoder = torch.nn.Module() + text_encoder.language_model = torch.nn.Linear(2, 2) + text_encoder.visual = torch.nn.Linear(2, 2) + infos = { + SubModelType.TextEncoder: SimpleNamespace( + model=text_encoder, compute_device=torch.device("cpu"), model_on_device=on_device + ), + SubModelType.Tokenizer: nullcontext(load_bundled_qwen3_vl_tokenizer()), + SubModelType.Processor: SimpleNamespace(model_on_device=lambda: nullcontext((None, None))), + } + node = QwenImage21TextEncoderInvocation.model_construct( + prompt="Combine the subjects of all reference images.", + qwen3_vl_encoder=Qwen3VLEncoderField(tokenizer=tokenizer, text_encoder=encoder, loras=[]), + reference_images=[ImageField(image_name=f"ref{i}.png") for i in range(num_references)], + ) + context = SimpleNamespace( + images=SimpleNamespace(get_pil=lambda _name: Image.new("RGB", (1024, 1024))), + models=SimpleNamespace( + get_config=lambda _identifier: SimpleNamespace( + base=BaseModelType.QwenImage21, type=ModelType.Main, format=ModelFormat.Diffusers + ), + load=lambda identifier: infos[identifier.submodel_type], + ), + util=SimpleNamespace(signal_progress=lambda _message: None), + ) + + with pytest.raises(_ReservationRecorded): + node.invoke(context) + + if num_references: + # Measured on an RTX 4090: four 1-megapixel references peak 5.44 GiB above the weights, and the allocator + # reserves about a fifth more. Less than that and the driver spills the encode to system memory. + assert reserved[0] >= 1.2 * 5.44 * 2**30 + else: + assert reserved[0] < 2**28 + + +class _VAE: + """Records what it encodes and answers with bf16 latents, as the real VAE runs in bf16.""" + + dtype = torch.bfloat16 + config = SimpleNamespace(latents_mean=[0.3] * LATENT_CHANNELS, latents_std=[1.7] * LATENT_CHANNELS) + + def __init__(self) -> None: + self.pixels: torch.Tensor | None = None + self.z: torch.Tensor | None = None + + def parameters(self): + return iter([torch.zeros(1, dtype=torch.bfloat16)]) + + def encode(self, pixels: torch.Tensor): + self.pixels = pixels + _, _, frames, height, width = pixels.shape + generator = torch.Generator().manual_seed(0) + self.z = torch.randn(1, LATENT_CHANNELS, frames, height // 16, width // 16, generator=generator).bfloat16() + return SimpleNamespace(latent_dist=SimpleNamespace(mode=lambda: self.z)) + + +def test_a_reference_encodes_at_its_reading_size_with_its_alpha_normalized_in_bf16(monkeypatch) -> None: + from invokeai.app.invocations.vae import qwen_image_2_1_image_to_latents as i2l_module + + monkeypatch.setattr(i2l_module, "patch_qwen_image_vae_tiling", lambda *_args: nullcontext()) + vae = _VAE() + saved: dict[str, torch.Tensor] = {} + + @contextmanager + def on_device(**_kwargs): + yield ({}, vae) + + context = SimpleNamespace( + images=SimpleNamespace(get_pil=lambda _name: Image.new("RGBA", (500, 300), (255, 0, 0, 64))), + models=SimpleNamespace( + load=lambda _identifier: SimpleNamespace( + compute_device=torch.device("cpu"), model=vae, model_on_device=on_device + ) + ), + config=SimpleNamespace(get=lambda: SimpleNamespace(force_tiled_decode=False, auto_tiled_decode=False)), + tensors=SimpleNamespace(save=lambda tensor: saved.update(latents=tensor) or "latents"), + ) + vae_field = VAEField( + vae=ModelIdentifierField(key="v", hash="h", name="v", base=BaseModelType.QwenImage21, type=ModelType.VAE) + ) + output = QwenImage21ImageToLatentsInvocation.model_construct( + image=ImageField(image_name="ref.png"), vae=vae_field, tiled=False, tile_size=0, reference=True + ).invoke(context) + + # 500x300 reads at 1312x800; the alpha is the image's own (64/255), not the canvas encode's 1. + assert (output.width, output.height) == (1312, 800) + assert vae.pixels is not None and vae.pixels.shape == (1, 4, 1, 800, 1312) + assert vae.pixels[0, 3].unique().tolist() == [torch.tensor(64 / 255 * 2 - 1).bfloat16().item()] + # Normalized in bf16, as the pipeline normalizes it -- which differs from normalizing in fp32. + mean, std = torch.tensor(0.3).bfloat16(), torch.tensor(1.7).bfloat16() + expected = ((vae.z - mean) / std).float()[:, :, 0] + assert torch.equal(saved["latents"], expected) + assert not torch.equal(expected, ((vae.z.float() - 0.3) / 1.7)[:, :, 0]) diff --git a/tests/backend/architectures/test_features.py b/tests/backend/architectures/test_features.py index 5b8aa98956e..65fae7c66d5 100644 --- a/tests/backend/architectures/test_features.py +++ b/tests/backend/architectures/test_features.py @@ -80,13 +80,20 @@ def test_control_kinds_match_the_frontend_policy(kind: ControlKind, expected: se assert declared == expected -def test_reference_images_and_the_one_variant_condition() -> None: - """Qwen-Image is the only base whose answer depends on the variant, so it is the only one with - `reference_images_require_variant` set.""" +def test_reference_images_and_the_conditions_on_them() -> None: + """Qwen-Image is the only base whose answer depends on the variant, and Qwen-Image-2.1 the only one whose + answer depends on the format.""" supported = { b.value for b in generative_bases() if (f := get(b, FeaturesFacet)) is not None and f.supports_reference_images } - assert supported == {"flux", "flux2", "sd-1", "sdxl", "qwen-image", "krea-2"} + assert supported == {"flux", "flux2", "sd-1", "sdxl", "qwen-image", "krea-2", "qwen-image-2-1"} + + by_format = { + b.value: f.reference_images_require_format + for b in generative_bases() + if (f := get(b, FeaturesFacet)) is not None and f.reference_images_require_format is not None + } + assert by_format == {"qwen-image-2-1": "diffusers"} conditional = { b.value diff --git a/tests/backend/model_manager/load/test_qwen3_vl_visual_tower.py b/tests/backend/model_manager/load/test_qwen3_vl_visual_tower.py index 55c6d116c70..c807f5c73ce 100644 --- a/tests/backend/model_manager/load/test_qwen3_vl_visual_tower.py +++ b/tests/backend/model_manager/load/test_qwen3_vl_visual_tower.py @@ -300,3 +300,24 @@ def test_the_folder_loader_drops_the_tower(tmp_path: Path) -> None: model = loader._load_model(config, SubModelType.TextEncoder) _assert_tower_gone(model) + + +def test_the_qwen_image_2_1_pipeline_encoder_keeps_the_tower(monkeypatch, tmp_path: Path) -> None: + """The one encoder that keeps it: Qwen-Image-2.1 reads reference images through it, and the node refuses an + encoder without it.""" + from transformers import Qwen3VLModel + + from invokeai.backend.model_manager.configs.main import Main_Diffusers_QwenImage21_Config + from invokeai.backend.model_manager.load.model_loaders.qwen_image_2_1 import QwenImage21DiffusersModel + from invokeai.backend.qwen_image_2_1.text_encoding import has_vision_tower + + encoder = _tiny_encoder() + config = Main_Diffusers_QwenImage21_Config.model_construct(path=str(tmp_path), repo_variant=None) + loader = object.__new__(QwenImage21DiffusersModel) + loader._apply_fp8_layerwise_casting = lambda model, _config, _submodel: model + monkeypatch.setattr(Qwen3VLModel, "from_pretrained", classmethod(lambda _cls, *_a, **_k: encoder)) + + model = loader._load_model(config, SubModelType.TextEncoder) + + assert has_vision_tower(model) + assert list(model.visual.parameters()) diff --git a/tests/backend/qwen_image_2_1/test_text_encoding.py b/tests/backend/qwen_image_2_1/test_text_encoding.py index 5797589414c..349af916949 100644 --- a/tests/backend/qwen_image_2_1/test_text_encoding.py +++ b/tests/backend/qwen_image_2_1/test_text_encoding.py @@ -1,13 +1,20 @@ -"""Qwen-Image-2.1's prompt encoding: the template the checkpoint was trained on, read before the final norm.""" +"""Qwen-Image-2.1's prompt encoding: the template the checkpoint was trained on, read before the final norm, and +the reference images an edit reads through the vision tower.""" + +import itertools import pytest import torch +from PIL import Image from invokeai.backend.qwen3_vl.qwen3_vl_assets import load_bundled_qwen3_vl_tokenizer from invokeai.backend.qwen_image_2_1.text_encoding import ( SYSTEM_BLOCK, encode_prompt, format_prompt, + has_vision_tower, + on_white, + sequence_length, system_prefix_length, ) from tests.backend.model_manager.load.qwen3vl_gguf_fixture import tiny_qwen3vl_config @@ -40,7 +47,7 @@ def test_the_embeddings_are_the_last_layer_before_the_final_norm(text_encoder, t language_model = text_encoder.language_model handle = language_model.norm.register_forward_hook(lambda m, args, out: captured.update(pre=args[0], post=out)) try: - embeds = encode_prompt(text_encoder, tokenizer, "a fox", torch.device("cpu")) + embeds, image_slots, grids = encode_prompt(text_encoder, tokenizer, "a fox", torch.device("cpu")) ids = tokenizer([format_prompt("a fox")], return_tensors="pt").input_ids text_encoder(input_ids=ids) finally: @@ -49,3 +56,89 @@ def test_the_embeddings_are_the_last_layer_before_the_final_norm(text_encoder, t drop = system_prefix_length(tokenizer) torch.testing.assert_close(embeds[0], captured["pre"][0, drop:]) assert not torch.allclose(embeds[0], captured["post"][0, drop:]) + assert image_slots is None and grids == () + + +def test_reference_placeholders_lead_the_prompt_as_the_pipeline_writes_them() -> None: + assert format_prompt("swap them", 2) == ( + SYSTEM_BLOCK + "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|> " + "<|vision_start|><|image_pad|><|vision_end|>swap them<|im_end|>\n<|im_start|>assistant\n" + ) + + +@pytest.fixture(scope="module") +def processor(tokenizer): + from transformers import Qwen2VLImageProcessor, Qwen3VLProcessor + from transformers.models.qwen3_vl.video_processing_qwen3_vl import Qwen3VLVideoProcessor + + # The released processor's geometry: 16px patches merged 2x2, so one slot per 32x32 pixels. + image_processor = Qwen2VLImageProcessor(patch_size=16, merge_size=2, temporal_patch_size=2) + return Qwen3VLProcessor( + image_processor=image_processor, tokenizer=tokenizer, video_processor=Qwen3VLVideoProcessor() + ) + + +def _runs(mask: torch.Tensor) -> list[int]: + return [len(list(run)) for is_slot, run in itertools.groupby(mask[0].tolist()) if is_slot] + + +def _encode(text_encoder, tokenizer, processor, references): + return encode_prompt( + text_encoder, tokenizer, "swap them", torch.device("cpu"), processor=processor, images=references + ) + + +def test_each_reference_fills_one_run_of_slots_per_32px_block(text_encoder, tokenizer, processor) -> None: + references = [Image.new("RGBA", (64, 96), (255, 0, 0, 128)), Image.new("RGB", (128, 64), (0, 255, 0))] + embeds, image_slots, grids = _encode(text_encoder, tokenizer, processor, references) + assert image_slots is not None and image_slots.shape == embeds.shape[:2] + # 64x96 is 3 rows of 2 blocks, 128x64 is 2 rows of 4, in the order given; the system turn is gone from both. + assert grids == ((3, 2), (2, 4)) + assert _runs(image_slots) == [6, 8] + + +def test_the_sequence_length_is_counted_before_encoding(text_encoder, tokenizer, processor) -> None: + # The working memory is reserved before the encoder runs, from this count; it has to be the length it reads. + references = [Image.new("RGB", (64, 96)), Image.new("RGB", (128, 64))] + embeds, _, _ = _encode(text_encoder, tokenizer, processor, references) + assert sequence_length(tokenizer, "swap them", references) == embeds.shape[1] + system_prefix_length(tokenizer) + + +def test_the_vision_tower_reads_the_pixels(text_encoder, tokenizer, processor) -> None: + red, green = Image.new("RGB", (64, 64), (255, 0, 0)), Image.new("RGB", (64, 64), (0, 255, 0)) + red_embeds, red_slots, _ = _encode(text_encoder, tokenizer, processor, [red]) + green_embeds, _, _ = _encode(text_encoder, tokenizer, processor, [green]) + # Same tokens either way; only the vision tower can tell the two apart. + assert not torch.allclose(red_embeds[red_slots], green_embeds[red_slots]) + + +def test_a_transparent_reference_reads_as_its_colors_over_white(text_encoder, tokenizer, processor) -> None: + # Fully transparent black must read as plain white, not as the black its RGB channels hold. + transparent = Image.new("RGBA", (64, 64), (0, 0, 0, 0)) + white = Image.new("RGB", (64, 64), (255, 255, 255)) + transparent_embeds, _, _ = _encode(text_encoder, tokenizer, processor, [transparent]) + white_embeds, _, _ = _encode(text_encoder, tokenizer, processor, [white]) + torch.testing.assert_close(transparent_embeds, white_embeds) + + +def test_references_need_the_vision_tower(tokenizer, processor) -> None: + from transformers import Qwen3VLModel + + from invokeai.backend.model_manager.util.qwen3_vl import drop_qwen3vl_visual_tower + + config = tiny_qwen3vl_config() + config.text_config.vocab_size = len(tokenizer) + encoder = Qwen3VLModel(config).float().eval() + # As the standalone loaders leave it: an Identity where the tower was. + drop_qwen3vl_visual_tower(encoder) + assert not has_vision_tower(encoder) + with pytest.raises(ValueError, match="vision tower"): + encode_prompt( + encoder, tokenizer, "a", torch.device("cpu"), processor=processor, images=[Image.new("RGB", (32, 32))] + ) + + +def test_transparency_is_flattened_over_white_for_the_vision_encoder() -> None: + half_red = Image.new("RGBA", (1, 1), (255, 0, 0, 128)) + assert on_white(half_red).getpixel((0, 0)) == (255, 127, 127) + assert on_white(Image.new("RGBA", (1, 1), (0, 0, 0, 0))).getpixel((0, 0)) == (255, 255, 255) diff --git a/tests/backend/qwen_image_2_1/test_vae.py b/tests/backend/qwen_image_2_1/test_vae.py index 69270995817..987523b3591 100644 --- a/tests/backend/qwen_image_2_1/test_vae.py +++ b/tests/backend/qwen_image_2_1/test_vae.py @@ -1,4 +1,5 @@ -"""Qwen-Image-2.1's RGBA VAE: how a decode becomes an image, what an encode reads, and the tile geometry it runs in.""" +"""Qwen-Image-2.1's RGBA VAE: how a decode becomes an image, what an encode reads -- a canvas image or an edit's +reference -- and the tile geometry it runs in.""" import numpy as np import pytest @@ -11,7 +12,9 @@ MIN_TILE_SIZE, SPATIAL_SCALE, choose_tile_size, + reference_size, to_image, + to_reference, to_vae_input, working_memory_bytes, ) @@ -142,3 +145,34 @@ def from_pretrained(_path, *, torch_dtype, **_kwargs): loader._load_model(SimpleNamespace(path=str(tmp_path)), SubModelType.VAE) assert loaded["dtype"] is expected + + +@pytest.mark.parametrize( + ("size", "expected"), + [ + # sqrt(1024^2 * 5/3) = 1322.0 -> 41.3 x 32 -> 1312; 1322.0 / (5/3) = 793.2 -> 24.8 x 32 -> 800. + ((500, 300), (1312, 800)), + ((1024, 1024), (1024, 1024)), + # sqrt(1024^2 * 3) = 1773.6 -> 1760; 591.2 -> 576. + ((3000, 1000), (1760, 576)), + ], +) +def test_a_reference_is_read_at_one_megapixel_on_the_32px_grid( + size: tuple[int, int], expected: tuple[int, int] +) -> None: + assert reference_size(*size) == expected + + +def test_a_reference_encodes_exactly_as_the_pipeline_preprocesses_it() -> None: + from diffusers.image_processor import VaeImageProcessor + + generator = torch.Generator().manual_seed(0) + pixels = torch.randint(0, 256, (300, 500, 4), generator=generator, dtype=torch.uint8).numpy() + image = Image.fromarray(pixels, mode="RGBA") + width, height = reference_size(*image.size) + expected = VaeImageProcessor(vae_scale_factor=16, vae_latent_channels=64).preprocess( + image, width=width, height=height + ) + + # The alpha is the reference's own, unlike a canvas encode. + assert torch.equal(to_vae_input(to_reference(image), keep_alpha=True), expected.float()) From d0f200e5f9a08b4f4f19771a93321c3ad3fd9bbe Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Thu, 22 Oct 2026 05:02:41 +0200 Subject: [PATCH 2/2] fix(qwen-image-2-1): load the vision tower only for edits and tighten edit pricing The tower is its own submodel, attached to the cached encoder only while it reads references; text prompts skip it. Memory estimates follow the dtype, the reference prefill and dedicated cards; the denoise records its cache mode. At most four references are accepted; a GGUF or single-file model explains why references are unavailable. --- .../03.Models/Local Models/qwen-image-2-1.mdx | 6 +- invokeai/app/api/routers/recall_parameters.py | 8 +- .../qwen_image_2_1/qwen_image_2_1_denoise.py | 127 +++++--- .../qwen_image_2_1_text_encoder.py | 72 +++-- .../model_manager/load/load_default.py | 11 +- .../load/model_loaders/qwen_image_2_1.py | 82 ++++- invokeai/backend/model_manager/taxonomy.py | 3 + .../backend/qwen_image_2_1/text_encoding.py | 69 ++++- invokeai/frontend/api/openapi.json | 9 +- invokeai/frontend/api/schema.ts | 14 +- .../frontend/webv2/public/locales/en.json | 1 + .../__snapshots__/generateGraphNodeTypes.json | 48 ++- .../core/baseGenerationPolicies.test.ts | 15 + .../generation/core/baseGenerationPolicies.ts | 22 ++ .../core/canvas/compileCanvasGraph.test.ts | 16 + .../features/generation/core/graph.test.ts | 27 +- .../src/features/generation/core/graph.ts | 12 +- .../generation/core/graphCoverage.testing.ts | 39 ++- .../features/generation/core/settings.test.ts | 40 +++ .../generation/ui/GenerateGuidanceSection.tsx | 9 +- .../GenerateReferenceImagesSection.tsx | 23 +- .../referenceImageFormatHint.browser.test.tsx | 113 +++++++ .../webv2/src/workbench/canvasMigration.ts | 1 - .../image-actions/imageRecall.test.ts | 42 +++ .../invocations/test_qwen_image_2_1_nodes.py | 292 ++++++++++++++++-- .../load/test_load_default_cpu_only.py | 10 + .../load/test_load_default_fp8.py | 12 + .../load/test_qwen3_vl_visual_tower.py | 96 +++++- .../qwen_image_2_1/test_text_encoding.py | 110 +++++++ 29 files changed, 1162 insertions(+), 167 deletions(-) create mode 100644 invokeai/frontend/webv2/src/features/generation/ui/reference-images/referenceImageFormatHint.browser.test.tsx diff --git a/docs/src/content/docs/Users Guide/03.Models/Local Models/qwen-image-2-1.mdx b/docs/src/content/docs/Users Guide/03.Models/Local Models/qwen-image-2-1.mdx index ecbb1190186..8ae96a8bc6f 100644 --- a/docs/src/content/docs/Users Guide/03.Models/Local Models/qwen-image-2-1.mdx +++ b/docs/src/content/docs/Users Guide/03.Models/Local Models/qwen-image-2-1.mdx @@ -1,7 +1,7 @@ --- title: Qwen-Image-2.1 description: Generate images with Qwen-Image-2.1 and its 8-step Turbo variant, from the Diffusers folders or a GGUF transformer, including transparent (RGBA) output. -lastUpdated: 2026-10-17 +lastUpdated: 2026-10-22 sidebar: order: 8 --- @@ -64,7 +64,9 @@ With a Qwen-Image-2.1 **Diffusers** model, add up to four images with **Add refe - **Diffusers models only.** Only the encoder that comes with the Diffusers model keeps the Qwen3-VL vision tower, which reads the images. A single-file or GGUF transformer with a standalone encoder does not offer reference images, and on a Diffusers model the references refuse to run while a standalone **Qwen3-VL Encoder** is selected under **Components**; clear it to edit. - **Size.** Each reference is read at about one megapixel at its own aspect ratio, whatever its size. The generated image takes the size set in the panel; **Use size** on a reference's thumbnail matches it. - **Transparency.** The encoder sees a transparent reference over white, as the model was trained; its latents keep the transparency. -- **Memory.** The vision tower adds about 1.1 GB to the Diffusers model's encoder, on every generation with it. Each reference adds about 4,100 tokens to the start of the sequence, which the model keeps for the whole generation: about 2 GB per reference, twice that with **CFG** above 1. Where that would take more than half of the GPU, Invoke reprocesses the references at every step instead of keeping them, which uses less memory but takes noticeably longer. Reading the references also needs encoder memory that grows quickly with their number: about 7 GB beyond its weights for four, so on a 24 GB card part of the encoder then runs from system memory. +- **At most four.** A graph or API request with more than four references is refused with an error. +- **Memory.** The encoder's vision tower, about 1.1 GB, is loaded only for generations with references and joins the encoder while it reads them; plain prompts never load it, and the encoder itself is not loaded again when you switch between edits and plain prompts. Each reference adds about 4,100 tokens to the start of the sequence, which the model keeps for the whole generation: about 2 GB per reference, twice that with **CFG** above 1. Where that would take more than half of the GPU, Invoke reprocesses the references at every step instead of keeping them, which uses less memory but takes noticeably longer. Reading the references also needs encoder memory that grows quickly with their number: about 7 GB beyond its weights for four, so on a 24 GB card part of the encoder then runs from system memory. +- **Reproducing an edit.** Keeping the references and reprocessing them give slightly different images from the same seed, and which one runs depends on the GPU's memory. The image's metadata records which one ran, as `qwen_image_2_1_prefix_kv_cache`: `true` when the references were kept, `false` when they were reprocessed. Generations without references record it too. ### Transparent images diff --git a/invokeai/app/api/routers/recall_parameters.py b/invokeai/app/api/routers/recall_parameters.py index 20caef0d1b5..f32e18bf008 100644 --- a/invokeai/app/api/routers/recall_parameters.py +++ b/invokeai/app/api/routers/recall_parameters.py @@ -65,10 +65,10 @@ class ReferenceImageRecallParameter(BaseModel): Used for reference images that feed directly into the main model rather than through a separate IP-Adapter / ControlNet model — for example - FLUX.2 Klein, FLUX Kontext, and Qwen Image Edit. The receiving frontend - picks the correct config type (``flux2_reference_image`` / - ``qwen_image_reference_image`` / ``flux_kontext_reference_image``) based - on the currently-selected main model. + FLUX.2 Klein, FLUX Kontext, Qwen Image Edit and Qwen-Image-2.1. The + receiving frontend picks the correct config type (``flux2_reference_image`` / + ``qwen_image_reference_image`` / ``qwen_image_2_1_reference_image`` / + ``flux_kontext_reference_image``) based on the currently-selected main model. """ image_name: str = Field(description="The filename of the reference image in outputs/images") diff --git a/invokeai/app/invocations/qwen_image_2_1/qwen_image_2_1_denoise.py b/invokeai/app/invocations/qwen_image_2_1/qwen_image_2_1_denoise.py index 7f69766c476..2fe3dfa4ae3 100644 --- a/invokeai/app/invocations/qwen_image_2_1/qwen_image_2_1_denoise.py +++ b/invokeai/app/invocations/qwen_image_2_1/qwen_image_2_1_denoise.py @@ -15,12 +15,13 @@ Input, InputField, LatentsField, + MetadataField, QwenImage21ConditioningField, WithBoard, WithMetadata, ) +from invokeai.app.invocations.metadata_linked import LatentsMetaOutput from invokeai.app.invocations.model import TransformerField -from invokeai.app.invocations.primitives import LatentsOutput from invokeai.app.services.shared.invocation_context import InvocationContext from invokeai.app.util.startup_utils import log_attention_backends from invokeai.backend.flux.sampling_utils import clip_timestep_schedule_fractional @@ -31,13 +32,16 @@ from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState from invokeai.backend.stable_diffusion.diffusion.conditioning_data import QwenImage21ConditioningInfo from invokeai.backend.util.attention import sdpa_score_matrix_bytes -from invokeai.backend.util.devices import TorchDevice +from invokeai.backend.util.devices import TorchDevice, device_has_dedicated_vram from invokeai.backend.util.logging import InvokeAILogger LATENT_CHANNELS = 64 VAE_SCALE_FACTOR = 16 # One vision-language slot of the transformer's joint sequence covers 2x2 latent tokens. TOKENS_PER_SLOT = 4 +# The metadata entry that records whether a run kept the prefix KV cache (True) or recomputed the prefix on every +# step (False). +PREFIX_CACHE_METADATA_KEY = "qwen_image_2_1_prefix_kv_cache" logger = InvokeAILogger.get_logger(__name__) @@ -82,17 +86,18 @@ def check_reference_slots( """Refuse references whose latents do not match the prompt's image slots, reference by reference. Each reference's latents fill one run of the prompt's slots, 2x2 latents per slot, at the slot grid the - encoder read it at. A mismatch means the prompt was encoded with other images, in another order or at another - size than the latents: the transformer would place one reference's latents in another's slots. + encoder read it at. This catches a different number of references, or one at a different size, than the + prompt was encoded with -- including references of different sizes in another order. It cannot tell apart + same-size references in another order, or a same-size image swapped for another: the slots carry no identity. """ expected = [(h // 2, w // 2) for _, h, w in reference_shapes] found = list(info.reference_grids) if found != expected or slot_runs(info.image_pad_mask) != [rows * columns for rows, columns in found]: raise ValueError( f"The {prompt} prompt was encoded with {len(found)} reference image(s) ({_describe_grids(found)}), but " - f"the denoise node received {len(expected)} reference latent(s) ({_describe_grids(expected)}). Encode " - "each reference with Image to Latents in reference mode, from the same images and in the same order " - "as the prompts." + f"the denoise node received {len(expected)} reference latent(s) ({_describe_grids(expected)}): the " + "counts or sizes differ. Encode each reference with Image to Latents in reference mode, from the same " + "images and in the same order as the prompts." ) @@ -102,8 +107,14 @@ def prefix_length(info: QwenImage21ConditioningInfo) -> int: return info.prompt_embeds.shape[1] + (TOKENS_PER_SLOT - 1) * slots -# K and V of one prefix token in all 32 blocks: 2 x 4096 x 32 x 2 bytes in bf16. -KV_BYTES_PER_PREFIX_TOKEN = 2 * 4096 * 32 * 2 +def kv_bytes_per_prefix_token(transformer: torch.nn.Module, dtype: torch.dtype) -> int: + """K and V of one prefix token in every block of `transformer`, in the compute dtype the caches are kept in. + + 2 x 4096 x 32 x 2 bytes for the released model in bf16 (32 heads of 128, 32 blocks): 0.5 MiB. + """ + config = transformer.config + hidden = config.num_attention_heads * config.attention_head_dim + return 2 * hidden * len(transformer.transformer_blocks) * dtype.itemsize def prefix_cache_fits(cache_bytes: int, device: torch.device) -> bool: @@ -112,11 +123,41 @@ def prefix_cache_fits(cache_bytes: int, device: torch.device) -> bool: They cannot be offloaded the way the transformer's weights can, so caches larger than half the card would leave too little for the model however it streams: four references at CFG above 1 hold ~16 GiB. Past that the prefix is recomputed every step instead -- slower, but within the memory of an ordinary run. Memory that is - not a dedicated card's (CPU, MPS) is not limited here. + not a dedicated card's (CPU, MPS, an integrated XPU) is not limited here, nor a card that reports no size. """ - if device.type != "cuda": + if not device_has_dedicated_vram(device): return True - return cache_bytes <= torch.cuda.get_device_properties(device).total_memory // 2 + backend = torch.xpu if device.type == "xpu" else torch.cuda + total = int(backend.get_device_properties(device).total_memory) + return total == 0 or cache_bytes <= total // 2 + + +def prefix_segments(info: QwenImage21ConditioningInfo) -> list[int]: + """The prefix's blocks in the joint sequence, in order: each run of text, and each reference widened to 2x2 + latents per slot. The prefill attends each block over every key up to its end.""" + if info.image_pad_mask is None: + return [info.prompt_embeds.shape[1]] + return [ + len(list(run)) * (TOKENS_PER_SLOT if is_slot else 1) + for is_slot, run in itertools.groupby(info.image_pad_mask[0].tolist()) + ] + + +def largest_attention_call(infos: list[QwenImage21ConditioningInfo], image_seq_len: int) -> int: + """Queries x keys of the largest attention call a run makes, over every prompt it runs. + + The prefill runs each prefix block over the keys before and in it, and the target over every key; a cached + step repeats only the target's call. With large references and a small target, the last reference's call is + the largest: four 1 MP references before a 512x512 target make it ~4x the target's. + """ + largest = 0 + for info in infos: + end = 0 + for length in prefix_segments(info): + end += length + largest = max(largest, length * end) + largest = max(largest, image_seq_len * (end + image_seq_len)) + return largest def pack_latents(latents: torch.Tensor) -> torch.Tensor: @@ -188,12 +229,20 @@ def validate_cfg_scale_is_finite(cls, value: float | list[float]) -> float | lis return value @torch.no_grad() - def invoke(self, context: InvocationContext) -> LatentsOutput: - latents = self._run_diffusion(context).detach().to("cpu") - name = context.tensors.save(tensor=latents) + def invoke(self, context: InvocationContext) -> LatentsMetaOutput: + latents, used_prefix_cache = self._run_diffusion(context) + name = context.tensors.save(tensor=latents.detach().to("cpu")) + # The two modes sample visibly differently in bf16, and which one runs depends on the card's memory: the + # image's metadata has to say, or nothing tells why one graph and seed give another image on another card. + # A graph routes this output to the node that saves the image. + metadata = dict(self.metadata.root) if self.metadata is not None else {} + metadata[PREFIX_CACHE_METADATA_KEY] = used_prefix_cache # `LatentsOutput.build` assumes 8x latents; these are 16x, so report the requested size directly. - return LatentsOutput( - latents=LatentsField(latents_name=name, seed=self.seed), width=self.width, height=self.height + return LatentsMetaOutput( + latents=LatentsField(latents_name=name, seed=self.seed), + width=self.width, + height=self.height, + metadata=MetadataField.model_validate(metadata), ) def _load_conditioning( @@ -245,7 +294,8 @@ def _prep_inpaint_mask(self, context: InvocationContext, latents: torch.Tensor) ) return mask.to(device=latents.device, dtype=latents.dtype) - def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: + def _run_diffusion(self, context: InvocationContext) -> tuple[torch.Tensor, bool]: + """The denoised latents, and whether the prefix KV cache ran (rather than the prefix being recomputed).""" from diffusers.models.transformers.transformer_qwenimage21 import QwenImage21KVCache if self.denoising_start >= self.denoising_end: @@ -350,7 +400,8 @@ def image_slot_mask(info: QwenImage21ConditioningInfo) -> torch.Tensor: neg_prefix = prefix_length(neg) if neg is not None else None # Text and reference keys and values do not depend on the step under `causal_condition`: the first step # prefills them, the rest compute the image tokens only -- when the caches fit. - cache_bytes = KV_BYTES_PER_PREFIX_TOKEN * (pos_prefix + (neg_prefix or 0)) + kv_bytes_per_token = kv_bytes_per_prefix_token(transformer_info.model, dtype) + cache_bytes = kv_bytes_per_token * (pos_prefix + (neg_prefix or 0)) use_cache = bool(getattr(transformer_info.model.config, "causal_condition", False)) if use_cache and not prefix_cache_fits(cache_bytes, device): use_cache = False @@ -359,10 +410,13 @@ def image_slot_mask(info: QwenImage21ConditioningInfo) -> torch.Tensor: "this GPU; recomputing the prefix every step instead, which is slower." ) working_memory = self._estimate_working_memory( - image_seq_len, pos_prefix, neg_prefix, do_cfg, use_cache, dtype.itemsize + image_seq_len, pos_prefix, neg_prefix, do_cfg, use_cache, dtype.itemsize, kv_bytes_per_token ) working_memory += self._attention_score_bytes( - transformer_info.model, image_seq_len, max(pos_prefix, neg_prefix or 0), device, dtype + transformer_info.model, + largest_attention_call([pos] if neg is None else [pos, neg], image_seq_len), + device, + dtype, ) working_memory += peak_dequant_transient_bytes(transformer_info.model, dtype) log_attention_backends(logger, device) @@ -436,7 +490,7 @@ def predict(embeds: torch.Tensor, img_mask: torch.Tensor, cache, mode, timestep) ) ) - return unpack_latents(latents, latent_height, latent_width).float() + return unpack_latents(latents, latent_height, latent_width).float(), use_cache @staticmethod def _estimate_working_memory( @@ -446,15 +500,16 @@ def _estimate_working_memory( do_cfg: bool, use_cache: bool, element_size: int, + kv_bytes_per_token: int, ) -> int: """Peak transformer activations, in bytes, for the model cache to keep free. Measured on an RTX 4090 in bf16 with the transformer fully resident: 0.55 GiB of activations at 1024x1024 (4096 image tokens) and 2.2 GiB at 2048x2048 (16384), about 0.14 MiB per token. 0.25 MiB per - token plus a 1 GiB base leaves room for the allocator and the prefix cache, which holds K and V of every - prefix token -- text and reference latents -- in all 32 blocks (~0.5 MiB per token, per guidance branch). - Both per-token terms are bf16 figures and scale with `element_size`, the compute dtype's width: a - float32 fallback needs twice as much. + token plus a 1 GiB base leaves room for the allocator. That term is a bf16 figure and scales with + `element_size`, the compute dtype's width: a float32 fallback needs twice as much. The prefix cache adds + `kv_bytes_per_token` (K and V in every block, already in that dtype: 0.5 MiB for the released model in + bf16) for every prefix token -- text and reference latents -- per guidance branch. """ mib, gib = 1024**2, 1024**3 width = element_size / 2 @@ -463,29 +518,29 @@ def _estimate_working_memory( if use_cache: branches = 2 if do_cfg else 1 prefix_tokens = branches * (pos_prefix if neg_prefix is None else prefix) - estimated += int(prefix_tokens * KV_BYTES_PER_PREFIX_TOKEN * width) + estimated += prefix_tokens * kv_bytes_per_token return estimated @staticmethod - def _attention_score_bytes( - transformer: object, image_seq_len: int, prefix: int, device: torch.device, dtype: torch.dtype - ) -> int: - """The score matrix attention materializes on a build without a fused kernel for it (0 where it has one).""" + def _attention_score_bytes(transformer: object, largest_call: int, device: torch.device, dtype: torch.dtype) -> int: + """The score matrix attention materializes on a build without a fused kernel for it (0 where it has one). + + `largest_call` is queries x keys of the largest call, from `largest_attention_call`. + """ config = getattr(transformer, "config", None) num_heads = getattr(config, "num_attention_heads", None) head_dim = getattr(config, "attention_head_dim", None) if not isinstance(num_heads, int) or not isinstance(head_dim, int): return 0 - # The largest call, the image queries over every key, goes unmasked both in the prefill and on decode steps: - # only the text segments carry a mask, and they are as short as the prompt. All of them go through diffusers' - # attention dispatch. Its score matrix is image x (prefix + image), which the square helper prices at the - # side with the same area. + # The largest call is a reference block's or the target's, both unmasked: only the text blocks carry a mask, + # and they are as short as the prompt. All of them go through diffusers' attention dispatch. The square + # helper prices the call at the side with the same area, rounded up. return sdpa_score_matrix_bytes( device=device, dtype=dtype, num_heads=num_heads, head_dim=head_dim, - seq_len=math.isqrt(image_seq_len * (image_seq_len + prefix)), + seq_len=math.isqrt(largest_call - 1) + 1 if largest_call > 0 else 0, via_diffusers_dispatch=True, ) diff --git a/invokeai/app/invocations/text_encoder/qwen_image_2_1_text_encoder.py b/invokeai/app/invocations/text_encoder/qwen_image_2_1_text_encoder.py index 7b59a30b17e..31b5384f4be 100644 --- a/invokeai/app/invocations/text_encoder/qwen_image_2_1_text_encoder.py +++ b/invokeai/app/invocations/text_encoder/qwen_image_2_1_text_encoder.py @@ -7,13 +7,15 @@ from invokeai.app.invocations.model import ModelIdentifierField, Qwen3VLEncoderField from invokeai.app.invocations.primitives import QwenImage21ConditioningOutput from invokeai.app.services.shared.invocation_context import InvocationContext +from invokeai.backend.architectures import FeaturesFacet, require from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelType, SubModelType from invokeai.backend.quantization.dequantizing_linear import peak_dequant_transient_bytes from invokeai.backend.qwen_image_2_1.text_encoding import ( + attached_vision_tower, encode_prompt, - has_vision_tower, sequence_length, + vision_attention_bytes, working_memory_bytes, ) from invokeai.backend.qwen_image_2_1.vae import to_reference @@ -51,39 +53,58 @@ class QwenImage21TextEncoderInvocation(BaseInvocation): ) reference_images: list[ImageField] = InputField( default=[], - description="Reference images to edit or draw from, in the order the denoise node receives their latents. " - "Needs a Qwen-Image-2.1 Diffusers model's own encoder.", + description="Reference images to edit or draw from, in the order the denoise node receives their latents, " + f"at most {require(BaseModelType.QwenImage21, FeaturesFacet).max_reference_images}. Needs a Qwen-Image-2.1 " + "Diffusers model's own encoder.", ) @torch.no_grad() def invoke(self, context: InvocationContext) -> QwenImage21ConditioningOutput: + max_references = require(BaseModelType.QwenImage21, FeaturesFacet).max_reference_images + if len(self.reference_images) > max_references: + raise ValueError( + f"Qwen-Image-2.1 reads at most {max_references} reference images; this prompt has " + f"{len(self.reference_images)}. Remove references until {max_references} remain." + ) images = [to_reference(context.images.get_pil(field.image_name)) for field in self.reference_images] - processor_id = self._processor_identifier(context) if images else None + tower_id, processor_id = self._edit_identifiers(context) if images else (None, None) # The encoder first: loading it evicts unlocked small entries from the RAM cache, which would otherwise drop - # the tokenizer and processor loaded before it. + # the tokenizer, tower and processor loaded after it. text_encoder_info = context.models.load(self.qwen3_vl_encoder.text_encoder) tokenizer_info = context.models.load(self.qwen3_vl_encoder.tokenizer) + tower_info = context.models.load(tower_id) if tower_id is not None else None processor_info = context.models.load(processor_id) if processor_id is not None else None + if tower_info is not None and tower_info.compute_device != text_encoder_info.compute_device: + raise RuntimeError( + f"The vision tower runs on {tower_info.compute_device} and the encoder on " + f"{text_encoder_info.compute_device}; an edit attaches one to the other, so they need one device." + ) + compute_dtype = TorchDevice.choose_bfloat16_safe_dtype(text_encoder_info.compute_device) # A GGUF or int8 encoder dequantizes one layer at a time; that transient is the working memory. - dequant_bytes = peak_dequant_transient_bytes( - text_encoder_info.model, TorchDevice.choose_bfloat16_safe_dtype(text_encoder_info.compute_device) + dequant_bytes = peak_dequant_transient_bytes(text_encoder_info.model, compute_dtype) + vision_bytes = ( + vision_attention_bytes(tower_info.model, images, tower_info.compute_device, compute_dtype) + if tower_info is not None + else 0 ) - # Before the encoder moves to the GPU: a graph can pair this pipeline's tokenizer with another encoder. - if images and not has_vision_tower(text_encoder_info.model): - raise ValueError(_NO_VISION_TOWER) context.util.signal_progress("Running Qwen3-VL text encoder") with tokenizer_info as tokenizer: # Reference images make the sequence long, and its attention grows with the square of it: room the # model cache has to make, or the driver spills the encode to system memory and it crawls. - working_memory = dequant_bytes + working_memory_bytes(sequence_length(tokenizer, self.prompt, images)) + num_tokens = sequence_length(tokenizer, self.prompt, images) + working_memory = dequant_bytes + working_memory_bytes(num_tokens, compute_dtype.itemsize, vision_bytes) + # The tower locks before the encoder, so the room the encoder reserves is still free once both are on + # the device. with ( processor_info.model_on_device() if processor_info is not None else nullcontext((None, None)) as ( _, processor, ), + tower_info.model_on_device() if tower_info is not None else nullcontext((None, None)) as (_, tower), text_encoder_info.model_on_device(working_mem_bytes=working_memory) as (_, text_encoder), + attached_vision_tower(text_encoder, tower) if tower is not None else nullcontext(), ): encoding = encode_prompt( text_encoder, @@ -102,14 +123,21 @@ def invoke(self, context: InvocationContext) -> QwenImage21ConditioningOutput: name = context.conditioning.save(ConditioningFieldData(conditionings=[conditioning])) return QwenImage21ConditioningOutput.build(name) - def _processor_identifier(self, context: InvocationContext) -> ModelIdentifierField: - """The image processor of the model the tokenizer came from, which has to be a Qwen-Image-2.1 pipeline.""" - tokenizer = self.qwen3_vl_encoder.tokenizer - config = context.models.get_config(tokenizer) - if ( - config.base is not BaseModelType.QwenImage21 - or config.type is not ModelType.Main - or config.format is not ModelFormat.Diffusers - ): - raise ValueError(_NO_VISION_TOWER) - return tokenizer.model_copy(update={"submodel_type": SubModelType.Processor}) + def _edit_identifiers(self, context: InvocationContext) -> tuple[ModelIdentifierField, ModelIdentifierField]: + """The vision tower and the image processor an edit adds, which only a Qwen-Image-2.1 pipeline brings. + + Checked on the configs, before anything loads: a graph can pair a standalone encoder with this pipeline's + tokenizer, or the other way round. + """ + for identifier in (self.qwen3_vl_encoder.tokenizer, self.qwen3_vl_encoder.text_encoder): + config = context.models.get_config(identifier) + if ( + config.base is not BaseModelType.QwenImage21 + or config.type is not ModelType.Main + or config.format is not ModelFormat.Diffusers + ): + raise ValueError(_NO_VISION_TOWER) + return ( + self.qwen3_vl_encoder.text_encoder.model_copy(update={"submodel_type": SubModelType.VisionTower}), + self.qwen3_vl_encoder.tokenizer.model_copy(update={"submodel_type": SubModelType.Processor}), + ) diff --git a/invokeai/backend/model_manager/load/load_default.py b/invokeai/backend/model_manager/load/load_default.py index 5f68261da37..1ba0800730f 100644 --- a/invokeai/backend/model_manager/load/load_default.py +++ b/invokeai/backend/model_manager/load/load_default.py @@ -51,6 +51,9 @@ # quality, which fp8 rounding degrades. Its tokenizer is listed for symmetry (it is not an nn.Module, # so casting is a no-op today only by accident). # +# The vision tower follows the text encoder it attaches to. It also could not take the cast as it stands: +# `Qwen3VLModel` casts the pixels to `visual.dtype`, which a cast tower reports as float8. +# # Unlike the model-kind rules in `fp8_capability.py`, these are about which *component* of a model is # being loaded, which no client can see and nothing outside this file asks about. _FP8_EXCLUDED_SUBMODEL_TYPES: frozenset[SubModelType] = frozenset( @@ -58,6 +61,7 @@ SubModelType.TextEncoder, SubModelType.TextEncoder2, SubModelType.TextEncoder3, + SubModelType.VisionTower, SubModelType.Tokenizer, SubModelType.Tokenizer2, SubModelType.Tokenizer3, @@ -342,7 +346,12 @@ def _get_execution_device( if hasattr(config, "default_settings") and config.default_settings is not None: if hasattr(config.default_settings, "cpu_only") and config.default_settings.cpu_only is True: # Only apply CPU execution to text encoder submodels - if submodel_type in [SubModelType.TextEncoder, SubModelType.TextEncoder2, SubModelType.TextEncoder3]: + if submodel_type in [ + SubModelType.TextEncoder, + SubModelType.TextEncoder2, + SubModelType.TextEncoder3, + SubModelType.VisionTower, + ]: return torch.device("cpu") # Check if this is a standalone text encoder config with cpu_only field (T5Encoder, Qwen3Encoder, etc.) diff --git a/invokeai/backend/model_manager/load/model_loaders/qwen_image_2_1.py b/invokeai/backend/model_manager/load/model_loaders/qwen_image_2_1.py index 08ff7d29996..549ee7f5be0 100644 --- a/invokeai/backend/model_manager/load/model_loaders/qwen_image_2_1.py +++ b/invokeai/backend/model_manager/load/model_loaders/qwen_image_2_1.py @@ -1,5 +1,7 @@ """Loaders for Qwen-Image-2.1 main models: diffusers pipelines, single-file and GGUF transformers.""" +import json +import math from pathlib import Path from typing import Any, Optional @@ -18,6 +20,7 @@ from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry from invokeai.backend.model_manager.load.model_loaders.generic_diffusers import GenericDiffusersLoader from invokeai.backend.model_manager.taxonomy import AnyModel, BaseModelType, ModelFormat, ModelType, SubModelType +from invokeai.backend.model_manager.util.qwen3_vl import drop_qwen3vl_visual_tower from invokeai.backend.quantization.fp8_scaled import ( attach_fp8_scales, cast_state_dict, @@ -66,10 +69,24 @@ class QwenImage21DiffusersModel(GenericDiffusersLoader): rather than the declared `Qwen3VLForConditionalGeneration`: the LM head is ~1.2 GiB the encoder never runs, and `Qwen3VLModel` is what the standalone Qwen3-VL loaders return, so the text encoder node sees one class. - Unlike those loaders, this one keeps the vision tower (~1.1 GiB): reference images reach the encoder only - through it, and only this pipeline's encoder carries one into Qwen-Image-2.1. + `TextEncoder` drops the vision tower, as the standalone loaders do: text prompts never reach it. + `VisionTower` is that tower alone (~1.1 GiB), read from the same folder for edits, which attach it to the + cached encoder while they encode. The language model's weights are cached once for both. """ + def get_size_fs( + self, config: AnyModelConfig, model_path: Path, submodel_type: Optional[SubModelType] = None + ) -> int: + if submodel_type is SubModelType.VisionTower: + # A few hundred tensors of the encoder's folder, not the folder: reserving its ~17 GB for the tower + # would evict models for nothing. Sized at the dtype they load in, whatever the file stores. + element_size = TorchDevice.choose_bfloat16_safe_dtype(TorchDevice.choose_torch_device()).itemsize + return sum( + math.prod(shape) * element_size + for _shard, _key, shape in _vision_tower_tensors(model_path / SubModelType.TextEncoder.value) + ) + return super().get_size_fs(config, model_path, submodel_type) + def _load_model(self, config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> AnyModel: if isinstance(config, Checkpoint_Config_Base): raise NotImplementedError( @@ -92,6 +109,11 @@ def _load_model(self, config: AnyModelConfig, submodel_type: Optional[SubModelTy dtype = self._torch_dtype_avoiding_float16() else: dtype = TorchDevice.choose_bfloat16_safe_dtype(TorchDevice.choose_torch_device()) + + if submodel_type is SubModelType.VisionTower: + tower = _load_vision_tower(model_path / SubModelType.TextEncoder.value, dtype) + return self._apply_fp8_layerwise_casting(tower, config, submodel_type) + repo_variant = config.repo_variant if isinstance(config, Diffusers_Config_Base) else None variant = repo_variant.value if repo_variant else None if submodel_type is SubModelType.TextEncoder: @@ -114,9 +136,65 @@ def _load_model(self, config: AnyModelConfig, submodel_type: Optional[SubModelTy else: raise + if submodel_type is SubModelType.TextEncoder: + # Text prompts never reach the vision tower; edits load it as `VisionTower` and attach it while they + # encode. + drop_qwen3vl_visual_tower(result) + return self._apply_fp8_layerwise_casting(result, config, submodel_type) +# The tower's tensors in an encoder folder: `model.visual.*` as `Qwen3VLForConditionalGeneration` saves them (the +# released pipeline), `visual.*` as `Qwen3VLModel` does. +_VISION_TOWER_PREFIXES = ("model.visual.", "visual.") + + +def _vision_tower_tensors(folder: Path) -> list[tuple[Path, str, list[int]]]: + """(shard, key, shape) of every vision-tower tensor in the folder, read from the headers. + + The shard index names the shards that hold them; without an index, every safetensors file in the folder. + """ + from safetensors import safe_open + + index = folder / "model.safetensors.index.json" + if index.exists(): + weight_map: dict[str, str] = json.loads(index.read_text(encoding="utf-8"))["weight_map"] + shards = sorted({folder / shard for key, shard in weight_map.items() if key.startswith(_VISION_TOWER_PREFIXES)}) + else: + shards = sorted(folder.glob("*.safetensors")) + tensors = [] + for shard in shards: + with safe_open(shard, framework="pt", device="cpu") as f: + for key in f.keys(): + if key.startswith(_VISION_TOWER_PREFIXES): + tensors.append((shard, key, f.get_slice(key).get_shape())) + if not tensors: + raise ValueError(f"{folder} holds no Qwen3-VL vision tower ({' or '.join(_VISION_TOWER_PREFIXES)} tensors).") + return tensors + + +def _load_vision_tower(folder: Path, dtype: torch.dtype) -> torch.nn.Module: + """The encoder's vision tower alone, as `Qwen3VLModel.visual` would hold it, without reading the language model.""" + from safetensors import safe_open + from transformers import AutoConfig + from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLVisionModel + + vision_config = AutoConfig.from_pretrained(folder, local_files_only=True).vision_config + with accelerate.init_empty_weights(): + tower = Qwen3VLVisionModel._from_config(vision_config, torch_dtype=dtype) + + state_dict: dict[str, torch.Tensor] = {} + tensors = _vision_tower_tensors(folder) + for shard in sorted({shard for shard, *_ in tensors}): + with safe_open(shard, framework="pt", device="cpu") as f: + for _shard, key, _shape in (t for t in tensors if t[0] == shard): + prefix = next(p for p in _VISION_TOWER_PREFIXES if key.startswith(p)) + state_dict[key[len(prefix) :]] = f.get_tensor(key).to(dtype) + load_state_dict_ignoring_extras(tower, state_dict, source=f"Qwen3-VL vision tower in {folder}", assign=True) + reject_incomplete_load(tower, what="Qwen3-VL vision tower") + return tower.eval() + + @ModelLoaderRegistry.register(base=BaseModelType.QwenImage21, type=ModelType.Main, format=ModelFormat.Checkpoint) class QwenImage21CheckpointModel(ModelLoader): """Loads a Qwen-Image-2.1 transformer from a single safetensors file. diff --git a/invokeai/backend/model_manager/taxonomy.py b/invokeai/backend/model_manager/taxonomy.py index 3a328b436ed..4a45c175c52 100644 --- a/invokeai/backend/model_manager/taxonomy.py +++ b/invokeai/backend/model_manager/taxonomy.py @@ -123,6 +123,9 @@ class SubModelType(str, Enum): TextEncoder = "text_encoder" TextEncoder2 = "text_encoder_2" TextEncoder3 = "text_encoder_3" + # A vision-language encoder's vision tower alone, which the text encoder is loaded without; edits attach it + # while they encode, so the language model is cached once. Qwen-Image-2.1 only. + VisionTower = "vision_tower" Tokenizer = "tokenizer" Tokenizer2 = "tokenizer_2" Tokenizer3 = "tokenizer_3" diff --git a/invokeai/backend/qwen_image_2_1/text_encoding.py b/invokeai/backend/qwen_image_2_1/text_encoding.py index a9f145ef78b..bb1a80c7b31 100644 --- a/invokeai/backend/qwen_image_2_1/text_encoding.py +++ b/invokeai/backend/qwen_image_2_1/text_encoding.py @@ -1,11 +1,14 @@ """Qwen-Image-2.1's prompt encoding, as `QwenImage21Pipeline._get_qwen_prompt_embeds` does it.""" -from collections.abc import Sequence +from collections.abc import Iterator, Sequence +from contextlib import contextmanager from typing import NamedTuple import torch from PIL import Image +from invokeai.backend.util.attention import sdpa_score_matrix_bytes + SYSTEM_PROMPT = "Comprehend and analyze the provided prompt." SYSTEM_BLOCK = f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n" # A raw string handed to the tokenizer, not `apply_chat_template`: the two tokenize differently, and the @@ -14,11 +17,14 @@ IMAGE_PAD = "<|image_pad|>" # Pixels per side of one vision token: 16px patches, merged 2x2. _PIXELS_PER_TOKEN = 32 +# The vision tower reads each reference's 16px patches before the 2x2 merge: four per vision token. +_PATCHES_PER_TOKEN = 4 -# What one encode needs above the weights for a sequence of T tokens, measured on an RTX 4090 (torch 2.13, +# What one encode needs above the weights for a sequence of T tokens, measured on an RTX 4090 in bf16 (torch 2.13, # transformers 5.10): ~286 B * T^2 + 224 KiB * T at peak. The quadratic term is the language model's attention, # which SDPA runs on its math kernel for grouped-query heads; four 1-megapixel references (~4.1k tokens) need -# ~5.4 GiB. The allocator reserves about a fifth more than it hands out. +# ~5.4 GiB. Both terms are activations, so they scale with the compute dtype's width. The allocator reserves about +# a fifth more than it hands out. _BYTES_PER_TOKEN_PAIR = 286 _BYTES_PER_TOKEN = 224 * 1024 _RESERVED_OVER_ALLOCATED = 1.25 @@ -26,6 +32,10 @@ def format_prompt(prompt: str, num_images: int = 0) -> str: """The prompt as the checkpoint was trained on it, with one vision placeholder per reference image first.""" + if num_images and IMAGE_PAD in prompt: + # The processor expands every placeholder into one image's tokens, so a literal one would claim an image + # that is not there. + raise ValueError(f"A prompt that reads reference images cannot contain {IMAGE_PAD}. Remove it from the prompt.") placeholders = " ".join(f"<|vision_start|>{IMAGE_PAD}<|vision_end|>" for i in range(1, num_images + 1)) # Qwen has no BOS token, so an empty prompt would leave the encoder nothing to read. return T2I_TEMPLATE.format(placeholders + (prompt or " ")) @@ -40,9 +50,35 @@ def sequence_length(tokenizer, prompt: str, images: Sequence[Image.Image] = ()) return text + sum((w // _PIXELS_PER_TOKEN) * (h // _PIXELS_PER_TOKEN) for w, h in (i.size for i in images)) -def working_memory_bytes(num_tokens: int) -> int: - """Device memory an encode of `num_tokens` tokens needs above the encoder's weights.""" - return int((_BYTES_PER_TOKEN_PAIR * num_tokens**2 + _BYTES_PER_TOKEN * num_tokens) * _RESERVED_OVER_ALLOCATED) +def vision_attention_bytes( + tower: torch.nn.Module, images: Sequence[Image.Image], device: torch.device, dtype: torch.dtype +) -> int: + """The score matrix the vision tower's attention materializes for the largest reference, 0 where it is fused. + + The tower attends within one image at a time, over its patches: four per vision token, ~4k for 1 megapixel. + """ + if not images: + return 0 + patches = max((w // _PIXELS_PER_TOKEN) * (h // _PIXELS_PER_TOKEN) for w, h in (i.size for i in images)) + return sdpa_score_matrix_bytes( + device=device, + dtype=dtype, + num_heads=tower.config.num_heads, + head_dim=tower.config.hidden_size // tower.config.num_heads, + seq_len=_PATCHES_PER_TOKEN * patches, + ) + + +def working_memory_bytes(num_tokens: int, element_size: int, vision_attention: int = 0) -> int: + """Device memory an encode of `num_tokens` tokens needs above the encoder's weights. + + `element_size` is the compute dtype's width. `vision_attention` is the tower's score matrix + (`vision_attention_bytes`): the tower runs, and frees it, before the language model starts, so only the larger + of the two attention peaks counts. + """ + width = element_size / 2 + attention = max(_BYTES_PER_TOKEN_PAIR * num_tokens**2 * width, vision_attention) + return int((attention + _BYTES_PER_TOKEN * num_tokens * width) * _RESERVED_OVER_ALLOCATED) def system_prefix_length(tokenizer) -> int: @@ -61,14 +97,31 @@ def _base_model(text_encoder: torch.nn.Module) -> torch.nn.Module: def has_vision_tower(text_encoder: torch.nn.Module) -> bool: - """Whether this encoder can read reference images: only the Qwen-Image-2.1 pipeline's keeps its vision tower. + """Whether this encoder can read reference images: whether a vision tower is attached (`attached_vision_tower`). - The standalone Qwen3-VL loaders replace the tower with an `Identity` rather than deleting it. + Every encoder loads with an `Identity` where the tower was. """ visual = getattr(_base_model(text_encoder), "visual", None) return isinstance(visual, torch.nn.Module) and not isinstance(visual, torch.nn.Identity) +@contextmanager +def attached_vision_tower(text_encoder: torch.nn.Module, tower: torch.nn.Module) -> Iterator[None]: + """Attach the pipeline's vision tower to its cached encoder for one encode. + + The encoder and the tower are separate model-cache entries, so the language model is cached once for text + prompts and edits. Both must be locked on the same device for the whole block. Whatever the block does, the + encoder leaves it with the module it had, so the cache never holds an encoder with a tower it does not own. + """ + base = _base_model(text_encoder) + previous = base.visual + base.visual = tower + try: + yield + finally: + base.visual = previous + + def on_white(image: Image.Image) -> Image.Image: """A reference as the vision encoder reads it: alpha composited over white, as the checkpoint was trained.""" if image.mode != "RGBA": diff --git a/invokeai/frontend/api/openapi.json b/invokeai/frontend/api/openapi.json index af354195dc7..ef1d9ea1be6 100644 --- a/invokeai/frontend/api/openapi.json +++ b/invokeai/frontend/api/openapi.json @@ -57786,7 +57786,7 @@ "$ref": "#/components/schemas/PromptTemplateOutput" }, "qwen_image_2_1_denoise": { - "$ref": "#/components/schemas/LatentsOutput" + "$ref": "#/components/schemas/LatentsMetaOutput" }, "qwen_image_2_1_i2l": { "$ref": "#/components/schemas/LatentsOutput" @@ -96401,7 +96401,7 @@ "type": "object", "version": "1.1.0", "output": { - "$ref": "#/components/schemas/LatentsOutput" + "$ref": "#/components/schemas/LatentsMetaOutput" } }, "QwenImage21ImageToLatentsInvocation": { @@ -96877,7 +96877,7 @@ }, "reference_images": { "default": [], - "description": "Reference images to edit or draw from, in the order the denoise node receives their latents. Needs a Qwen-Image-2.1 Diffusers model's own encoder.", + "description": "Reference images to edit or draw from, in the order the denoise node receives their latents, at most 4. Needs a Qwen-Image-2.1 Diffusers model's own encoder.", "field_kind": "input", "input": "any", "items": { @@ -99706,7 +99706,7 @@ "type": "object", "required": ["image_name"], "title": "ReferenceImageRecallParameter", - "description": "Global reference-image configuration for recall.\n\nUsed for reference images that feed directly into the main model rather\nthan through a separate IP-Adapter / ControlNet model \u2014 for example\nFLUX.2 Klein, FLUX Kontext, and Qwen Image Edit. The receiving frontend\npicks the correct config type (``flux2_reference_image`` /\n``qwen_image_reference_image`` / ``flux_kontext_reference_image``) based\non the currently-selected main model." + "description": "Global reference-image configuration for recall.\n\nUsed for reference images that feed directly into the main model rather\nthan through a separate IP-Adapter / ControlNet model \u2014 for example\nFLUX.2 Klein, FLUX Kontext, Qwen Image Edit and Qwen-Image-2.1. The\nreceiving frontend picks the correct config type (``flux2_reference_image`` /\n``qwen_image_reference_image`` / ``qwen_image_2_1_reference_image`` /\n``flux_kontext_reference_image``) based on the currently-selected main model." }, "RemoteModelFile": { "properties": { @@ -106081,6 +106081,7 @@ "text_encoder", "text_encoder_2", "text_encoder_3", + "vision_tower", "tokenizer", "tokenizer_2", "tokenizer_3", diff --git a/invokeai/frontend/api/schema.ts b/invokeai/frontend/api/schema.ts index cba788848f4..a97d0d0aeba 100644 --- a/invokeai/frontend/api/schema.ts +++ b/invokeai/frontend/api/schema.ts @@ -22998,7 +22998,7 @@ export type components = { pidi_edge_detection: components["schemas"]["ImageOutput"]; prompt_from_file: components["schemas"]["StringCollectionOutput"]; prompt_template: components["schemas"]["PromptTemplateOutput"]; - qwen_image_2_1_denoise: components["schemas"]["LatentsOutput"]; + qwen_image_2_1_denoise: components["schemas"]["LatentsMetaOutput"]; qwen_image_2_1_i2l: components["schemas"]["LatentsOutput"]; qwen_image_2_1_l2i: components["schemas"]["ImageOutput"]; qwen_image_2_1_model_loader: components["schemas"]["QwenImage21ModelLoaderOutput"]; @@ -41892,7 +41892,7 @@ export type components = { qwen3_vl_encoder?: components["schemas"]["Qwen3VLEncoderField"] | null; /** * Reference Images - * @description Reference images to edit or draw from, in the order the denoise node receives their latents. Needs a Qwen-Image-2.1 Diffusers model's own encoder. + * @description Reference images to edit or draw from, in the order the denoise node receives their latents, at most 4. Needs a Qwen-Image-2.1 Diffusers model's own encoder. * @default [] */ reference_images?: components["schemas"]["ImageField"][]; @@ -43354,10 +43354,10 @@ export type components = { * * Used for reference images that feed directly into the main model rather * than through a separate IP-Adapter / ControlNet model — for example - * FLUX.2 Klein, FLUX Kontext, and Qwen Image Edit. The receiving frontend - * picks the correct config type (``flux2_reference_image`` / - * ``qwen_image_reference_image`` / ``flux_kontext_reference_image``) based - * on the currently-selected main model. + * FLUX.2 Klein, FLUX Kontext, Qwen Image Edit and Qwen-Image-2.1. The + * receiving frontend picks the correct config type (``flux2_reference_image`` / + * ``qwen_image_reference_image`` / ``qwen_image_2_1_reference_image`` / + * ``flux_kontext_reference_image``) based on the currently-selected main model. */ ReferenceImageRecallParameter: { /** @@ -46586,7 +46586,7 @@ export type components = { * @description Submodel type. * @enum {string} */ - SubModelType: "unet" | "transformer" | "transformer_2" | "text_encoder" | "text_encoder_2" | "text_encoder_3" | "tokenizer" | "tokenizer_2" | "tokenizer_3" | "processor" | "pe" | "pe_tokenizer" | "vae" | "vae_decoder" | "vae_encoder" | "audio_vae" | "vocoder" | "connectors" | "latent_upsampler" | "temporal_latent_upsampler" | "scheduler" | "safety_checker"; + SubModelType: "unet" | "transformer" | "transformer_2" | "text_encoder" | "text_encoder_2" | "text_encoder_3" | "vision_tower" | "tokenizer" | "tokenizer_2" | "tokenizer_3" | "processor" | "pe" | "pe_tokenizer" | "vae" | "vae_decoder" | "vae_encoder" | "audio_vae" | "vocoder" | "connectors" | "latent_upsampler" | "temporal_latent_upsampler" | "scheduler" | "safety_checker"; /** SubmodelDefinition */ SubmodelDefinition: { /** Path Or Prefix */ diff --git a/invokeai/frontend/webv2/public/locales/en.json b/invokeai/frontend/webv2/public/locales/en.json index 6b04ac9322e..be79fdeb608 100644 --- a/invokeai/frontend/webv2/public/locales/en.json +++ b/invokeai/frontend/webv2/public/locales/en.json @@ -2832,6 +2832,7 @@ "referenceImage": "Reference Image", "referenceImageModel": "Reference image model", "referenceImages": "Reference images", + "referenceImagesNeedFormat": "Reference images need a {{format}} {{base}} model.", "referenceImagesUnsupported": "Reference images are not supported by the selected model.", "removeConcept": "Remove concept", "removeConceptNamed": "Remove {{name}}", diff --git a/invokeai/frontend/webv2/src/features/generation/core/__snapshots__/generateGraphNodeTypes.json b/invokeai/frontend/webv2/src/features/generation/core/__snapshots__/generateGraphNodeTypes.json index 6d79322aa91..80aac9dfcd9 100644 --- a/invokeai/frontend/webv2/src/features/generation/core/__snapshots__/generateGraphNodeTypes.json +++ b/invokeai/frontend/webv2/src/features/generation/core/__snapshots__/generateGraphNodeTypes.json @@ -265,11 +265,14 @@ ] }, "nodeTypes": [ + "collect", "core_metadata", + "image", "img_nsfw", "img_watermark", "integer", "qwen_image_2_1_denoise", + "qwen_image_2_1_i2l", "qwen_image_2_1_l2i", "qwen_image_2_1_model_loader", "qwen_image_2_1_text_encoder", @@ -1252,6 +1255,21 @@ "conditioning" ] }, + "image": { + "inputs": [], + "literalInputs": { + "image": [], + "is_intermediate": [ + true + ], + "use_cache": [ + true + ] + }, + "outputs": [ + "image" + ] + }, "img_nsfw": { "inputs": [ "image" @@ -1463,13 +1481,17 @@ }, "qwen_image_2_1_denoise": { "inputs": [ + "metadata", + "negative_conditioning", "positive_conditioning", + "reference_latents", "seed", "transformer" ], "literalInputs": { "cfg_scale": [ - 1 + 1, + 4 ], "height": [ 1024 @@ -1487,6 +1509,27 @@ 1024 ] }, + "outputs": [ + "latents", + "metadata" + ] + }, + "qwen_image_2_1_i2l": { + "inputs": [ + "image", + "vae" + ], + "literalInputs": { + "is_intermediate": [ + true + ], + "reference": [ + true + ], + "use_cache": [ + true + ] + }, "outputs": [ "latents" ] @@ -1532,7 +1575,8 @@ "qwen_image_2_1_text_encoder": { "inputs": [ "prompt", - "qwen3_vl_encoder" + "qwen3_vl_encoder", + "reference_images" ], "literalInputs": { "is_intermediate": [ diff --git a/invokeai/frontend/webv2/src/features/generation/core/baseGenerationPolicies.test.ts b/invokeai/frontend/webv2/src/features/generation/core/baseGenerationPolicies.test.ts index d0d799bc87d..80b62638d1b 100644 --- a/invokeai/frontend/webv2/src/features/generation/core/baseGenerationPolicies.test.ts +++ b/invokeai/frontend/webv2/src/features/generation/core/baseGenerationPolicies.test.ts @@ -29,6 +29,7 @@ import { getGenerationUiPolicy, getGenerationValidationReasons, getMaxReferenceImages, + getMissingReferenceImageFormat, getPromptPolicy, getSchedulerOptions, getSettingsWithModelDefaults, @@ -1158,6 +1159,20 @@ describe('component policies', () => { expect(isReferenceImageSupported(createModel('qwen-image-2-1', { format: 'checkpoint' }))).toBe(false); }); + // The UI names the missing format only when switching to it would enable references. + it('names the format a model lacks for reference images only when the format alone blocks them', () => { + expect(getMissingReferenceImageFormat(createModel('qwen-image-2-1', { format: 'gguf_quantized' }))).toBe( + 'diffusers' + ); + expect(getMissingReferenceImageFormat(createModel('qwen-image-2-1', { format: 'checkpoint' }))).toBe('diffusers'); + expect(getMissingReferenceImageFormat(createModel('qwen-image-2-1', { format: 'diffusers' }))).toBeNull(); + // Qwen-Image needs the Edit variant, not a format; CogView4 takes no references at all. + expect(getMissingReferenceImageFormat(createModel('qwen-image', { format: 'gguf_quantized' }))).toBeNull(); + expect(getMissingReferenceImageFormat(createModel('cogview4', { format: 'gguf_quantized' }))).toBeNull(); + expect(getMissingReferenceImageFormat(externalModel)).toBeNull(); + expect(getMissingReferenceImageFormat(undefined)).toBeNull(); + }); + it('derives the reference image limit from external provider capabilities', () => { expect(getMaxReferenceImages(createModel('sdxl'))).toBe(5); expect(getMaxReferenceImages(createModel('cogview4'))).toBe(0); diff --git a/invokeai/frontend/webv2/src/features/generation/core/baseGenerationPolicies.ts b/invokeai/frontend/webv2/src/features/generation/core/baseGenerationPolicies.ts index a97cc615c6d..149421d0d20 100644 --- a/invokeai/frontend/webv2/src/features/generation/core/baseGenerationPolicies.ts +++ b/invokeai/frontend/webv2/src/features/generation/core/baseGenerationPolicies.ts @@ -1333,6 +1333,28 @@ export const isReferenceImageSupported = (model: GenerateModelConfig | undefined ); }; +/** + * The model format the selected base needs for reference images, when the model's format is all that keeps them + * unavailable, so the UI can say what to switch to instead of hiding the section without a reason. + */ +export const getMissingReferenceImageFormat = (model: GenerateModelConfig | undefined): string | null => { + if (!model || model.type === 'external_image_generator' || isReferenceImageSupported(model)) { + return null; + } + + const features = getArchitectureFeatures(model.base); + + if ( + !features || + features.max_reference_images <= 0 || + (features.reference_images_require_variant !== null && model.variant !== features.reference_images_require_variant) + ) { + return null; + } + + return features.reference_images_require_format; +}; + export const getMaxReferenceImages = (model: GenerateModelConfig | undefined): number => { if (!model || !isReferenceImageSupported(model)) { return 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..5821f83fb2a 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 @@ -712,6 +712,22 @@ describe('compileCanvasGraph — inpaint per base', () => { expect(metaEdge?.destination.node_id).toBe('canvas_output'); }); + // Qwen-Image-2.1's denoise adds its cache mode to the record, so the record still passes through it once the blend + // takes over the output. + it('routes Qwen-Image-2.1 inpaint metadata through the denoise onto the final blend output', () => { + const { backendGraph } = compile(qwenImage21Model, 'inpaint'); + const metadata = getNodeByType(backendGraph, 'core_metadata'); + const metadataEdges = backendGraph.edges.filter( + (edge) => edge.source.field === 'metadata' || edge.destination.field === 'metadata' + ); + + expect(backendGraph.nodes.canvas_output?.type).toBe('invokeai_img_blend'); + expect(metadataEdges.map(({ destination, source }) => [source.node_id, destination.node_id])).toEqual([ + [metadata?.id, 'denoise_latents'], + ['denoise_latents', 'canvas_output'], + ]); + }); + it('outputs only the generated mask region with transparency when requested', () => { const { backendGraph } = compile(sd1Model, 'inpaint', { outputOnlyMaskedRegions: true }); diff --git a/invokeai/frontend/webv2/src/features/generation/core/graph.test.ts b/invokeai/frontend/webv2/src/features/generation/core/graph.test.ts index 625e857b683..679b87ffac9 100644 --- a/invokeai/frontend/webv2/src/features/generation/core/graph.test.ts +++ b/invokeai/frontend/webv2/src/features/generation/core/graph.test.ts @@ -813,18 +813,18 @@ describe('compileGenerateGraph', () => { expect(latentImages).toEqual(['a.png', 'b.png']); }); - it('refuses references with a standalone encoder, which has no vision tower', () => { - const encoder: ComponentModelConfig = { - base: 'any', - key: 'qwen3-vl-8b', - name: 'Qwen3-VL 8B', - type: 'qwen3_vl_encoder', - variant: 'qwen3_vl_8b', - }; + // The denoise records whether it reused the prompt's attention cache, which changes the image in bf16. + it('passes the image metadata through the denoise to the decoder that saves the image', () => { + const graph = compile(qwenImage21Model); + const metadata = getNodeByType(graph, 'core_metadata'); + const metadataEdges = graph.edges.filter( + (edge) => edge.source.field === 'metadata' || edge.destination.field === 'metadata' + ); - expect(() => - compile(qwenImage21Model, { qwen3VLEncoderModel: encoder, referenceImages: references('a.png') }) - ).toThrow(/vision tower/); + expect(metadataEdges.map(({ destination, source }) => [source.node_id, destination.node_id])).toEqual([ + [metadata?.id, 'denoise_latents'], + ['denoise_latents', 'canvas_output'], + ]); }); }); }); @@ -1818,7 +1818,10 @@ describe('output filters', () => { node_id: 'canvas_unfiltered_output', }); expect(getEdge(graph, 'canvas_output', 'image')?.source).toEqual({ field: 'image', node_id: 'canvas_img_nsfw' }); - expect(getEdge(graph, 'canvas_output', 'metadata')?.source.node_id).toBe(getNodeByType(graph, 'core_metadata')?.id); + // Qwen-Image-2.1's denoise adds its cache mode to the record on the way, so the record reaches the result from it. + const metadataSource = + base === 'qwen-image-2-1' ? 'denoise_latents' : (getNodeByType(graph, 'core_metadata')?.id ?? ''); + expect(getEdge(graph, 'canvas_output', 'metadata')?.source).toEqual({ field: 'metadata', node_id: metadataSource }); expect(getEdge(graph, 'canvas_unfiltered_output', 'metadata')).toBeUndefined(); }); diff --git a/invokeai/frontend/webv2/src/features/generation/core/graph.ts b/invokeai/frontend/webv2/src/features/generation/core/graph.ts index fcf2563b428..c4b0151a212 100644 --- a/invokeai/frontend/webv2/src/features/generation/core/graph.ts +++ b/invokeai/frontend/webv2/src/features/generation/core/graph.ts @@ -25,7 +25,6 @@ import { getFlux2DiffusersComponentSource, getGenerationDimensions, getGenerationValidationReasons, - QWEN_IMAGE_21_STANDALONE_ENCODER_REASON, type SupportedGenerateBase, } from './baseGenerationPolicies'; import { @@ -1442,13 +1441,15 @@ const buildQwenImage21Graph = ( addQwenImage21ReferenceImages(graph, settings, { conditionings: negCond ? [posCond, negCond] : [posCond], denoise, - hasStandaloneEncoder: Boolean(qwen3VlEncoderModel), modelLoader, }); - addMetadata(graph, output, settings, model, 'qwen_image_2_1_txt2img', projectSettings, { + // The denoise adds the prefix-cache mode it ran to the record, which changes the image, so the metadata passes + // through it on the way to the decoder. + addMetadata(graph, denoise, settings, model, 'qwen_image_2_1_txt2img', projectSettings, { qwen3_vl_encoder: qwen3VlEncoderModel ?? undefined, vae: vaeModel ?? undefined, }); + addEdge(graph, denoise, 'metadata', output, 'metadata'); addReferenceImageMetadata(graph, output, settings); return graph; @@ -1465,7 +1466,6 @@ const addQwenImage21ReferenceImages = ( nodes: { conditionings: BackendInvocationContract[]; denoise: BackendInvocationContract; - hasStandaloneEncoder: boolean; modelLoader: BackendInvocationContract; } ): void => { @@ -1473,10 +1473,6 @@ const addQwenImage21ReferenceImages = ( if (referenceImages.length === 0) { return; } - // A backstop: validation reports this before Invoke. - if (nodes.hasStandaloneEncoder) { - throw new Error(QWEN_IMAGE_21_STANDALONE_ENCODER_REASON); - } let imageCollect: BackendInvocationContract | null = null; let latentsCollect: BackendInvocationContract | null = null; diff --git a/invokeai/frontend/webv2/src/features/generation/core/graphCoverage.testing.ts b/invokeai/frontend/webv2/src/features/generation/core/graphCoverage.testing.ts index 81e4ec21c27..885e623419b 100644 --- a/invokeai/frontend/webv2/src/features/generation/core/graphCoverage.testing.ts +++ b/invokeai/frontend/webv2/src/features/generation/core/graphCoverage.testing.ts @@ -6,11 +6,18 @@ import type { GenerateComponentValueKey, SupportedGenerateBase, } from './baseGenerationPolicies'; -import type { GenerateSettings, MainModelConfig, ModelIdentifierConfig } from './types'; +import type { + GenerateReferenceImage, + GenerateReferenceImageAsset, + GenerateSettings, + MainModelConfig, + ModelIdentifierConfig, +} from './types'; import { getComponentSectionPolicy, getDefaultGenerateSettings, + isReferenceImageSupported, SUPPORTED_GENERATE_BASES, } from './baseGenerationPolicies'; import { getIsPidSupportedBase } from './pid'; @@ -210,6 +217,23 @@ export const fullyFilledSettingsFor = (base: SupportedGenerateBase, shape: Model return { filled: [...allFilled].sort(), model, settings }; }; +/** References for the families whose builders wire them through their own nodes rather than an adapter. */ +const tunedReferenceImagesFor = ( + model: MainModelConfig, + image: GenerateReferenceImageAsset +): GenerateReferenceImage[] | undefined => { + if (model.base === 'krea-2') { + return [{ config: { image, styleStrength: 1, type: 'krea2_reference_image' }, id: 'style', isEnabled: true }]; + } + + // Only the Diffusers pipeline's own encoder reads references; other shapes compile without them. + if (model.base === 'qwen-image-2-1' && isReferenceImageSupported(model)) { + return [{ config: { image, type: 'qwen_image_2_1_reference_image' }, id: 'edit', isEnabled: true }]; + } + + return undefined; +}; + /** * The optional per-family tuning and output filters, all switched on, so the backend contract also covers the nodes * and literal values a default recipe never emits. Settings a family ignores are harmless to carry. @@ -224,21 +248,14 @@ export const tunedSettingsFor = (base: SupportedGenerateBase, shape: ModelShape) projectSettings: { useCpuNoise: true, useNsfwChecker: true, useWatermark: true }, settings: { ...settings, + // CFG above 1 adds the negative prompt, which Qwen-Image-2.1 also encodes with the references. + cfgScale: base === 'qwen-image-2-1' ? 4 : settings.cfgScale, fluxDypeExponent: 1, fluxDypePreset: 'manual' as const, fluxDypeScale: 3, qwenImageEncoderQuantization: 'nf4' as const, qwenImageShift: 3, - referenceImages: - base === 'krea-2' - ? [ - { - config: { image, styleStrength: 1, type: 'krea2_reference_image' as const }, - id: 'style', - isEnabled: true, - }, - ] - : settings.referenceImages, + referenceImages: tunedReferenceImagesFor(model, image) ?? settings.referenceImages, zImageSeedVarianceEnabled: true, zImageSeedVarianceStrength: 0.2, zImageShift: 3, diff --git a/invokeai/frontend/webv2/src/features/generation/core/settings.test.ts b/invokeai/frontend/webv2/src/features/generation/core/settings.test.ts index f5d373e07d4..0f81cc4707a 100644 --- a/invokeai/frontend/webv2/src/features/generation/core/settings.test.ts +++ b/invokeai/frontend/webv2/src/features/generation/core/settings.test.ts @@ -191,6 +191,46 @@ describe('normalizeGenerateSettings', () => { ).toBe(false); }); + // Qwen-Image-2.1 references are their own flavor, so a stored one must not fold into Qwen-Image Edit's. + it('keeps persisted Qwen-Image-2.1 reference images as their own flavor', () => { + const normalized = normalizeGenerateSettings({ + ...legacyStoredValues, + referenceImages: [ + { + id: 'edit-1', + isEnabled: false, + config: { + type: 'qwen_image_2_1_reference_image', + image: { + height: 640, + imageName: 'edit.png', + imageUrl: '/api/v1/images/i/edit.png/full', + queuedAt: '2026-01-01T00:00:00.000Z', + sourceQueueItemId: 'backend-gallery', + thumbnailUrl: '/api/v1/images/i/edit.png/thumbnail', + width: 480, + }, + }, + }, + { id: 'empty', isEnabled: true, config: { type: 'qwen_image_2_1_reference_image', image: null } }, + { id: 'broken', isEnabled: true, config: { type: 'qwen_image_2_1_reference_image', image: { width: 1 } } }, + ], + }); + + expect(normalized?.referenceImages).toEqual([ + { + config: { + image: { original: { image: { height: 640, image_name: 'edit.png', width: 480 } } }, + type: 'qwen_image_2_1_reference_image', + }, + id: 'edit-1', + isEnabled: false, + }, + { config: { image: null, type: 'qwen_image_2_1_reference_image' }, id: 'empty', isEnabled: true }, + ]); + expect(normalized && isGenerateSettings(normalized)).toBe(true); + }); + it('keeps more than five valid persisted reference images for provider-specific caps', () => { const referenceImages = Array.from({ length: 6 }, (_, index) => ({ id: `ref-${index}`, diff --git a/invokeai/frontend/webv2/src/features/generation/ui/GenerateGuidanceSection.tsx b/invokeai/frontend/webv2/src/features/generation/ui/GenerateGuidanceSection.tsx index 5bc8723152e..4043d293809 100644 --- a/invokeai/frontend/webv2/src/features/generation/ui/GenerateGuidanceSection.tsx +++ b/invokeai/frontend/webv2/src/features/generation/ui/GenerateGuidanceSection.tsx @@ -2,7 +2,10 @@ import type { GenerationModelCatalogItem as ModelConfig } from '@features/genera import type { GenerateModelConfig, GenerateSettings, LoraModelConfig } from '@features/generation/core/types'; import { Badge, Separator, Stack } from '@chakra-ui/react'; -import { isReferenceImageSupported } from '@features/generation/core/baseGenerationPolicies'; +import { + getMissingReferenceImageFormat, + isReferenceImageSupported, +} from '@features/generation/core/baseGenerationPolicies'; import { type GenerateDraft, type GenerateSettingsUpdate, @@ -85,7 +88,9 @@ export const GenerateGuidanceSection = ({ onCommit={onReferenceCommit} onCommitImmediate={onCommitImmediate} /> - {referenceImagesSupported || referenceImageCount > 0 ? : null} + {referenceImagesSupported || referenceImageCount > 0 || getMissingReferenceImageFormat(selectedModel) ? ( + + ) : null} = { diffusers: 'Diffusers' }; const IMAGE_ONLY = ['image'] as const; interface GenerateReferenceImagesContentProps { @@ -50,8 +53,9 @@ export const GenerateReferenceImagesContent = ({ selectedModel, }: GenerateReferenceImagesContentProps) => { const { t } = useTranslation(); - const { gallery } = useGenerationUi(); + const { gallery, models: modelUi } = useGenerationUi(); const isSupported = isReferenceImageSupported(selectedModel); + const missingFormat = getMissingReferenceImageFormat(selectedModel); const maxReferenceImages = getMaxReferenceImages(selectedModel); const canAdd = referenceImages.length < maxReferenceImages; const { acceptsActiveDrag, isOver, setNodeRef } = useGalleryImageDroppable({ @@ -194,8 +198,21 @@ export const GenerateReferenceImagesContent = ({ }, }); + // A model that only lacks the right format says which one would take references. + const unsupportedMessage = + selectedModel && missingFormat + ? t('widgets.generate.referenceImagesNeedFormat', { + base: modelUi.getBaseLabel(selectedModel.base), + format: MODEL_FORMAT_LABELS[missingFormat] ?? missingFormat, + }) + : t('widgets.generate.referenceImagesUnsupported'); + if (!isSupported && referenceImageCount === 0) { - return null; + return missingFormat ? ( + + {unsupportedMessage} + + ) : null; } // Unsupported-model references remain read-only with removal as recovery. @@ -203,7 +220,7 @@ export const GenerateReferenceImagesContent = ({ return ( - {t('widgets.generate.referenceImagesUnsupported')} + {unsupportedMessage}