Class ResAttBlock

model/transformer_utils.py:256–299

Attention block.

        return attn_output, None

class ResAttBlock(nn.Module):
    """
    Attention block.
    """
    def __init__(self, d_model, n_head, window_size=None, drop_path_rate=0.0):
        super().__init__()
        self.attn = MultiHeadAttention(d_model, d_model, d_model, d_model, n_head)
        self.layernorm1 = LayerNorm(d_model)
        self.mlp = nn.Sequential(OrderedDict([
            ("c_fc", nn.Linear(d_model, d_model * 4, bias=False)),
            ("silu", nn.SiLU(inplace=True)),
            ("c_proj", nn.Linear(d_model * 4, d_model, bias=False))
        ]))
        self.layernorm2 = LayerNorm(d_model)
        self.window_size = window_size

    def attention(self, x, index):
        attn_mask = None
        if self.window_size is not None:
            l = x.shape[1]
            assert l % self.window_size == 0
            if index % 2 == 0:
                x = rearrange(x, 'b (p w) c -> (b p) w c', w=self.window_size)
                x = self.attn(x, x, x, need_weights=False, attn_mask=attn_mask)[0] 
                x = rearrange(x, '(b l) w c -> b (l w) c', l=l//self.window_size, w=self.window_size)
            else:
                x = torch.roll(x, shifts=self.window_size//2, dims=1)
                x = rearrange(x, 'b (p w) c -> (b p) w c', w=self.window_size)
                x = self.attn(x, x, x, need_weights=False, attn_mask=attn_mask)[0] 
                x = rearrange(x, '(b l) w c -> b (l w) c', l=l//self.window_size, w=self.window_size)
                x = torch.roll(x, shifts=-self.window_size//2, dims=1)
        else:
            x = self.attn(x, x, x, need_weights=False, attn_mask=attn_mask)[0]
        return x

    def forward(self, x, index, condition=None):
        # no condition in encoder, its a dummy argument
        y = self.layernorm1(x)
        y = self.attention(y, index)
        x = x.type(torch.float32) + y  # residual in fp32
        y = self.layernorm2(x)
        y = self.mlp(y)
        x = x.type(torch.float32) + y  # residual in fp32
        return x

class ConditionalResAttBlock(nn.Module):
    def __init__(self, d_model, n_head, window_size=None, drop_path_rate=0.0):