From 62bfa5a23efb41b9685e5946ced9aa134d2ff3fa Mon Sep 17 00:00:00 2001
From: Ricardo-M-L <69202550+Ricardo-M-L@users.noreply.github.com>
Date: Tue, 21 Apr 2026 19:45:10 +0800
Subject: [PATCH] fix(freeu): run FFT in float32 for float16 inputs to avoid
 ComplexHalf (#13503)

* fix(freeu): run FFT in float32 for float16 inputs to avoid ComplexHalf

`fourier_filter` already upcasts `bfloat16` inputs to `float32` before
calling `torch.fft.fftn`, because PyTorch's FFT does not support bf16.
The same is true for `float16`: depending on the PyTorch version,
`fftn` either

- produces the experimental `torch.complex32` (ComplexHalf) dtype and
  emits a `UserWarning: ComplexHalf support is experimental`, or
- raises `RuntimeError: Unsupported dtype Half` outright.

Both paths were reachable from FreeU with half-precision models
(e.g. `sd-turbo` + `fp16` + `enable_freeu`) as reported in #12504.

Extend the existing upcast branch to cover `float16` too. The function
already downcasts the result back to `x_in.dtype` at the end, so the
externally observable dtype is unchanged.

Closes #12504.

* Address review: generalize upcast to non-float32 + fix ruff F821

- Apply @sayakpaul's suggestion: use `elif x.dtype != torch.float32:`
  so any non-float32 dtype (bf16, fp16, and future half-precision
  dtypes) is upcast to float32 before the FFT.
- Drop the `"torch.Tensor"` return annotation on the test helper
  that triggered ruff F821 in CI (torch is imported inside the
  method body, not at module scope).
---
 src/diffusers/utils/torch_utils.py |  6 +++--
 tests/others/test_utils.py         | 43 ++++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/src/diffusers/utils/torch_utils.py b/src/diffusers/utils/torch_utils.py
index a73ad4acf3..07036a4ee0 100644
--- a/src/diffusers/utils/torch_utils.py
+++ b/src/diffusers/utils/torch_utils.py
@@ -223,8 +223,10 @@ def fourier_filter(x_in: "torch.Tensor", threshold: int, scale: int) -> "torch.T
     # 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:
+    # fftn does not support bfloat16, and produces the experimental ComplexHalf
+    # dtype (torch.complex32) when given float16, which is numerically unstable
+    # and triggers a UserWarning. Upcast any non-float32 dtype to float32.
+    elif x.dtype != torch.float32:
         x = x.to(dtype=torch.float32)
 
     # FFT
diff --git a/tests/others/test_utils.py b/tests/others/test_utils.py
index bb06563863..7b445e3a21 100755
--- a/tests/others/test_utils.py
+++ b/tests/others/test_utils.py
@@ -204,6 +204,49 @@ class DeprecateTester(unittest.TestCase):
         ), f"Expected deprecation message substring not found, got: {messages}"
 
 
+class FourierFilterTester(unittest.TestCase):
+    """Tests for :func:`diffusers.utils.torch_utils.fourier_filter` (FreeU helper)."""
+
+    def _run_without_complexhalf_warning(self, dtype):
+        import torch
+
+        from diffusers.utils.torch_utils import fourier_filter
+
+        x = torch.randn(1, 4, 32, 32, dtype=dtype)
+        with warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter("always")
+            out = fourier_filter(x, threshold=1, scale=0.5)
+
+        messages = [str(w.message) for w in caught]
+        assert not any("ComplexHalf" in m for m in messages), (
+            f"Unexpected ComplexHalf warning emitted by fourier_filter: {messages}"
+        )
+        return out
+
+    def test_fourier_filter_float16_no_complexhalf_warning(self):
+        import torch
+
+        out = self._run_without_complexhalf_warning(torch.float16)
+        assert out.dtype == torch.float16
+
+    def test_fourier_filter_bfloat16_no_complexhalf_warning(self):
+        import torch
+
+        out = self._run_without_complexhalf_warning(torch.bfloat16)
+        assert out.dtype == torch.bfloat16
+
+    def test_fourier_filter_preserves_dtype_and_shape(self):
+        import torch
+
+        from diffusers.utils.torch_utils import fourier_filter
+
+        for dtype in (torch.float32, torch.float16, torch.bfloat16):
+            x = torch.randn(2, 3, 16, 16, dtype=dtype)
+            out = fourier_filter(x, threshold=1, scale=0.5)
+            assert out.dtype == dtype
+            assert out.shape == x.shape
+
+
 # Copied from https://github.com/huggingface/transformers/blob/main/tests/utils/test_expectations.py
 class ExpectationsTester(unittest.TestCase):
     def test_expectations(self):