Create losses.py

losses/losses.py

@@ -0,0 +1,463 @@
+import torch
+import wandb
+import cv2
+import torch.nn.functional as F
+import numpy as np
+from facenet_pytorch import MTCNN
+from torchvision import transforms
+from dreamsim import dreamsim
+from einops import rearrange
+import kornia.augmentation as K
+import lpips
+
+from pretrained_models.arcface import Backbone
+from utils.vis_utils import add_text_to_image
+from utils.utils import extract_faces_and_landmarks
+import clip
+
+
+class Loss():
+    """
+    General purpose loss class. 
+    Mainly handles dtype and visualize_every_k.
+    keeps current iteration of loss, mainly for visualization purposes.
+    """
+    def __init__(self, visualize_every_k=-1, dtype=torch.float32, accelerator=None, **kwargs):
+        self.visualize_every_k = visualize_every_k
+        self.iteration = -1
+        self.dtype=dtype
+        self.accelerator = accelerator
+        
+    def __call__(self, **kwargs):
+        self.iteration += 1
+        return self.forward(**kwargs)
+
+
+class L1Loss(Loss):
+    """
+    Simple L1 loss between predicted_pixel_values and pixel_values
+    
+    Args:
+        predicted_pixel_values (torch.Tensor): The predicted pixel values using 1 step LCM and the VAE decoder.
+        encoder_pixel_values (torch.Tesnor): The input image to the encoder
+    """
+    def forward(
+        self, 
+        predict: torch.Tensor,
+        target: torch.Tensor,
+        **kwargs
+    ) -> torch.Tensor:
+        return F.l1_loss(predict, target, reduction="mean")
+
+
+class DreamSIMLoss(Loss):
+    """DreamSIM loss between predicted_pixel_values and pixel_values.
+    DreamSIM is similar to LPIPS (https://dreamsim-nights.github.io/) but is trained on more human defined similarity dataset
+    DreamSIM expects an RGB image of size 224x224 and values between 0 and 1. So we need to normalize the input images to 0-1 range and resize them to 224x224.
+    Args:
+        predicted_pixel_values (torch.Tensor): The predicted pixel values using 1 step LCM and the VAE decoder.
+        encoder_pixel_values (torch.Tesnor): The input image to the encoder
+    """
+    def __init__(self, device: str='cuda:0', **kwargs):
+        super().__init__(**kwargs)
+        self.model, _ = dreamsim(pretrained=True, device=device)
+        self.model.to(dtype=self.dtype, device=device)
+        self.model = self.accelerator.prepare(self.model)
+        self.transforms = transforms.Compose([
+            transforms.Lambda(lambda x: (x + 1) / 2),
+            transforms.Resize((224, 224), interpolation=transforms.InterpolationMode.BICUBIC)])
+
+    def forward(
+        self,
+        predicted_pixel_values: torch.Tensor,
+        encoder_pixel_values: torch.Tensor,
+        **kwargs,
+    ) -> torch.Tensor:
+        predicted_pixel_values.to(dtype=self.dtype)
+        encoder_pixel_values.to(dtype=self.dtype)
+        return self.model(self.transforms(predicted_pixel_values), self.transforms(encoder_pixel_values)).mean()
+
+
+class LPIPSLoss(Loss):
+    """LPIPS loss between predicted_pixel_values and pixel_values.
+    Args:
+        predicted_pixel_values (torch.Tensor): The predicted pixel values using 1 step LCM and the VAE decoder.
+        encoder_pixel_values (torch.Tesnor): The input image to the encoder
+    """
+    def __init__(self,  **kwargs):
+        super().__init__(**kwargs)
+        self.model = lpips.LPIPS(net='vgg')
+        self.model.to(dtype=self.dtype, device=self.accelerator.device)
+        self.model = self.accelerator.prepare(self.model)
+
+    def forward(self, predict, target, **kwargs):
+        predict.to(dtype=self.dtype)
+        target.to(dtype=self.dtype)
+        return self.model(predict, target).mean()
+
+
+class LCMVisualization(Loss):
+    """Dummy loss used to visualize the LCM outputs
+    Args:
+        predicted_pixel_values (torch.Tensor): The predicted pixel values using 1 step LCM and the VAE decoder.
+        pixel_values (torch.Tensor): The input image to the decoder
+        encoder_pixel_values (torch.Tesnor): The input image to the encoder
+    """
+    def forward(
+        self, 
+        predicted_pixel_values: torch.Tensor,
+        pixel_values: torch.Tensor,
+        encoder_pixel_values: torch.Tensor,
+        timesteps: torch.Tensor,
+        **kwargs,
+    ) -> None:
+        if self.visualize_every_k > 0 and self.iteration % self.visualize_every_k == 0:
+            predicted_pixel_values = rearrange(predicted_pixel_values, "n c h w -> (n h) w c").detach().cpu().numpy()
+            pixel_values = rearrange(pixel_values, "n c h w -> (n h) w c").detach().cpu().numpy()
+            encoder_pixel_values = rearrange(encoder_pixel_values, "n c h w -> (n h) w c").detach().cpu().numpy()
+            image = np.hstack([encoder_pixel_values, pixel_values, predicted_pixel_values])
+            for tracker in self.accelerator.trackers:
+                if tracker.name == 'wandb':
+                    tracker.log({"TrainVisualization": wandb.Image(image, caption=f"Encoder Input Image, Decoder Input Image, Predicted LCM Image. Timesteps {timesteps.cpu().tolist()}")})
+        return torch.tensor(0.0)
+
+
+class L2Loss(Loss):
+    """
+    Regular diffusion loss between predicted noise and target noise.
+    Args:
+        predicted_noise (torch.Tensor): noise predicted by the diffusion model
+        target_noise (torch.Tensor): actual noise added to the image.
+    """
+    def forward(
+        self,
+        predict: torch.Tensor,
+        target: torch.Tensor,
+        weights: torch.Tensor = None,
+        **kwargs
+    ) -> torch.Tensor:
+        if weights is not None:
+            loss = (predict.float() - target.float()).pow(2) * weights
+            return loss.mean()
+        return F.mse_loss(predict.float(), target.float(), reduction="mean")
+
+
+class HuberLoss(Loss):
+    """Huber loss between predicted_pixel_values and pixel_values.
+    Args:
+        predicted_pixel_values (torch.Tensor): The predicted pixel values using 1 step LCM and the VAE decoder.
+        encoder_pixel_values (torch.Tesnor): The input image to the encoder
+    """
+    def __init__(self, huber_c=0.001, **kwargs):
+        super().__init__(**kwargs)
+        self.huber_c = huber_c
+
+    def forward(
+        self,
+        predict: torch.Tensor,
+        target: torch.Tensor,
+        weights: torch.Tensor = None,
+        **kwargs
+    ) -> torch.Tensor:
+        loss = torch.sqrt((predict.float() - target.float()) ** 2 + self.huber_c**2) - self.huber_c
+        if weights is not None:
+            return (loss * weights).mean()
+        return loss.mean()
+
+
+class WeightedNoiseLoss(Loss):
+    """
+    Weighted diffusion loss between predicted noise and target noise.
+    Args:
+        predicted_noise (torch.Tensor): noise predicted by the diffusion model
+        target_noise (torch.Tensor): actual noise added to the image.
+        loss_batch_weights (torch.Tensor): weighting for each batch item. Can be used to e.g. zero-out loss for InstantID training if keypoint extraction fails.
+    """
+    def forward(
+        self,
+        predict: torch.Tensor,
+        target: torch.Tensor,
+        weights,
+        **kwargs
+    ) -> torch.Tensor:
+        return F.mse_loss(predict.float() * weights, target.float() * weights, reduction="mean")
+
+
+class IDLoss(Loss):
+    """
+    Use pretrained facenet model to extract features from the face of the predicted image and target image.
+    Facenet expects 112x112 images, so we crop the face using MTCNN and resize it to 112x112.
+    Then we use the cosine similarity between the features to calculate the loss. (The cosine similarity is 1 - cosine distance).
+    Also notice that the outputs of facenet are normalized so the dot product is the same as cosine distance.
+    """
+    def __init__(self, pretrained_arcface_path: str, skip_not_found=True, **kwargs):
+        super().__init__(**kwargs)
+        assert pretrained_arcface_path is not None, "please pass `pretrained_arcface_path` in the losses config. You can download the pretrained model from "\
+            "https://drive.google.com/file/d/1KW7bjndL3QG3sxBbZxreGHigcCCpsDgn/view?usp=sharing"
+        self.mtcnn = MTCNN(device=self.accelerator.device)
+        self.mtcnn.forward = self.mtcnn.detect
+        self.facenet_input_size = 112  # Has to be 112, can't find weights for 224 size.
+        self.facenet = Backbone(input_size=112, num_layers=50, drop_ratio=0.6, mode='ir_se')
+        self.facenet.load_state_dict(torch.load(pretrained_arcface_path))
+        self.face_pool = torch.nn.AdaptiveAvgPool2d((self.facenet_input_size, self.facenet_input_size))
+        self.facenet.requires_grad_(False)
+        self.facenet.eval()
+        self.facenet.to(device=self.accelerator.device, dtype=self.dtype)  # not implemented for half precision
+        self.face_pool.to(device=self.accelerator.device, dtype=self.dtype)  # not implemented for half precision
+        self.visualization_resize = transforms.Resize((self.facenet_input_size, self.facenet_input_size), interpolation=transforms.InterpolationMode.BICUBIC)
+        self.reference_facial_points = np.array([[38.29459953, 51.69630051],
+                                                 [72.53179932, 51.50139999],
+                                                 [56.02519989, 71.73660278],
+                                                 [41.54930115, 92.3655014],
+                                                 [70.72990036, 92.20410156]
+                                                 ])  # Original points are 112 * 96 added 8 to the x axis to make it 112 * 112
+        self.facenet, self.face_pool, self.mtcnn = self.accelerator.prepare(self.facenet, self.face_pool, self.mtcnn)
+
+        self.skip_not_found = skip_not_found
+    
+    def extract_feats(self, x: torch.Tensor):
+        """
+        Extract features from the face of the image using facenet model.
+        """
+        x = self.face_pool(x)
+        x_feats = self.facenet(x)
+
+        return x_feats
+
+    def forward(
+        self, 
+        predicted_pixel_values: torch.Tensor,
+        encoder_pixel_values: torch.Tensor,
+        timesteps: torch.Tensor,
+        **kwargs
+    ):
+        encoder_pixel_values = encoder_pixel_values.to(dtype=self.dtype)
+        predicted_pixel_values = predicted_pixel_values.to(dtype=self.dtype)
+
+        predicted_pixel_values_face, predicted_invalid_indices = extract_faces_and_landmarks(predicted_pixel_values, mtcnn=self.mtcnn)
+        with torch.no_grad():
+            encoder_pixel_values_face, source_invalid_indices = extract_faces_and_landmarks(encoder_pixel_values, mtcnn=self.mtcnn)
+        
+        if self.skip_not_found:
+            valid_indices = []
+            for i in range(predicted_pixel_values.shape[0]):
+                if i not in predicted_invalid_indices and i not in source_invalid_indices:
+                    valid_indices.append(i)
+        else:
+            valid_indices = list(range(predicted_pixel_values))
+            
+        valid_indices = torch.tensor(valid_indices).to(device=predicted_pixel_values.device)
+
+        if len(valid_indices) == 0:
+            loss =  (predicted_pixel_values_face * 0.0).mean()  # It's done this way so the `backwards` will delete the computation graph of the predicted_pixel_values.
+            if self.visualize_every_k > 0 and self.iteration % self.visualize_every_k == 0:
+                self.visualize(predicted_pixel_values, encoder_pixel_values, predicted_pixel_values_face, encoder_pixel_values_face, timesteps, valid_indices, loss)
+            return loss
+
+        with torch.no_grad():
+            pixel_values_feats = self.extract_feats(encoder_pixel_values_face[valid_indices])
+            
+        predicted_pixel_values_feats = self.extract_feats(predicted_pixel_values_face[valid_indices])
+        loss = 1 - torch.einsum("bi,bi->b", pixel_values_feats, predicted_pixel_values_feats)
+
+        if self.visualize_every_k > 0 and self.iteration % self.visualize_every_k == 0:
+            self.visualize(predicted_pixel_values, encoder_pixel_values, predicted_pixel_values_face, encoder_pixel_values_face, timesteps, valid_indices, loss)
+        return loss.mean()
+    
+    def visualize(
+        self,
+        predicted_pixel_values: torch.Tensor,
+        encoder_pixel_values: torch.Tensor,
+        predicted_pixel_values_face: torch.Tensor,
+        encoder_pixel_values_face: torch.Tensor,
+        timesteps: torch.Tensor,
+        valid_indices: torch.Tensor,
+        loss: torch.Tensor,
+    ) -> None:
+        small_predicted_pixel_values = (rearrange(self.visualization_resize(predicted_pixel_values), "n c h w -> (n h) w c").detach().cpu().numpy())
+        small_pixle_values = rearrange(self.visualization_resize(encoder_pixel_values), "n c h w -> (n h) w c").detach().cpu().numpy() 
+        small_predicted_pixel_values_face = rearrange(self.visualization_resize(predicted_pixel_values_face), "n c h w -> (n h) w c").detach().cpu().numpy()
+        small_pixle_values_face = rearrange(self.visualization_resize(encoder_pixel_values_face), "n c h w -> (n h) w c").detach().cpu().numpy()
+        
+        small_predicted_pixel_values = add_text_to_image(((small_predicted_pixel_values * 0.5 + 0.5) * 255).astype(np.uint8), "Pred Images", add_below=False)
+        small_pixle_values = add_text_to_image(((small_pixle_values * 0.5 + 0.5) * 255).astype(np.uint8), "Target Images", add_below=False)
+        small_predicted_pixel_values_face = add_text_to_image(((small_predicted_pixel_values_face * 0.5 + 0.5) * 255).astype(np.uint8), "Pred Faces", add_below=False)
+        small_pixle_values_face = add_text_to_image(((small_pixle_values_face * 0.5 + 0.5) * 255).astype(np.uint8), "Target Faces", add_below=False)
+
+
+        final_image = np.hstack([small_predicted_pixel_values, small_pixle_values, small_predicted_pixel_values_face, small_pixle_values_face])
+        for tracker in self.accelerator.trackers:
+            if tracker.name == 'wandb':
+                tracker.log({"IDLoss Visualization": wandb.Image(final_image, caption=f"loss: {loss.cpu().tolist()} timesteps: {timesteps.cpu().tolist()}, valid_indices: {valid_indices.cpu().tolist()}")})
+
+
+class ImageAugmentations(torch.nn.Module):
+    # Standard image augmentations used for CLIP loss to discourage adversarial outputs.
+    def __init__(self, output_size, augmentations_number, p=0.7):
+        super().__init__()
+        self.output_size = output_size
+        self.augmentations_number = augmentations_number
+
+        self.augmentations = torch.nn.Sequential(
+            K.RandomAffine(degrees=15, translate=0.1, p=p, padding_mode="border"),  # type: ignore
+            K.RandomPerspective(0.7, p=p),
+        )
+
+        self.avg_pool = torch.nn.AdaptiveAvgPool2d((self.output_size, self.output_size))
+
+        self.device = None
+
+    def forward(self, input):
+        """Extents the input batch with augmentations
+        If the input is consists of images [I1, I2] the extended augmented output
+        will be [I1_resized, I2_resized, I1_aug1, I2_aug1, I1_aug2, I2_aug2 ...]
+        Args:
+            input ([type]): input batch of shape [batch, C, H, W]
+        Returns:
+            updated batch: of shape [batch * augmentations_number, C, H, W]
+        """
+        # We want to multiply the number of images in the batch in contrast to regular augmantations
+        # that do not change the number of samples in the batch)
+        resized_images = self.avg_pool(input)
+        resized_images = torch.tile(resized_images, dims=(self.augmentations_number, 1, 1, 1))
+
+        batch_size = input.shape[0]
+        # We want at least one non augmented image
+        non_augmented_batch = resized_images[:batch_size]
+        augmented_batch = self.augmentations(resized_images[batch_size:])
+        updated_batch = torch.cat([non_augmented_batch, augmented_batch], dim=0)
+
+        return updated_batch
+
+
+class CLIPLoss(Loss):
+    def __init__(self, augmentations_number: int = 4, **kwargs):
+        super().__init__(**kwargs)
+
+        self.clip_model, clip_preprocess = clip.load("ViT-B/16", device=self.accelerator.device, jit=False)
+
+        self.clip_model.device = None
+
+        self.clip_model.eval().requires_grad_(False)
+        
+        self.preprocess = transforms.Compose([transforms.Normalize(mean=[-1.0, -1.0, -1.0], std=[2.0, 2.0, 2.0])] + # Un-normalize from [-1.0, 1.0] (SD output) to [0, 1].
+                                              clip_preprocess.transforms[:2] +                                      # to match CLIP input scale assumptions
+                                              clip_preprocess.transforms[4:])                                       # + skip convert PIL to tensor
+
+        self.clip_size = self.clip_model.visual.input_resolution
+
+        self.clip_normalize = transforms.Normalize(
+            mean=[0.48145466, 0.4578275, 0.40821073], std=[0.26862954, 0.26130258, 0.27577711]
+        )
+
+        self.image_augmentations = ImageAugmentations(output_size=self.clip_size,
+                                                      augmentations_number=augmentations_number)
+        
+        self.clip_model, self.image_augmentations = self.accelerator.prepare(self.clip_model, self.image_augmentations)
+
+    def forward(self, decoder_prompts, predicted_pixel_values: torch.Tensor, **kwargs) -> torch.Tensor:
+
+        if not isinstance(decoder_prompts, list):
+            decoder_prompts = [decoder_prompts]
+
+        tokens = clip.tokenize(decoder_prompts).to(predicted_pixel_values.device)
+        image  = self.preprocess(predicted_pixel_values)
+
+        logits_per_image, _ = self.clip_model(image, tokens)
+
+        logits_per_image = torch.diagonal(logits_per_image)
+
+        return (1. - logits_per_image / 100).mean()
+
+
+class DINOLoss(Loss):
+    def __init__(
+            self,
+            dino_model,
+            dino_preprocess,
+            output_hidden_states: bool = False,
+            center_momentum: float = 0.9,
+            student_temp: float = 0.1,
+            teacher_temp: float = 0.04,
+            warmup_teacher_temp: float = 0.04,
+            warmup_teacher_temp_epochs: int = 30,
+            **kwargs):
+        super().__init__(**kwargs)
+
+        self.dino_model = dino_model
+        self.output_hidden_states = output_hidden_states
+        self.rescale_factor = dino_preprocess.rescale_factor
+
+        # Un-normalize from [-1.0, 1.0] (SD output) to [0, 1].
+        self.preprocess = transforms.Compose(
+            [
+                transforms.Normalize(mean=[-1.0, -1.0, -1.0], std=[2.0, 2.0, 2.0]),
+                transforms.Resize(size=256),
+                transforms.CenterCrop(size=(224, 224)),
+                transforms.Normalize(mean=dino_preprocess.image_mean, std=dino_preprocess.image_std)
+            ]
+        )
+
+        self.student_temp = student_temp
+        self.teacher_temp = teacher_temp
+        self.center_momentum = center_momentum
+        self.center = torch.zeros(1, 257, 1024).to(self.accelerator.device, dtype=self.dtype)
+
+        # TODO: add temp, now fixed to 0.04
+        # we apply a warm up for the teacher temperature because
+        # a too high temperature makes the training instable at the beginning
+        # self.teacher_temp_schedule = np.concatenate((
+        #     np.linspace(warmup_teacher_temp,
+        #                 teacher_temp, warmup_teacher_temp_epochs),
+        #     np.ones(nepochs - warmup_teacher_temp_epochs) * teacher_temp
+        # ))
+
+        self.dino_model = self.accelerator.prepare(self.dino_model)
+
+    def forward(
+            self,
+            target: torch.Tensor,
+            predict: torch.Tensor,
+            weights: torch.Tensor = None,
+            **kwargs) -> torch.Tensor:
+
+        predict = self.preprocess(predict)
+        target = self.preprocess(target)
+
+        encoder_input = torch.cat([target, predict]).to(self.dino_model.device, dtype=self.dino_model.dtype)
+
+        if self.output_hidden_states:
+            raise ValueError("Output hidden states not supported for DINO loss.")
+            image_enc_hidden_states = self.dino_model(encoder_input, output_hidden_states=True).hidden_states[-2]
+        else:
+            image_enc_hidden_states = self.dino_model(encoder_input).last_hidden_state
+
+        teacher_output, student_output = image_enc_hidden_states.chunk(2, dim=0)         # [B, 257, 1024]
+
+        student_out = student_output.float() / self.student_temp
+
+        # teacher centering and sharpening
+        # temp = self.teacher_temp_schedule[epoch]
+        temp = self.teacher_temp
+        teacher_out = F.softmax((teacher_output.float() - self.center) / temp, dim=-1)
+        teacher_out = teacher_out.detach()
+
+        loss = torch.sum(-teacher_out * F.log_softmax(student_out, dim=-1), dim=-1, keepdim=True)
+        # self.update_center(teacher_output)
+
+        if weights is not None:
+            loss = loss * weights
+            return loss.mean()
+        return loss.mean()
+
+    @torch.no_grad()
+    def update_center(self, teacher_output):
+        """
+        Update center used for teacher output.
+        """
+        batch_center = torch.sum(teacher_output, dim=0, keepdim=True)
+        self.accelerator.reduce(batch_center, reduction="sum")
+        batch_center = batch_center / (len(teacher_output) * self.accelerator.num_processes)
+
+        # ema update
+        self.center = self.center * self.center_momentum + batch_center * (1 - self.center_momentum)