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
})