Update src/diffusers/utils/torch_utils.py

Author: hlky

src/diffusers/utils/torch_utils.py

@@ -98,7 +98,12 @@ def fourier_filter(x_in: "torch.Tensor", threshold: int, scale: int) -> "torch.T
     """
     x = x_in
     B, C, H, W = x.shape
-    x = x.to(dtype=torch.float32)
+    # Non-power of 2 images must be float32
+    if (W & (W - 1)) != 0 or (H & (H - 1)) != 0:
+        x = x.to(dtype=torch.float32)
+    # fftn does not support bfloat16
+    elif x.dtype == torch.bfloat16:
+        x = x.to(dtype=torch.float32)
 
     # FFT
     x_freq = fftn(x, dim=(-2, -1))