Update ldm/modules/diffusionmodules/flag_large_dit.py

ldm/modules/diffusionmodules/flag_large_dit.py

@@ -240,13 +240,14 @@ class TxtFlagLargeDiT(nn.Module):
         theta = theta * ntk_factor
 
         print(f"theta {theta} rope scaling {rope_scaling_factor} ntk {ntk_factor}")
-
-        freqs = 1.0 / (theta ** (
+        if torch.cuda.is_available():
+            freqs = 1.0 / (theta ** (
             torch.arange(0, dim, 2)[: (dim // 2)].float().cuda() / dim
         ))
-        # freqs = 1.0 / (theta ** (
-        #     torch.arange(0, dim, 2)[: (dim // 2)].float() / dim
-        # ))
+        else:
+            freqs = 1.0 / (theta ** (
+                torch.arange(0, dim, 2)[: (dim // 2)].float() / dim
+            ))
         t = torch.arange(end, device=freqs.device, dtype=torch.float)  # type: ignore
         t = t / rope_scaling_factor
         freqs = torch.outer(t, freqs).float()  # type: ignore