Add match_target_res for Qwen Image 2.1 so reference images are scaled to the target's pixel area with their own aspect ratio, consistently across text-embedding caching, training, and sampling, and key the embedding cache on it.

extensions_built_in/diffusion_models/qwen_image_2/qwen_image_2.py

@@ -132,11 +132,10 @@ class QwenImage2Model(BaseModel):
         # embeddings as vision tokens and their latents into the sequence.
         self.encode_control_in_text_embeddings = True
         self.has_multiple_control_images = True
-        # References are bucket-resized to the target size by the dataloader.
-        # The DiT reads the image-slot layout from row 0 of the batch, so every
-        # sample must contribute the same slot count -- raw (per-image aspect)
-        # references would break that.
-        self.use_raw_control_images = False
+        # References keep their own size/aspect; prepare_condition_image budgets them
+        # identically in the cache and training paths so slot counts always agree.
+        # Batch > 1 requires same-size references (checked in encode_condition_images).
+        self.use_raw_control_images = True
 
     @staticmethod
     def get_train_scheduler():
@@ -294,13 +293,38 @@ class QwenImage2Model(BaseModel):
     # ------------------------------------------------------------------
     @property
     def control_image_max_pixels(self) -> int:
-        """Pixel budget a reference image is shrunk to fit. A reference that is
-        already smaller keeps its size (only snapped to the 32 px grid), so a
-        bucket-resized training reference stays at the target resolution."""
+        """Pixel budget a reference image is shrunk to fit when no target size is
+        known (blank/static prompts). A smaller reference keeps its size, only
+        snapped to the 32 px grid."""
         return int(
             self.model_config.model_kwargs.get("control_image_max_pixels", 1024 * 1024)
         )
 
