updates

src/common/layers.py

@@ -361,7 +361,7 @@ class ChebyKANLayer(nn.Module):
     """
     https://github.com/SynodicMonth/ChebyKAN/blob/main/ChebyKANLayer.py
     """
-    def __init__(self, in_features, out_features, degree=5):
+    def __init__(self, in_features, out_features, degree=8):
         super(ChebyKANLayer, self).__init__()
         self.inputdim = in_features
         self.outdim = out_features
@@ -396,7 +396,7 @@ class ChebyKANLayer(nn.Module):
 
 
 class UpscaleBlockChebyKAN(nn.Module):
-    def __init__(self, in_features=4, hidden_dim = 32, layers_count=4, upscale_factor=1, input_max_value=255, output_max_value=255):
+    def __init__(self, in_features=4, hidden_dim = 32, layers_count=4, upscale_factor=1, input_max_value=255, output_max_value=255, degree=8):
         super(UpscaleBlockChebyKAN, self).__init__()   
         assert layers_count > 0     
         self.upscale_factor = upscale_factor 
@@ -405,7 +405,7 @@ class UpscaleBlockChebyKAN(nn.Module):
         
         self.linear_projections = []
         for i in range(layers_count):
-            self.linear_projections.append(ChebyKANLayer(in_features=hidden_dim, out_features=hidden_dim))
+            self.linear_projections.append(ChebyKANLayer(in_features=hidden_dim, out_features=hidden_dim, degree=degree))
         self.linear_projections = nn.ModuleList(self.linear_projections)          
 
         self.project_channels = nn.Linear(in_features=hidden_dim, out_features=upscale_factor * upscale_factor, bias=True)

src/models/hdbnet.py

@@ -11,10 +11,11 @@ from itertools import cycle
 from models.base import SRNetBase
 
 class HDBNet(SRNetBase):
-    def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
+    def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4, rotations = 4):
         super(HDBNet, self).__init__()
         assert scale == 4
         self.scale = scale 
+        self.rotations = rotations
         self.stage1_3H = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=255, output_max_value=255)
         self.stage1_3D = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=255, output_max_value=255)
         self.stage1_3B = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=255, output_max_value=255)
@@ -31,9 +32,7 @@ class HDBNet(SRNetBase):
         self._extract_pattern_3D = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,1],[2,2]], center=[0,0], window_size=3)
         self._extract_pattern_3B = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,2],[2,1]], center=[0,0], window_size=3)
         self._extract_pattern_2H = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1]], center=[0,0], window_size=2)
-        self._extract_pattern_2D = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,1]], center=[0,0], window_size=2)
-
-        self.rotations = 4
+        self._extract_pattern_2D = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,1]], center=[0,0], window_size=2)        
 
     def forward_stage(self, x, scale, percieve_pattern, stage):
         b,c,h,w = x.shape
@@ -113,6 +112,102 @@ class HDBNet(SRNetBase):
             return F.mse_loss(pred/255, target/255)
         return loss_fn
 
