| from . import gaussian_diffusion as gd_orig |
| from . import gaussian_diffusion_dual as gd_dual |
| |
|
|
|
|
| def create_diffusion( |
| timestep_respacing, |
| noise_schedule="linear", |
| use_kl=False, |
| sigma_small=False, |
| predict_xstart=False, |
| learn_sigma=True, |
| rescale_learned_sigmas=False, |
| diffusion_steps=1000, |
| dual=False |
| ): |
| if dual: |
| print("Using DUAL diffusion") |
| from .respace_dual import SpacedDiffusion, space_timesteps |
| gd_module = gd_dual |
| else: |
| print("Using SINGLE diffusion") |
| from .respace import SpacedDiffusion, space_timesteps |
| gd_module = gd_orig |
|
|
| betas = gd_module.get_named_beta_schedule(noise_schedule, diffusion_steps) |
| |
| if use_kl: |
| loss_type = gd_module.LossType.RESCALED_KL |
| elif rescale_learned_sigmas: |
| loss_type = gd_module.LossType.RESCALED_MSE |
| else: |
| loss_type = gd_module.LossType.MSE |
| if timestep_respacing is None or timestep_respacing == "": |
| timestep_respacing = [diffusion_steps] |
| return SpacedDiffusion( |
| use_timesteps=space_timesteps(diffusion_steps, timestep_respacing), |
| betas=betas, |
| model_mean_type=( |
| gd_module.ModelMeanType.EPSILON if not predict_xstart else gd_module.ModelMeanType.START_X |
| ), |
| model_var_type=( |
| ( |
| gd_module.ModelVarType.FIXED_LARGE |
| if not sigma_small |
| else gd_module.ModelVarType.FIXED_SMALL |
| ) |
| if not learn_sigma |
| else gd_module.ModelVarType.LEARNED_RANGE |
| ), |
| loss_type=loss_type |
| |
| ) |
|
|
