diff --git a/extensions_built_in/diffusion_models/qwen_image_2/qwen_image_2.py b/extensions_built_in/diffusion_models/qwen_image_2/qwen_image_2.py index 54f6488..32e0770 100644 --- a/extensions_built_in/diffusion_models/qwen_image_2/qwen_image_2.py +++ b/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( diff --git a/extensions_built_in/diffusion_models/qwen_image_2/src/pipeline.py b/extensions_built_in/diffusion_models/qwen_image_2/src/pipeline.py index 3a3fe67..710a554 100644 --- a/extensions_built_in/diffusion_models/qwen_image_2/src/pipeline.py +++ b/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) diff --git a/extensions_built_in/sd_trainer/SDTrainer.py b/extensions_built_in/sd_trainer/SDTrainer.py index 69537d4..d464d3d 100644 --- a/extensions_built_in/sd_trainer/SDTrainer.py +++ b/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, diff --git a/toolkit/data_loader.py b/toolkit/data_loader.py index 1825e65..3f96183 100644 --- a/toolkit/data_loader.py +++ b/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, diff --git a/toolkit/data_transfer_object/data_loader.py b/toolkit/data_transfer_object/data_loader.py index 0521885..7ddc3ce 100644 --- a/toolkit/data_transfer_object/data_loader.py +++ b/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, "") diff --git a/toolkit/dataloader_mixins.py b/toolkit/dataloader_mixins.py index 5302aad..b4cd1d0 100644 --- a/toolkit/dataloader_mixins.py +++ b/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 ( diff --git a/toolkit/models/base_model.py b/toolkit/models/base_model.py index 35a28d5..29256fd 100644 --- a/toolkit/models/base_model.py +++ b/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) diff --git a/toolkit/stable_diffusion_model.py b/toolkit/stable_diffusion_model.py index 354dc34..24c8043 100644 --- a/toolkit/stable_diffusion_model.py +++ b/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