(
timestep_respacing,
noise_schedule="linear",
use_kl=False,
sigma_small=False,
predict_xstart=False,
learn_sigma=True,
rescale_learned_sigmas=False,
diffusion_steps=1000
)
| 8 | |
| 9 | |
| 10 | def create_diffusion( |
| 11 | timestep_respacing, |
| 12 | noise_schedule="linear", |
| 13 | use_kl=False, |
| 14 | sigma_small=False, |
| 15 | predict_xstart=False, |
| 16 | learn_sigma=True, |
| 17 | rescale_learned_sigmas=False, |
| 18 | diffusion_steps=1000 |
| 19 | ): |
| 20 | betas = gd.get_named_beta_schedule(noise_schedule, diffusion_steps) |
| 21 | if use_kl: |
| 22 | loss_type = gd.LossType.RESCALED_KL |
| 23 | elif rescale_learned_sigmas: |
| 24 | loss_type = gd.LossType.RESCALED_MSE |
| 25 | else: |
| 26 | loss_type = gd.LossType.MSE |
| 27 | if timestep_respacing is None or timestep_respacing == "": |
| 28 | timestep_respacing = [diffusion_steps] |
| 29 | return SpacedDiffusion( |
| 30 | use_timesteps=space_timesteps(diffusion_steps, timestep_respacing), |
| 31 | betas=betas, |
| 32 | model_mean_type=( |
| 33 | gd.ModelMeanType.EPSILON if not predict_xstart else gd.ModelMeanType.START_X |
| 34 | ), |
| 35 | model_var_type=( |
| 36 | ( |
| 37 | gd.ModelVarType.FIXED_LARGE |
| 38 | if not sigma_small |
| 39 | else gd.ModelVarType.FIXED_SMALL |
| 40 | ) |
| 41 | if not learn_sigma |
| 42 | else gd.ModelVarType.LEARNED_RANGE |
| 43 | ), |
| 44 | loss_type=loss_type |
| 45 | # rescale_timesteps=rescale_timesteps, |
| 46 | ) |
no test coverage detected