Project-MONAI · Warvito · Nov 30, 2022 · Nov 8, 2022 · Nov 8, 2022 · Nov 8, 2022
diff --git a/generative/networks/nets/diffusion_model_unet.py b/generative/networks/nets/diffusion_model_unet.py
diff --git a/generative/utils/__init__.py b/generative/utils/__init__.py
@@ -10,4 +10,3 @@
 # limitations under the License.
 
 from .enums import AdversarialIterationEvents, AdversarialKeys
-from .misc import default, exists, extract
diff --git a/generative/utils/misc.py b/generative/utils/misc.py
diff --git a/tests/test_diffusion_inferer.py b/tests/test_diffusion_inferer.py
@@ -24,12 +24,11 @@
             "spatial_dims": 2,
             "in_channels": 1,
             "out_channels": 1,
-            "model_channels": 8,
+            "num_channels": [8],
             "norm_num_groups": 8,
-            "attention_resolutions": [1],
+            "attention_levels": [True],
             "num_res_blocks": 1,
-            "channel_mult": [1],
-            "num_heads": 1,
+            "num_head_channels": 8,
         },
         (2, 1, 8, 8),
     ],
@@ -38,12 +37,11 @@
             "spatial_dims": 3,
             "in_channels": 1,
             "out_channels": 1,
-            "model_channels": 8,
+            "num_channels": [8],
             "norm_num_groups": 8,
-            "attention_resolutions": [1],
+            "attention_levels": [True],
             "num_res_blocks": 1,
-            "channel_mult": [1],
-            "num_heads": 1,
+            "num_head_channels": 8,
         },
         (2, 1, 8, 8, 8),
     ],
Original file line number	Diff line number	Diff line change
Expand Up		@@ -10,4 +10,3 @@
		# limitations under the License.

		from .enums import AdversarialIterationEvents, AdversarialKeys
		from .misc import default, exists, extract