-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
8131a33
commit 2ea67cc
Showing
13 changed files
with
549 additions
and
5 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,9 @@ | ||
"""It holds some constants""" | ||
|
||
WIDTH = 512 | ||
HEIGHT = 512 | ||
LATENTS_WIDTH = WIDTH // 8 | ||
LATENTS_HEIGHT = HEIGHT // 8 | ||
|
||
ALLOW_CUDA = False | ||
ALLOW_MPS = False |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,77 @@ | ||
"""Attention mechanisms""" | ||
|
||
import math | ||
import torch | ||
from torch import nn | ||
from torch.nn import functional as F | ||
|
||
|
||
class SelfAttention(nn.Module): | ||
"""Self Attention""" | ||
|
||
def __init__(self, n_heads: int, | ||
d_embed: int, | ||
in_proj_bias: bool = True, | ||
out_proj_bias: bool = True | ||
): | ||
super().__init__() | ||
# This combines the Wq, Wk and Wv matrices into one matrix | ||
self.in_proj = nn.Linear(d_embed, 3 * d_embed, bias=in_proj_bias) | ||
# This one represents the Wo matrix | ||
self.out_proj = nn.Linear(d_embed, d_embed, bias=out_proj_bias) | ||
self.n_heads = n_heads | ||
self.d_head = d_embed // n_heads | ||
|
||
def forward(self, x: torch.Tensor, | ||
causal_mask: bool = False) -> torch.Tensor: | ||
"""Foward method""" | ||
# (Batch_Size, Seq_Len, Dim) | ||
input_shape = x.shape | ||
|
||
# (Batch_Size, Seq_Len, Dim) | ||
batch_size, sequence_length, _ = input_shape | ||
|
||
# (Batch_Size, Seq_Len, H, Dim / H) | ||
inter_shape = (batch_size, sequence_length, self.n_heads, self.d_head) | ||
|
||
# (Batch_Size, Seq_Len, Dim) -> (Batch_Size, Seq_Len, Dim * 3) | ||
qkv: torch.Tensor = self.in_proj(x) | ||
|
||
# -> 3 tensor of shape (Batch_Size, Seq_Len, Dim) | ||
q, k, v = qkv.chunk(3, dim=-1) | ||
|
||
# (Batch_Size, Seq_Len, Dim) -> | ||
# (Batch_Size, Seq_Len, H, Dim / H) -> | ||
# (Batch_Size, H, Seq_Len, Dim / H) | ||
q = q.view(inter_shape).transpose(1, 2) | ||
k = k.view(inter_shape).transpose(1, 2) | ||
v = v.view(inter_shape).transpose(1, 2) | ||
|
||
# (Batch_Size, H, Seq_Len, Dim / H) @ | ||
# (Batch_Size, H, Dim / H, Seq_Len) -> (Batch_Size, H, Seq_Len, Seq_Len) | ||
weight = q @ k.transpose(-1, -2) / math.sqrt(self.d_head) | ||
|
||
if causal_mask: | ||
# Mask where the upper triangle (above the principal diagonal) is 1 | ||
mask = torch.ones_like(weight, dtype=torch.bool).triu(1) | ||
# Fill the upper triangle with -inf | ||
weight.masked_fill_(mask, -torch.inf) | ||
|
||
# (Batch_Size, H, Seq_Len, Seq_Len) -> (Batch_Size, H, Seq_Len, Seq_Len) | ||
weight = F.softmax(weight, dim=-1) | ||
|
||
# (Batch_Size, H, Seq_Len, Seq_Len) @ | ||
# (Batch_Size, H, Seq_Len, Dim / H) -> (Batch_Size, H, Seq_Len, Dim / H) | ||
output = weight @ v | ||
|
||
# (Batch_Size, H, Seq_Len, Dim / H) -> (Batch_Size, Seq_Len, H, Dim / H) | ||
output = output.transpose(1, 2) | ||
|
||
# (Batch_Size, Seq_Len, H, Dim / H) -> (Batch_Size, Seq_Len, Dim) | ||
output = output.reshape(input_shape) | ||
|
||
# (Batch_Size, Seq_Len, Dim) -> (Batch_Size, Seq_Len, Dim) | ||
output = self.out_proj(output) | ||
|
||
# (Batch_Size, Seq_Len, Dim) | ||
return output |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,107 @@ | ||
"""VAE Decoder Pytorch Module""" | ||
|
||
import torch | ||
from torch import nn | ||
from src.models.vae_blocks import VAEResidualBlock, VAEAttentionBlock | ||
|
||
|
||
class VAEDecoder(nn.Module): | ||
"""VAE Decoder""" | ||
|
||
def __init__(self): | ||
super().__init__() | ||
self.layers = nn.ModuleList( | ||
[ | ||
# (Batch_Size, 4, Height / 8, Width / 8) -> (Batch_Size, 4, Height / 8, Width / 8) | ||
nn.Conv2d(4, 4, kernel_size=1, padding=0), | ||
|
||
# (Batch_Size, 4, Height / 8, Width / 8) -> (Batch_Size, 512, Height / 8, Width / 8) | ||
nn.Conv2d(4, 512, kernel_size=3, padding=1), | ||
|
||
# (Batch_Size, 512, Height / 8, Width / 8) -> (Batch_Size, 512, Height / 8, Width / 8) | ||
VAEResidualBlock(512, 512), | ||
|
||
# (Batch_Size, 512, Height / 8, Width / 8) -> (Batch_Size, 512, Height / 8, Width / 8) | ||
VAEAttentionBlock(512), | ||
|
||
# (Batch_Size, 512, Height / 8, Width / 8) -> (Batch_Size, 512, Height / 8, Width / 8) | ||
VAEResidualBlock(512, 512), | ||
|
||
# (Batch_Size, 512, Height / 8, Width / 8) -> (Batch_Size, 512, Height / 8, Width / 8) | ||
VAEResidualBlock(512, 512), | ||
|
||
# (Batch_Size, 512, Height / 8, Width / 8) -> (Batch_Size, 512, Height / 8, Width / 8) | ||
VAEResidualBlock(512, 512), | ||
|
||
# (Batch_Size, 512, Height / 8, Width / 8) -> (Batch_Size, 512, Height / 8, Width / 8) | ||
VAEResidualBlock(512, 512), | ||
|
||
# Repeats the rows and columns of the data by scale_factor. | ||
# (Batch_Size, 512, Height / 8, Width / 8) -> (Batch_Size, 512, Height / 4, Width / 4) | ||
nn.Upsample(scale_factor=2), | ||
|
||
# (Batch_Size, 512, Height / 4, Width / 4) -> (Batch_Size, 512, Height / 4, Width / 4) | ||
nn.Conv2d(512, 512, kernel_size=3, padding=1), | ||
|
||
# (Batch_Size, 512, Height / 4, Width / 4) -> (Batch_Size, 512, Height / 4, Width / 4) | ||
VAEResidualBlock(512, 512), | ||
|
||
# (Batch_Size, 512, Height / 4, Width / 4) -> (Batch_Size, 512, Height / 4, Width / 4) | ||
VAEResidualBlock(512, 512), | ||
|
||
# (Batch_Size, 512, Height / 4, Width / 4) -> (Batch_Size, 512, Height / 4, Width / 4) | ||
VAEResidualBlock(512, 512), | ||
|
||
# (Batch_Size, 512, Height / 4, Width / 4) -> (Batch_Size, 512, Height / 2, Width / 2) | ||
nn.Upsample(scale_factor=2), | ||
|
||
# (Batch_Size, 512, Height / 2, Width / 2) -> (Batch_Size, 512, Height / 2, Width / 2) | ||
nn.Conv2d(512, 512, kernel_size=3, padding=1), | ||
|
||
# (Batch_Size, 512, Height / 2, Width / 2) -> (Batch_Size, 256, Height / 2, Width / 2) | ||
VAEResidualBlock(512, 256), | ||
|
||
# (Batch_Size, 256, Height / 2, Width / 2) -> (Batch_Size, 256, Height / 2, Width / 2) | ||
VAEResidualBlock(256, 256), | ||
|
||
# (Batch_Size, 256, Height / 2, Width / 2) -> (Batch_Size, 256, Height / 2, Width / 2) | ||
VAEResidualBlock(256, 256), | ||
|
||
# (Batch_Size, 256, Height / 2, Width / 2) -> (Batch_Size, 256, Height, Width) | ||
nn.Upsample(scale_factor=2), | ||
|
||
# (Batch_Size, 256, Height, Width) -> (Batch_Size, 256, Height, Width) | ||
nn.Conv2d(256, 256, kernel_size=3, padding=1), | ||
|
||
# (Batch_Size, 256, Height, Width) -> (Batch_Size, 128, Height, Width) | ||
VAEResidualBlock(256, 128), | ||
|
||
# (Batch_Size, 128, Height, Width) -> (Batch_Size, 128, Height, Width) | ||
VAEResidualBlock(128, 128), | ||
|
||
# (Batch_Size, 128, Height, Width) -> (Batch_Size, 128, Height, Width) | ||
VAEResidualBlock(128, 128), | ||
|
||
# (Batch_Size, 128, Height, Width) -> (Batch_Size, 128, Height, Width) | ||
nn.GroupNorm(32, 128), | ||
|
||
# (Batch_Size, 128, Height, Width) -> (Batch_Size, 128, Height, Width) | ||
nn.SiLU(), | ||
|
||
# (Batch_Size, 128, Height, Width) -> (Batch_Size, 3, Height, Width) | ||
nn.Conv2d(128, 3, kernel_size=3, padding=1), | ||
] | ||
) | ||
|
||
def forward(self, x: torch.Tensor) -> torch.Tensor: | ||
"""Forward method""" | ||
# x: (Batch_Size, 4, Height / 8, Width / 8) | ||
# Remove the scaling added by the Encoder. | ||
x = x / 0.18215 | ||
|
||
for layer in self.layers: | ||
|
||
x = layer(x) | ||
|
||
# (Batch_Size, 3, Height, Width) | ||
return x |
Oops, something went wrong.