VictorMylle/Kohya-ss-sd-scripts
1
0
Fork 1
mirror of https://github.com/kohya-ss/sd-scripts.git synced 2026-04-08 22:35:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix to work cond mask and alpha mask

Browse Source
This commit is contained in:
Kohya S
2024-05-26 22:01:37 +09:00
parent da6fea3d97
commit e8cfd4ba1d
3 changed files with 17 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
library/custom_train_functions.py
View File

@@ -484,9 +484,11 @@ def apply_masked_loss(loss, batch):
# conditioning image is -1 to 1. we need to convert it to 0 to 1
mask_image = batch["conditioning_images"].to(dtype=loss.dtype)[:, 0].unsqueeze(1) # use R channel
mask_image = mask_image / 2 + 0.5
# print(f"conditioning_image: {mask_image.shape}")
elif "alpha_masks" in batch and batch["alpha_masks"] is not None:
# alpha mask is 0 to 1
mask_image = batch["alpha_masks"].to(dtype=loss.dtype)
mask_image = batch["alpha_masks"].to(dtype=loss.dtype).unsqueeze(1) # add channel dimension
# print(f"mask_image: {mask_image.shape}, {mask_image.mean()}")
else:
return loss