def __init__(self,
in_size,
out_size,
cond_size=512,
depth=2,
add_noise=False,
kernel_size=1,
activation=None):
super().__init__()
self.add_noise = add_noise
self.activation = activation or nn.LeakyReLU(0.2)
self.conds = nn.ModuleList([
lr_equal(nn.Linear(cond_size, in_size, bias=True))
for idx in range(depth)
])
with torch.no_grad():
for cond in self.conds:
cond.module.bias.zero_().add_(1.0)
self.blocks = nn.ModuleList([
WeightDemodulation(nn.Conv2d(in_size, in_size, kernel_size))
for idx in range(depth - 1)
] + [WeightDemodulation(nn.Conv2d(in_size, out_size, kernel_size))])
self._noise_scale = []
for idx in range(depth):
self._noise_scale.append(idx)
size = in_size if idx < depth - 1 else out_size
setattr(self, f"noise_scale_{idx}",
nn.Parameter(torch.zeros(1, requires_grad=True)))
self.rgb = lr_equal(nn.Conv2d(out_size, 3, 1))
def __init__(self,
in_channels,
base_channels=128,
cond_size=512,
channel_factors=None,
block_depth=2,
activation=None,
add_noise=True):
super().__init__()
self.block_depth = block_depth
self.activation = activation or nn.LeakyReLU(0.2)
self.project_in = lr_equal(
nn.Conv2d(in_channels,
base_channels * channel_factors[0],
1,
bias=False))
self.blocks = nn.ModuleList([
IndependentStyleGAN2Block(in_factor * base_channels,
out_factor * base_channels,
cond_size=cond_size,
depth=block_depth,
activation=activation,
kernel_size=3,
add_noise=add_noise)
for idx, (in_factor, out_factor) in enumerate(
zip(channel_factors[:-1], channel_factors[1:]))
])
def __init__(self,
in_size=512,
cond_size=512,
hidden_size=128,
depth=3,
lr_scale=0.01,
activation=None):
super().__init__()
self.project_in = lr_equal(nn.Conv2d(in_size, hidden_size, 1),
lr_scale=lr_scale)
self.project_out = lr_equal(nn.Conv2d(hidden_size, cond_size, 1),
lr_scale=lr_scale)
self.blocks = nn.ModuleList([
lr_equal(nn.Conv2d(hidden_size, hidden_size, 1), lr_scale=lr_scale)
for idx in range(depth - 2)
])
self.activation = activation or nn.LeakyReLU(0.2)
def __init__(self,
in_size=None,
base_channels=128,
cond_size=512,
channel_factors=None,
block_depth=2,
activation=None):
super().__init__(base_channels=base_channels,
cond_size=cond_size,
channel_factors=channel_factors,
block_depth=block_depth,
activation=activation)
if in_size is None:
in_size = base_channels * channel_factors[0]
self.project_in = lr_equal(
nn.Conv2d(in_size + base_channels * channel_factors[0],
base_channels * channel_factors[0], 1))
def __init__(self,
base_channels=128,
cond_size=512,
channel_factors=None,
block_depth=2,
scale=50.0,
target_size=128,
activation=None,
add_noise=True):
super().__init__()
self.target_size = target_size
self.channel_factors = channel_factors
self.base_channels = base_channels
self.block_depth = block_depth
self.project = ProjectCoordinates(cond_size=cond_size,
in_size=2,
out_size=cond_size,
scale=scale)
self.project_local = LocalMapping(cond_size=cond_size)
self.levels = nn.ModuleList([
LevelStyleGAN2GeneratorBlock(2 * cond_size,
base_channels=base_channels,
cond_size=cond_size,
channel_factors=cf,
block_depth=block_depth,
activation=activation,
add_noise=add_noise)
for cf in channel_factors
])
self.projections = nn.ModuleList([
lr_equal(
nn.Conv2d(cf[-1] * base_channels + 3, cond_size, 1,
bias=False)) for cf in channel_factors
])
coordinate_size = self.input_size(0)
self.inputs = nn.Parameter(
torch.randn(cond_size,
coordinate_size,
coordinate_size,
requires_grad=True))
def __init__(self,
base_channels=128,
cond_size=512,
channel_factors=None,
block_depth=2,
scale=50.0,
activation=None,
add_noise=True):
super().__init__()
self.channel_factors = channel_factors
self.base_channels = base_channels
self.block_depth = block_depth
self.project = ProjectCoordinates(cond_size=cond_size,
in_size=2,
out_size=cond_size,
scale=scale)
self.project_local = LocalMapping(cond_size=cond_size)
self.combine = lr_equal(
nn.Conv2d(2 * cond_size, base_channels * channel_factors[0], 1))
self.blocks = nn.ModuleList([
IndependentStyleGAN2Block(in_factor * base_channels,
out_factor * base_channels,
cond_size=cond_size,
depth=block_depth,
activation=activation,
kernel_size=3,
add_noise=add_noise)
for idx, (in_factor, out_factor) in enumerate(
zip(channel_factors[:-1], channel_factors[1:]))
])
self.zeros = nn.ModuleList([
ReZero(3) for idx, (in_factor, out_factor) in enumerate(
zip(channel_factors[:-1], channel_factors[1:]))
])