Add Minimal Implementation of Masked Weight Loss

Author: AI-Casanova

library/train_util.py

@@ -68,6 +68,8 @@ def __init__(self, image_key: str, num_repeats: int, caption: str, is_reg: bool,
     self.latents_flipped: torch.Tensor = None
     self.latents_npz: str = None
     self.latents_npz_flipped: str = None
+    self.mask: np.ndarray = None
+    self.mask_flipped: np.ndarray = None
 
 
 class BucketManager():
@@ -467,9 +469,12 @@ def shuffle_buckets(self):
 
   def load_image(self, image_path):
     image = Image.open(image_path)
-    if not image.mode == "RGB":
-      image = image.convert("RGB")
+    # if not image.mode == "RGB":
+      # image = image.convert("RGB")
+    if not image.mode == "RGBA":
+      image = image.convert("RGBA")
     img = np.array(image, np.uint8)
+    # alpha_channel = np.array(image, np.uint8)[:,:,-1]
     return img
 
   def trim_and_resize_if_required(self, image, reso, resized_size):
@@ -508,16 +513,19 @@ def cache_latents(self, vae):
 
       image = self.load_image(info.absolute_path)
       image = self.trim_and_resize_if_required(image, info.bucket_reso, info.resized_size)
-
+      mask = image[:,:,-1] #grab alpha channel
+      image = image[:,:,:3] #drop alpha channel
       img_tensor = self.image_transforms(image)
       img_tensor = img_tensor.unsqueeze(0).to(device=vae.device, dtype=vae.dtype)
       info.latents = vae.encode(img_tensor).latent_dist.sample().squeeze(0).to("cpu")
+      info.mask = mask/255
 
       if self.flip_aug:
         image = image[:, ::-1].copy()     # cannot convert to Tensor without copy
         img_tensor = self.image_transforms(image)
         img_tensor = img_tensor.unsqueeze(0).to(device=vae.device, dtype=vae.dtype)
         info.latents_flipped = vae.encode(img_tensor).latent_dist.sample().squeeze(0).to("cpu")
+        info.mask_flipped = mask[::-1]/255
 
   def get_image_size(self, image_path):
     image = Image.open(image_path)
@@ -606,14 +614,17 @@ def __getitem__(self, index):
     input_ids_list = []
     latents_list = []
     images = []
+    masks = []
 
     for image_key in bucket[image_index:image_index + bucket_batch_size]:
       image_info = self.image_data[image_key]
       loss_weights.append(self.prior_loss_weight if image_info.is_reg else 1.0)
 
       # image/latentsを処理する
       if image_info.latents is not None:
-        latents = image_info.latents if not self.flip_aug or random.random() < .5 else image_info.latents_flipped
+        rand_flip = random.random()
+        latents = image_info.latents if not self.flip_aug or rand_flip < .5 else image_info.latents_flipped
+        mask = image_info.mask if not self.flip_aug or rand_flip < .5 else image_info.mask_flipped
         image = None
       elif image_info.latents_npz is not None:
         latents = self.load_latents_from_npz(image_info, self.flip_aug and random.random() >= .5)
@@ -622,6 +633,8 @@ def __getitem__(self, index):
       else:
         # 画像を読み込み、必要ならcropする
         img, face_cx, face_cy, face_w, face_h = self.load_image_with_face_info(image_info.absolute_path)
+        mask = img[:,:,-1] #grab alpha channel
+        img = img[:,:,:3] #drop alpha channel
         im_h, im_w = img.shape[0:2]
 
         if self.enable_bucket:
@@ -647,7 +660,8 @@ def __getitem__(self, index):
 
         latents = None
         image = self.image_transforms(img)      # -1.0~1.0のtorch.Tensorになる
-
+        mask = (self.image_transforms(mask)+1)*.5
+      masks.append(torch.tensor(mask))
       images.append(image)
       latents_list.append(latents)
 
@@ -672,7 +686,7 @@ def __getitem__(self, index):
     else:
       images = None
     example['images'] = images
-
+    example['masks'] = torch.stack(masks) if masks[0] is not None else None
     example['latents'] = torch.stack(latents_list) if latents_list[0] is not None else None
 
     if self.debug_dataset:
@@ -1494,6 +1508,8 @@ def add_dataset_arguments(parser: argparse.ArgumentParser, support_dreambooth: b
                       help="steps of resolution for buckets, divisible by 8 is recommended / bucketの解像度の単位、8で割り切れる値を推奨します")
   parser.add_argument("--bucket_no_upscale", action="store_true",
                       help="make bucket for each image without upscaling / 画像を拡大せずbucketを作成します")
+  parser.add_argument("--masked_loss", action="store_true",
+                      help="Enable Masked Loss from Alpha Channel")
 
   if support_caption_dropout:
     # Textual Inversion はcaptionのdropoutをsupportしない
@@ -2059,4 +2075,4 @@ def __getitem__(self, idx):
     return (tensor_pil, img_path)
 
 
-# endregion
+# endregion

train_network.py

@@ -1,5 +1,6 @@
 from torch.cuda.amp import autocast
 from torch.nn.parallel import DistributedDataParallel as DDP
+import torch.nn.functional
 import importlib
 import argparse
 import gc
@@ -377,6 +378,24 @@ def train(args):
           target = noise_scheduler.get_velocity(latents, noise, timesteps)
         else:
           target = noise
+          
+        if args.masked_loss and batch['masks'] is not None:
+            
+          mask = (
+              batch['masks']
+              .to(noise_pred.device)
+              .reshape(
+                  noise_pred.shape[0], 1, noise_pred.shape[2] * 8, noise_pred.shape[3] * 8
+              )
+          )
+          # resize to match model_pred
+          mask = torch.nn.functional.interpolate(
+              mask.float(),
+              size=noise_pred.shape[-2:],
+              mode="nearest",
+          )
+          noise_pred = noise_pred * mask
+          target = target * mask
 
         loss = torch.nn.functional.mse_loss(noise_pred.float(), target.float(), reduction="none")
         loss = loss.mean([1, 2, 3])