(
image_size,
classifier_use_fp16,
classifier_width,
classifier_depth,
classifier_attention_resolutions,
classifier_use_scale_shift_norm,
classifier_resblock_updown,
classifier_pool,
)
| 236 | |
| 237 | |
| 238 | def create_classifier( |
| 239 | image_size, |
| 240 | classifier_use_fp16, |
| 241 | classifier_width, |
| 242 | classifier_depth, |
| 243 | classifier_attention_resolutions, |
| 244 | classifier_use_scale_shift_norm, |
| 245 | classifier_resblock_updown, |
| 246 | classifier_pool, |
| 247 | ): |
| 248 | if image_size == 512: |
| 249 | channel_mult = (0.5, 1, 1, 2, 2, 4, 4) |
| 250 | elif image_size == 256: |
| 251 | channel_mult = (1, 1, 2, 2, 4, 4) |
| 252 | elif image_size == 128: |
| 253 | channel_mult = (1, 1, 2, 3, 4) |
| 254 | elif image_size == 64: |
| 255 | channel_mult = (1, 2, 3, 4) |
| 256 | else: |
| 257 | raise ValueError(f"unsupported image size: {image_size}") |
| 258 | |
| 259 | attention_ds = [] |
| 260 | for res in classifier_attention_resolutions.split(","): |
| 261 | attention_ds.append(image_size // int(res)) |
| 262 | |
| 263 | return EncoderUNetModel( |
| 264 | image_size=image_size, |
| 265 | in_channels=3, |
| 266 | model_channels=classifier_width, |
| 267 | out_channels=1000, |
| 268 | num_res_blocks=classifier_depth, |
| 269 | attention_resolutions=tuple(attention_ds), |
| 270 | channel_mult=channel_mult, |
| 271 | use_fp16=classifier_use_fp16, |
| 272 | num_head_channels=64, |
| 273 | use_scale_shift_norm=classifier_use_scale_shift_norm, |
| 274 | resblock_updown=classifier_resblock_updown, |
| 275 | pool=classifier_pool, |
| 276 | ) |
| 277 | |
| 278 | |
| 279 | def sr_model_and_diffusion_defaults(): |
no test coverage detected