From de5fcf6fe322dde5ea119f0491dada9945f8a649 Mon Sep 17 00:00:00 2001 From: Ricardo-M-L <69202550+Ricardo-M-L@users.noreply.github.com> Date: Sat, 16 May 2026 02:11:54 +0800 Subject: [PATCH] fix(randn_tensor): compare device.type, not torch.device, when suppressing MPS info log (#13508) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(randn_tensor): compare device.type, not torch.device to str, when suppressing MPS info log When a CPU generator is passed with an MPS target, randn_tensor intentionally skips the 'generator was on cpu, tensor will be moved to ' info log — MPS doesn't support device-side generators, so the suggestion to create one on MPS would be misleading. The guard was written as `if device != "mps"`, but a few lines earlier `device` is coerced to a `torch.device` object, and `torch.device("mps") == "mps"` is False (torch.device's __eq__ with a string returns NotImplemented, falling back to identity — they're different types). Result: the guard is effectively always True, so MPS users get the spurious log whenever they pass a CPU generator — the opposite of the documented intent. Fix: compare `device.type` (a str) against "mps". Added a regression test in tests/others/test_utils.py that exercises both the MPS and non-MPS paths via `assertLogs` on the diffusers logger. * refactor: use CaptureLogger instead of assertLogs in randn_tensor test Co-Authored-By: Claude Opus 4.7 --------- Co-authored-by: Claude Opus 4.7 Co-authored-by: YiYi Xu --- src/diffusers/utils/torch_utils.py | 2 +- tests/others/test_utils.py | 37 ++++++++++++++++++++++++++++++ 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/src/diffusers/utils/torch_utils.py b/src/diffusers/utils/torch_utils.py index 07036a4ee0..c314a8609b 100644 --- a/src/diffusers/utils/torch_utils.py +++ b/src/diffusers/utils/torch_utils.py @@ -173,7 +173,7 @@ def randn_tensor( gen_device_type = generator.device.type if not isinstance(generator, list) else generator[0].device.type if gen_device_type != device.type and gen_device_type == "cpu": rand_device = "cpu" - if device != "mps": + if device.type != "mps": logger.info( f"The passed generator was created on 'cpu' even though a tensor on {device} was expected." f" Tensors will be created on 'cpu' and then moved to {device}. Note that one can probably" diff --git a/tests/others/test_utils.py b/tests/others/test_utils.py index 7b445e3a21..4600f5f371 100755 --- a/tests/others/test_utils.py +++ b/tests/others/test_utils.py @@ -247,6 +247,43 @@ class FourierFilterTester(unittest.TestCase): assert out.shape == x.shape +class RandnTensorTester(unittest.TestCase): + """Tests for :func:`diffusers.utils.torch_utils.randn_tensor`.""" + + def test_mps_suppresses_cpu_generator_info_log(self): + import torch + + from diffusers.utils import logging as diffusers_logging + from diffusers.utils import torch_utils + + from ..testing_utils import CaptureLogger + + gen = torch.Generator(device="cpu") + diffusers_logging.set_verbosity_info() + + def _capture(target_device): + with CaptureLogger(torch_utils.logger) as cl: + try: + torch_utils.randn_tensor((1, 2), generator=gen, device=target_device, dtype=torch.float32) + except (AssertionError, RuntimeError): + pass + return cl.out + + mps_out = _capture("mps") + self.assertNotIn( + "moved to", + mps_out, + f"MPS target should not emit the CPU-fallback info log, got: {mps_out}", + ) + + cuda_out = _capture("cuda") + self.assertIn( + "moved to", + cuda_out, + f"Non-MPS target should still emit the CPU-fallback info log, got: {cuda_out}", + ) + + # Copied from https://github.com/huggingface/transformers/blob/main/tests/utils/test_expectations.py class ExpectationsTester(unittest.TestCase): def test_expectations(self):