Update src/model.py

src/model.py

@@ -9,20 +9,54 @@ class B(nn.Module):
         s.s = nn.Conv2d(n_i, n_o, 1, bias=False) if n_i != n_o else nn.Identity()
         s.f = nn.ReLU()
     def forward(s, x): return s.f(s.c(x) + s.s(x))
-def E(lc=4):
+
+import torch
+import torch as th
+import torch.nn as nn
+import torch.nn.functional as F
+
+def conv(n_in, n_out, **kwargs):
+    return nn.Conv2d(n_in, n_out, 3, padding=1, **kwargs)
+
+class Clamp(nn.Module):
+    def forward(self, x):
+        return torch.tanh(x / 3) * 3
+
+class Block(nn.Module):
+    def __init__(self, n_in, n_out):
+        super().__init__()
+        self.conv = nn.Sequential(conv(n_in, n_out), nn.ReLU(), conv(n_out, n_out), nn.ReLU(), conv(n_out, n_out))
+        self.skip = nn.Conv2d(n_in, n_out, 1, bias=False) if n_in != n_out else nn.Identity()
+        self.fuse = nn.ReLU()
+    def forward(self, x):
+        return self.fuse(self.conv(x) + self.skip(x))
+
+def E(latent_channels=4):
     return nn.Sequential(
-        cv(3, 64), B(64, 64), cv(64, 64, stride=2, bias=False), B(64, 64), B(64, 64), B(64, 64),
-        cv(64, 64, stride=2, bias=False), B(64, 64), B(64, 64), B(64, 64),
-        cv(64, 64, stride=2, bias=False), B(64, 64), B(64, 64), B(64, 64),
-        cv(64, lc),
+        conv(3, 64), Block(64, 64),
+        conv(64, 64, stride=2, bias=False), Block(64, 64), Block(64, 64), Block(64, 64),
+        conv(64, 64, stride=2, bias=False), Block(64, 64), Block(64, 64), Block(64, 64),
+        conv(64, 64, stride=2, bias=False), Block(64, 64), Block(64, 64), Block(64, 64),
+        conv(64, latent_channels),
     )
-def D(lc=16):
+
+def D(latent_channels=16):  # Adjusted to match expected input channels
     return nn.Sequential(
-        C(), cv(lc, 48), nn.ReLU(), B(48, 48), B(48, 48), nn.Upsample(scale_factor=2),
-        cv(48, 48, bias=False), B(48, 48), B(48, 48), nn.Upsample(scale_factor=2),
-        cv(48, 48, bias=False), B(48, 48), nn.Upsample(scale_factor=2),
-        cv(48, 48, bias=False), B(48, 48), cv(48, 3),
+        Clamp(), 
+        conv(latent_channels, 48),  # Reduced from 64 to 48 channels
+        nn.ReLU(),
+        Block(48, 48), Block(48, 48),  # Reduced number of blocks
+        nn.Upsample(scale_factor=2), conv(48, 48, bias=False),
+        Block(48, 48), Block(48, 48),  # Reduced number of blocks
+        nn.Upsample(scale_factor=2), conv(48, 48, bias=False),
+        Block(48, 48),  # Further reduction in blocks
+        nn.Upsample(scale_factor=2), conv(48, 48, bias=False),
+        Block(48, 48), 
+        conv(48, 3),  # Final convolution to output channels
     )
+
+
+    
 class M(nn.Module):
     lm, ls = 3, 0.5
     def __init__(s, ep="encoder.pth", dp="decoder.pth", lc=None):