Method forward

modules/layers.py:333–346

(self, x: Tensor, vec: Tensor, pe: Tensor, mask: Tensor = None)

        self.backend = backend

    def forward(self, x: Tensor, vec: Tensor, pe: Tensor, mask: Tensor = None) -> Tensor:
        mod, _ = self.modulation(vec)
        x_mod = (1 + mod.scale) * self.pre_norm(x) + mod.shift
        qkv, mlp = torch.split(self.linear1(x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1)

        q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
        q, k = self.norm(q, k, v)
        if mask is not None:
            mask = repeat(mask, 'B L S->  B H L S', H=self.num_heads)
        # compute attention
        attn = attention(q, k, v, pe=pe, mask = mask, backend=self.backend)
        # compute activation in mlp stream, cat again and run second linear layer
        output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2))
        return x + mod.gate * output


class DoubleStreamBlockC(DoubleStreamBlock):