add more test cases

Signed-off-by: Pengfei Guo <pengfeig@nvidia.com>

Author: guopengf

monai/apps/generation/maisi/networks/controlnet_maisi.py

@@ -168,11 +168,11 @@ def _apply_mid_block(self, emb, context, h):
 
     def _apply_controlnet_blocks(self, h, down_block_res_samples):
         # 6. Control net blocks
-    controlnet_down_block_res_samples = []
-    for down_block_res_sample, controlnet_block in zip(down_block_res_samples, self.controlnet_down_blocks):
-        down_block_res_sample = controlnet_block(down_block_res_sample)
-        controlnet_down_block_res_samples.append(down_block_res_sample)
+        controlnet_down_block_res_samples = []
+        for down_block_res_sample, controlnet_block in zip(down_block_res_samples, self.controlnet_down_blocks):
+            down_block_res_sample = controlnet_block(down_block_res_sample)
+            controlnet_down_block_res_samples.append(down_block_res_sample)
 
-    mid_block_res_sample = self.controlnet_mid_block(h)
+        mid_block_res_sample = self.controlnet_mid_block(h)
 
-    return controlnet_down_block_res_samples, mid_block_res_sample
+        return controlnet_down_block_res_samples, mid_block_res_sample

tests/test_controlnet_maisi.py

@@ -60,6 +60,71 @@
     ],
 ]
 
+TEST_CASES_CONDITIONAL = [
+    [
+        {
+            "spatial_dims": 2,
+            "in_channels": 1,
+            "num_res_blocks": 1,
+            "num_channels": (8, 8, 8),
+            "attention_levels": (False, False, True),
+            "num_head_channels": 8,
+            "norm_num_groups": 8,
+            "conditioning_embedding_in_channels": 1,
+            "conditioning_embedding_num_channels": (8, 8),
+            "use_checkpointing": False,
+            "with_conditioning": True,
+            "cross_attention_dim": 2,
+        },
+        6,
+        (1, 8, 4, 4),
+    ],
+    [
+        {
+            "spatial_dims": 3,
+            "in_channels": 1,
+            "num_res_blocks": 1,
+            "num_channels": (8, 8, 8),
+            "attention_levels": (False, False, True),
+            "num_head_channels": 8,
+            "norm_num_groups": 8,
+            "conditioning_embedding_in_channels": 1,
+            "conditioning_embedding_num_channels": (8, 8),
+            "use_checkpointing": True,
+            "with_conditioning": True,
+            "cross_attention_dim": 2,
+        },
+        6,
+        (1, 8, 4, 4, 4),
+    ],
+]
+
+TEST_CASES_ERROR = [
+    [
+        {"spatial_dims": 2, "in_channels": 1, "with_conditioning": True, "cross_attention_dim": None},
+        "ControlNet expects dimension of the cross-attention conditioning "
+        "(cross_attention_dim) when using with_conditioning.",
+    ],
+    [
+        {"spatial_dims": 2, "in_channels": 1, "with_conditioning": False, "cross_attention_dim": 2},
+        "ControlNet expects with_conditioning=True when specifying the cross_attention_dim.",
+    ],
+    [
+        {"spatial_dims": 2, "in_channels": 1, "num_channels": (8, 16), "norm_num_groups": 16},
+        "ControlNet expects all num_channels being multiple of norm_num_groups",
+    ],
+    [
+        {
+            "spatial_dims": 2,
+            "in_channels": 1,
+            "num_channels": (8, 16),
+            "attention_levels": (True,),
+            "norm_num_groups": 8,
+        },
+        "ControlNet expects num_channels being same size of attention_levels",
+    ],
+]
+
 
 @skipUnless(has_generative, "monai-generative required")
 class TestControlNet(unittest.TestCase):
@@ -76,6 +141,27 @@ def test_shape_unconditioned_models(self, input_param, expected_num_down_blocks_
             self.assertEqual(len(result[0]), expected_num_down_blocks_residuals)
             self.assertEqual(result[1].shape, expected_shape)
 
+    @parameterized.expand(TEST_CASES_CONDITIONAL)
+    def test_shape_conditioned_models(self, input_param, expected_num_down_blocks_residuals, expected_shape):
+        net = ControlNetMaisi(**input_param)
+        with eval_mode(net):
+            x = torch.rand((1, 1, 16, 16)) if input_param["spatial_dims"] == 2 else torch.rand((1, 1, 16, 16, 16))
+            timesteps = torch.randint(0, 1000, (1,)).long()
+            controlnet_cond = (
+                torch.rand((1, 1, 32, 32)) if input_param["spatial_dims"] == 2 else torch.rand((1, 1, 32, 32, 32))
+            )
+            context = torch.randn((1, 1, input_param["cross_attention_dim"]))
+            result = net.forward(x, timesteps, controlnet_cond, context=context)
+            self.assertEqual(len(result[0]), expected_num_down_blocks_residuals)
+            self.assertEqual(result[1].shape, expected_shape)
+
+    @parameterized.expand(TEST_CASES_ERROR)
+    def test_error_input(self, input_param, expected_error):
+        with self.assertRaises(ValueError) as context:  # output shape too small
+            _ = ControlNetMaisi(**input_param)
+        runtime_error = context.exception
+        self.assertEqual(str(runtime_error), expected_error)
+
 
 if __name__ == "__main__":
     unittest.main()