The half UNet model with attention and timestep embedding. For usage, see UNet.
| 786 | |
| 787 | |
| 788 | class EncoderUNetModel(nn.Module): |
| 789 | """ |
| 790 | The half UNet model with attention and timestep embedding. |
| 791 | For usage, see UNet. |
| 792 | """ |
| 793 | |
| 794 | def __init__( |
| 795 | self, |
| 796 | image_size, |
| 797 | in_channels, |
| 798 | model_channels, |
| 799 | out_channels, |
| 800 | num_res_blocks, |
| 801 | attention_resolutions, |
| 802 | dropout=0, |
| 803 | channel_mult=(1, 2, 4, 8), |
| 804 | conv_resample=True, |
| 805 | dims=2, |
| 806 | use_checkpoint=False, |
| 807 | use_fp16=False, |
| 808 | num_heads=1, |
| 809 | num_head_channels=-1, |
| 810 | num_heads_upsample=-1, |
| 811 | use_scale_shift_norm=False, |
| 812 | resblock_updown=False, |
| 813 | use_new_attention_order=False, |
| 814 | pool="adaptive", |
| 815 | *args, |
| 816 | **kwargs |
| 817 | ): |
| 818 | super().__init__() |
| 819 | |
| 820 | if num_heads_upsample == -1: |
| 821 | num_heads_upsample = num_heads |
| 822 | |
| 823 | self.in_channels = in_channels |
| 824 | self.model_channels = model_channels |
| 825 | self.out_channels = out_channels |
| 826 | self.num_res_blocks = num_res_blocks |
| 827 | self.attention_resolutions = attention_resolutions |
| 828 | self.dropout = dropout |
| 829 | self.channel_mult = channel_mult |
| 830 | self.conv_resample = conv_resample |
| 831 | self.use_checkpoint = use_checkpoint |
| 832 | self.dtype = th.float16 if use_fp16 else th.float32 |
| 833 | self.num_heads = num_heads |
| 834 | self.num_head_channels = num_head_channels |
| 835 | self.num_heads_upsample = num_heads_upsample |
| 836 | |
| 837 | time_embed_dim = model_channels * 4 |
| 838 | self.time_embed = nn.Sequential( |
| 839 | linear(model_channels, time_embed_dim), |
| 840 | nn.SiLU(), |
| 841 | linear(time_embed_dim, time_embed_dim), |
| 842 | ) |
| 843 | |
| 844 | self.input_blocks = nn.ModuleList( |
| 845 | [ |
nothing calls this directly
no outgoing calls
no test coverage detected