Spaces:
Configuration error
Configuration error
File size: 3,217 Bytes
cfdc687 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 |
import torch
from modules.base import BaseModule
from modules.linear_modulation import FeatureWiseAffine
from modules.interpolation import InterpolationBlock
from modules.layers import Conv1dWithInitialization
class BasicModulationBlock(BaseModule):
"""
Linear modulation part of UBlock, represented by sequence of the following layers:
- Feature-wise Affine
- LReLU
- 3x1 Conv
"""
def __init__(self, n_channels, dilation):
super(BasicModulationBlock, self).__init__()
self.featurewise_affine = FeatureWiseAffine()
self.leaky_relu = torch.nn.LeakyReLU(0.2)
self.convolution = Conv1dWithInitialization(
in_channels=n_channels,
out_channels=n_channels,
kernel_size=3,
stride=1,
padding=dilation,
dilation=dilation
)
def forward(self, x, scale, shift):
outputs = self.featurewise_affine(x, scale, shift)
outputs = self.leaky_relu(outputs)
outputs = self.convolution(outputs)
return outputs
class UpsamplingBlock(BaseModule):
def __init__(self, in_channels, out_channels, factor, dilations):
super(UpsamplingBlock, self).__init__()
self.first_block_main_branch = torch.nn.ModuleDict({
'upsampling': torch.nn.Sequential(*[
torch.nn.LeakyReLU(0.2),
InterpolationBlock(
scale_factor=factor,
mode='linear',
align_corners=False
),
Conv1dWithInitialization(
in_channels=in_channels,
out_channels=out_channels,
kernel_size=3,
stride=1,
padding=dilations[0],
dilation=dilations[0]
),
torch.nn.LeakyReLU(0.2)
]),
'modulation': BasicModulationBlock(
out_channels, dilation=dilations[1]
)
})
self.first_block_residual_branch = torch.nn.Sequential(*[
InterpolationBlock(
scale_factor=factor,
mode='linear',
align_corners=False
),
Conv1dWithInitialization(
in_channels=in_channels,
out_channels=out_channels,
kernel_size=1,
stride=1
)
])
self.second_block_main_branch = torch.nn.ModuleDict({
f'modulation_{idx}': BasicModulationBlock(
out_channels, dilation=dilations[2 + idx]
) for idx in range(2)
})
def forward(self, x, scale, shift):
# First upsampling residual block
outputs = self.first_block_main_branch['upsampling'](x)
outputs = self.first_block_main_branch['modulation'](outputs, scale, shift)
outputs = outputs + self.first_block_residual_branch(x)
# Second residual block
residual = self.second_block_main_branch['modulation_0'](outputs, scale, shift)
outputs = outputs + self.second_block_main_branch['modulation_1'](residual, scale, shift)
return outputs
|