| model_config.update({ | |
| 'attention_resolutions': '32, 16, 8', | |
| 'class_cond': False, | |
| 'rescale_timesteps': True, | |
| 'image_size': 256, | |
| 'learn_sigma': True, | |
| 'noise_schedule': 'linear', | |
| 'num_channels': 256, | |
| 'num_head_channels': 64, | |
| 'num_res_blocks': 2, | |
| 'resblock_updown': True, | |
| 'use_checkpoint': use_checkpoint, | |
| 'use_fp16': True, | |
| 'use_scale_shift_norm': True | |
| }) |