(
image_size,
class_cond,
learn_sigma,
sigma_small,
num_channels,
num_res_blocks,
num_heads,
num_heads_upsample,
attention_resolutions,
dropout,
diffusion_steps,
noise_schedule,
timestep_respacing,
use_kl,
predict_xstart,
rescale_timesteps,
rescale_learned_sigmas,
use_checkpoint,
use_scale_shift_norm,
)
| 36 | |
| 37 | |
| 38 | def create_model_and_diffusion( |
| 39 | image_size, |
| 40 | class_cond, |
| 41 | learn_sigma, |
| 42 | sigma_small, |
| 43 | num_channels, |
| 44 | num_res_blocks, |
| 45 | num_heads, |
| 46 | num_heads_upsample, |
| 47 | attention_resolutions, |
| 48 | dropout, |
| 49 | diffusion_steps, |
| 50 | noise_schedule, |
| 51 | timestep_respacing, |
| 52 | use_kl, |
| 53 | predict_xstart, |
| 54 | rescale_timesteps, |
| 55 | rescale_learned_sigmas, |
| 56 | use_checkpoint, |
| 57 | use_scale_shift_norm, |
| 58 | ): |
| 59 | model = create_model( |
| 60 | image_size, |
| 61 | num_channels, |
| 62 | num_res_blocks, |
| 63 | learn_sigma=learn_sigma, |
| 64 | class_cond=class_cond, |
| 65 | use_checkpoint=use_checkpoint, |
| 66 | attention_resolutions=attention_resolutions, |
| 67 | num_heads=num_heads, |
| 68 | num_heads_upsample=num_heads_upsample, |
| 69 | use_scale_shift_norm=use_scale_shift_norm, |
| 70 | dropout=dropout, |
| 71 | ) |
| 72 | diffusion = create_gaussian_diffusion( |
| 73 | steps=diffusion_steps, |
| 74 | learn_sigma=learn_sigma, |
| 75 | sigma_small=sigma_small, |
| 76 | noise_schedule=noise_schedule, |
| 77 | use_kl=use_kl, |
| 78 | predict_xstart=predict_xstart, |
| 79 | rescale_timesteps=rescale_timesteps, |
| 80 | rescale_learned_sigmas=rescale_learned_sigmas, |
| 81 | timestep_respacing=timestep_respacing, |
| 82 | ) |
| 83 | return model, diffusion |
| 84 | |
| 85 | |
| 86 | def create_model( |
no test coverage detected