(beta_type: str = None,
use_dpm_solver: bool = False)
| 91 | |
| 92 | |
| 93 | def get_scheduler(beta_type: str = None, |
| 94 | use_dpm_solver: bool = False) -> DDIMScheduler: |
| 95 | ddim_kwargs = dict( |
| 96 | num_train_timesteps=1000, |
| 97 | beta_start=0.00085, |
| 98 | beta_end=0.012, |
| 99 | beta_schedule='scaled_linear', |
| 100 | clip_sample=False, |
| 101 | set_alpha_to_one=True, |
| 102 | steps_offset=1, |
| 103 | rescale_betas_zero_snr=True, |
| 104 | timestep_spacing='trailing', |
| 105 | prediction_type='v_prediction', |
| 106 | ) |
| 107 | if 'cogvideo' not in beta_type: |
| 108 | return DDIMScheduler(**ddim_kwargs) |
| 109 | else: |
| 110 | snr_shift_scale = 3.0 if '2b' in beta_type else 1.0 |
| 111 | ddim_kwargs['steps_offset'] = 0 |
| 112 | ddim_kwargs['snr_shift_scale'] = snr_shift_scale |
| 113 | if use_dpm_solver: |
| 114 | return CogVideoXDPMScheduler(**ddim_kwargs) |
| 115 | else: |
| 116 | return CogVideoXDDIMScheduler(**ddim_kwargs) |
| 117 | |
| 118 | |
| 119 | def get_smooth_dynamic_cfg_scale_list(cfg_scale, num_inference_steps): |
nothing calls this directly
no test coverage detected