The full UNet model with attention and timestep embedding. :param in_channels: channels in the input Tensor. :param model_channels: base channel count for the model. :param out_channels: channels in the output Tensor. :param num_res_blocks: number of residual blocks per downsamp
| 423 | |
| 424 | |
| 425 | class UNetModel(nn.Module): |
| 426 | """ |
| 427 | The full UNet model with attention and timestep embedding. |
| 428 | :param in_channels: channels in the input Tensor. |
| 429 | :param model_channels: base channel count for the model. |
| 430 | :param out_channels: channels in the output Tensor. |
| 431 | :param num_res_blocks: number of residual blocks per downsample. |
| 432 | :param attention_resolutions: a collection of downsample rates at which |
| 433 | attention will take place. May be a set, list, or tuple. |
| 434 | For example, if this contains 4, then at 4x downsampling, attention |
| 435 | will be used. |
| 436 | :param dropout: the dropout probability. |
| 437 | :param channel_mult: channel multiplier for each level of the UNet. |
| 438 | :param conv_resample: if True, use learned convolutions for upsampling and |
| 439 | downsampling. |
| 440 | :param dims: determines if the signal is 1D, 2D, or 3D. |
| 441 | :param num_classes: if specified (as an int), then this model will be |
| 442 | class-conditional with `num_classes` classes. |
| 443 | :param use_checkpoint: use gradient checkpointing to reduce memory usage. |
| 444 | :param num_heads: the number of attention heads in each attention layer. |
| 445 | :param num_heads_channels: if specified, ignore num_heads and instead use |
| 446 | a fixed channel width per attention head. |
| 447 | :param num_heads_upsample: works with num_heads to set a different number |
| 448 | of heads for upsampling. Deprecated. |
| 449 | :param use_scale_shift_norm: use a FiLM-like conditioning mechanism. |
| 450 | :param resblock_updown: use residual blocks for up/downsampling. |
| 451 | :param use_new_attention_order: use a different attention pattern for potentially |
| 452 | increased efficiency. |
| 453 | """ |
| 454 | def __init__( |
| 455 | self, |
| 456 | # ------------------- DEFAULTS FROM Stable Diffusion 2-1 |
| 457 | image_size=32, |
| 458 | in_channels=4, |
| 459 | out_channels=4, |
| 460 | model_channels=320, |
| 461 | attention_resolutions=[ 4, 2, 1 ], |
| 462 | num_res_blocks=2, |
| 463 | channel_mult=[ 1, 2, 4, 4 ], |
| 464 | num_heads=-1, |
| 465 | num_head_channels=64, |
| 466 | use_spatial_transformer=True, # custom transformer support |
| 467 | use_linear_in_transformer=True, |
| 468 | transformer_depth=1, # custom transformer support |
| 469 | context_dim=1024, # custom transformer support |
| 470 | legacy=False, |
| 471 | load_from_ckpt=None, |
| 472 | # ---------------------------------------------------- |
| 473 | dropout=0, |
| 474 | conv_resample=True, |
| 475 | dims=2, |
| 476 | num_classes=None, |
| 477 | use_checkpoint=False, |
| 478 | use_fp16=False, |
| 479 | use_bf16=False, |
| 480 | num_heads_upsample=-1, |
| 481 | use_scale_shift_norm=False, |
| 482 | resblock_updown=False, |