masked_loss_blur.patch

Subject: [PATCH] add code for latent blur test
add masked loss implementation
---
Index: library/train_util.py
IDEA additional info:
Subsystem: com.intellij.openapi.diff.impl.patch.CharsetEP
<+>UTF-8
===================================================================
diff --git a/library/train_util.py b/library/train_util.py
--- a/library/train_util.py	(revision 89c30334016a27f1d890466151a61a46d8d1545e)
+++ b/library/train_util.py	(revision 4fcdeeff1486e3cc88f2c932d59d762b9e53457c)
@@ -100,6 +100,10 @@
         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
+        self.loss_start_step: float = None
+        self.loss_end_step: float = None
 
 
 class BucketManager:
@@ -697,11 +701,32 @@
         random.shuffle(self.buckets_indices)
         self.bucket_manager.shuffle()
 
+    def load_mask(self, path):
+        try:
+            p = pathlib.Path(path)
+            mask_path = os.path.join(p.parent, 'mask', p.stem + '.png')
+            mask = np.array(Image.open(mask_path))
+            if len(mask.shape) > 2 and mask.max() <= 255:
+                print(mask.shape)
+                return np.array(Image.open(mask_path).convert("L"))
+            elif len(mask.shape) == 2 and mask.max() > 255:
+                print(mask.max())
+                return mask // (((2 ** 16) - 1) // 255)
+            elif len(mask.shape) == 2 and mask.max() <= 255:
+                return mask
+            else:
+                print(f"{mask_path} has invalid mask format: Defaulting to no mask")
+                return np.ones_like(np.array(Image.open(path).convert("L"))) * 255
+        except:
+            print(f"{mask_path} not found: Defaulting to no mask")
+        return np.ones_like(np.array(Image.open(path).convert("L"))) * 255
+
     def load_image(self, image_path):
         image = Image.open(image_path)
-        if not image.mode == "RGB":
-            image = image.convert("RGB")
+        if not image.mode == "RGBA":
+            image = image.convert("RGBA")
         img = np.array(image, np.uint8)
+        img[..., -1] = self.load_mask(image_path)
         return img
 
     def trim_and_resize_if_required(self, subset: BaseSubset, image, reso, resized_size):
@@ -802,31 +827,37 @@
         # iterate batches
         for batch in tqdm(batches, smoothing=1, total=len(batches)):
             images = []
+            masks = []
             for info in batch:
                 image = self.load_image(info.absolute_path)
                 image = self.trim_and_resize_if_required(subset, image, info.bucket_reso, info.resized_size)
+                mask = image[:, :, -1] / 255  # grab alpha channel
+                image = image[:, :, :3]  # drop alpha channel
                 image = self.image_transforms(image)
                 images.append(image)
+                masks.append(mask)
 
             img_tensors = torch.stack(images, dim=0)
             img_tensors = img_tensors.to(device=vae.device, dtype=vae.dtype)
 
             latents = vae.encode(img_tensors).latent_dist.sample().to("cpu")
 
-            for info, latent in zip(batch, latents):
+            for info, latent, mask in zip(batch, latents, masks):
                 if cache_to_disk:
                     np.savez(info.latents_npz, latent.float().numpy())
                 else:
                     info.latents = latent
+                    info.mask = mask
 
             if subset.flip_aug:
                 img_tensors = torch.flip(img_tensors, dims=[3])
                 latents = vae.encode(img_tensors).latent_dist.sample().to("cpu")
-                for info, latent in zip(batch, latents):
+                for info, latent, mask in zip(batch, latents, masks):
                     if cache_to_disk:
                         np.savez(info.latents_npz_flipped, latent.float().numpy())
                     else:
                         info.latents_flipped = latent
+                        info.mask_flipped = mask
 
     def get_image_size(self, image_path):
         image = Image.open(image_path)
@@ -913,15 +944,20 @@
         input_ids_list = []
         latents_list = []
         images = []
+        masks = []
+        loss_start_steps = []
+        loss_end_steps = []
 
         for image_key in bucket[image_index : image_index + bucket_batch_size]:
             image_info = self.image_data[image_key]
             subset = self.image_to_subset[image_key]
-            loss_weights.append(self.prior_loss_weight if image_info.is_reg else 1.0)
+            loss_weights.append(self.prior_loss_weight if image_info.is_reg else 1)
 
             # image/latentsを処理する
             if image_info.latents is not None:  # cache_latents=Trueの場合
