From 8b8a426ba90daa67f406cfa343291ae1eb8f8cf6 Mon Sep 17 00:00:00 2001 From: Priyansh Date: Wed, 30 Sep 2026 17:40:35 +0000 Subject: [PATCH 1/2] Add test_default_dtype_float64 for UniPCMultistepScheduler Covers the float64 default-dtype path that fails with a dtype mismatch in torch.linalg.solve. Modeled on test_fp16_support. --- tests/schedulers/test_scheduler_unipc.py | 35 ++++++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/tests/schedulers/test_scheduler_unipc.py b/tests/schedulers/test_scheduler_unipc.py index ac7e1d3f88b4..e46db0706ed3 100644 --- a/tests/schedulers/test_scheduler_unipc.py +++ b/tests/schedulers/test_scheduler_unipc.py @@ -257,6 +257,41 @@ def test_fp16_support(self): assert sample.dtype == torch.float16 + def test_default_dtype_float64(self): + """The unit entry of ``rks`` must take the sample's dtype, not the default dtype. + + ``torch.ones(())`` takes the default dtype, so under + ``torch.set_default_dtype(torch.float64)`` the stacked ``R`` was promoted + to float64 while ``b`` stayed float32 and ``torch.linalg.solve`` raised. + """ + default_dtype = torch.get_default_dtype() + torch.set_default_dtype(torch.float64) + try: + for order in [1, 2, 3]: + for solver_type in ["bh1", "bh2"]: + for prediction_type in ["epsilon", "sample", "v_prediction"]: + scheduler_class = self.scheduler_classes[0] + scheduler_config = self.get_scheduler_config( + prediction_type=prediction_type, + solver_order=order, + solver_type=solver_type, + ) + scheduler = scheduler_class(**scheduler_config) + + num_inference_steps = 10 + model = self.dummy_model() + sample = self.dummy_sample_deter.to(torch.float64) + scheduler.set_timesteps(num_inference_steps) + + for i, t in enumerate(scheduler.timesteps): + residual = model(sample, t) + sample = scheduler.step(residual, t, sample).prev_sample + + assert sample.dtype == torch.float64 + assert torch.isfinite(sample).all() + finally: + torch.set_default_dtype(default_dtype) + def test_full_loop_with_noise(self): scheduler_class = self.scheduler_classes[0] scheduler_config = self.get_scheduler_config() From 6b3715350e914d68085d0e0964a1fbe985c95454 Mon Sep 17 00:00:00 2001 From: Priyansh Date: Wed, 30 Sep 2026 17:40:40 +0000 Subject: [PATCH 2/2] Fix UniPCMultistepScheduler under a float64 default dtype multistep_uni_p_bh_update and multistep_uni_c_bh_update built rks from the float32 sigmas and appended torch.ones(()), which takes the default dtype. Under torch.set_default_dtype(torch.float64) torch.stack promoted R to float64 while b stayed float32, so torch.linalg.solve raised: RuntimeError: linalg.solve: Expected A and B to have the same dtype, but found A of type Double and B of type Float instead Use torch.ones_like(h) so the unit entry has the dtype and device of the other rks entries. Under the default float32 this is the same tensor as before, so existing behaviour is unchanged. Fixes #14888 --- src/diffusers/schedulers/scheduling_unipc_multistep.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/diffusers/schedulers/scheduling_unipc_multistep.py b/src/diffusers/schedulers/scheduling_unipc_multistep.py index 5c2cbcc13ff1..e8f5af16fbe4 100644 --- a/src/diffusers/schedulers/scheduling_unipc_multistep.py +++ b/src/diffusers/schedulers/scheduling_unipc_multistep.py @@ -903,7 +903,7 @@ def multistep_uni_p_bh_update( rks.append(rk) D1s.append((mi - m0) / rk) - rks.append(torch.ones((), device=device)) + rks.append(torch.ones_like(h)) rks = torch.stack(rks) R = [] @@ -1038,7 +1038,7 @@ def multistep_uni_c_bh_update( rks.append(rk) D1s.append((mi - m0) / rk) - rks.append(torch.ones((), device=device)) + rks.append(torch.ones_like(h)) rks = torch.stack(rks) R = []