ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-04-28 02:01:29 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add ability to set transparent color for control images

Browse Source
This commit is contained in:
Jaret Burkett
2025-09-02 11:08:44 -06:00
parent 85dcae6e2b
commit f699f4be5f
3 changed files with 16 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
toolkit/dataloader_mixins.py
View File

@@ -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}")