Class ConditionalResAttBlock

model/transformer_utils.py:301–360

        return x

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

        self.self_attn = MultiHeadAttention(d_model, d_model, d_model, d_model, n_head)
        self.self_attn_ln = LayerNorm(d_model)

        self.cross_attn = MultiHeadAttention(d_model, d_model, d_model, d_model, n_head)
        self.cross_attn_ln = 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.mlp_ln = LayerNorm(d_model)

        self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0. else nn.Identity()

    def window_attention(self, x, attn_layer, index):
        attn_mask = None
        if self.window_size is not None:
            l = x.shape[1]
            assert l % self.window_size == 0, "Sequence length must be divisible by window size"
            if index % 2 == 0:
                # Even index: split into windows without shifting.
                x = rearrange(x, 'b (p w) c -> (b p) w c', w=self.window_size)
                x = attn_layer(x, x, x, need_weights=False, attn_mask=attn_mask)[0]
                x = rearrange(x, '(b p) w c -> b (p w) c', p=l // self.window_size)
            else:
                # Odd index: roll by half a window, then split.
                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 = attn_layer(x, x, x, need_weights=False, attn_mask=attn_mask)[0]
                x = rearrange(x, '(b p) w c -> b (p w) c', p=l // self.window_size)
                x = torch.roll(x, shifts=-self.window_size // 2, dims=1)
        else:
            x = attn_layer(x, x, x, need_weights=False, attn_mask=attn_mask)[0]
        return x

    def forward(self, x, index, condition):
        residual = x.type(torch.float32)
        x_ln = self.self_attn_ln(x)
        x2 = self.window_attention(x_ln, self.self_attn, index)
        x = residual + self.drop_path(x2)

        x = rearrange(x, 'b (p n) d -> (b p) n d', p=2)  # split back to frame_0 and frame_1
        residual = x.type(torch.float32)
        x_ln = self.cross_attn_ln(x)
        x2 = self.cross_attn(x_ln, condition, condition, need_weights=False)[0]
        x = residual + self.drop_path(x2)
        x = rearrange(x, '(b p) n d -> b (p n) d', p=2)  # combine frame_0 and frame_1

        residual = x.type(torch.float32)
        x_ln = self.mlp_ln(x)
        x2 = self.mlp(x_ln)
        x = residual + self.drop_path(x2)