pe.api.image.improved_diffusion_lib.unet module

This code contains minor edits from the original code at https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/unet.py and https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/script_util.py to avoid calling self.input_blocks.parameters() in the original code, which is not supported by DataParallel.

pe.api.image.improved_diffusion_lib.unet.create_model(image_size, num_channels, num_res_blocks, learn_sigma, class_cond, use_checkpoint, attention_resolutions, num_heads, num_heads_upsample, use_scale_shift_norm, dropout)