Upload stae.py with huggingface_hub

stae.py

@@ -0,0 +1,266 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class RMSNorm2d(nn.Module):
+	def __init__(self, channels, eps=1e-8, affine=True):
+		super().__init__()
+		self.eps = eps
+		self.affine = affine
+		if affine:
+			self.weight = nn.Parameter(torch.ones(channels))
+		else:
+			self.register_parameter("weight", None)
+
+	def forward(self, x):
+		norm = x.pow(2).mean(dim=1, keepdim=True).add(self.eps).rsqrt()
+		x = x * norm
+		if self.affine:
+			x = x * self.weight[:, None, None]
+		return x
+
+class ConvMlp(nn.Module):
+	def __init__(self, in_features, hidden_features=None, out_features=None):
+		super().__init__()
+		self.model = nn.Sequential(
+			nn.Conv2d(in_channels=in_features, out_channels=hidden_features, kernel_size=1),
+			nn.GELU(), 
+			nn.Conv2d(in_channels=hidden_features, out_channels=out_features, kernel_size=1),
+		)
+
+	def forward(self, x):
+		return self.model(x)
+
+import torch
+import torch.nn as nn
+class GegluMlp(nn.Module):
+	def __init__(self, hidden_dim, out_dim=None):
+		super().__init__()
+
+		if(out_dim is None):
+			out_dim = hidden_dim
+		self.conv_up = nn.Conv2d(hidden_dim, hidden_dim * 4, kernel_size=1)
+		self.conv_down = nn.Conv2d(hidden_dim * 2, out_dim, kernel_size=1)
+		self.activation = nn.GELU(approximate="tanh")
+
+	def forward(self, x):
+		x = self.conv_up(x)
+		x_gate, x_act = torch.chunk(x, 2, dim=1)
+		x = self.activation(x_act) * x_gate
+		x = self.conv_down(x)
+		
+		return x
+
+class EncoderBlock(nn.Module):
+	def __init__(self, channels):
+		super().__init__()
+		self.norm = RMSNorm2d(channels)
+		hidden_dim = channels
+
+		self.mlp = GegluMlp(hidden_dim)
+	   
+	def forward(self, x):
+		norm = self.norm(x)
+		mlp_out = self.mlp(norm)
+		x = x + mlp_out
+
+		return x
+
+class DecoderBlock(nn.Module):
+	def __init__(self, channels):
+		super().__init__()
+		self.norm = RMSNorm2d(channels)
+
+		self.mlp = nn.Sequential(
+			nn.Conv2d(channels, channels, kernel_size=1),
+			nn.GELU(approximate="tanh"),
+			nn.Conv2d(channels, channels, kernel_size=3, padding=1),
+		)
+		
+	def forward(self, x):
+		norm = self.norm(x)
+		mlp_out = self.mlp(norm)
+		x = x + mlp_out
+
+		return x
+
+class StupidEncoder(nn.Module):
+	def __init__(self,
+				 hidden_dim,
+				 in_channels,
+				 out_channels,
+				 patch_size,
+				 num_blocks):
+		super().__init__()
+
+		self.initial = nn.Sequential(
+			nn.Conv2d(in_channels, hidden_dim, patch_size, padding=0, stride=patch_size),
+		)
+
+		self.blocks = nn.ModuleList(EncoderBlock(hidden_dim) for _ in range(num_blocks))
+		self.out = ConvMlp(hidden_dim, hidden_dim, out_channels)
+
+	def forward(self, x):
+		x = self.initial(x)
+
+		for block in self.blocks:
+			x = block(x)
+
+		x = self.out(x)
+		return x
+
+class NerfHead(nn.Module):
+	def __init__(self, patch_dim, mlp_dim):
+		super().__init__()
+		self.mlp_dim = mlp_dim
+		self.param_gen = nn.Linear(patch_dim, self.mlp_dim*self.mlp_dim*2)
+		self.norm = nn.RMSNorm(self.mlp_dim)
+
+	def forward(self, pixels, patches):
+		bs = pixels.shape[0]
+		params = self.param_gen(patches)
+		layer1, layer2 = params.chunk(2, dim=-1)
+		layer1 = layer1.view(bs, self.mlp_dim, self.mlp_dim)
+		layer2 = layer2.view(bs, self.mlp_dim, self.mlp_dim)
+
+		layer1 = torch.nn.functional.normalize(layer1, dim=-2)
+
+		res_x = pixels
+		pixels = self.norm(pixels)
+		pixels = torch.bmm(pixels, layer1)
+		pixels = torch.nn.functional.silu(pixels)
+		pixels = torch.bmm(pixels, layer2)
+		pixels = pixels + res_x
+		return pixels
+
+class StupidDecoder(nn.Module):
+	def __init__(self,
+				 hidden_dim,
+				 in_channels,
+				 out_channels,
+				 patch_size,
+				 num_blocks,
+				 nerf_blocks,
+				 mlp_dim):
+		super().__init__()
+		
+		self.out_channels = out_channels
+
+		self.patch_size = patch_size
+		self.conv_in = ConvMlp(in_channels, hidden_dim, hidden_dim)
+		self.blocks = []
+		for _ in range(num_blocks):
+			self.blocks.append(DecoderBlock(hidden_dim))
+			self.blocks.append(EncoderBlock(hidden_dim))
+		self.blocks = nn.ModuleList(self.blocks)
+
+		self.nerf = nn.ModuleList(NerfHead(hidden_dim, mlp_dim) for _ in range(nerf_blocks))
+		self.positions = nn.Parameter(torch.randn(1, self.patch_size**2, mlp_dim))
+		self.last = nn.Linear(mlp_dim, self.out_channels)
+
+	def forward(self, x):
+		B, C, H, W = x.shape
+		x = self.conv_in(x)
+		for block in self.blocks:
+			x = block(x)
+
+		patches = x.flatten(2).transpose(1,2) # B C H W -> B (HW) C 
+		patch_count = H*W
+		total_len = x.shape[0] * patch_count
+		patches = patches.reshape(total_len, -1)
+		x = self.positions.repeat(total_len, 1, 1)
+
+		for block in self.nerf:
+			x = block(x, patches) # B * patch_count, ps*ps, C
+		x = self.last(x)
+		x = x.transpose(1,2) # [B * patch_count, ps*ps, C] -> [B*patch_count, C, ps*ps]
+		x = x.reshape(B, patch_count, -1) # [B*patch_count, C, ps*ps] -> [B, patch_count, ps*ps*3]
+		x = x.transpose(1,2) # [B, patch_count, ps*ps*3] -> [B, ps*ps*3, patch_count]
+		x = torch.nn.functional.fold(x.contiguous(),
+									 (H*self.patch_size, W*self.patch_size),
+									 kernel_size=self.patch_size,
+									 stride=self.patch_size)
+
+		return x
+
+class SimpleStupidDecoder(nn.Module):
+	def __init__(self,
+				 hidden_dim,
+				 in_channels,
+				 out_channels,
+				 patch_size,
+				 num_blocks):
+		super().__init__()
+		
+		self.out_channels = out_channels
+		self.patch_size = patch_size
+
+		self.conv_in = ConvMlp(in_channels, hidden_dim, hidden_dim)
+		self.blocks = nn.ModuleList(DecoderBlock(hidden_dim) for _ in range(num_blocks))
+
+		self.last = nn.Sequential(
+			ConvMlp(hidden_dim, hidden_dim, out_channels * patch_size * patch_size),
+			nn.PixelShuffle(patch_size)
+		)
+
+	def forward(self, x):
+		x = self.conv_in(x)
+		for block in self.blocks:
+			x = block(x)
+
+		return self.last(x)
+
+class StupidAE(nn.Module):
+	def __init__(self):
+		super().__init__()
+		
+		self.encoder = nn.Sequential(
+			StupidEncoder(in_channels=3, out_channels=32, hidden_dim=512, patch_size=8, num_blocks=2),
+			StupidEncoder(in_channels=32, out_channels=256, hidden_dim=1024, patch_size=4, num_blocks=2),
+		)
+
+		self.decoder = nn.Sequential(
+			StupidDecoder(in_channels=256, out_channels=32, hidden_dim=1024, patch_size=8, num_blocks=2, nerf_blocks=1, mlp_dim=128),
+			StupidDecoder(in_channels=32, out_channels=3, hidden_dim=512, patch_size=4, num_blocks=2, nerf_blocks=1, mlp_dim=32)
+		)
+
+		self.semantic_decoder = GegluMlp(256, 768)
+	
+	@torch.compile(mode="default")
+	def encode(self, x):
+		return self.encoder(x)
+	
+	@torch.compile(mode="default")
+	def decode(self, x):
+		return self.decoder(x)
+	
+	def decode_from_tokens(self, tokens, H, W):
+		tokens = tokens * 1.28
+		results = []
+		downsample_factor = 32 
+		batch_size = tokens.shape[0]
+
+		for i in range(batch_size):
+			h = int(H[i])
+			w = int(W[i])
+
+			h_lat = h // downsample_factor
+			w_lat = w // downsample_factor
+			num_tokens = h_lat * w_lat
+
+			# Достаем токены для текущей картинки: [Num_Tokens, C]
+			t = tokens[i, :num_tokens] 
+
+			# Решейп в формат сверток [1, C, H_lat, W_lat]
+			t = t.transpose(0, 1).view(1, -1, h_lat, w_lat)
+
+			# Декодируем
+			img = self.decoder(t).squeeze(0) * 0.5 + 0.5
+			results.append(img)
+
+		return results
+	
+	def forward(self, x):
+		x = self.encode(x)
+		x = self.decode(x)
+		return x