Fix output channel calculation in PixelShuffle2DUpBlock

src/virtual_stain_flow/models/blocks/up_down_blocks.py

@@ -287,6 +287,7 @@ def __init__(
         self,
         in_channels: int,
         out_channels: Optional[int] = None,
+        preserve_channels: bool = False,
         **kwargs
     ):
         """
@@ -304,6 +305,8 @@ def __init__(
         # as the pixel shuffle operation merely rearranges the channels
         # to the spatial dimensions
         out_channels = in_channels
+        if not preserve_channels:
+            out_channels = out_channels // (scale_factor ** spatial_dims)
 
         super().__init__(
             in_channels=in_channels,

src/virtual_stain_flow/models/stages.py

@@ -176,15 +176,13 @@ def out_channels(self) -> int:
     def out_h(self, in_h: int) -> int:
         _out_h = in_h
         for block in [self.in_block, self.comp_block]:
-            if isinstance(block, Conv2DDownBlock):
-                _out_h = block.out_h(_out_h)
+            _out_h = block.out_h(_out_h)
         return _out_h
-    
+
     def out_w(self, in_w: int) -> int:
         _out_w = in_w
         for block in [self.in_block, self.comp_block]:
-            if isinstance(block, Conv2DDownBlock):
-                _out_w = block.out_w(_out_w)
+            _out_w = block.out_w(_out_w)
         return _out_w
 
 """

src/virtual_stain_flow/models/unext.py

@@ -47,7 +47,8 @@ def __init__(
         decoder_up_block: Literal['pixelshuffle', 'convt'] = 'pixelshuffle',
         decoder_compute_block: Literal['convnext', 'conv2d'] = 'convnext',
         act_type: ActivationType = 'sigmoid',
-        _num_units: Union[List[int], int] = 2
+        _num_units: Union[List[int], int] = 2,
+        _pixel_shuffle_preserve_channels: bool = False,
     ):
         """
         Initializes the ConvNeXtUNet model.
@@ -98,13 +99,16 @@ def __init__(
 
         if decoder_up_block == 'pixelshuffle':
             in_block_handles = [PixelShuffle2DUpBlock] * (depth - 1)
+            in_block_kwargs = [{'preserve_channels': _pixel_shuffle_preserve_channels}] * (depth - 1)
         elif decoder_up_block == 'convt':
             in_block_handles = [ConvTrans2DUpBlock] * (depth - 1)
+            in_block_kwargs = [{'norm_type': 'layer'}] * (depth - 1)
         else:
             raise ValueError(
                 f"Unsupported decoder_up_block: {decoder_up_block!r}. "
                 "Expected 'pixelshuffle' or 'convt'."
             )
+        self._pixel_shuffle_preserve_channels = _pixel_shuffle_preserve_channels
         self._decoder_up_block = decoder_up_block
 
         if decoder_compute_block == 'convnext':
@@ -138,7 +142,7 @@ def __init__(
             encoder_feature_map_channels=convnextv2_model.feature_info.channels(),
             # use convolutional up-sampling blocks
             in_block_handles=in_block_handles,
-            in_block_kwargs=[{'norm_type': 'layer'}] * (depth - 1),
+            in_block_kwargs=in_block_kwargs,
             comp_block_handles=comp_block_handles,
             comp_block_kwargs=comp_block_kwargs,
         )
@@ -194,6 +198,7 @@ def to_config(self) -> Dict[str, Any]:
                 "decoder_compute_block": self._decoder_compute_block,
                 "act_type": self._act_type,
                 "_num_units": self._num_units_cfg,
+                "_pixel_shuffle_preserve_channels": self._pixel_shuffle_preserve_channels,
             },
         }
 
@@ -205,5 +210,8 @@ def from_config(cls, config: Dict[str, Any]) -> "ConvNeXtUNet":
         """
 
         init_cfg = config.get("init", config)
+        if "_pixel_shuffle_preserve_channels" not in init_cfg:
+            # For backward compatibility with configs that don't have this key
+            init_cfg["_pixel_shuffle_preserve_channels"] = False
 
         return cls(**init_cfg)

tests/models/test_up_down_blocks.py

@@ -23,8 +23,10 @@ class TestUpDownBlocks:
 			(MaxPool2DDownBlock, {"out_channels": 8}, 3, 3, 0.5),
 			(ConvTrans2DUpBlock, {}, 4, 2, 2),
 			(ConvTrans2DUpBlock, {"out_channels": 3}, 4, 3, 2),
-			(PixelShuffle2DUpBlock, {}, 4, 4, 2),
-			(PixelShuffle2DUpBlock, {"out_channels": 8}, 4, 4, 2),
+			(PixelShuffle2DUpBlock, {}, 4, 1, 2),
+			(PixelShuffle2DUpBlock, {"out_channels": 8}, 4, 1, 2),
+			(PixelShuffle2DUpBlock, {"preserve_channels": True}, 4, 4, 2),
+			(PixelShuffle2DUpBlock, {"preserve_channels": True, "out_channels": 8}, 4, 4, 2),
 			(Bilinear2DUpsampleBlock, {}, 3, 3, 2),
 			(Bilinear2DUpsampleBlock, {"out_channels": 8}, 3, 3, 2),
 		],