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

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",