Files
automatic/modules/sd_samplers_diffusers.py
T
2023-07-02 14:04:54 -04:00

38 lines
2.7 KiB
Python

from diffusers import (
DDIMScheduler,
DDPMScheduler,
DEISMultistepScheduler,
DPMSolverMultistepScheduler,
DPMSolverSinglestepScheduler,
EulerAncestralDiscreteScheduler,
EulerDiscreteScheduler,
HeunDiscreteScheduler,
KDPM2DiscreteScheduler,
PNDMScheduler,
UniPCMultistepScheduler,
)
from modules import sd_samplers_common
samplers_data_diffusers = [
sd_samplers_common.SamplerData('UniPC', lambda model: DiffusionSampler('UniPC', UniPCMultistepScheduler, model), [], {}),
sd_samplers_common.SamplerData('DDIM', lambda model: DiffusionSampler('DDIM', DDIMScheduler, model), [], {}),
sd_samplers_common.SamplerData('DDPM', lambda model: DiffusionSampler('DDPM', DDPMScheduler, model), [], {}),
sd_samplers_common.SamplerData('DEIS', lambda model: DiffusionSampler('DEIS', DEISMultistepScheduler, model), [], {}),
sd_samplers_common.SamplerData('DPM++ 2M', lambda model: DiffusionSampler('DPM++ 2M', DPMSolverMultistepScheduler, model), [], {}),
sd_samplers_common.SamplerData('DPM++ 1S', lambda model: DiffusionSampler('DPM++ 1S', DPMSolverSinglestepScheduler, model), [], {}),
sd_samplers_common.SamplerData('DPM++ 2M SDE', lambda model: DiffusionSampler('DPM++ 2M SDE', DPMSolverMultistepScheduler, model, algorithm_type="sde-dpmsolver++"), [], {}),
sd_samplers_common.SamplerData('DPM++ 2M Karras', lambda model: DiffusionSampler('DPM++ 2M Karras', DPMSolverMultistepScheduler, model, use_karras_sigmas=True), [], {}),
sd_samplers_common.SamplerData('DPM++ 1S Karras', lambda model: DiffusionSampler('DPM++ 1S Karras', DPMSolverSinglestepScheduler, model, use_karras_sigmas=True), [], {}),
sd_samplers_common.SamplerData('DPM++ 2M SDE Karras', lambda model: DiffusionSampler('DPM++ 2M SDE Karras', DPMSolverMultistepScheduler, model, use_karras_sigmas=True, algorithm_type="sde-dpmsolver++"), [], {}),
sd_samplers_common.SamplerData('Euler', lambda model: DiffusionSampler('Euler', EulerDiscreteScheduler, model), [], {}),
sd_samplers_common.SamplerData('Euler a', lambda model: DiffusionSampler('Euler a', EulerAncestralDiscreteScheduler, model), [], {}),
sd_samplers_common.SamplerData('Heun', lambda model: DiffusionSampler('Heun', HeunDiscreteScheduler, model), [], {}),
sd_samplers_common.SamplerData('DPM2++ 2M', lambda model: DiffusionSampler('KDPM2', KDPM2DiscreteScheduler, model), [], {}),
sd_samplers_common.SamplerData('PNDM', lambda model: DiffusionSampler('PNDM', PNDMScheduler, model), [], {}),
]
class DiffusionSampler:
def __init__(self, name, constructor, sd_model, **kwargs):
self.sampler = constructor.from_pretrained(sd_model, subfolder="scheduler", **kwargs)
self.sampler.name = name