(
image_size,
num_channels,
num_res_blocks,
channel_mult="",
learn_sigma=False,
class_cond=False,
rep_cond=False,
rep_dim=256,
use_checkpoint=False,
attention_resolutions="16",
num_heads=1,
num_head_channels=-1,
num_heads_upsample=-1,
use_scale_shift_norm=False,
dropout=0,
resblock_updown=False,
use_fp16=False,
use_new_attention_order=False,
)
| 134 | |
| 135 | |
| 136 | def create_model( |
| 137 | image_size, |
| 138 | num_channels, |
| 139 | num_res_blocks, |
| 140 | channel_mult="", |
| 141 | learn_sigma=False, |
| 142 | class_cond=False, |
| 143 | rep_cond=False, |
| 144 | rep_dim=256, |
| 145 | use_checkpoint=False, |
| 146 | attention_resolutions="16", |
| 147 | num_heads=1, |
| 148 | num_head_channels=-1, |
| 149 | num_heads_upsample=-1, |
| 150 | use_scale_shift_norm=False, |
| 151 | dropout=0, |
| 152 | resblock_updown=False, |
| 153 | use_fp16=False, |
| 154 | use_new_attention_order=False, |
| 155 | ): |
| 156 | if channel_mult == "": |
| 157 | if image_size == 512: |
| 158 | channel_mult = (0.5, 1, 1, 2, 2, 4, 4) |
| 159 | elif image_size == 256: |
| 160 | channel_mult = (1, 1, 2, 2, 4, 4) |
| 161 | elif image_size == 128: |
| 162 | channel_mult = (1, 1, 2, 3, 4) |
| 163 | elif image_size == 64: |
| 164 | channel_mult = (1, 2, 3, 4) |
| 165 | else: |
| 166 | raise ValueError(f"unsupported image size: {image_size}") |
| 167 | else: |
| 168 | channel_mult = tuple(int(ch_mult) for ch_mult in channel_mult.split(",")) |
| 169 | |
| 170 | attention_ds = [] |
| 171 | for res in attention_resolutions.split(","): |
| 172 | attention_ds.append(image_size // int(res)) |
| 173 | |
| 174 | return UNetModel( |
| 175 | image_size=image_size, |
| 176 | in_channels=3, |
| 177 | model_channels=num_channels, |
| 178 | out_channels=(3 if not learn_sigma else 6), |
| 179 | num_res_blocks=num_res_blocks, |
| 180 | attention_resolutions=tuple(attention_ds), |
| 181 | dropout=dropout, |
| 182 | channel_mult=channel_mult, |
| 183 | num_classes=(NUM_CLASSES if class_cond else None), |
| 184 | rep_cond=rep_cond, |
| 185 | rep_dim=rep_dim, |
| 186 | use_checkpoint=use_checkpoint, |
| 187 | use_fp16=use_fp16, |
| 188 | num_heads=num_heads, |
| 189 | num_head_channels=num_head_channels, |
| 190 | num_heads_upsample=num_heads_upsample, |
| 191 | use_scale_shift_norm=use_scale_shift_norm, |
| 192 | resblock_updown=resblock_updown, |
| 193 | use_new_attention_order=use_new_attention_order, |
no test coverage detected