(
scheduler,
num_inference_steps=None,
device=None,
timesteps=None,
sigmas=None,
**kwargs,
)
| 27 | |
| 28 | |
| 29 | def retrieve_timesteps( |
| 30 | scheduler, |
| 31 | num_inference_steps=None, |
| 32 | device=None, |
| 33 | timesteps=None, |
| 34 | sigmas=None, |
| 35 | **kwargs, |
| 36 | ): |
| 37 | if timesteps is not None and sigmas is not None: |
| 38 | raise ValueError( |
| 39 | "Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values" |
| 40 | ) |
| 41 | if timesteps is not None: |
| 42 | accepts_timesteps = "timesteps" in set( |
| 43 | inspect.signature(scheduler.set_timesteps).parameters.keys()) |
| 44 | if not accepts_timesteps: |
| 45 | raise ValueError( |
| 46 | f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" |
| 47 | f" timestep schedules. Please check whether you are using the correct scheduler." |
| 48 | ) |
| 49 | scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) |
| 50 | timesteps = scheduler.timesteps |
| 51 | num_inference_steps = len(timesteps) |
| 52 | elif sigmas is not None: |
| 53 | accept_sigmas = "sigmas" in set( |
| 54 | inspect.signature(scheduler.set_timesteps).parameters.keys()) |
| 55 | if not accept_sigmas: |
| 56 | raise ValueError( |
| 57 | f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" |
| 58 | f" sigmas schedules. Please check whether you are using the correct scheduler." |
| 59 | ) |
| 60 | scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) |
| 61 | timesteps = scheduler.timesteps |
| 62 | num_inference_steps = len(timesteps) |
| 63 | else: |
| 64 | scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) |
| 65 | timesteps = scheduler.timesteps |
| 66 | return timesteps, num_inference_steps |
| 67 | |
| 68 | |
| 69 | class FlowDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin): |
nothing calls this directly
no test coverage detected