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 1d12e63..54f6488 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 @@ -206,11 +206,11 @@ class QwenImage2Model(BaseModel): self.print_and_status_update("Model Loaded") # ------------------------------------------------------------------ - # VAE. The latents are RGBA; toolkit images are RGB, so encode pads an - # opaque alpha channel and decode drops it again (unless model_kwargs.rgba). + # VAE. The latents are RGBA. Images without alpha get an opaque one on + # encode, and decode drops it again unless RGBA output is on. # ------------------------------------------------------------------ @property - def output_rgba(self) -> bool: + def load_rgba(self) -> bool: return bool(self.model_config.model_kwargs.get("rgba", False)) def _latent_stats(self, device, dtype): @@ -242,7 +242,7 @@ class QwenImage2Model(BaseModel): def decode_latents(self, latents: torch.Tensor, device=None, dtype=None): images = self._decode_rgba(latents, device=device, dtype=dtype) - if not self.output_rgba: + if not self.load_rgba: images = images[:, :3] return images @@ -277,7 +277,7 @@ class QwenImage2Model(BaseModel): return images.squeeze(2).to(device, dtype=dtype) def decode_to_images(self, latents: torch.Tensor) -> List[Image.Image]: - """Decode to PIL, keeping the alpha channel when model_kwargs.rgba is set.""" + """Decode to PIL, keeping the alpha channel when load_rgba is set.""" return [ self.image_tensor_to_pil(image) for image in self.decode_latents(latents) ] @@ -521,9 +521,12 @@ class QwenImage2Model(BaseModel): ] condition_images = None if paths: + # same channels the dataloader gives training references, so a + # transparent reference behaves the same way in both + mode = "RGBA" if self.load_rgba else "RGB" tensors = [ torch.from_numpy( - np.array(Image.open(path).convert("RGBA"), dtype=np.float32) / 255.0 + np.array(Image.open(path).convert(mode), dtype=np.float32) / 255.0 ) .permute(2, 0, 1) .unsqueeze(0) diff --git a/extensions_built_in/diffusion_models/ui.tsx b/extensions_built_in/diffusion_models/ui.tsx index fa51ba5..13f5d5d 100644 --- a/extensions_built_in/diffusion_models/ui.tsx +++ b/extensions_built_in/diffusion_models/ui.tsx @@ -556,6 +556,13 @@ export const AI_TOOLKIT_UI_MODELS: ModelArch[] = [ "config.process[0].model.qtype": ["convrot8", "qfloat8"], "config.process[0].model.qtype_te": ["convrot8", "qfloat8"], "config.process[0].sample.guidance_scale": [3.0, 4.0], + // the VAE is RGBA: images load, encode and decode with their alpha + "config.process[0].model.model_kwargs": [ + { + rgba: false, + }, + {}, + ], }, disableSections: ["network.conv", "train.unload_text_encoder"], // one model: it edits when the dataset has control paths, and is plain @@ -566,6 +573,52 @@ export const AI_TOOLKIT_UI_MODELS: ModelArch[] = [ "model.low_vram", "model.layer_offloading", ], + customModelSelectOptions: [ + { + type: "checkbox", + label: "Transparency (RGBA)", + getValue: (config: JobConfig) => + config?.config?.process?.[0]?.model?.model_kwargs?.rgba ?? false, + onChange: ( + value: boolean, + config: JobConfig, + setJobConfig: (value: any, key: string) => void, + ) => { + const kwargs = { + ...(config?.config?.process?.[0]?.model?.model_kwargs ?? {}), + }; + if (value) { + kwargs.rgba = true; + } else { + delete kwargs.rgba; + } + setJobConfig(kwargs, "config.process[0].model.model_kwargs"); + }, + doc: { + title: "Transparency (RGBA)", + description: ( +
+

+ This model's VAE is natively RGBA, so alpha can be carried + end to end. When on, dataset images and reference images load + with their alpha channel, the VAE encodes all four channels, and + samples are saved as PNGs with their transparency intact. +

+

+ Images with no alpha of their own get a fully opaque one, so a + mixed dataset is fine. Turn this off to train and sample flat + RGB: alpha is dropped on the way in and added back as opaque. +

+

+ Changing this re-caches latents, and it cannot be combined with + a dataset using alpha_mask, which consumes the + alpha channel as a loss mask instead. +

