sgm/modules/diffusionmodules/loss.py

from typing import List, Optional, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
from omegaconf import ListConfig
from torchvision.utils import save_image
from ...util import append_dims, instantiate_from_config


class StandardDiffusionLoss(nn.Module):
    def __init__(
        self,
        sigma_sampler_config,
        type="l2",
        offset_noise_level=0.0,
        batch2model_keys: Optional[Union[str, List[str], ListConfig]] = None,
    ):
        super().__init__()

        assert type in ["l2", "l1"]

        self.sigma_sampler = instantiate_from_config(sigma_sampler_config)

        self.type = type
        self.offset_noise_level = offset_noise_level

        if not batch2model_keys:
            batch2model_keys = []

        if isinstance(batch2model_keys, str):
            batch2model_keys = [batch2model_keys]

        self.batch2model_keys = set(batch2model_keys)

    def __call__(self, network, denoiser, conditioner, input, batch, *args, **kwarg):
        cond = conditioner(batch)
        additional_model_inputs = {
            key: batch[key] for key in self.batch2model_keys.intersection(batch)
        }

        sigmas = self.sigma_sampler(input.shape[0]).to(input.device)
        noise = torch.randn_like(input)
        if self.offset_noise_level > 0.0:
            noise = noise + self.offset_noise_level * append_dims(
                torch.randn(input.shape[0], device=input.device), input.ndim
            )
        noised_input = input + noise * append_dims(sigmas, input.ndim)
        model_output = denoiser(
            network, noised_input, sigmas, cond, **additional_model_inputs
        )
        w = append_dims(denoiser.w(sigmas), input.ndim)

        loss = self.get_diff_loss(model_output, input, w)
        loss = loss.mean()
        loss_dict = {"loss": loss}

        return loss, loss_dict

    def get_diff_loss(self, model_output, target, w):
        if self.type == "l2":
            return torch.mean(
                (w * (model_output - target) ** 2).reshape(target.shape[0], -1), 1
            )
        elif self.type == "l1":
            return torch.mean(
                (w * (model_output - target).abs()).reshape(target.shape[0], -1), 1
            )


