Method __init__

diffsynth/models/stepvideo_vae.py:467–476

(self,
        in_channels
    )

class AttnBlock(nn.Module):
    def __init__(self,
        in_channels
    ):
        super().__init__()

        self.norm = BaseGroupNorm(num_groups=32, num_channels=in_channels)
        self.q        = CausalConvChannelLast(in_channels, in_channels, kernel_size=1)
        self.k        = CausalConvChannelLast(in_channels, in_channels, kernel_size=1)
        self.v        = CausalConvChannelLast(in_channels, in_channels, kernel_size=1)
        self.proj_out = CausalConvChannelLast(in_channels, in_channels, kernel_size=1)

    def attention(self, x, is_init=True):
        x = self.norm(x, act_silu=False, channel_last=True)