src/aiunn/upsampler/aiunn.py

@@ -15,17 +15,16 @@ class aiuNN(AIIA):
         # Pass the unified base configuration using the new parameter.
         self.config = aiuNNConfig(base_config=base_model.config)
         
-        self.upsample = nn.Upsample(
-            scale_factor=self.config.upsample_scale,
-            mode=self.config.upsample_mode,
-            align_corners=self.config.upsample_align_corners
-        )
-        # Conversion layer: change from hidden size channels to number of channels from the config.
-        self.to_rgb = nn.Conv2d(
+        # Enhanced approach
+        scale_factor = self.config.upsample_scale
+        out_channels = self.base_model.config.num_channels * (scale_factor ** 2)
+        self.pixel_shuffle_conv = nn.Conv2d(
             in_channels=self.base_model.config.hidden_size,
-            out_channels=self.base_model.config.num_channels,
-            kernel_size=1
+            out_channels=out_channels,
+            kernel_size=self.base_model.config.kernel_size,
+            padding=1
         )
+        self.pixel_shuffle = nn.PixelShuffle(scale_factor)
 
 
     def forward(self, x):