Spaces:
Running
on
Zero
Running
on
Zero
| # Modified from OpenAI's diffusion repos | |
| # GLIDE: https://github.com/openai/glide-text2im/blob/main/glide_text2im/gaussian_diffusion.py | |
| # ADM: https://github.com/openai/guided-diffusion/blob/main/guided_diffusion | |
| # IDDPM: https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py | |
| from diffusion.model.respace import SpacedDiffusion, space_timesteps | |
| from .model import gaussian_diffusion as gd | |
| def IDDPM( | |
| timestep_respacing, | |
| noise_schedule="linear", | |
| use_kl=False, | |
| sigma_small=False, | |
| predict_xstart=False, | |
| learn_sigma=True, | |
| pred_sigma=True, | |
| rescale_learned_sigmas=False, | |
| diffusion_steps=1000, | |
| snr=False, | |
| return_startx=False, | |
| ): | |
| betas = gd.get_named_beta_schedule(noise_schedule, diffusion_steps) | |
| if use_kl: | |
| loss_type = gd.LossType.RESCALED_KL | |
| elif rescale_learned_sigmas: | |
| loss_type = gd.LossType.RESCALED_MSE | |
| else: | |
| loss_type = gd.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.ModelMeanType.EPSILON if not predict_xstart else gd.ModelMeanType.START_X | |
| ), | |
| model_var_type=( | |
| (( | |
| gd.ModelVarType.FIXED_LARGE | |
| if not sigma_small | |
| else gd.ModelVarType.FIXED_SMALL | |
| ) | |
| if not learn_sigma | |
| else gd.ModelVarType.LEARNED_RANGE | |
| ) | |
| if pred_sigma | |
| else None | |
| ), | |
| loss_type=loss_type, | |
| snr=snr, | |
| return_startx=return_startx, | |
| # rescale_timesteps=rescale_timesteps, | |
| ) |