Update arch check for SD (NVIDIA#9783)

* Update arch check for SD Signed-off-by: Jaemin Choi <[email protected]> * Apply isort and black reformatting Signed-off-by: minitu <[email protected]> --------- Signed-off-by: Jaemin Choi <[email protected]> Signed-off-by: minitu <[email protected]> Co-authored-by: Jaemin Choi <[email protected]> Co-authored-by: minitu <[email protected]> Signed-off-by: Boxiang Wang <[email protected]>
BoxiangW · Jul 30, 2024 · 0db4d5c · 0db4d5c
1 parent f54331f
commit 0db4d5c
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 9 deletions.
diff --git a/nemo/collections/multimodal/modules/imagen/diffusionmodules/blocks.py b/nemo/collections/multimodal/modules/imagen/diffusionmodules/blocks.py
@@ -58,10 +58,9 @@ def check_cuda():
     dprops = th.cuda.get_device_properties(cur_device)
 
     is_sm75 = dprops.major == 7 and dprops.minor == 5
-    is_sm8x = dprops.major == 8 and dprops.minor >= 0
-    is_sm90 = dprops.major == 9 and dprops.minor >= 0
+    is_sm8x_or_later = dprops.major >= 8
 
-    return is_sm8x or is_sm75 or is_sm90
+    return is_sm75 or is_sm8x_or_later
 
 
 try:
@@ -154,7 +153,9 @@ def __init__(
         self.use_scale_shift_norm = use_scale_shift_norm
 
         self.in_layers = nn.Sequential(
-            normalization(channels), nn.SiLU(), conv_nd(dims, channels, self.out_channels, 3, padding=1),
+            normalization(channels),
+            nn.SiLU(),
+            conv_nd(dims, channels, self.out_channels, 3, padding=1),
         )
 
         self.updown = up or down
@@ -173,7 +174,11 @@ def __init__(
             self.h_upd = self.x_upd = nn.Identity()
 
         self.emb_layers = nn.Sequential(
-            nn.SiLU(), linear(emb_channels, 2 * self.out_channels if use_scale_shift_norm else self.out_channels,),
+            nn.SiLU(),
+            linear(
+                emb_channels,
+                2 * self.out_channels if use_scale_shift_norm else self.out_channels,
+            ),
         )
         self.out_layers = nn.Sequential(
             normalization(self.out_channels),
@@ -263,7 +268,11 @@ def __init__(
         )
 
         self.emb_layers = nn.Sequential(
-            nn.SiLU(), nn.Linear(emb_channels, 2 * out_channels if use_scale_shift_norm else out_channels,),
+            nn.SiLU(),
+            nn.Linear(
+                emb_channels,
+                2 * out_channels if use_scale_shift_norm else out_channels,
+            ),
         )
 
         self.out_layers = nn.Sequential(

diff --git a/nemo/collections/multimodal/modules/stable_diffusion/attention.py b/nemo/collections/multimodal/modules/stable_diffusion/attention.py
@@ -56,10 +56,9 @@ def check_cuda():
     dprops = torch.cuda.get_device_properties(cur_device)
 
     is_sm75 = dprops.major == 7 and dprops.minor == 5
-    is_sm8x = dprops.major == 8 and dprops.minor >= 0
-    is_sm90 = dprops.major == 9 and dprops.minor >= 0
+    is_sm8x_or_later = dprops.major >= 8
 
-    return is_sm8x or is_sm75 or is_sm90
+    return is_sm75 or is_sm8x_or_later
 
 
 try: