Add support for training with an alpha channel for Qwen Image 2.1

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
+    # VAE. The latents are RGBA. Images without alpha get an opaque one on
-    # opaque alpha channel and decode drops it again (unless model_kwargs.rgba).
+    # 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)

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: (
+            <div className="space-y-2">
+              <p>
+                This model&apos;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.
+              </p>
+              <p>
+                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.
+              </p>
+              <p>
+                Changing this re-caches latents, and it cannot be combined with
+                a dataset using <code>alpha_mask</code>, which consumes the
+                alpha channel as a loss mask instead.
+              </p>
+            </div>
+          ),
+        },
+      },
+    ],
   },
   {
     name: "hidream",

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

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(

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
+        # Convert RGB to BGR, leaving any alpha channel where it is
-        open_cv_image = open_cv_image[:, :, ::-1].copy()
+        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
+        # convert to rgb, keeping alpha so it rides the same spatial transform
-        img = img.convert('RGB')
+        img = img.convert('RGBA' if has_alpha else 'RGB')
 
         open_cv_image = np.array(img)
-        # Convert RGB to BGR
+        # Convert RGB to BGR, leaving any alpha channel where it is
-        open_cv_image = open_cv_image[:, :, ::-1].copy()
+        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):

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'

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):

version.py

@@ -1 +1 @@
-VERSION = "0.13.18"
+VERSION = "0.13.19"