Add V2 model and improved training code

upscaler.py

@@ -3,23 +3,40 @@ import torch.nn as nn
 import numpy as np
 
 class LatentUpscaler(nn.Module):
-	def __init__(self, fac):
-		super().__init__()
-
-		module_list = [
-			nn.Conv2d(4, 64, kernel_size=5, padding=2),
+	def head(self):
+		return [
+			nn.Conv2d(self.chan, self.size, kernel_size=self.krn, padding=self.pad),
 			nn.ReLU(),
-			nn.Upsample(scale_factor=fac, mode="nearest"), # bicubic was blurry
+			nn.Upsample(scale_factor=self.fac, mode="nearest"),
 			nn.ReLU(),
-			nn.Conv2d(64, 64, kernel_size=7, padding=3),
-			nn.ReLU(),
-			nn.Conv2d(64, 64, kernel_size=7, padding=3),
-			nn.ReLU(),
-			nn.Conv2d(64, 32, kernel_size=7, padding=3),
-			nn.ReLU(),
-			nn.Conv2d(32, 4, kernel_size=5, padding=2),
 		]
-		self.sequential = nn.Sequential(*module_list)
+	def core(self):
+		layers = []
+		for _ in range(self.depth):
+			layers += [
+				nn.Conv2d(self.size, self.size, kernel_size=self.krn, padding=self.pad),
+				nn.ReLU(),
+			]
+		return layers
+	def tail(self):
+		return [
+			nn.Conv2d(self.size, self.chan, kernel_size=self.krn, padding=self.pad),
+		]
+
+	def __init__(self, fac, depth=16):
+		super().__init__()
+		self.size = 64      # Conv2d size
+		self.chan = 4       # in/out channels
+		self.depth = depth  # no. of layers
+		self.fac = fac      # scale factor
+		self.krn = 3        # kernel size
+		self.pad = 1        # padding
+
+		self.sequential = nn.Sequential(
+			*self.head(),
+			*self.core(),
+			*self.tail(),
+		)
 
 	def forward(self, x: torch.Tensor) -> torch.Tensor:
 		return self.sequential(x)