Fix cpu_only schedulers(unipc)

This commit is contained in:
Sergey Borisov
2023-08-14 05:14:05 +03:00
parent 58d5c61c79
commit 409e5d01ba

View File

@ -326,10 +326,11 @@ class DenoiseLatentsInvocation(BaseInvocation):
# original idea by https://github.com/AmericanPresidentJimmyCarter # original idea by https://github.com/AmericanPresidentJimmyCarter
# TODO: research more for second order schedulers timesteps # TODO: research more for second order schedulers timesteps
def init_scheduler(self, scheduler, device, steps, denoising_start, denoising_end): def init_scheduler(self, scheduler, device, steps, denoising_start, denoising_end):
if scheduler.config.get("cpu_only", False):
device = torch.device("cpu")
num_inference_steps = steps num_inference_steps = steps
if scheduler.config.get("cpu_only", False):
scheduler.set_timesteps(num_inference_steps, device="cpu")
timesteps = scheduler.timesteps.to(device=device)
else:
scheduler.set_timesteps(num_inference_steps, device=device) scheduler.set_timesteps(num_inference_steps, device=device)
timesteps = scheduler.timesteps timesteps = scheduler.timesteps