diff --git a/src/diffusers/schedulers/scheduling_unipc_multistep.py b/src/diffusers/schedulers/scheduling_unipc_multistep.py index 5c2cbcc13ff1..d2867be838eb 100644 --- a/src/diffusers/schedulers/scheduling_unipc_multistep.py +++ b/src/diffusers/schedulers/scheduling_unipc_multistep.py @@ -275,6 +275,11 @@ def __init__( else: raise NotImplementedError(f"{solver_type} is not implemented for {self.__class__}") + if not predict_x0 and final_sigmas_type == "zero": + raise ValueError( + "`final_sigmas_type` zero is not supported for `predict_x0=False`. Please choose `sigma_min` instead." + ) + self.predict_x0 = predict_x0 # setable values self.num_inference_steps = None @@ -903,7 +908,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 = [] @@ -945,16 +950,16 @@ def multistep_uni_p_bh_update( x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0 if D1s is not None: pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s) + x_t = x_t_ - alpha_t * B_h * pred_res else: - pred_res = 0 - x_t = x_t_ - alpha_t * B_h * pred_res + x_t = x_t_ else: x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0 if D1s is not None: pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s) + x_t = x_t_ - sigma_t * B_h * pred_res else: - pred_res = 0 - x_t = x_t_ - sigma_t * B_h * pred_res + x_t = x_t_ x_t = x_t.to(x.dtype) return x_t @@ -1038,7 +1043,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 = [] @@ -1205,7 +1210,9 @@ def step( self.model_outputs[-1] = model_output_convert self.timestep_list[-1] = timestep - if self.config.lower_order_final: + if self.config.final_sigmas_type == "zero" and self.step_index == len(self.timesteps) - 1: + this_order = 1 + elif self.config.lower_order_final: this_order = min(self.config.solver_order, len(self.timesteps) - self.step_index) else: this_order = self.config.solver_order diff --git a/tests/schedulers/test_scheduler_unipc.py b/tests/schedulers/test_scheduler_unipc.py index ac7e1d3f88b4..52d96883951c 100644 --- a/tests/schedulers/test_scheduler_unipc.py +++ b/tests/schedulers/test_scheduler_unipc.py @@ -284,6 +284,61 @@ def test_full_loop_with_noise(self): assert abs(result_sum.item() - 315.5757) < 1e-2, f" expected result sum 315.5757, but get {result_sum}" assert abs(result_mean.item() - 0.4109) < 1e-3, f" expected result mean 0.4109, but get {result_mean}" + def test_last_step_to_sigma_zero_is_the_data_prediction(self): + scheduler_class = self.scheduler_classes[0] + model = self.dummy_model() + + for solver_type in ["bh1", "bh2"]: + for order in [1, 2, 3]: + for lower_order_final in [True, False]: + scheduler_config = self.get_scheduler_config( + solver_type=solver_type, + solver_order=order, + lower_order_final=lower_order_final, + final_sigmas_type="zero", + ) + scheduler = scheduler_class(**scheduler_config) + scheduler.set_timesteps(10) + + sample = self.dummy_sample_deter.clone() + for t in scheduler.timesteps[:-1]: + sample = scheduler.step(model(sample, t), t, sample).prev_sample + + t_last = scheduler.timesteps[-1] + last_residual = model(sample, t_last) + expected = scheduler.convert_model_output(last_residual, sample=sample) + prev_sample = scheduler.step(last_residual, t_last, sample).prev_sample + + assert torch.isfinite(prev_sample).all() + assert torch.allclose(prev_sample, expected, atol=1e-6) + + def test_predict_x0_false_requires_final_sigmas_type_sigma_min(self): + scheduler_class = self.scheduler_classes[0] + scheduler_config = self.get_scheduler_config(predict_x0=False, final_sigmas_type="zero") + with self.assertRaises(ValueError): + scheduler_class(**scheduler_config) + + def test_default_dtype_float64(self): + scheduler_class = self.scheduler_classes[0] + scheduler_config = self.get_scheduler_config(solver_order=2) + scheduler = scheduler_class(**scheduler_config) + scheduler.set_timesteps(10) + + model = self.dummy_model() + sample = self.dummy_sample_deter.clone() + + prev_dtype = torch.get_default_dtype() + torch.set_default_dtype(torch.float64) + try: + for t in scheduler.timesteps: + residual = model(sample, t) + sample = scheduler.step(residual, t, sample).prev_sample + finally: + torch.set_default_dtype(prev_dtype) + + assert sample.dtype == torch.float32 + assert torch.isfinite(sample).all() + class UniPCMultistepScheduler1DTest(UniPCMultistepSchedulerTest): @property