File size: 690 Bytes
bb2108c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 |
import math
import numpy as np
import torch
def set_timesteps_patched(self, num_inference_steps: int, device = None):
self.num_inference_steps = num_inference_steps
ramp = np.linspace(0, 1, self.num_inference_steps)
sigmas = torch.linspace(math.log(self.config.sigma_min), math.log(self.config.sigma_max), len(ramp)).exp().flip(0)
sigmas = (sigmas).to(dtype=torch.float32, device=device)
self.timesteps = self.precondition_noise(sigmas)
self.sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)])
self._step_index = None
self._begin_index = None
self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication
|