-                latents = image_info.latents if not subset.flip_aug or random.random() < 0.5 else image_info.latents_flipped
+                rand_flip = random.random()
+                latents = image_info.latents if not subset.flip_aug or rand_flip < .5 else image_info.latents_flipped
+                mask = image_info.mask if not subset.flip_aug or rand_flip < .5 else image_info.mask_flipped
                 image = None
             elif image_info.latents_npz is not None:  # FineTuningDatasetまたはcache_latents_to_disk=Trueの場合
                 latents = self.load_latents_from_npz(image_info, subset.flip_aug and random.random() >= 0.5)
@@ -930,6 +966,8 @@
             else:
                 # 画像を読み込み、必要ならcropする
                 img, face_cx, face_cy, face_w, face_h = self.load_image_with_face_info(subset, image_info.absolute_path)
+                mask = img[:, :, -1] / 255  # grab alpha channel
+                img = img[:, :, :3]  # drop alpha channel
                 im_h, im_w = img.shape[0:2]
 
                 if self.enable_bucket:
@@ -960,9 +998,13 @@
 
                 latents = None
                 image = self.image_transforms(img)  # -1.0~1.0のtorch.Tensorになる
+                mask = torch.from_numpy(mask)
 
             images.append(image)
+            masks.append(torch.tensor(mask))
             latents_list.append(latents)
+            loss_start_steps.append(image_info.loss_start_step)
+            loss_end_steps.append(image_info.loss_end_step)
 
             caption = self.process_caption(subset, image_info.caption)
             if self.XTI_layers:
@@ -998,9 +1040,10 @@
         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
         example["captions"] = captions
+        example["loss_ranges"] = list(zip(loss_start_steps, loss_end_steps))
 
         if self.debug_dataset:
             example["image_keys"] = bucket[image_index : image_index + self.batch_size]
@@ -1266,6 +1309,8 @@
 
                 image_info = ImageInfo(image_key, subset.num_repeats, caption, False, abs_path)
                 image_info.image_size = img_md.get("train_resolution")
+                image_info.loss_start_step = img_md.get("loss_start") or 0.0
+                image_info.loss_end_step = img_md.get("loss_end") or 1.0
 
                 if not subset.color_aug and not subset.random_crop:
                     # if npz exists, use them
@@ -2388,6 +2433,8 @@
         "--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 Mask File")
+
     parser.add_argument(
         "--token_warmup_min",
         type=int,
Index: train_network.py
IDEA additional info:
Subsystem: com.intellij.openapi.diff.impl.patch.CharsetEP
<+>UTF-8
===================================================================
diff --git a/train_network.py b/train_network.py
--- a/train_network.py	(revision 89c30334016a27f1d890466151a61a46d8d1545e)
+++ b/train_network.py	(revision 4fcdeeff1486e3cc88f2c932d59d762b9e53457c)
@@ -11,6 +11,8 @@
 
 from tqdm import tqdm
 import torch
+import torch.nn.functional
+import torchvision
 from accelerate.utils import set_seed
 from diffusers import DDPMScheduler
 
@@ -644,6 +646,7 @@
                 # Sample a random timestep for each image
                 timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (b_size,), device=latents.device)
                 timesteps = timesteps.long()
+
                 # Add noise to the latents according to the noise magnitude at each timestep
                 # (this is the forward diffusion process)
                 noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
@@ -658,6 +661,34 @@
                 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 noise_pred
+                    mask = torch.nn.functional.interpolate(
+                        mask.float(),
+                        size=noise_pred.shape[-2:],
+                        mode="nearest",
+                    )
+
+                    steps = timesteps / noise_scheduler.config.num_train_timesteps
+                    decay = 1 / 0.2
+
+                    a = steps.reshape(steps.shape + (1,) * 3) - mask
+                    b = torch.exp2(-decay * a)
+                    c = torch.minimum(torch.ones_like(mask), b)
+
+                    mask = torch.where(mask > 0.0, c, 0.0)
+                    mask = mask / torch.sqrt(mask.mean())
+
+                    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])
 
@@ -721,6 +752,7 @@
             loss_total += current_loss
             avr_loss = loss_total / len(loss_list)
             logs = {"loss": avr_loss}  # , "lr": lr_scheduler.get_last_lr()[0]}
+
             progress_bar.set_postfix(**logs)
 
             if args.scale_weight_norms: