mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-22 08:16:08 +00:00
Delete animation/SadTalker directory
This commit is contained in:
@@ -1,140 +0,0 @@
|
||||
import torch.nn as nn
|
||||
import torch
|
||||
|
||||
|
||||
class ResidualConv(nn.Module):
|
||||
def __init__(self, input_dim, output_dim, stride, padding):
|
||||
super(ResidualConv, self).__init__()
|
||||
|
||||
self.conv_block = nn.Sequential(
|
||||
nn.BatchNorm2d(input_dim),
|
||||
nn.ReLU(),
|
||||
nn.Conv2d(
|
||||
input_dim, output_dim, kernel_size=3, stride=stride, padding=padding
|
||||
),
|
||||
nn.BatchNorm2d(output_dim),
|
||||
nn.ReLU(),
|
||||
nn.Conv2d(output_dim, output_dim, kernel_size=3, padding=1),
|
||||
)
|
||||
self.conv_skip = nn.Sequential(
|
||||
nn.Conv2d(input_dim, output_dim, kernel_size=3, stride=stride, padding=1),
|
||||
nn.BatchNorm2d(output_dim),
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
|
||||
return self.conv_block(x) + self.conv_skip(x)
|
||||
|
||||
|
||||
class Upsample(nn.Module):
|
||||
def __init__(self, input_dim, output_dim, kernel, stride):
|
||||
super(Upsample, self).__init__()
|
||||
|
||||
self.upsample = nn.ConvTranspose2d(
|
||||
input_dim, output_dim, kernel_size=kernel, stride=stride
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return self.upsample(x)
|
||||
|
||||
|
||||
class Squeeze_Excite_Block(nn.Module):
|
||||
def __init__(self, channel, reduction=16):
|
||||
super(Squeeze_Excite_Block, self).__init__()
|
||||
self.avg_pool = nn.AdaptiveAvgPool2d(1)
|
||||
self.fc = nn.Sequential(
|
||||
nn.Linear(channel, channel // reduction, bias=False),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Linear(channel // reduction, channel, bias=False),
|
||||
nn.Sigmoid(),
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
b, c, _, _ = x.size()
|
||||
y = self.avg_pool(x).view(b, c)
|
||||
y = self.fc(y).view(b, c, 1, 1)
|
||||
return x * y.expand_as(x)
|
||||
|
||||
|
||||
class ASPP(nn.Module):
|
||||
def __init__(self, in_dims, out_dims, rate=[6, 12, 18]):
|
||||
super(ASPP, self).__init__()
|
||||
|
||||
self.aspp_block1 = nn.Sequential(
|
||||
nn.Conv2d(
|
||||
in_dims, out_dims, 3, stride=1, padding=rate[0], dilation=rate[0]
|
||||
),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.BatchNorm2d(out_dims),
|
||||
)
|
||||
self.aspp_block2 = nn.Sequential(
|
||||
nn.Conv2d(
|
||||
in_dims, out_dims, 3, stride=1, padding=rate[1], dilation=rate[1]
|
||||
),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.BatchNorm2d(out_dims),
|
||||
)
|
||||
self.aspp_block3 = nn.Sequential(
|
||||
nn.Conv2d(
|
||||
in_dims, out_dims, 3, stride=1, padding=rate[2], dilation=rate[2]
|
||||
),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.BatchNorm2d(out_dims),
|
||||
)
|
||||
|
||||
self.output = nn.Conv2d(len(rate) * out_dims, out_dims, 1)
|
||||
self._init_weights()
|
||||
|
||||
def forward(self, x):
|
||||
x1 = self.aspp_block1(x)
|
||||
x2 = self.aspp_block2(x)
|
||||
x3 = self.aspp_block3(x)
|
||||
out = torch.cat([x1, x2, x3], dim=1)
|
||||
return self.output(out)
|
||||
|
||||
def _init_weights(self):
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
nn.init.kaiming_normal_(m.weight)
|
||||
elif isinstance(m, nn.BatchNorm2d):
|
||||
m.weight.data.fill_(1)
|
||||
m.bias.data.zero_()
|
||||
|
||||
|
||||
class Upsample_(nn.Module):
|
||||
def __init__(self, scale=2):
|
||||
super(Upsample_, self).__init__()
|
||||
|
||||
self.upsample = nn.Upsample(mode="bilinear", scale_factor=scale)
|
||||
|
||||
def forward(self, x):
|
||||
return self.upsample(x)
|
||||
|
||||
|
||||
class AttentionBlock(nn.Module):
|
||||
def __init__(self, input_encoder, input_decoder, output_dim):
|
||||
super(AttentionBlock, self).__init__()
|
||||
|
||||
self.conv_encoder = nn.Sequential(
|
||||
nn.BatchNorm2d(input_encoder),
|
||||
nn.ReLU(),
|
||||
nn.Conv2d(input_encoder, output_dim, 3, padding=1),
|
||||
nn.MaxPool2d(2, 2),
|
||||
)
|
||||
|
||||
self.conv_decoder = nn.Sequential(
|
||||
nn.BatchNorm2d(input_decoder),
|
||||
nn.ReLU(),
|
||||
nn.Conv2d(input_decoder, output_dim, 3, padding=1),
|
||||
)
|
||||
|
||||
self.conv_attn = nn.Sequential(
|
||||
nn.BatchNorm2d(output_dim),
|
||||
nn.ReLU(),
|
||||
nn.Conv2d(output_dim, 1, 1),
|
||||
)
|
||||
|
||||
def forward(self, x1, x2):
|
||||
out = self.conv_encoder(x1) + self.conv_decoder(x2)
|
||||
out = self.conv_attn(out)
|
||||
return out * x2
|
||||
Reference in New Issue
Block a user