+    @property
+    def match_target_res(self) -> bool:
+        """References are scaled to the target's pixel area (own aspect kept).
+        model_kwargs.match_target_res: false -> control_image_max_pixels cap only."""
+        return bool(self.model_config.model_kwargs.get("match_target_res", True))
+
+    @property
+    def text_embedding_uses_target_size(self) -> bool:
+        # the dataloader adds the item's bucket size to the text-embedding cache key
+        return self.match_target_res
+
+    def get_text_embedding_space_version(self) -> str:
+        # reference sizing changes the vision tokens; keep caches from different rules apart
+        rule = "match" if self.match_target_res else f"cap{self.control_image_max_pixels}"
+        return f"{self.text_embedding_space_version}_ref{rule}"
+
+    def _target_pixels(self, target_size) -> Optional[int]:
+        """`(width, height)` -> pixel area on the 32 px grid, or None. Floors the
+        same way generate_single_image does so the TE and VAE passes agree."""
+        if target_size is None:
+            return None
+        divisor = self.get_bucket_divisibility()
+        width, height = target_size
+        return int(width // divisor * divisor) * int(height // divisor * divisor)
+
     def _normalize_control_images(self, control_images, batch_size: int) -> List[List]:
         """Any of the shapes the toolkit hands over -> one list per batch item.
 
@@ -327,10 +351,14 @@ class QwenImage2Model(BaseModel):
         return control_images
 
     def _prepare_control_images(
-        self, control_images: List[List[torch.Tensor]]
+        self,
+        control_images: List[List[torch.Tensor]],
+        target_pixels: Optional[int] = None,
     ) -> List[List[torch.Tensor]]:
-        """Put every reference on the 32 px grid, as `(1, C, H, W)` in [0, 1]."""
-        budget = self.control_image_max_pixels
+        """Put every reference on the 32 px grid, as `(1, C, H, W)` in [0, 1].
+        With match_target_res and a known target, scale each to the target's area."""
+        match = self.match_target_res and target_pixels is not None
+        budget = target_pixels if match else self.control_image_max_pixels
         prepared = []
         for sample in control_images:
             images = []
@@ -338,7 +366,9 @@ class QwenImage2Model(BaseModel):
                 if image.dim() == 3:
                     image = image.unsqueeze(0)
                 images.append(
-                    prepare_condition_image(image.to(self.device_torch), budget)
+                    prepare_condition_image(
+                        image.to(self.device_torch), budget, match=match
+                    )
                 )
             prepared.append(images)
         return prepared
@@ -378,7 +408,9 @@ class QwenImage2Model(BaseModel):
     # ------------------------------------------------------------------
     # Prompts
     # ------------------------------------------------------------------
-    def get_prompt_embeds(self, prompt, control_images=None) -> AdvancedPromptEmbeds:
+    def get_prompt_embeds(
+        self, prompt, control_images=None, target_size=None
+    ) -> AdvancedPromptEmbeds:
         if isinstance(prompt, str):
             prompt = [prompt]
         if self.text_encoder[0].device != self.device_torch:
@@ -387,7 +419,8 @@ class QwenImage2Model(BaseModel):
         images = None
         if control_images is not None:
             samples = self._prepare_control_images(
-                self._normalize_control_images(control_images, len(prompt))
+                self._normalize_control_images(control_images, len(prompt)),
+                target_pixels=self._target_pixels(target_size),
             )
             images = [[tensor_to_pil(image) for image in sample] for sample in samples]
 
@@ -446,8 +479,14 @@ class QwenImage2Model(BaseModel):
                 control = batch.control_tensor_list
                 if control is None:
                     control = batch.control_tensor
+                # same area the dataloader cached the prompt against (bucket crop)
+                target_pixels = self._target_pixels((
+                    latent_model_input.shape[3] * VAE_SCALE_FACTOR,
+                    latent_model_input.shape[2] * VAE_SCALE_FACTOR,
+                ))
                 samples = self._prepare_control_images(
-                    self._normalize_control_images(control, batch_size)
+                    self._normalize_control_images(control, batch_size),
+                    target_pixels=target_pixels,
                 )
                 condition_latents, condition_shapes = self.encode_condition_images(
                     samples
@@ -533,7 +572,8 @@ class QwenImage2Model(BaseModel):
                 for path in paths
             ]
             condition_images = self._prepare_control_images(
-                self._normalize_control_images([tensors], 1)
+                self._normalize_control_images([tensors], 1),
+                target_pixels=self._target_pixels((gen_config.width, gen_config.height)),
             )
 
         return pipeline(

extensions_built_in/diffusion_models/qwen_image_2/src/pipeline.py

@@ -54,17 +54,20 @@ def calculate_dimensions(target_area: int, ratio: float) -> tuple[int, int]:
     return _snap(width), _snap(width / ratio)
 
 
-def prepare_condition_image(image: torch.Tensor, max_pixels: int) -> torch.Tensor:
+def prepare_condition_image(
+    image: torch.Tensor, max_pixels: int, match: bool = False
+) -> torch.Tensor:
     """Put one `(1, C, H, W)` condition image on the 32 px grid, shrinking it
-    first if it is over the pixel budget.
+    first if it is over the pixel budget. With `match`, scale it to the budget's
+    area in both directions (own aspect kept) instead of only shrinking.
 
-    Deterministic in the image alone, on purpose: the Qwen3-VL pass and the VAE
-    pass happen in different calls (`get_prompt_embeds` and
+    Deterministic in the image and budget alone, on purpose: the Qwen3-VL pass
+    and the VAE pass happen in different calls (`get_prompt_embeds` and
     `get_noise_prediction`) and must agree on the size to the pixel, or the slot
     count and the latent-token count disagree and the sequence is rejected.
     """
     height, width = image.shape[2], image.shape[3]
-    if height * width > max_pixels:
+    if match or height * width > max_pixels:
         new_width, new_height = calculate_dimensions(max_pixels, width / height)
     else:
         new_width, new_height = _snap(width), _snap(height)

extensions_built_in/sd_trainer/SDTrainer.py

@@ -122,6 +122,16 @@ class SDTrainer(BaseSDTrainProcess):
     def before_model_load(self):
         pass
 
+    def get_batch_target_size(self, batch: 'DataLoaderBatchDTO'):
+        # (width, height) of the bucket the batch was cropped to; what the noisy
+        # latents will be, so reference sizing can match it
+        item = batch.file_items[0]
+        width = getattr(item, 'crop_width', None)
+        height = getattr(item, 'crop_height', None)
+        if width is None or height is None:
+            return None
+        return (width, height)
+
     def get_blank_control_image(self):
         # noise instead of a black image so the fallback does not read as a
         # meaningful (solid black) reference
@@ -242,13 +252,16 @@ class SDTrainer(BaseSDTrainProcess):
                         ctrl_img = ctrl_img_list[0] if len(ctrl_img_list) > 0 else None
                     
                     
+                    target_size = (gen_img_config.width, gen_img_config.height)
                     positive = self.sd.encode_prompt(
                         gen_img_config.prompt,
-                        control_images=ctrl_img
+                        control_images=ctrl_img,
+                        target_size=target_size,
                     ).to('cpu')
                     negative = self.sd.encode_prompt(
                         gen_img_config.negative_prompt,
-                        control_images=ctrl_img
+                        control_images=ctrl_img,
+                        target_size=target_size,
                     ).to('cpu')
                 else:
                     positive = self.sd.encode_prompt(gen_img_config.prompt).to('cpu')
@@ -1325,12 +1338,13 @@ class SDTrainer(BaseSDTrainProcess):
                     prompt_kwargs = {}
                     if self.sd.encode_control_in_text_embeddings and batch.control_tensor is not None:
                         prompt_kwargs['control_images'] = batch.control_tensor.to(self.sd.device_torch, dtype=self.sd.torch_dtype)
+                        prompt_kwargs['target_size'] = self.get_batch_target_size(batch)
                     embeds_to_use = self.sd.encode_prompt(
                         prompt_list,
-                        long_prompts=self.do_long_prompts).to(
+                        long_prompts=self.do_long_prompts,
+                        **prompt_kwargs).to(
                         self.device_torch,
-                        dtype=dtype,
-                        **prompt_kwargs
+                        dtype=dtype
                     ).detach()
 
             # dont use network on this
@@ -1739,6 +1753,7 @@ class SDTrainer(BaseSDTrainProcess):
                     prompt_kwargs = {}
                     if self.sd.encode_control_in_text_embeddings and batch.control_tensor is not None:
                         prompt_kwargs['control_images'] = batch.control_tensor.to(self.sd.device_torch, dtype=self.sd.torch_dtype)
+                        prompt_kwargs['target_size'] = self.get_batch_target_size(batch)
                     if self.train_config.unload_text_encoder or self.is_caching_text_embeddings:
                         with torch.set_grad_enabled(False):
                             if batch.prompt_embeds is not None:
@@ -1818,6 +1833,7 @@ class SDTrainer(BaseSDTrainProcess):
                                 self.adapter.is_unconditional_run = False
                             if self.sd.encode_control_in_text_embeddings and batch.control_tensor_list is not None:
                                 prompt_kwargs['control_images'] = batch.control_tensor_list
+                                prompt_kwargs['target_size'] = self.get_batch_target_size(batch)
                             conditional_embeds = self.sd.encode_prompt(
                                 conditioned_prompts, prompt_2,
                                 dropout_prob=self.train_config.prompt_dropout_prob,

toolkit/data_loader.py

@@ -530,6 +530,7 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
                     encode_first_frame_in_text_embeddings=getattr(self.sd, 'encode_first_frame_in_text_embeddings', False) if self.sd else False,
                     dopsd_self_ref=getattr(self.sd, 'dopsd_self_ref', False) if self.sd else False,
                     text_embedding_space_version=text_embedding_space_version,
+                    text_embedding_uses_target_size=getattr(self.sd, 'text_embedding_uses_target_size', False) if self.sd else False,
                     te_padding_side=self.sd.te_padding_side if self.sd else "right",
                     latent_space_version=latent_space_version,
                     temporal_compression=temporal_compression,

toolkit/data_transfer_object/data_loader.py

@@ -91,6 +91,8 @@ class FileItemDTO(
         self.te_padding_side = kwargs.get("te_padding_side", "right")
         self.latent_space_version = kwargs.get("latent_space_version", "sd1")
         self.text_embedding_space_version = kwargs.get("text_embedding_space_version", "sd1")
+        # model sizes references against the item's bucket, so the cache key must carry it
+        self.text_embedding_uses_target_size = kwargs.get("text_embedding_uses_target_size", False)
         if dataset_root is not None:
             # remove dataset root from path
             file_key = self.path.replace(dataset_root, "")

toolkit/dataloader_mixins.py

@@ -2250,6 +2250,8 @@ class TextEmbeddingFileItemDTOMixin:
         # if we have a control image, cache the path
         if self.encode_control_in_text_embeddings and self.control_path is not None:
             item["control_path"] = self.control_path
+            if getattr(self, 'text_embedding_uses_target_size', False) and getattr(self, 'crop_width', None):
+                item["control_target_size"] = [self.crop_width, self.crop_height]
         if self.encode_control_in_text_embeddings and getattr(self, 'control_video_paths', None):
             item["control_videos"] = sorted(self.control_video_paths)
             # v2: reference-video vision blocks are no longer resampled by the
@@ -2506,6 +2508,10 @@ class TextEmbeddingCachingMixin:
                             ctrl_img = ctrl_img_list[0]
                         else:
                             ctrl_img = ctrl_img_list
+                        # the bucket the item trains at, so references can be sized against it
+                        target_size = None
+                        if getattr(file_item, 'crop_width', None) and getattr(file_item, 'crop_height', None):
+                            target_size = (file_item.crop_width, file_item.crop_height)
                         for path, caption in encode_targets:
                             if path in dropout_target_paths:
                                 # dropout embeds are plain text. Only fall back to the
@@ -2513,9 +2519,11 @@ class TextEmbeddingCachingMixin:
                                 try:
                                     prompt_embeds: PromptEmbeds = self.sd.encode_prompt(caption)
                                 except Exception:
-                                    prompt_embeds: PromptEmbeds = self.sd.encode_prompt(caption, control_images=ctrl_img)
+                                    prompt_embeds: PromptEmbeds = self.sd.encode_prompt(
+                                        caption, control_images=ctrl_img, target_size=target_size)
                             else:
-                                prompt_embeds: PromptEmbeds = self.sd.encode_prompt(caption, control_images=ctrl_img)
+                                prompt_embeds: PromptEmbeds = self.sd.encode_prompt(
+                                    caption, control_images=ctrl_img, target_size=target_size)
                             prompt_embeds.save(path)
                             del prompt_embeds
                     elif (

toolkit/models/base_model.py

@@ -681,19 +681,21 @@ class BaseModel:
                         if isinstance(self.adapter, CustomAdapter):
                             self.adapter.is_unconditional_run = False
                         conditional_embeds = self.encode_prompt(
-                            gen_config.prompt, 
-                            gen_config.prompt_2, 
+                            gen_config.prompt,
+                            gen_config.prompt_2,
                             force_all=True,
-                            control_images=ctrl_img
+                            control_images=ctrl_img,
+                            target_size=(gen_config.width, gen_config.height),
                         )
 
                         if isinstance(self.adapter, CustomAdapter):
                             self.adapter.is_unconditional_run = True
                         unconditional_embeds = self.encode_prompt(
-                            gen_config.negative_prompt, 
-                            gen_config.negative_prompt_2, 
+                            gen_config.negative_prompt,
+                            gen_config.negative_prompt_2,
                             force_all=True,
-                            control_images=ctrl_img
+                            control_images=ctrl_img,
+                            target_size=(gen_config.width, gen_config.height),
                         )
                         if isinstance(self.adapter, CustomAdapter):
                             self.adapter.is_unconditional_run = False
@@ -1183,6 +1185,7 @@ class BaseModel:
             max_length=None,
             dropout_prob=0.0,
             control_images=None,
+            target_size=None,
     ) -> PromptEmbeds:
         # sd1.5 embeddings are (bs, 77, 768)
         prompt = prompt
@@ -1194,7 +1197,11 @@ class BaseModel:
             prompt2 = [prompt2]
         # if control_images in the signature, pass it. This keep from breaking plugins
         if self.encode_control_in_text_embeddings:
-            return self.get_prompt_embeds(prompt, control_images=control_images)
+            kwargs = {"control_images": control_images}
+            # target (width, height) only for models that size references against it
+            if target_size is not None and "target_size" in inspect.signature(self.get_prompt_embeds).parameters:
+                kwargs["target_size"] = target_size
+            return self.get_prompt_embeds(prompt, **kwargs)
 
         return self.get_prompt_embeds(prompt)
 

toolkit/stable_diffusion_model.py

@@ -2485,6 +2485,7 @@ class StableDiffusion:
             max_length=None,
             dropout_prob=0.0,
             control_images=None,
+            target_size=None,  # accepted for parity with BaseModel; legacy archs ignore it
     ) -> PromptEmbeds:
         # sd1.5 embeddings are (bs, 77, 768)
         prompt = prompt