+class HDBNetv2(SRNetBase):
+    def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4, rotations = 4):
+        super(HDBNetv2, self).__init__()
+        assert scale == 4
+        self.scale = scale 
+        self.rotations = rotations
+        self.stage1_3H = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=255, output_max_value=255)
+        self.stage1_3D = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=255, output_max_value=255)
+        self.stage1_3B = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=255, output_max_value=255)
+        self.stage1_2H = layers.UpscaleBlock(in_features=2, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=15, output_max_value=255)
+        self.stage1_2D = layers.UpscaleBlock(in_features=2, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=15, output_max_value=255)
+
+        self.stage2_3H = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=255, output_max_value=255)
+        self.stage2_3D = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=255, output_max_value=255)
+        self.stage2_3B = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=255, output_max_value=255)
+        self.stage2_2H = layers.UpscaleBlock(in_features=2, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=15, output_max_value=255)
+        self.stage2_2D = layers.UpscaleBlock(in_features=2, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=15, output_max_value=255)
+
+        self._extract_pattern_3H = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[0,2]], center=[0,0], window_size=3)
+        self._extract_pattern_3D = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,1],[2,2]], center=[0,0], window_size=3)
+        self._extract_pattern_3B = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,2],[2,1]], center=[0,0], window_size=3)
+        self._extract_pattern_2H = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1]], center=[0,0], window_size=2)
+        self._extract_pattern_2D = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,1]], center=[0,0], window_size=2)        
+
+    def forward(self, x, config=None):
+        b,c,h,w = x.shape
+        x = x.reshape(b*c, 1, h, w)
+        lsb = x % 16
+        msb = x - lsb     
+        output_msb = torch.zeros([b*c, 1, h*2, w*2], dtype=x.dtype, device=x.device)
+        output_lsb = torch.zeros([b*c, 1, h*2, w*2], dtype=x.dtype, device=x.device)
+        for rotations_count in range(self.rotations):
+           rotated_msb = torch.rot90(msb, k=rotations_count, dims=[2, 3])
+           rotated_lsb = torch.rot90(lsb, k=rotations_count, dims=[2, 3])
+           output_msbt = self.forward_stage(rotated_msb, 2, self._extract_pattern_3H, self.stage1_3H) + \
+                         self.forward_stage(rotated_msb, 2, self._extract_pattern_3D, self.stage1_3D) + \
+                         self.forward_stage(rotated_msb, 2, self._extract_pattern_3B, self.stage1_3B)
+           output_lsbt = self.forward_stage(rotated_lsb, 2, self._extract_pattern_2H, self.stage1_2H) + \
+                         self.forward_stage(rotated_lsb, 2, self._extract_pattern_2D, self.stage1_2D)
+           if not config is None and config.current_iter % config.display_step == 0:
+                config.writer.add_histogram('s1_output_lsb', output_lsb.detach().cpu().numpy(), config.current_iter)
+                config.writer.add_histogram('s1_output_msb', output_msb.detach().cpu().numpy(), config.current_iter) 
+           output_msb += torch.rot90(output_msbt, k=-rotations_count, dims=[2, 3])
+           output_lsb += torch.rot90(output_lsbt, k=-rotations_count, dims=[2, 3])
+        output_msb /= self.rotations*3
+        output_lsb /= self.rotations*2
+        output = output_msb + output_lsb
+        x = output.clamp(0, 255)
+        lsb = x % 16
+        msb = x - lsb
+        output_msb = torch.zeros([b*c, 1, h*2*2, w*2*2], dtype=x.dtype, device=x.device)
+        output_lsb = torch.zeros([b*c, 1, h*2*2, w*2*2], dtype=x.dtype, device=x.device)
+        for rotations_count in range(self.rotations):
+           rotated_msb = torch.rot90(msb, k=rotations_count, dims=[2, 3])
+           rotated_lsb = torch.rot90(lsb, k=rotations_count, dims=[2, 3])
+           output_msbt = self.forward_stage(rotated_msb, 2, self._extract_pattern_3H, self.stage2_3H) + \
+                         self.forward_stage(rotated_msb, 2, self._extract_pattern_3D, self.stage2_3D) + \
+                         self.forward_stage(rotated_msb, 2, self._extract_pattern_3B, self.stage2_3B)
+           output_lsbt = self.forward_stage(rotated_lsb, 2, self._extract_pattern_2H, self.stage2_2H) + \
+                         self.forward_stage(rotated_lsb, 2, self._extract_pattern_2D, self.stage2_2D)
+           if not config is None and config.current_iter % config.display_step == 0:
+                config.writer.add_histogram('s2_output_lsb', output_lsb.detach().cpu().numpy(), config.current_iter)
+                config.writer.add_histogram('s2_output_msb', output_msb.detach().cpu().numpy(), config.current_iter) 
+           output_msb += torch.rot90(output_msbt, k=-rotations_count, dims=[2, 3])
+           output_lsb += torch.rot90(output_lsbt, k=-rotations_count, dims=[2, 3])
+        output_msb /= self.rotations*3
+        output_lsb /= self.rotations*2
+        output = output_msb + output_lsb
+        x = output.clamp(0, 255)
+        x = x.reshape(b, c, h*self.scale, w*self.scale)
+        return x
+
+    def get_lut_model(self, quantization_interval=16, batch_size=2**10):
+        stage1_3H = lut.transfer_3_input_SxS_output(self.stage1_3H, quantization_interval=quantization_interval, batch_size=batch_size)
+        stage1_3D = lut.transfer_3_input_SxS_output(self.stage1_3D, quantization_interval=quantization_interval, batch_size=batch_size)
+        stage1_3B = lut.transfer_3_input_SxS_output(self.stage1_3B, quantization_interval=quantization_interval, batch_size=batch_size)
+        stage1_2H = lut.transfer_2_input_SxS_output(self.stage1_2H, quantization_interval=quantization_interval, batch_size=batch_size)
+        stage1_2D = lut.transfer_2_input_SxS_output(self.stage1_2D, quantization_interval=quantization_interval, batch_size=batch_size)
+
+        stage2_3H = lut.transfer_3_input_SxS_output(self.stage2_3H, quantization_interval=quantization_interval, batch_size=batch_size)
+        stage2_3D = lut.transfer_3_input_SxS_output(self.stage2_3D, quantization_interval=quantization_interval, batch_size=batch_size)
+        stage2_3B = lut.transfer_3_input_SxS_output(self.stage2_3B, quantization_interval=quantization_interval, batch_size=batch_size)
+        stage2_2H = lut.transfer_2_input_SxS_output(self.stage2_2H, quantization_interval=quantization_interval, batch_size=batch_size)
+        stage2_2D = lut.transfer_2_input_SxS_output(self.stage2_2D, quantization_interval=quantization_interval, batch_size=batch_size)
+
+        lut_model = hdblut.HDBLut.init_from_numpy(
+            stage1_3H, stage1_3D, stage1_3B, stage1_2H, stage1_2D, 
+            stage2_3H, stage2_3D, stage2_3B, stage2_2H, stage2_2D
+        )
+        return lut_model
+
+    def get_loss_fn(self):
+        def loss_fn(pred, target):
+            return F.mse_loss(pred/255, target/255)
+        return loss_fn
+
 
 class HDBLNet(nn.Module):
     def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):

