Spaces:
Paused
Paused
File size: 1,548 Bytes
ef6c3c2 |
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 |
import torch
import torch.nn as nn
import torch.nn.functional as F
from .decoder import VAE_AttentionBlock, VAE_ResidualBlock
class VAE_Encoder(nn.Sequential):
def __init__(self):
super().__init__(
nn.Conv2d(3, 128, kernel_size=3, padding=1),
VAE_ResidualBlock(128, 128),
VAE_ResidualBlock(128, 128),
nn.Conv2d(128, 128, kernel_size=3, stride=2, padding=0),
VAE_ResidualBlock(128, 256),
VAE_ResidualBlock(256, 256),
nn.Conv2d(256, 256, kernel_size=3, stride=2, padding=0),
VAE_ResidualBlock(256, 512),
VAE_ResidualBlock(512, 512),
nn.Conv2d(512, 512, kernel_size=3, stride=2, padding=0),
VAE_ResidualBlock(512, 512),
VAE_ResidualBlock(512, 512),
VAE_ResidualBlock(512, 512),
VAE_AttentionBlock(512),
VAE_ResidualBlock(512, 512),
nn.GroupNorm(32, 512),
nn.SiLU(),
nn.Conv2d(512, 8, kernel_size=3, padding=1),
nn.Conv2d(8, 8, kernel_size=1, padding=0),
)
def forward(self, x, noise):
for module in self:
if getattr(module, 'stride', None) == (2, 2):
x = F.pad(x, (0, 1, 0, 1))
x = module(x)
mean, log_variance = torch.chunk(x, 2, dim=1)
log_variance = torch.clamp(log_variance, -30, 20)
variance = log_variance.exp()
stdev = variance.sqrt()
x = mean + stdev * noise
x *= 0.18215
return x
|