Change num_res_channels and num_channels to Sequence[int] | int

.pre-commit-config.yaml

@@ -12,6 +12,7 @@ repos:
     rev: v4.4.0
     hooks:
       - id: end-of-file-fixer
+        exclude: ^tutorials/
       - id: trailing-whitespace
       - id: check-yaml
       - id: check-docstring-first
@@ -24,7 +25,6 @@ repos:
       - id: forbid-new-submodules
       - id: pretty-format-json
         args: ['--autofix', '--no-sort-keys', '--indent=4']
-      - id: end-of-file-fixer
       - id: mixed-line-ending
 
   - repo: https://github.com/asottile/pyupgrade

generative/networks/nets/vqvae.py

@@ -17,6 +17,7 @@
 import torch.nn as nn
 from monai.networks.blocks import Convolution
 from monai.networks.layers import Act
+from monai.utils import ensure_tuple_rep
 
 from generative.networks.layers.vector_quantizer import EMAQuantizer, VectorQuantizer
 
@@ -138,8 +139,8 @@ def __init__(
             (2, 4, 1, 1, 0),
         ),
         num_res_layers: int = 3,
-        num_channels: Sequence[int] = (96, 96, 192),
-        num_res_channels: Sequence[int] = (96, 96, 192),
+        num_channels: Sequence[int] | int = (96, 96, 192),
+        num_res_channels: Sequence[int] | int = (96, 96, 192),
         num_embeddings: int = 32,
         embedding_dim: int = 64,
         embedding_init: str = "normal",
@@ -158,6 +159,11 @@ def __init__(
         self.out_channels = out_channels
         self.spatial_dims = spatial_dims
 
+        if isinstance(num_channels, int):
+            num_channels = ensure_tuple_rep(num_channels, num_levels)
+        if isinstance(num_res_channels, int):
+            num_res_channels = ensure_tuple_rep(num_res_channels, num_levels)
+
         assert (
             num_levels == len(downsample_parameters)
             and num_levels == len(upsample_parameters)

tests/test_vqvae.py

@@ -21,6 +21,56 @@
 from tests.utils import test_script_save
 
 TEST_CASES = [
+    [
+        {
+            "spatial_dims": 2,
+            "in_channels": 1,
+            "out_channels": 1,
+            "num_levels": 2,
+            "downsample_parameters": [(2, 4, 1, 1)] * 2,
+            "upsample_parameters": [(2, 4, 1, 1, 0)] * 2,
+            "num_res_layers": 1,
+            "num_channels": 8,
+            "num_res_channels": [8, 8],
+            "num_embeddings": 16,
+            "embedding_dim": 8,
+            "embedding_init": "normal",
+            "commitment_cost": 0.25,
+            "decay": 0.5,
+            "epsilon": 1e-5,
+            "adn_ordering": "NDA",
+            "dropout": 0.1,
+            "act": "RELU",
+            "output_act": None,
+        },
+        (1, 1, 16, 16),
+        (1, 1, 16, 16),
+    ],
+    [
+        {
+            "spatial_dims": 2,
+            "in_channels": 1,
+            "out_channels": 1,
+            "num_levels": 2,
+            "downsample_parameters": [(2, 4, 1, 1)] * 2,
+            "upsample_parameters": [(2, 4, 1, 1, 0)] * 2,
+            "num_res_layers": 1,
+            "num_channels": 8,
+            "num_res_channels": 8,
+            "num_embeddings": 16,
+            "embedding_dim": 8,
+            "embedding_init": "normal",
+            "commitment_cost": 0.25,
+            "decay": 0.5,
+            "epsilon": 1e-5,
+            "adn_ordering": "NDA",
+            "dropout": 0.1,
+            "act": "RELU",
+            "output_act": None,
+        },
+        (1, 1, 16, 16),
+        (1, 1, 16, 16),
+    ],
     [
         {
             "spatial_dims": 2,