Coverage for widitapp/diffusion/__init__.py: 100.00%

1# Modified from OpenAI's diffusion repos 

2# GLIDE: https://github.com/openai/glide-text2im/blob/main/glide_text2im/gaussian_diffusion.py 

3# ADM: https://github.com/openai/guided-diffusion/blob/main/guided_diffusion 

4# IDDPM: https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py 

5 

6from . import gaussian_diffusion as gd 

7from .respace import SpacedDiffusion, space_timesteps 

8 

9 

10def 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 )