Upload stae_pixel.py with huggingface_hub

stae_pixel.py

@@ -0,0 +1,314 @@
+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, cond=None):
+		x = self.initial(x)
+
+		if(cond is None):
+			for block in self.blocks:
+				x = block(x)
+		else:
+			cond = cond.chunk(len(self.blocks), dim=1)
+			for block, cond in zip(self.blocks, cond):
+				x = block(x) + cond
+
+		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 NerfEmbedder(nn.Module):
+	def __init__(self, in_channels, hidden_size_input, max_freqs):
+		super().__init__()
+		self.max_freqs = max_freqs
+		self.hidden_size_input = hidden_size_input
+		self.embedder = nn.Sequential(
+			nn.Linear(in_channels+max_freqs**2, hidden_size_input, bias=True),
+		)
+		self.positions = nn.Parameter(torch.randn(1, 16**2, max_freqs**2))
+
+
+	def forward(self, inputs):
+		B, P2, C = inputs.shape
+		
+		dct = self.positions
+		dct = dct.repeat(B, 1, 1)
+		inputs = torch.cat([inputs, dct], dim=-1)
+		inputs = self.embedder(inputs)
+		return inputs
+
+
+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.last = nn.Linear(mlp_dim, self.out_channels)
+		self.x_embedder = NerfEmbedder(3, mlp_dim, 8)
+
+	def forward(self, x, x_orig, cond=None):
+		B, C, H, W = x.shape
+		x = self.conv_in(x)
+		if(cond is None):
+			for block in self.blocks:
+				x = block(x)
+		else:
+			cond = cond.chunk(len(self.blocks), dim=1)
+			for block, cond in zip(self.blocks, cond):
+				add, scale = cond.chunk(2, dim=1)
+				x = (block(x) + add) * (1 + scale)
+
+		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 = torch.nn.functional.unfold(x_orig, kernel_size=self.patch_size, stride=self.patch_size).transpose(1, 2)
+		x = x.reshape(total_len, 3, self.patch_size ** 2 )
+		x = x.transpose(1, 2)
+		x = self.x_embedder(x)
+
+		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 Upsampler(nn.Module):
+	def __init__(self,
+				hidden_dim,
+				nerf_blocks,
+				mlp_dim,
+				patch_size,
+				out_channels):
+		super().__init__()
+		
+		self.patch_size = patch_size
+		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, out_channels)
+
+	def forward(self, x):
+		B, C, H, W = x.shape
+
+		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
+
+def weights_init_zeros(m):
+	if hasattr(m, 'weight') and m.weight is not None:
+		nn.init.constant_(m.weight, 0)
+	if hasattr(m, 'bias') and m.bias is not None:
+		nn.init.constant_(m.bias, 0)
+
+class StupidAE(nn.Module):
+	def __init__(self):
+		super().__init__()
+		
+		self.real_encoder = nn.Sequential(
+			StupidEncoder(in_channels=3, out_channels=32, hidden_dim=512, patch_size=8, num_blocks=1),
+			StupidEncoder(in_channels=32, out_channels=256, hidden_dim=1024, patch_size=4, num_blocks=2),
+			StupidEncoder(in_channels=256, out_channels=1024, hidden_dim=1024, patch_size=2, num_blocks=2),
+			Upsampler(1024, 1, 128, 4, 16)
+		)
+
+		encoder_dim = 1024
+		num_encoder_blocks = 1
+		self.encoder_proj = nn.Sequential(
+			nn.Conv2d(16, 1024, kernel_size=3, stride=1, padding=1),
+			nn.GELU(),
+			nn.Conv2d(1024, 24 * 1024, kernel_size=1, stride=1)
+		)
+
+		self.encoder_proj[2].apply(weights_init_zeros)
+
+		self.encoder = nn.Sequential(
+			StupidEncoder(in_channels=3, out_channels=512, hidden_dim=512, patch_size=8, num_blocks=1),
+			StupidEncoder(in_channels=512, out_channels=1024, hidden_dim=encoder_dim, patch_size=2, num_blocks=num_encoder_blocks),
+		)
+
+		self.decoder = StupidDecoder(in_channels=1024, out_channels=3, hidden_dim=1024, patch_size=16, num_blocks=6, nerf_blocks=2, mlp_dim=96)
+
+		# self.encoder.requires_grad_(False)
+		# self.decoder.requires_grad_(False)
+		
+		# self.real_encoder.requires_grad_(False)
+
+	@torch.compile(mode="default")
+	def encode(self, x):
+		return self.real_encoder(x)
+
+	@torch.compile(mode="default")
+	def forward(self, x, cond=None):
+		x_orig = x 
+
+		x = self.encoder(x)
+
+		if(cond is not None):
+			projected = self.encoder_proj(cond)
+			x = self.decoder(x, x_orig, projected)
+		else:
+			x = self.decoder(x, x_orig)
+
+
+		return x