q-future/Co-Instruct

1import math2import torch3import torch.nn as nn4import torch.nn.functional as F5from torch.nn import init as init6from torch.nn.modules.batchnorm import _BatchNorm7 8from insir_models.nafnet_utils import Local_Base, LayerNorm2d9from insir_models.nafnet import SimpleGate, NAFBlock10 11 12class ICB(nn.Module):13    """14    Instruction Condition Block (ICB)15    Paper Section 3.316    """17 18    def __init__(self, feature_dim, text_dim=768):19        super(ICB, self).__init__()20        self.fc    = nn.Linear(text_dim, feature_dim)21        self.block = NAFBlock(feature_dim)22        self.beta  = nn.Parameter(torch.zeros((1, feature_dim, 1, 1)), requires_grad=True)23        self.gamma = nn.Parameter(torch.zeros((1, feature_dim, 1, 1)), requires_grad=True)24 25    def forward(self, x, text_embedding):26        gating_factors = torch.sigmoid(self.fc(text_embedding))27        gating_factors = gating_factors.unsqueeze(-1).unsqueeze(-1)28 29        f = x * self.gamma + self.beta  # 1) learned feature scaling/modulation30        f = f * gating_factors          # 2) (soft) feature routing based on text31        f = self.block(f)               # 3) block feature enhancement32        return f + x33 34 35class InstructIR(nn.Module):36    """37    InstructIR model using NAFNet (ECCV 2022) as backbone.38    The model takes as input an RGB image and a text embedding (encoded instruction).39    Described in Paper Section 3.340    """41 42    def __init__(self, img_channel=3, width=16, middle_blk_num=1, enc_blk_nums=[], dec_blk_nums=[], txtdim=768):43        super().__init__()44 45        self.intro  = nn.Conv2d(in_channels=img_channel, out_channels=width, kernel_size=3, padding=1, stride=1, groups=1,46                              bias=True)47        self.ending = nn.Conv2d(in_channels=width, out_channels=img_channel, kernel_size=3, padding=1, stride=1, groups=1,48                              bias=True)49 50        self.encoders    = nn.ModuleList()51        self.decoders    = nn.ModuleList()52        self.middle_blks = nn.ModuleList()53        self.ups         = nn.ModuleList()54        self.downs       = nn.ModuleList()55        self.enc_cond    = nn.ModuleList()56        self.dec_cond    = nn.ModuleList()57 58        chan = width59        for num in enc_blk_nums:60            self.encoders.append(61                nn.Sequential(62                    *[NAFBlock(chan) for _ in range(num)]63                )64            )65            66            self.enc_cond.append(ICB(chan, txtdim))67 68            self.downs.append(69                nn.Conv2d(chan, 2*chan, 2, 2)70            )71            chan = chan * 272 73        self.middle_blks = nn.Sequential(74                *[NAFBlock(chan) for _ in range(middle_blk_num)]75            )76 77        for num in dec_blk_nums:78            self.ups.append(79                nn.Sequential(80                    nn.Conv2d(chan, chan * 2, 1, bias=False),81                    nn.PixelShuffle(2)82                )83            )84            chan = chan // 285            self.decoders.append(86                nn.Sequential(87                    *[NAFBlock(chan) for _ in range(num)]88                )89            )90            # Add text embedding as modulation91            self.dec_cond.append(ICB(chan, txtdim))92 93        self.padder_size = 2 ** len(self.encoders)94 95    def forward(self, inp, txtembd):96        B, C, H, W = inp.shape97        inp = self.check_image_size(inp)98 99        x = self.intro(inp)100        encs = []101 102        for encoder, enc_mod, down in zip(self.encoders, self.enc_cond, self.downs):103            x = encoder(x)104            x = enc_mod(x, txtembd)105            encs.append(x)106            x = down(x)107 108        x = self.middle_blks(x)109 110        for decoder, up, enc_skip, dec_mod in zip(self.decoders, self.ups, encs[::-1], self.dec_cond):111            x = up(x)112            x = x + enc_skip113            x = decoder(x)114            x = dec_mod(x, txtembd)115 116        x = self.ending(x)117        x = x + inp118 119        return x[:, :, :H, :W]120 121    def check_image_size(self, x):122        _, _, h, w = x.size()123        mod_pad_h = (self.padder_size - h % self.padder_size) % self.padder_size124        mod_pad_w = (self.padder_size - w % self.padder_size) % self.padder_size125        x = F.pad(x, (0, mod_pad_w, 0, mod_pad_h))126        return x127 128 129def create_model(input_channels = 3, width = 32, enc_blks = [2, 2, 4, 8], middle_blk_num = 12, dec_blks = [2, 2, 2, 2], txtdim=768):130 131    net = InstructIR(img_channel=input_channels, width=width, middle_blk_num=middle_blk_num,132                      enc_blk_nums=enc_blks, dec_blk_nums=dec_blks, txtdim=txtdim)133 134    return net