File size: 1,485 Bytes
8eca2ce
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
from diffusers import DPMSolverMultistepScheduler, DPMSolverSinglestepScheduler, FlowMatchEulerDiscreteScheduler

schedulers = {"Euler": (FlowMatchEulerDiscreteScheduler,
                        {}),
              "DPM++ 2M": (DPMSolverMultistepScheduler,
                           {}),
              "DPM++ 2M Karras": (DPMSolverMultistepScheduler,
                                  {"use_karras_sigmas": True}),
              "DPM++ 2M SDE": (DPMSolverMultistepScheduler,
                               {"algorithm_type": "sde-dpmsolver++"},
                               ),
              "DPM++ 2M SDE Karras": (DPMSolverMultistepScheduler,
                                      {"use_karras_sigmas": True,
                                       "algorithm_type": "sde-dpmsolver++"},
                                      ),
              "DPM++ SDE": (DPMSolverSinglestepScheduler,
                            {}),
              "DPM++ SDE Karras": (DPMSolverSinglestepScheduler,
                                   {"use_karras_sigmas": True}),
              }


def apply_scheduler(name, pipeline):
    # Check if the scheduler name exists in the schedulers dictionary
    if name in schedulers:
        scheduler_class, config = schedulers[name]
        # Update the pipeline's scheduler with the selected scheduler class and
        # configuration
        pipeline.scheduler = scheduler_class.from_config(
            pipeline.scheduler.config, **config)
    return pipeline