Spaces:
Runtime error
Runtime error
File size: 1,798 Bytes
cab0202 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 |
from diffusers import (
DDIMInverseScheduler,
DDIMScheduler,
DDPMScheduler,
DEISMultistepScheduler,
DPMSolverMultistepInverseScheduler,
DPMSolverMultistepScheduler,
DPMSolverSinglestepScheduler,
EulerAncestralDiscreteScheduler,
EulerDiscreteScheduler,
HeunDiscreteScheduler,
IPNDMScheduler,
KarrasVeScheduler,
KDPM2AncestralDiscreteScheduler,
KDPM2DiscreteScheduler,
PNDMScheduler,
RePaintScheduler,
SchedulerMixin,
ScoreSdeVeScheduler,
UnCLIPScheduler,
UniPCMultistepScheduler,
VQDiffusionScheduler,
)
SCHEDULER_MAPPING = {
"DDIM": DDIMScheduler,
"DDIMInverse": DDIMInverseScheduler,
"DDPMScheduler": DDPMScheduler,
"DEISMultistep": DEISMultistepScheduler,
"DPMSolverMultistepInverse": DPMSolverMultistepInverseScheduler,
"DPMSolverMultistep": DPMSolverMultistepScheduler,
"DPMSolverSinglestep": DPMSolverSinglestepScheduler,
"EulerAncestralDiscrete": EulerAncestralDiscreteScheduler,
"EulerDiscrete": EulerDiscreteScheduler,
"HeunDiscrete": HeunDiscreteScheduler,
"IPNDMScheduler": IPNDMScheduler,
"KarrasVe": KarrasVeScheduler,
"KDPM2AncestralDiscrete": KDPM2AncestralDiscreteScheduler,
"KDPM2Discrete": KDPM2DiscreteScheduler,
"PNDMScheduler": PNDMScheduler,
"RePaint": RePaintScheduler,
"ScoreSdeVe": ScoreSdeVeScheduler,
"UnCLIP": UnCLIPScheduler,
"UniPCMultistep": UniPCMultistepScheduler,
"VQDiffusion": VQDiffusionScheduler,
}
def get_scheduler(pipe, scheduler):
if scheduler in SCHEDULER_MAPPING:
SchedulerClass = SCHEDULER_MAPPING[scheduler]
pipe.scheduler = SchedulerClass.from_config(pipe.scheduler.config)
else:
raise ValueError(f"Invalid scheduler name {scheduler}")
return pipe
|