automatic/modules/sd_samplers_diffusors.py

46 lines
2.8 KiB
Python

from diffusers import (
DDIMScheduler,
DDPMScheduler,
DEISMultistepScheduler,
DPMSolverMultistepScheduler,
EulerAncestralDiscreteScheduler,
EulerDiscreteScheduler,
HeunDiscreteScheduler,
IPNDMScheduler,
KDPM2AncestralDiscreteScheduler,
PNDMScheduler,
UniPCMultistepScheduler,
# KarrasVeScheduler,
# RePaintScheduler,
# ScoreSdeVeScheduler,
# UnCLIPScheduler,
# VQDiffusionScheduler,
)
from modules import sd_samplers_common
# scheduler = diffusers.UniPCMultistepScheduler.from_pretrained(shared.cmd_opts.ckpt, subfolder="scheduler")
samplers_data_diffusors = [
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('DDPMS', lambda model: DiffusionSampler('DDPMS', DDPMScheduler, model), [], {}),
sd_samplers_common.SamplerData('DEIS', lambda model: DiffusionSampler('DEIS', DEISMultistepScheduler, model), [], {}),
sd_samplers_common.SamplerData('DPMSolver', lambda model: DiffusionSampler('DPMSolver', DPMSolverMultistepScheduler, model), [], {}),
sd_samplers_common.SamplerData('Euler', lambda model: DiffusionSampler('Euler', EulerDiscreteScheduler, model), [], {}),
sd_samplers_common.SamplerData('EulerAncestral', lambda model: DiffusionSampler('EulerAncestral', EulerAncestralDiscreteScheduler, model), [], {}),
sd_samplers_common.SamplerData('Heun', lambda model: DiffusionSampler('Heun', HeunDiscreteScheduler, model), [], {}),
sd_samplers_common.SamplerData('IPNDM', lambda model: DiffusionSampler('IPNDM', IPNDMScheduler, model), [], {}),
sd_samplers_common.SamplerData('KDPM2Ancestral', lambda model: DiffusionSampler('KDPM2Ancestral', KDPM2AncestralDiscreteScheduler, model), [], {}),
sd_samplers_common.SamplerData('PNDMS', lambda model: DiffusionSampler('PNDMS', PNDMScheduler, model), [], {}),
# sd_samplers_common.SamplerData('KarrasVe', lambda model: DiffusionSampler('KarrasVe', KarrasVeScheduler, model), [], {}),
# sd_samplers_common.SamplerData('RePaint', lambda model: DiffusionSampler('RePaint', RePaintScheduler, model), [], {}),
# sd_samplers_common.SamplerData('ScoreSdeVe', lambda model: DiffusionSampler('ScoreSdeVe', ScoreSdeVeScheduler, model), [], {}),
# sd_samplers_common.SamplerData('UnCLIP', lambda model: DiffusionSampler('UnCLIP', UnCLIPScheduler, model), [], {}),
# sd_samplers_common.SamplerData('VQDiffusion', lambda model: DiffusionSampler('VQDiffusion', VQDiffusionScheduler, model), [], {}),
]
class DiffusionSampler:
def __init__(self, name, constructor, sd_model):
self.sampler = constructor.from_pretrained(sd_model, subfolder="scheduler")
self.sampler.name = name