+
+ ), + }, + }, + ], }, { name: "hidream", diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py index 4976e8d..df6b86d 100644 --- a/toolkit/config_modules.py +++ b/toolkit/config_modules.py @@ -1362,6 +1362,10 @@ class GenerateImageConfig: if self.output_ext == 'mp3': add_album_artwork(audio_path) else: + if image.mode == 'RGBA' and self.output_ext not in ['png', 'webp']: + # jpg cannot carry alpha, and dropping it silently would hide + # the transparency an RGBA model just generated + self.output_ext = 'png' # TODO save image gen header info for A1111 and us, our seeds probably wont match image.save(self.get_image_path(count, max_count)) # do prompt file diff --git a/toolkit/data_transfer_object/data_loader.py b/toolkit/data_transfer_object/data_loader.py index 86d4a13..0521885 100644 --- a/toolkit/data_transfer_object/data_loader.py +++ b/toolkit/data_transfer_object/data_loader.py @@ -76,6 +76,8 @@ class FileItemDTO( if _sd is not None and hasattr(_sd, "get_frame_count_snapper") else None ) + # models with an RGBA VAE want the alpha channel loaded and encoded + self.load_rgba = _sd is not None and _sd.load_rgba size_database = kwargs.get("size_database", {}) dataset_root = kwargs.get("dataset_root", None) self.encode_control_in_text_embeddings = kwargs.get( diff --git a/toolkit/dataloader_mixins.py b/toolkit/dataloader_mixins.py index e108ec5..5302aad 100644 --- a/toolkit/dataloader_mixins.py +++ b/toolkit/dataloader_mixins.py @@ -931,7 +931,13 @@ class ImageProcessingDTOMixin: np_img = np_img[:, :, :3] img = Image.fromarray(np_img) - img = img.convert('RGB') + # load_rgba keeps a 4th channel for models with an RGBA VAE. alpha_mask + # consumes the alpha itself, so the two are mutually exclusive. + if self.load_rgba and not self.use_alpha_as_mask: + # sources without alpha get a fully opaque one + img = img.convert('RGBA') + else: + img = img.convert('RGB') w, h = img.size if w > h and self.scale_to_width < self.scale_to_height: # throw error, they should match @@ -1182,7 +1188,11 @@ class ControlFileItemDTOMixin: img = Image.open(control_path) img = exif_transpose(img) - if img.mode in ("RGBA", "LA"): + if self.load_rgba: + # keep the alpha instead of flattening it; sources without + # one get a fully opaque alpha + img = img.convert("RGBA") + elif img.mode in ("RGBA", "LA"): # Create a background with the specified transparent color transparent_color = tuple(self.dataset_config.control_transparent_color) background = Image.new("RGB", img.size, transparent_color) @@ -1514,9 +1524,13 @@ class AugmentationFileItemDTOMixin: # save the original tensor self.unaugmented_tensor = transforms.ToTensor()(img) if transform is None else transform(img) + has_alpha = img.mode == 'RGBA' open_cv_image = np.array(img) - # Convert RGB to BGR - open_cv_image = open_cv_image[:, :, ::-1].copy() + # Convert RGB to BGR, leaving any alpha channel where it is + if has_alpha: + open_cv_image = open_cv_image[:, :, [2, 1, 0, 3]].copy() + else: + open_cv_image = open_cv_image[:, :, ::-1].copy() # apply augmentations transformed = self.aug_transform(image=open_cv_image) @@ -1533,7 +1547,7 @@ class AugmentationFileItemDTOMixin: self.aug_replay_spatial_transforms = augmented_params # convert back to RGB tensor - augmented = cv2.cvtColor(augmented, cv2.COLOR_BGR2RGB) + augmented = cv2.cvtColor(augmented, cv2.COLOR_BGRA2RGBA if has_alpha else cv2.COLOR_BGR2RGB) # convert to PIL image augmented = Image.fromarray(augmented) @@ -1550,20 +1564,24 @@ class AugmentationFileItemDTOMixin: # save colorspace to convert back to colorspace = img.mode + has_alpha = colorspace == 'RGBA' - # convert to rgb - img = img.convert('RGB') + # convert to rgb, keeping alpha so it rides the same spatial transform + img = img.convert('RGBA' if has_alpha else 'RGB') open_cv_image = np.array(img) - # Convert RGB to BGR - open_cv_image = open_cv_image[:, :, ::-1].copy() + # Convert RGB to BGR, leaving any alpha channel where it is + if has_alpha: + open_cv_image = open_cv_image[:, :, [2, 1, 0, 3]].copy() + else: + open_cv_image = open_cv_image[:, :, ::-1].copy() # Replay transforms transformed = A.ReplayCompose.replay(self.aug_replay_spatial_transforms, image=open_cv_image) augmented = transformed["image"] # convert back to RGB tensor - augmented = cv2.cvtColor(augmented, cv2.COLOR_BGR2RGB) + augmented = cv2.cvtColor(augmented, cv2.COLOR_BGRA2RGBA if has_alpha else cv2.COLOR_BGR2RGB) # convert to PIL image augmented = Image.fromarray(augmented) @@ -1859,6 +1877,10 @@ class LatentCachingFileItemDTOMixin: if self.dataset_config.cache_tensors_to_disk: # tensor is stored in the cache file, invalidate caches made without it item["cache_tensors_to_disk"] = True + if self.load_rgba: + # the encoded alpha changes the latent; only added when on so caches + # made before this existed stay valid for models that do not use it + item["load_rgba"] = True return item def get_latent_path(self: 'FileItemDTO', recalculate=False): diff --git a/toolkit/models/base_model.py b/toolkit/models/base_model.py index 0abb3ae..28e81ee 100644 --- a/toolkit/models/base_model.py +++ b/toolkit/models/base_model.py @@ -256,6 +256,14 @@ class BaseModel: def is_ssd(self): return self.arch == 'ssd' + @property + def load_rgba(self) -> bool: + """Images keep an alpha channel end to end: the dataloader loads them + as RGBA (opaque alpha when the source has none), the VAE encodes four + channels, and decoded samples keep the alpha. Only models with an RGBA + VAE override this.""" + return False + @property def is_v3(self): return self.arch == 'sd3' diff --git a/toolkit/stable_diffusion_model.py b/toolkit/stable_diffusion_model.py index 60660ac..354dc34 100644 --- a/toolkit/stable_diffusion_model.py +++ b/toolkit/stable_diffusion_model.py @@ -301,6 +301,11 @@ class StableDiffusion: @property def is_ssd(self): return self.arch == 'ssd' + + @property + def load_rgba(self) -> bool: + # no legacy arch has an RGBA VAE + return False @property def is_v3(self): diff --git a/version.py b/version.py index 29cfb35..761bf74 100644 --- a/version.py +++ b/version.py @@ -1 +1 @@ -VERSION = "0.13.18" +VERSION = "0.13.19"