| 770 | # --------------------------------------------------------------------------------------------------- |
| 771 | |
| 772 | class STAttentionBlock2(nn.Module): |
| 773 | def __init__( |
| 774 | self, |
| 775 | channels, |
| 776 | num_heads=1, |
| 777 | num_head_channels=-1, |
| 778 | use_checkpoint=False, # not used, only used in ResBlock |
| 779 | use_new_attention_order=False, # QKVAttention or QKVAttentionLegacy |
| 780 | temporal_length=16, # used in relative positional representation. |
| 781 | image_length=8, # used for image-video joint training. |
| 782 | use_relative_position=False, # whether use relative positional representation in temporal attention. |
| 783 | img_video_joint_train=False, |
| 784 | # norm_type="groupnorm", |
| 785 | attn_norm_type="group", |
| 786 | use_tempoal_causal_attn=False, |
| 787 | ): |
| 788 | """ |
| 789 | version 1: guided_diffusion implemented version |
| 790 | version 2: remove args input argument |
| 791 | """ |
| 792 | super().__init__() |
| 793 | |
| 794 | if num_head_channels == -1: |
| 795 | self.num_heads = num_heads |
| 796 | else: |
| 797 | assert ( |
| 798 | channels % num_head_channels == 0 |
| 799 | ), f"q,k,v channels {channels} is not divisible by num_head_channels {num_head_channels}" |
| 800 | self.num_heads = channels // num_head_channels |
| 801 | self.use_checkpoint = use_checkpoint |
| 802 | |
| 803 | self.temporal_length = temporal_length |
| 804 | self.image_length = image_length |
| 805 | self.use_relative_position = use_relative_position |
| 806 | self.img_video_joint_train = img_video_joint_train |
| 807 | self.attn_norm_type = attn_norm_type |
| 808 | assert(self.attn_norm_type in ["group", "no_norm"]) |
| 809 | self.use_tempoal_causal_attn = use_tempoal_causal_attn |
| 810 | |
| 811 | if self.attn_norm_type == "group": |
| 812 | self.norm_s = normalization(channels) |
| 813 | self.norm_t = normalization(channels) |
| 814 | |
| 815 | self.qkv_s = conv_nd(1, channels, channels * 3, 1) |
| 816 | self.qkv_t = conv_nd(1, channels, channels * 3, 1) |
| 817 | |
| 818 | if self.img_video_joint_train: |
| 819 | mask = th.ones([1, temporal_length+image_length, temporal_length+image_length]) |
| 820 | mask[:, temporal_length:, :] = 0 |
| 821 | mask[:, :, temporal_length:] = 0 |
| 822 | self.register_buffer("mask", mask) |
| 823 | else: |
| 824 | self.mask = None |
| 825 | |
| 826 | if use_new_attention_order: |
| 827 | # split qkv before split heads |
| 828 | self.attention_s = QKVAttention(self.num_heads) |
| 829 | self.attention_t = QKVAttention(self.num_heads) |
nothing calls this directly
no outgoing calls
no test coverage detected