DL-Art-School/codes/models/diffusion/rrdb_diffusion.py

import torch
import torch.nn as nn
import torch.nn.functional as F

from models.arch_util import ConvGnLelu, default_init_weights, make_layer
from models.diffusion.nn import timestep_embedding
from trainer.networks import register_model
from utils.util import checkpoint


# Conditionally uses torch's checkpoint functionality if it is enabled in the opt file.


class ResidualDenseBlock(nn.Module):
    """Residual Dense Block.

    Used in RRDB block in ESRGAN.

    Args:
        mid_channels (int): Channel number of intermediate features.
        growth_channels (int): Channels for each growth.
    """

    def __init__(self, mid_channels=64, growth_channels=32, embedding=False, init_weight=.1):
        super(ResidualDenseBlock, self).__init__()
        self.embedding = embedding
        if embedding:
            self.first_conv = ConvGnLelu(mid_channels, mid_channels, activation=True, norm=False, bias=True)
            self.emb_layers = nn.Sequential(
                nn.SiLU(),
                nn.Linear(
                    mid_channels*4,
                    mid_channels,
                ),
            )
        for i in range(5):
            out_channels = mid_channels if i == 4 else growth_channels
            self.add_module(
                f'conv{i + 1}',
                nn.Conv2d(mid_channels + i * growth_channels, out_channels, 3,
                          1, 1))
        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
        for i in range(5):
            default_init_weights(getattr(self, f'conv{i + 1}'), init_weight)
        default_init_weights(self.conv5, 0)

        self.normalize = nn.GroupNorm(num_groups=8, num_channels=mid_channels)

    def forward(self, x, emb):
        """Forward function.

        Args:
            x (Tensor): Input tensor with shape (n, c, h, w).

        Returns:
            Tensor: Forward results.
        """
        if self.embedding:
            x0 = self.first_conv(x)
            emb_out = self.emb_layers(emb).type(x0.dtype)
            while len(emb_out.shape) < len(x0.shape):
                emb_out = emb_out[..., None]
            x0 = x0 + emb_out
        else:
            x0 = x
        x1 = self.lrelu(self.conv1(x0))
        x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1)))
        x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1)))
        x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1)))
        x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))

        return self.normalize(x5 * .2 + x)


class RRDB(nn.Module):
    """Residual in Residual Dense Block.

    Used in RRDB-Net in ESRGAN.

    Args:
        mid_channels (int): Channel number of intermediate features.
        growth_channels (int): Channels for each growth.
    """

    def __init__(self, mid_channels, growth_channels=32):
        super(RRDB, self).__init__()
        self.rdb1 = ResidualDenseBlock(mid_channels, growth_channels, embedding=True)
        self.rdb2 = ResidualDenseBlock(mid_channels, growth_channels)
        self.rdb3 = ResidualDenseBlock(mid_channels, growth_channels)
        self.normalize = nn.GroupNorm(num_groups=8, num_channels=mid_channels)
        self.residual_mult = nn.Parameter(torch.FloatTensor([.1]))

    def forward(self, x, emb):
        """Forward function.

        Args:
            x (Tensor): Input tensor with shape (n, c, h, w).

        Returns:
            Tensor: Forward results.
        """
        out = self.rdb1(x, emb)
        out = self.rdb2(out, emb)
        out = self.rdb3(out, emb)

        return self.normalize(out * self.residual_mult + x)


class RRDBNet(nn.Module):
    """Networks consisting of Residual in Residual Dense Block, which is used
    in ESRGAN.

    ESRGAN: Enhanced Super-Resolution Generative Adversarial Networks.
    Currently, it supports x4 upsampling scale factor.

    Args:
        in_channels (int): Channel number of inputs.
        out_channels (int): Channel number of outputs.
        mid_channels (int): Channel number of intermediate features.
            Default: 64
        num_blocks (int): Block number in the trunk network. Defaults: 23
        growth_channels (int): Channels for each growth. Default: 32.
    """

    def __init__(self,
                 in_channels,
                 out_channels,
                 mid_channels=64,
                 num_blocks=23,
                 growth_channels=32,
                 body_block=RRDB,
                 ):
        super(RRDBNet, self).__init__()
        self.num_blocks = num_blocks
        self.in_channels = in_channels
        self.mid_channels = mid_channels

        # The diffusion RRDB starts with a full resolution image and downsamples into a .25 working space
        self.input_block = ConvGnLelu(in_channels, mid_channels, kernel_size=7, stride=1, activation=True, norm=False, bias=True)
        self.down1 = ConvGnLelu(mid_channels, mid_channels, kernel_size=3, stride=2, activation=True, norm=False, bias=True)
        self.down2 = ConvGnLelu(mid_channels, mid_channels, kernel_size=3, stride=2, activation=True, norm=False, bias=True)

        # Guided diffusion uses a time embedding.
        time_embed_dim = mid_channels * 4
        self.time_embed = nn.Sequential(
            nn.Linear(mid_channels, time_embed_dim),
            nn.SiLU(),
            nn.Linear(time_embed_dim, time_embed_dim),
        )

        self.body = make_layer(
            body_block,
            num_blocks,
            mid_channels=mid_channels,
            growth_channels=growth_channels)

        self.conv_body = nn.Conv2d(self.mid_channels, self.mid_channels, 3, 1, 1)
        # upsample
        self.conv_up1 = nn.Conv2d(self.mid_channels, self.mid_channels, 3, 1, 1)
        self.conv_up2 = nn.Conv2d(self.mid_channels*2, self.mid_channels, 3, 1, 1)
        self.conv_up3 = None
        self.conv_hr = nn.Conv2d(self.mid_channels*2, self.mid_channels, 3, 1, 1)
        self.conv_last = nn.Conv2d(self.mid_channels, out_channels, 3, 1, 1)

        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
        self.normalize = nn.GroupNorm(num_groups=8, num_channels=self.mid_channels)

        for m in [
            self.conv_body, self.conv_up1,
            self.conv_up2, self.conv_hr
        ]:
            if m is not None:
                default_init_weights(m, 1.0)
        default_init_weights(self.conv_last, 0)

    def forward(self, x, timesteps, low_res=None):
        emb = self.time_embed(timestep_embedding(timesteps, self.mid_channels))

        _, _, new_height, new_width = x.shape
        upsampled = F.interpolate(low_res, (new_height, new_width), mode="bilinear")
        x = torch.cat([x, upsampled], dim=1)

        d1 = self.input_block(x)
        d2 = self.down1(d1)
        feat = self.down2(d2)
        for bl in self.body:
            feat = checkpoint(bl, feat, emb)
        feat = feat[:, :self.mid_channels]
        feat = self.conv_body(feat)

        # upsample
        out = torch.cat([self.lrelu(
            self.normalize(self.conv_up1(F.interpolate(feat, scale_factor=2, mode='nearest')))),
            d2], dim=1)
        out = torch.cat([self.lrelu(
            self.normalize(self.conv_up2(F.interpolate(out, scale_factor=2, mode='nearest')))),
            d1], dim=1)
        out = self.conv_last(self.normalize(self.lrelu(self.conv_hr(out))))

        return out


@register_model
def register_rrdb_diffusion(opt_net, opt):
    return RRDBNet(**opt_net['args'])


if __name__ == '__main__':
    model = RRDBNet(6,6)
    x = torch.randn(1,3,128,128)
    l = torch.randn(1,3,32,32)
    t = torch.LongTensor([555])
    y = model(x, t, l)
    print(y.shape, y.mean(), y.std(), y.min(), y.max())