class FullLoss(StandardDiffusionLoss):

    def __init__(
        self,
        seq_len=12,
        kernel_size=3,
        gaussian_sigma=0.5,
        min_attn_size=16,
        lambda_local_loss=0.0,
        lambda_ocr_loss=0.0,
        ocr_enabled = False,
        predictor_config = None,
        *args, **kwarg
    ):
        super().__init__(*args, **kwarg)

        self.gaussian_kernel_size = kernel_size
        gaussian_kernel = self.get_gaussian_kernel(kernel_size=self.gaussian_kernel_size, sigma=gaussian_sigma, out_channels=seq_len)
        self.register_buffer("g_kernel", gaussian_kernel.requires_grad_(False))

        self.min_attn_size = min_attn_size
        self.lambda_local_loss = lambda_local_loss
        self.lambda_ocr_loss = lambda_ocr_loss

        self.ocr_enabled = ocr_enabled
        if ocr_enabled:
            self.predictor = instantiate_from_config(predictor_config)
    
    def get_gaussian_kernel(self, kernel_size=3, sigma=1, out_channels=3):
        # Create a x, y coordinate grid of shape (kernel_size, kernel_size, 2)
        x_coord = torch.arange(kernel_size)
        x_grid = x_coord.repeat(kernel_size).view(kernel_size, kernel_size)
        y_grid = x_grid.t()
        xy_grid = torch.stack([x_grid, y_grid], dim=-1).float()

        mean = (kernel_size - 1)/2.
        variance = sigma**2.

        # Calculate the 2-dimensional gaussian kernel which is
        # the product of two gaussian distributions for two different
        # variables (in this case called x and y)
        gaussian_kernel = (1./(2.*torch.pi*variance)) *\
                        torch.exp(
                            -torch.sum((xy_grid - mean)**2., dim=-1) /\
                            (2*variance)
                        )

        # Make sure sum of values in gaussian kernel equals 1.
        gaussian_kernel = gaussian_kernel / torch.sum(gaussian_kernel)

        # Reshape to 2d depthwise convolutional weight
        gaussian_kernel = gaussian_kernel.view(1, 1, kernel_size, kernel_size)
        gaussian_kernel = gaussian_kernel.tile(out_channels, 1, 1, 1)
        
        return gaussian_kernel

    def __call__(self, network, denoiser, conditioner, input, batch, first_stage_model, scaler):

        cond = conditioner(batch)

        sigmas = self.sigma_sampler(input.shape[0]).to(input.device)
        noise = torch.randn_like(input)
        if self.offset_noise_level > 0.0:
            noise = noise + self.offset_noise_level * append_dims(
                torch.randn(input.shape[0], device=input.device), input.ndim
            )

        noised_input = input + noise * append_dims(sigmas, input.ndim)
        model_output = denoiser(network, noised_input, sigmas, cond)
        w = append_dims(denoiser.w(sigmas), input.ndim)

        diff_loss = self.get_diff_loss(model_output, input, w)
        local_loss = self.get_local_loss(network.diffusion_model.attn_map_cache, batch["seg"], batch["seg_mask"])
        diff_loss = diff_loss.mean()
        local_loss = local_loss.mean()

        if self.ocr_enabled:
            ocr_loss = self.get_ocr_loss(model_output, batch["r_bbox"], batch["label"], first_stage_model, scaler)
            ocr_loss = ocr_loss.mean()

        loss = diff_loss + self.lambda_local_loss * local_loss
        if self.ocr_enabled:
            loss += self.lambda_ocr_loss * ocr_loss

        loss_dict = {
            "loss/diff_loss": diff_loss,
            "loss/local_loss": local_loss,
            "loss/full_loss": loss
        }

        if self.ocr_enabled:
            loss_dict["loss/ocr_loss"] = ocr_loss

        return loss, loss_dict
    
    def get_ocr_loss(self, model_output, r_bbox, label, first_stage_model, scaler):

        model_output = 1 / scaler * model_output
        model_output_decoded = first_stage_model.decode(model_output)
        model_output_crops = []
        
        for i, bbox in enumerate(r_bbox):
            m_top, m_bottom, m_left, m_right = bbox
            model_output_crops.append(model_output_decoded[i, :, m_top:m_bottom, m_left:m_right])

        loss = self.predictor.calc_loss(model_output_crops, label)

        return loss

    def get_min_local_loss(self, attn_map_cache, mask, seg_mask):

        loss = 0
        count = 0

        for item in attn_map_cache:

            heads = item["heads"]
            size = item["size"]
            attn_map = item["attn_map"]

            if size < self.min_attn_size: continue

            seg_l = seg_mask.shape[1]

            bh, n, l = attn_map.shape # bh: batch size * heads / n : pixel length(h*w) / l: token length
            attn_map = attn_map.reshape((-1, heads, n, l)) # b, h, n, l
            
            assert seg_l <= l
            attn_map = attn_map[..., :seg_l]
            attn_map = attn_map.permute(0, 1, 3, 2) # b, h, l, n
            attn_map = attn_map.mean(dim = 1) # b, l, n

            attn_map = attn_map.reshape((-1, seg_l, size, size)) # b, l, s, s
            attn_map = F.conv2d(attn_map, self.g_kernel, padding = self.gaussian_kernel_size//2, groups=seg_l) # gaussian blur on each channel
            attn_map = attn_map.reshape((-1, seg_l, n)) # b, l, n
            
            mask_map = F.interpolate(mask, (size, size))
            mask_map = mask_map.tile((1, seg_l, 1, 1))
            mask_map = mask_map.reshape((-1, seg_l, n)) # b, l, n

            p_loss = (mask_map * attn_map).max(dim = -1)[0] # b, l
            p_loss = p_loss + (1 - seg_mask) # b, l
            p_loss = p_loss.min(dim = -1)[0] # b,

            loss += -p_loss
            count += 1

        loss = loss / count

        return loss

    def get_local_loss(self, attn_map_cache, seg, seg_mask):

        loss = 0
        count = 0

        for item in attn_map_cache:

            heads = item["heads"]
            size = item["size"]
            attn_map = item["attn_map"]

            if size < self.min_attn_size: continue

            seg_l = seg_mask.shape[1]

            bh, n, l = attn_map.shape # bh: batch size * heads / n : pixel length(h*w) / l: token length
            attn_map = attn_map.reshape((-1, heads, n, l)) # b, h, n, l
            
            assert seg_l <= l
            attn_map = attn_map[..., :seg_l]
            attn_map = attn_map.permute(0, 1, 3, 2) # b, h, l, n
            attn_map = attn_map.mean(dim = 1) # b, l, n

            attn_map = attn_map.reshape((-1, seg_l, size, size)) # b, l, s, s
            attn_map = F.conv2d(attn_map, self.g_kernel, padding = self.gaussian_kernel_size//2, groups=seg_l) # gaussian blur on each channel
            attn_map = attn_map.reshape((-1, seg_l, n)) # b, l, n

            seg_map = F.interpolate(seg, (size, size))
            seg_map = seg_map.reshape((-1, seg_l, n)) # b, l, n
            n_seg_map = 1 - seg_map

            p_loss = (seg_map * attn_map).max(dim = -1)[0] # b, l
            n_loss = (n_seg_map * attn_map).max(dim = -1)[0] # b, l

            p_loss = p_loss * seg_mask # b, l
            n_loss = n_loss * seg_mask # b, l

            p_loss = p_loss.sum(dim = -1) / seg_mask.sum(dim = -1) # b,
            n_loss = n_loss.sum(dim = -1) / seg_mask.sum(dim = -1) # b,

            f_loss = n_loss - p_loss # b,
            loss += f_loss
            count += 1

        loss = loss / count

        return loss