95 lines
3.1 KiB
Python
95 lines
3.1 KiB
Python
from tinygrad import Tensor, nn
|
|
|
|
class gen:
|
|
def __init__(self, input_channels=1, height=128, width=216, latent_dim=1024):
|
|
self.height = height
|
|
self.width = width
|
|
self.latent_dim = latent_dim
|
|
|
|
self.w = width // 8
|
|
self.h = height // 8
|
|
self.flattened_size = 256 * self.h * self.w
|
|
|
|
self.num_tokens = 16
|
|
self.dim_per_token = self.latent_dim // self.num_tokens
|
|
|
|
|
|
self.e1 = nn.Conv2d(input_channels, 64, kernel_size=3, stride=2, padding=1)
|
|
self.e2 = nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1)
|
|
self.e3 = nn.Conv2d(128,256, kernel_size=3,stride=2,padding=1)
|
|
|
|
|
|
self.el = nn.Linear(self.flattened_size, self.latent_dim)
|
|
|
|
self.q = nn.Linear(self.dim_per_token,self.dim_per_token)
|
|
self.k = nn.Linear(self.dim_per_token,self.dim_per_token)
|
|
self.v = nn.Linear(self.dim_per_token,self.dim_per_token)
|
|
self.norm1 = nn.LayerNorm(self.dim_per_token)
|
|
|
|
ffn_dim = self.dim_per_token * 4
|
|
self.ffn1 = nn.Linear(self.dim_per_token, ffn_dim)
|
|
self.ffn2 = nn.Linear(ffn_dim, self.dim_per_token)
|
|
self.norm2 = nn.LayerNorm(self.dim_per_token)
|
|
|
|
self.dl = nn.Linear(self.latent_dim, self.flattened_size)
|
|
|
|
self.d1 = nn.ConvTranspose2d(256,128,kernel_size=3,stride=2,padding=1,output_padding=1)
|
|
self.d2 = nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1, output_padding=1)
|
|
self.d3 = nn.ConvTranspose2d(64, input_channels, kernel_size=3, stride=2, padding=1, output_padding=1)
|
|
|
|
def __call__(self, x: Tensor) -> Tensor:
|
|
y, shape = self.encode(x)
|
|
z = self.atten(y)
|
|
return self.decode(z, shape)
|
|
|
|
def encode(self, x: Tensor):
|
|
x = self.e1(x).leakyrelu()
|
|
x = self.e2(x).leakyrelu()
|
|
x = self.e3(x).leakyrelu()
|
|
b, c, h, w = x.shape
|
|
|
|
flattened_size = c * h * w
|
|
x = x.reshape(shape=(b, flattened_size))
|
|
z = self.el(x)
|
|
|
|
# reshape to multi-token: (batch, num_tokens, dim_per_token)
|
|
z = z.reshape(shape=(b, self.num_tokens, self.dim_per_token))
|
|
return z, (c, h, w)
|
|
|
|
def atten(self, x: Tensor):
|
|
q = self.q(x)
|
|
k = self.k(x)
|
|
v = self.v(x)
|
|
attn = q.scaled_dot_product_attention(k, v)
|
|
x = self.norm1(x+attn)
|
|
|
|
ffn = self.ffn1(x).relu()
|
|
ffn = self.ffn2(ffn)
|
|
x = self.norm2(x+ffn)
|
|
|
|
return x
|
|
|
|
def decode(self, z: Tensor, shape):
|
|
z = z.reshape(shape=(z.shape[0], -1))
|
|
x = self.dl(z).leakyrelu()
|
|
x = x.reshape(shape=(-1, 256, self.h, self.w))
|
|
x = self.d1(x).leakyrelu()
|
|
x = self.d2(x).leakyrelu()
|
|
x = self.d3(x).sigmoid()
|
|
|
|
# Crop or pad to match input size
|
|
out_h, out_w = x.shape[2], x.shape[3]
|
|
if out_h > self.height:
|
|
x = x[:, :, :self.height, :]
|
|
elif out_h < self.height:
|
|
pad_h = self.height - out_h
|
|
x = x.pad2d((0, 0, 0, pad_h))
|
|
|
|
if out_w > self.width:
|
|
x = x[:, :, :, :self.width]
|
|
elif out_w < self.width:
|
|
pad_w = self.width - out_w
|
|
x = x.pad2d((0, pad_w, 0, 0))
|
|
|
|
return x
|