src/models/sdynet.py

@@ -7,9 +7,9 @@ from common import lut
 from common import layers
 from pathlib import Path
 from . import sdylut
+from models.base import SRNetBase
 
-
-class SDYNetx1(nn.Module):
+class SDYNetx1(SRNetBase):
     def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
         super(SDYNetx1, self).__init__()
         self.scale = scale   
@@ -20,16 +20,6 @@ class SDYNetx1(nn.Module):
         self.stage1_D = layers.UpscaleBlock(hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=scale)
         self.stage1_Y = layers.UpscaleBlock(hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=scale)
 
-    def forward_stage(self, x, scale, percieve_pattern, stage):
-        b,c,h,w = x.shape
-        x = percieve_pattern(x)   
-        x = stage(x)
-        x = round_func(x)
-        x = x.reshape(b, c, h, w, scale, scale)
-        x = x.permute(0,1,2,4,3,5)
-        x = x.reshape(b, c, h*scale, w*scale)
-        return x
-
     def forward(self, x, config=None):
         b,c,h,w = x.shape
         x = x.reshape(b*c, 1, h, w)
@@ -39,7 +29,6 @@ class SDYNetx1(nn.Module):
         output += self.forward_stage(x, self.scale, self._extract_pattern_Y, self.stage1_Y)
         output /= 3
         x = output
-        x = round_func(x)
         x = x.reshape(b, c, h*self.scale, w*self.scale)
         return x
 
@@ -56,7 +45,7 @@ class SDYNetx1(nn.Module):
         return loss_fn
         
 
-class SDYNetx2(nn.Module):
+class SDYNetx2(SRNetBase):
     def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
         super(SDYNetx2, self).__init__()
         self.scale = scale   
@@ -70,16 +59,6 @@ class SDYNetx2(nn.Module):
         self.stage2_D = layers.UpscaleBlock(hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=scale)
         self.stage2_Y = layers.UpscaleBlock(hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=scale)
 
-    def forward_stage(self, x, scale, percieve_pattern, stage):
-        b,c,h,w = x.shape
-        x = percieve_pattern(x)   
-        x = stage(x)
-        x = round_func(x)
-        x = x.reshape(b, c, h, w, scale, scale)
-        x = x.permute(0,1,2,4,3,5)
-        x = x.reshape(b, c, h*scale, w*scale)
-        return x
-
     def forward(self, x, config=None):
         b,c,h,w = x.shape
         x = x.reshape(b*c, 1, h, w)
@@ -89,14 +68,12 @@ class SDYNetx2(nn.Module):
         output += self.forward_stage(x, 1, self._extract_pattern_Y, self.stage1_Y)
         output /= 3
         x = output
-        x = round_func(x)
         output = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device)
         output += self.forward_stage(x, self.scale, self._extract_pattern_S, self.stage2_S)
         output += self.forward_stage(x, self.scale, self._extract_pattern_D, self.stage2_D)
         output += self.forward_stage(x, self.scale, self._extract_pattern_Y, self.stage2_Y)
         output /= 3
         x = output
-        x = round_func(x)
         x = x.reshape(b, c, h*self.scale, w*self.scale)
         return x
 
@@ -115,7 +92,7 @@ class SDYNetx2(nn.Module):
             return F.mse_loss(pred/255, target/255)
         return loss_fn
 
-class SDYNetx3(nn.Module):
+class SDYNetx3(SRNetBase):
     def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
         super(SDYNetx3, self).__init__()
         self.scale = scale   
