Add ability to set transparent color for control images

toolkit/config_modules.py

@@ -814,6 +814,9 @@ class DatasetConfig:
         self.control_path: Union[str,List[str]] = kwargs.get('control_path', None)  # depth maps, etc
         if self.control_path == '':
             self.control_path = None
+        
+        # color for transparent reigon of control images with transparency
+        self.control_transparent_color: List[int] = kwargs.get('control_transparent_color', [0, 0, 0])
         # inpaint images should be webp/png images with alpha channel. The alpha 0 (invisible) section will
         # be the part conditioned to be inpainted. The alpha 1 (visible) section will be the part that is ignored
         self.inpaint_path: Union[str,List[str]] = kwargs.get('inpaint_path', None)

toolkit/dataloader_mixins.py

@@ -876,8 +876,19 @@ class ControlFileItemDTOMixin:
         
         for control_path in control_path_list:
             try:
-                img = Image.open(control_path).convert('RGB')
+                img = Image.open(control_path)
                 img = exif_transpose(img)
+
+                if 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)
+                    # Paste the image on top using its alpha channel as mask
+                    background.paste(img, mask=img.getchannel("A"))
+                    img = background
+                else:
+                    # Already no alpha channel
+                    img = img.convert("RGB")
             except Exception as e:
                 print_acc(f"Error: {e}")
                 print_acc(f"Error loading image: {control_path}")