import torch
import torch.nn as nn
import torch.nn.functional as F
import math
import copy
from einops import rearrangeclass SinusoidalPosEmb(nn.Module):def __init__(self, dim):super().__init__()self.dim = dimdef forward(self, x):device = x.devicehalf_dim = self.dim // 2emb = math.log(10000) / (half_dim - 1)emb = torch.exp(torch.arange(half_dim, device=device) * -emb)emb = x[:, None] * emb[None, :]emb = torch.cat((emb.sin(), emb.cos()), dim=-1)return embclass single_conv(nn.Module):def __init__(self, in_ch, out_ch):super(single_conv, self).__init__()self.conv = nn.Sequential(nn.Conv2d(in_ch, out_ch, 5, padding=2),nn.ReLU(inplace=True))def forward(self, x):return self.conv(x)class up(nn.Module):def __init__(self, in_ch):super(up, self).__init__()self.up = nn.ConvTranspose2d(in_ch, in_ch // 2, 2, stride=2)self.conv_concat = nn.Conv2d(in_ch, in_ch // 2, 3, 1, 1)self.relu = nn.ReLU(inplace=True)def forward(self, x1, x2):x1 = self.up(x1)# input is CHWdiffY = x2.size()[2] - x1.size()[2]diffX = x2.size()[3] - x1.size()[3]x1 = F.pad(x1, (diffX // 2, diffX - diffX // 2,diffY // 2, diffY - diffY // 2))x = torch.cat((x1, x2), dim=1)x = self.conv_concat(x)x = self.relu(x)return x# diff
import torch
import torch.nn as nnclass EMA(nn.Module):def __init__(self, channels, c2=None, factor=32):super(EMA, self).__init__()self.groups = factorassert channels // self.groups > 0self.softmax = nn.Softmax(-1)self.agp = nn.AdaptiveAvgPool2d((1, 1))self.pool_h = nn.AdaptiveAvgPool2d((None, 1))self.pool_w = nn.AdaptiveAvgPool2d((1, None))self.gn = nn.GroupNorm(channels // self.groups, channels // self.groups)self.conv1x1 = nn.Conv2d(channels // self.groups, channels // self.groups, kernel_size=5, stride=1, padding=2)self.conv1x1_3 = nn.Conv2d(channels // self.groups, channels // self.groups, kernel_size=3, stride=1, padding=1)self.conv1x1_c = nn.Conv2d(channels // self.groups*2, channels // self.groups, kernel_size=5, stride=1, padding=2)self.conv3x3 = nn.Conv2d(channels // self.groups, channels // self.groups, kernel_size=5, stride=1, padding=2)self.conv3x3_3 = nn.Conv2d(channels // self.groups, channels // self.groups, kernel_size=3, stride=1, padding=1)self.conv3x3_c = nn.Conv2d(channels // self.groups*2, channels // self.groups, kernel_size=5, stride=1, padding=2)def forward(self, x):b, c, h, w = x.size()group_x = x.reshape(b * self.groups, -1, h, w) # b*g,c//g,h,wx_h = self.pool_h(group_x)x_w = self.pool_w(group_x).permute(0, 1, 3, 2)hw_5 = self.conv1x1(torch.cat([x_h, x_w], dim=2))hw_3 = self.conv1x1_3(torch.cat([x_h, x_w], dim=2))hw = torch.cat([hw_5, hw_3], dim=1)hw = self.conv1x1_c(hw)x_h, x_w = torch.split(hw, [h, w], dim=2)x1 = self.gn(group_x * x_h.sigmoid() * x_w.permute(0, 1, 3, 2).sigmoid())x2_5 = self.conv3x3(group_x)x2_3 = self.conv3x3_3(group_x)x2 = torch.cat([x2_5, x2_3], dim=1)x2 = self.conv3x3_c(x2)x11 = self.softmax(self.agp(x1).reshape(b * self.groups, -1, 1).permute(0, 2, 1))x12 = x2.reshape(b * self.groups, c // self.groups, -1) # b*g, c//g, hwx21 = self.softmax(self.agp(x2).reshape(b * self.groups, -1, 1).permute(0, 2, 1))x22 = x1.reshape(b * self.groups, c // self.groups, -1) # b*g, c//g, hwweights = (torch.matmul(x11, x12) + torch.matmul(x21, x22)).reshape(b * self.groups, 1, h, w)return (group_x * weights.sigmoid()).reshape(b, c, h, w)# diff
class down(nn.Module):def __init__(self, in_channels, dilation_rates=(1, 2, 3)):super(down, self).__init__()self.conv = nn.ModuleList([nn.Sequential(nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=d, dilation=d, bias=False),nn.ReLU(inplace=True))for d in dilation_rates])self.conv_2 = nn.Conv2d(in_channels * 3, in_channels, 3, 1, 1)self.relu = nn.ReLU(inplace=True)def forward(self, x):outputs = [conv(x) for conv in self.conv]out = torch.cat(outputs, dim=1)out = self.conv_2(out)out = self.relu(out)return outclass outconv(nn.Module):def __init__(self, in_ch, out_ch):super(outconv, self).__init__()self.conv = nn.Conv2d(in_ch, out_ch, 1)def forward(self, x):x = self.conv(x)return xclass adjust_net(nn.Module):def __init__(self, out_channels=64, middle_channels=32):super(adjust_net, self).__init__()self.model = nn.Sequential(nn.Conv2d(2, middle_channels, 3, padding=1),nn.ReLU(inplace=True),nn.AvgPool2d(2),nn.Conv2d(middle_channels, middle_channels * 2, 3, padding=1),nn.ReLU(inplace=True),nn.AvgPool2d(2),nn.Conv2d(middle_channels * 2, middle_channels * 4, 3, padding=1),nn.ReLU(inplace=True),nn.AvgPool2d(2),nn.Conv2d(middle_channels * 4, out_channels * 2, 1, padding=0))def forward(self, x):out = self.model(x)out = F.adaptive_avg_pool2d(out, (1, 1))out1 = out[:, :out.shape[1] // 2]out2 = out[:, out.shape[1] // 2:]return out1, out2# The architecture of U-Net refers to "Toward Convolutional Blind Denoising of Real Photographs",
# official MATLAB implementation: https://github.com/GuoShi28/CBDNet.
# unofficial PyTorch implementation: https://github.com/IDKiro/CBDNet-pytorch/tree/master.
# We improved it by adding time step embedding and EMM module, while removing the noise estimation network.
class UNet(nn.Module):def __init__(self, in_channels=2, out_channels=1):super(UNet, self).__init__()dim = 32self.time_mlp = nn.Sequential(SinusoidalPosEmb(dim),nn.Linear(dim, dim * 4),nn.GELU(),nn.Linear(dim * 4, dim))self.inc = nn.Sequential(single_conv(in_channels, 64),single_conv(64, 64))self.down1 = down(64)self.mlp1 = nn.Sequential(nn.GELU(),nn.Linear(dim, 64))self.adjust1 = adjust_net(64)self.conv1 = nn.Sequential(single_conv(64, 128),single_conv(128, 128),single_conv(128, 128))self.down2 = down(128)self.mlp2 = nn.Sequential(nn.GELU(),nn.Linear(dim, 128))self.adjust2 = adjust_net(128)self.conv2 = nn.Sequential(single_conv(128, 256),single_conv(256, 256),single_conv(256, 256),single_conv(256, 256),single_conv(256, 256),single_conv(256, 256))self.up1 = up(256)self.mlp3 = nn.Sequential(nn.GELU(),nn.Linear(dim, 128))self.adjust3 = adjust_net(128)self.conv3 = nn.Sequential(single_conv(128, 128),single_conv(128, 128),single_conv(128, 128))self.up2 = up(128)self.mlp4 = nn.Sequential(nn.GELU(),nn.Linear(dim, 64))self.adjust4 = adjust_net(64)self.conv4 = nn.Sequential(single_conv(64, 64),single_conv(64, 64))self.outc = outconv(64, out_channels)self.ema1 = EMA(channels=128)self.ema2 = EMA(channels=256)self.ema3 = EMA(channels=128)self.ema4 = EMA(channels=64)def forward(self, x, t, x_adjust, adjust):inx = self.inc(x)time_emb = self.time_mlp(t)down1 = self.down1(inx)condition1 = self.mlp1(time_emb)b, c = condition1.shapecondition1 = rearrange(condition1, 'b c -> b c 1 1')if adjust:gamma1, beta1 = self.adjust1(x_adjust)down1 = down1 + gamma1 * condition1 + beta1else:down1 = down1 + condition1conv1 = self.conv1(down1)# diff: 通过EMA模块处理conv1的输出conv1 = self.ema1(conv1)down2 = self.down2(conv1)condition2 = self.mlp2(time_emb)b, c = condition2.shapecondition2 = rearrange(condition2, 'b c -> b c 1 1')if adjust:gamma2, beta2 = self.adjust2(x_adjust)down2 = down2 + gamma2 * condition2 + beta2else:down2 = down2 + condition2conv2 = self.conv2(down2)# diff: 通过EMA模块处理conv2的输出conv2 = self.ema2(conv2)up1 = self.up1(conv2, conv1)condition3 = self.mlp3(time_emb)b, c = condition3.shapecondition3 = rearrange(condition3, 'b c -> b c 1 1')if adjust:gamma3, beta3 = self.adjust3(x_adjust)up1 = up1 + gamma3 * condition3 + beta3else:up1 = up1 + condition3conv3 = self.conv3(up1)conv3 = self.ema3(conv3)up2 = self.up2(conv3, inx)condition4 = self.mlp4(time_emb)b, c = condition4.shapecondition4 = rearrange(condition4, 'b c -> b c 1 1')if adjust:gamma4, beta4 = self.adjust4(x_adjust)up2 = up2 + gamma4 * condition4 + beta4else:up2 = up2 + condition4conv4 = self.conv4(up2)conv4 = self.ema4(conv4)out = self.outc(conv4)return outclass Network(nn.Module):def __init__(self, in_channels=3, out_channels=1, context=True):super(Network, self).__init__()self.unet = UNet(in_channels=in_channels, out_channels=out_channels)self.context = contextdef forward(self, x, t, y, x_end, adjust=True):if self.context:x_middle = x[:, 1].unsqueeze(1)else:x_middle = xx_adjust = torch.cat((y, x_end), dim=1)out = self.unet(x, t, x_adjust, adjust=adjust) + x_middlereturn out# WeightNet of the one-shot learning framework
class WeightNet(nn.Module):def __init__(self, weight_num=10):super(WeightNet, self).__init__()init = torch.ones([1, weight_num, 1, 1]) / weight_numself.weights = nn.Parameter(init)def forward(self, x):weights = F.softmax(self.weights, 1)out = weights * xout = out.sum(dim=1, keepdim=True)return out, weights# Main function to test the model
def main():# Set random seed for reproducibilitytorch.manual_seed(0)# Instantiate the modelmodel = Network(in_channels=2, out_channels=1, context=True)# Set the model to evaluation mode (since we're debugging dimensions)model.eval()# Define input dimensionsbatch_size = 2num_channels = 2 # Input channels for xheight = 64width = 64# Create sample inputsx = torch.randn(batch_size, num_channels, height, width)t = torch.randint(0, 1000, (batch_size,))y = torch.randn(batch_size, 1, height, width)x_end = torch.randn(batch_size, 1, height, width)# Pass the inputs through the modelwith torch.no_grad():output = model(x, t, y, x_end, adjust=True)# Print the output shapeprint(f"Final output shape: {output.shape}")if __name__ == "__main__":main()
