From 18ef5684c01a0bb56f457e31a058023220fae18f Mon Sep 17 00:00:00 2001 From: Laurent Date: Tue, 1 Oct 2024 08:12:51 +0000 Subject: [PATCH] add missing remove_noise method from the DPMSolver class --- .../foundationals/latent_diffusion/solvers/dpm.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/src/refiners/foundationals/latent_diffusion/solvers/dpm.py b/src/refiners/foundationals/latent_diffusion/solvers/dpm.py index 2ec2de4..b2296dc 100644 --- a/src/refiners/foundationals/latent_diffusion/solvers/dpm.py +++ b/src/refiners/foundationals/latent_diffusion/solvers/dpm.py @@ -191,6 +191,16 @@ class DPMSolver(Solver): noised_x = cumulative_scale_factors * x + noise_stds * noise return noised_x + def remove_noise(self, x: torch.Tensor, noise: torch.Tensor, step: int) -> torch.Tensor: + """Remove noise from the input tensor using the current step of the diffusion process. + + See [`Solver.remove_noise`][refiners.foundationals.latent_diffusion.solvers.solver.Solver.remove_noise] for more details. + """ + cumulative_scale_factors = self.cumulative_scale_factors[step] + noise_stds = self.noise_std[step] + denoised_x = (x - noise_stds * noise) / cumulative_scale_factors + return denoised_x + def _solver_tensors_from_sigmas(self, sigmas: torch.Tensor) -> SolverTensors: """Generate the tensors from the sigmas.""" cumulative_scale_factors = 1 / torch.sqrt(sigmas**2 + 1)