Finegrain/refiners
1
0
Fork 0
mirror of https://github.com/finegrain-ai/refiners.git synced 2024-11-12 16:18:22 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix precision of DPM solver in bfloat16

Browse source
This commit is contained in:
Pierre Chapuis 2024-09-25 21:46:58 +02:00
parent bdba91312b
commit 74ce42f923
No known key found for this signature in database
1 changed files with 2 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
src/refiners/foundationals/latent_diffusion/solvers/dpm.py
View file

@ -77,7 +77,7 @@ class DPMSolver(Solver):
first_inference_step=first_inference_step,
params=params,
device=device,
dtype=dtype,
dtype=torch.float64, # compute constants precisely
)
self.estimated_data = deque([torch.tensor([])] * 2, maxlen=2)
self.last_step_first_order = last_step_first_order
@ -89,6 +89,7 @@ class DPMSolver(Solver):
self.sigmas
)
self.timesteps = self._timesteps_from_sigmas(sigmas)
self.to(dtype=dtype)
def rebuild(
self: "DPMSolver",
@ -131,12 +132,9 @@ class DPMSolver(Solver):
case NoiseSchedule.KARRAS:
rho = 7
case None:
if sigmas.dtype == torch.bfloat16:
sigmas = sigmas.to(torch.float32)
return torch.tensor(
np.interp(self.timesteps.cpu(), np.arange(0, len(sigmas)), sigmas.cpu()),
device=self.device,
dtype=self.dtype,
)
linear_schedule = torch.linspace(0, 1, steps=self.num_inference_steps, device=self.device)