mirror of
https://github.com/huggingface/diffusers.git
synced 2026-10-23 09:09:01 +03:00
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).
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user