The half UNet model with attention and timestep embedding. For usage, see UNet.
| 692 | |
| 693 | |
| 694 | class EncoderUNetModel(nn.Module): |
| 695 | """ |
| 696 | The half UNet model with attention and timestep embedding. |
| 697 | |
| 698 | For usage, see UNet. |
| 699 | """ |
| 700 | |
| 701 | def __init__( |
| 702 | self, |
| 703 | image_size, |
| 704 | in_channels, |
| 705 | model_channels, |
| 706 | out_channels, |
| 707 | num_res_blocks, |
| 708 | attention_resolutions, |
| 709 | dropout=0, |
| 710 | channel_mult=(1, 2, 4, 8), |
| 711 | conv_resample=True, |
| 712 | dims=2, |
| 713 | use_checkpoint=False, |
| 714 | use_fp16=False, |
| 715 | num_heads=1, |
| 716 | num_head_channels=-1, |
| 717 | num_heads_upsample=-1, |
| 718 | use_scale_shift_norm=False, |
| 719 | resblock_updown=False, |
| 720 | use_new_attention_order=False, |
| 721 | pool="adaptive", |
| 722 | ): |
| 723 | super().__init__() |
| 724 | |
| 725 | if num_heads_upsample == -1: |
| 726 | num_heads_upsample = num_heads |
| 727 | |
| 728 | self.in_channels = in_channels |
| 729 | self.model_channels = model_channels |
| 730 | self.out_channels = out_channels |
| 731 | self.num_res_blocks = num_res_blocks |
| 732 | self.attention_resolutions = attention_resolutions |
| 733 | self.dropout = dropout |
| 734 | self.channel_mult = channel_mult |
| 735 | self.conv_resample = conv_resample |
| 736 | self.use_checkpoint = use_checkpoint |
| 737 | self.dtype = th.float16 if use_fp16 else th.float32 |
| 738 | self.num_heads = num_heads |
| 739 | self.num_head_channels = num_head_channels |
| 740 | self.num_heads_upsample = num_heads_upsample |
| 741 | |
| 742 | time_embed_dim = model_channels * 4 |
| 743 | self.time_embed = nn.Sequential( |
| 744 | linear(model_channels, time_embed_dim), |
| 745 | nn.SiLU(), |
| 746 | linear(time_embed_dim, time_embed_dim), |
| 747 | ) |
| 748 | |
| 749 | ch = int(channel_mult[0] * model_channels) |
| 750 | self.input_blocks = nn.ModuleList( |
| 751 | [TimestepEmbedSequential(conv_nd(dims, in_channels, ch, 3, padding=1))] |