(
*,
steps=1000,
learn_sigma=False,
sigma_small=False,
noise_schedule="linear",
use_kl=False,
predict_xstart=False,
rescale_timesteps=False,
rescale_learned_sigmas=False,
timestep_respacing="",
)
| 20 | |
| 21 | |
| 22 | def create_gaussian_diffusion( |
| 23 | *, |
| 24 | steps=1000, |
| 25 | learn_sigma=False, |
| 26 | sigma_small=False, |
| 27 | noise_schedule="linear", |
| 28 | use_kl=False, |
| 29 | predict_xstart=False, |
| 30 | rescale_timesteps=False, |
| 31 | rescale_learned_sigmas=False, |
| 32 | timestep_respacing="", |
| 33 | ): |
| 34 | betas = gd.get_named_beta_schedule(noise_schedule, steps) |
| 35 | if use_kl: |
| 36 | loss_type = gd.LossType.RESCALED_KL |
| 37 | elif rescale_learned_sigmas: |
| 38 | loss_type = gd.LossType.RESCALED_MSE |
| 39 | else: |
| 40 | loss_type = gd.LossType.MSE |
| 41 | if not timestep_respacing: |
| 42 | timestep_respacing = [steps] |
| 43 | return SpacedDiffusion( |
| 44 | use_timesteps=space_timesteps(steps, timestep_respacing), |
| 45 | betas=betas, |
| 46 | model_mean_type=( |
| 47 | gd.ModelMeanType.EPSILON if not predict_xstart else gd.ModelMeanType.START_X |
| 48 | ), |
| 49 | model_var_type=( |
| 50 | ( |
| 51 | gd.ModelVarType.FIXED_LARGE |
| 52 | if not sigma_small |
| 53 | else gd.ModelVarType.FIXED_SMALL |
| 54 | ) |
| 55 | if not learn_sigma |
| 56 | else gd.ModelVarType.LEARNED_RANGE |
| 57 | ), |
| 58 | loss_type=loss_type, |
| 59 | rescale_timesteps=rescale_timesteps, |
| 60 | ) |
no test coverage detected