@@ -132,16 +109,6 @@ class SDYNetx3(nn.Module):
         self.stage3_D = layers.UpscaleBlock(hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=scale)
         self.stage3_Y = layers.UpscaleBlock(hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=scale)
 
-    def forward_stage(self, x, scale, percieve_pattern, stage):
-        b,c,h,w = x.shape
-        x = percieve_pattern(x)   
-        x = stage(x)
-        x = round_func(x)
-        x = x.reshape(b, c, h, w, scale, scale)
-        x = x.permute(0,1,2,4,3,5)
-        x = x.reshape(b, c, h*scale, w*scale)
-        return x
-
     def forward(self, x, config=None):
         b,c,h,w = x.shape
         x = x.reshape(b*c, 1, h, w)
@@ -151,21 +118,18 @@ class SDYNetx3(nn.Module):
         output += self.forward_stage(x, 1, self._extract_pattern_Y, self.stage1_Y)
         output /= 3
         x = output
-        x = round_func(x)
         output = torch.zeros([b*c, 1, h, w], dtype=x.dtype, device=x.device)
         output += self.forward_stage(x, 1, self._extract_pattern_S, self.stage2_S)
         output += self.forward_stage(x, 1, self._extract_pattern_D, self.stage2_D)
         output += self.forward_stage(x, 1, self._extract_pattern_Y, self.stage2_Y)
         output /= 3
         x = output
-        x = round_func(x)
         output = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device)
         output += self.forward_stage(x, self.scale, self._extract_pattern_S, self.stage3_S)
         output += self.forward_stage(x, self.scale, self._extract_pattern_D, self.stage3_D)
         output += self.forward_stage(x, self.scale, self._extract_pattern_Y, self.stage3_Y)
         output /= 3
         x = output
-        x = round_func(x)
         x = x.reshape(b, c, h*self.scale, w*self.scale)
         return x
 
@@ -222,7 +186,6 @@ class SDYNetR90x1(nn.Module):
             output += torch.rot90(self.forward_stage(rotated, self.scale, self._extract_pattern_Y, self.stage1_Y), k=-rotations_count, dims=[-2, -1])
         output /= 4*3
         x = output
-        x = round_func(x)
         x = x.reshape(b, c, h*self.scale, w*self.scale)
         return x
 
@@ -238,7 +201,7 @@ class SDYNetR90x1(nn.Module):
             return F.mse_loss(pred/255, target/255)
         return loss_fn
 
-class SDYNetR90x2(nn.Module):
+class SDYNetR90x2(SRNetBase):
     def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
         super(SDYNetR90x2, self).__init__()
         self.scale = scale   
@@ -252,16 +215,6 @@ class SDYNetR90x2(nn.Module):
         self.stage2_D = layers.UpscaleBlock(hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=scale)
         self.stage2_Y = layers.UpscaleBlock(hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=scale)
 
-    def forward_stage(self, x, scale, percieve_pattern, stage):
-        b,c,h,w = x.shape
-        x = percieve_pattern(x)   
-        x = stage(x)
-        x = round_func(x)
-        x = x.reshape(b, c, h, w, scale, scale)
-        x = x.permute(0,1,2,4,3,5)
-        x = x.reshape(b, c, h*scale, w*scale)
-        return x
-
     def forward(self, x, config=None):
         b,c,h,w = x.shape
         x = x.view(b*c, 1, h, w)
@@ -275,7 +228,7 @@ class SDYNetR90x2(nn.Module):
             output_1 += torch.rot90(self.forward_stage(rotated, 1, self._extract_pattern_D, self.stage1_D), k=-rotations_count, dims=[-2, -1])
             output_1 += torch.rot90(self.forward_stage(rotated, 1, self._extract_pattern_Y, self.stage1_Y), k=-rotations_count, dims=[-2, -1])              
         output_1 /= 4*3
-        x = round_func(output_1)
+        x = output_1
         output_2 = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device)
         output_2 += self.forward_stage(x, self.scale, self._extract_pattern_S, self.stage2_S)
         output_2 += self.forward_stage(x, self.scale, self._extract_pattern_D, self.stage2_D)
@@ -286,7 +239,7 @@ class SDYNetR90x2(nn.Module):
             output_2 += torch.rot90(self.forward_stage(rotated, self.scale, self._extract_pattern_D, self.stage2_D), k=-rotations_count, dims=[-2, -1])
             output_2 += torch.rot90(self.forward_stage(rotated, self.scale, self._extract_pattern_Y, self.stage2_Y), k=-rotations_count, dims=[-2, -1])
         output_2 /= 4*3
-        x = round_func(output_2)
+        x = output_2
         x = x.view(b, c, h*self.scale, w*self.scale)
         return x