diffusers: upgrade to diffusers 0.10, add Heun scheduler

This commit is contained in:
Kevin Turner 2022-12-08 13:02:47 -08:00
parent 30a8d4c2b3
commit 9199d698f8
3 changed files with 10 additions and 4 deletions

View File

@ -4,7 +4,7 @@
--trusted-host https://download.pytorch.org
accelerate~=0.14
albumentations
diffusers[torch]~=0.9
diffusers[torch]~=0.10
einops
eventlet
flask_cors

View File

@ -17,6 +17,7 @@ import skimage
import torch
import transformers
from PIL import Image, ImageOps
from diffusers import HeunDiscreteScheduler
from diffusers.pipeline_utils import DiffusionPipeline
from diffusers.schedulers.scheduling_ddim import DDIMScheduler
from diffusers.schedulers.scheduling_dpmsolver_multistep import DPMSolverMultistepScheduler
@ -1008,12 +1009,17 @@ class Generate:
scheduler_map = dict(
ddim=DDIMScheduler,
dpmpp_2=DPMSolverMultistepScheduler,
ipndm=IPNDMScheduler,
# DPMSolverMultistepScheduler is technically not `k_` anything, as it is neither
# the k-diffusers implementation nor included in EDM (Karras 2022), but we can
# provide an alias for compatibility.
k_dpmpp_2=DPMSolverMultistepScheduler,
k_euler=EulerDiscreteScheduler,
k_euler_a=EulerAncestralDiscreteScheduler,
k_heun=HeunDiscreteScheduler,
k_lms=LMSDiscreteScheduler,
plms=PNDMScheduler,
k_dpmpp_2=DPMSolverMultistepScheduler,
)
if self.sampler_name in scheduler_map:

View File

@ -407,7 +407,7 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline):
batch_size = initial_latents.size(0)
img2img_pipeline = StableDiffusionImg2ImgPipeline(**self.components)
img2img_pipeline.scheduler.set_timesteps(num_inference_steps, device=device)
timesteps = img2img_pipeline.get_timesteps(num_inference_steps, strength, device=device)
timesteps, _ = img2img_pipeline.get_timesteps(num_inference_steps, strength, device=device)
latent_timestep = timesteps[:1].repeat(batch_size)
latents = self.noise_latents_for_time(initial_latents, latent_timestep, noise_func=noise_func)
@ -454,7 +454,7 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline):
img2img_pipeline = StableDiffusionImg2ImgPipeline(**self.components)
img2img_pipeline.scheduler.set_timesteps(num_inference_steps, device=device)
timesteps = img2img_pipeline.get_timesteps(num_inference_steps, strength, device=device)
timesteps, _ = img2img_pipeline.get_timesteps(num_inference_steps, strength, device=device)
assert img2img_pipeline.scheduler is self.scheduler