Class DoubleStreamBlockD

modules/layers.py:408–492

A DiT block with parallel linear layers as described in https://arxiv.org/abs/2302.05442 and adapted modulation interface.

class DoubleStreamBlockD(DoubleStreamBlock):
    """
    A DiT block with parallel linear layers as described in
    https://arxiv.org/abs/2302.05442 and adapted modulation interface.
    """

    def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float,
                 qkv_bias: bool = False, backend='pytorch'):
        super().__init__(hidden_size, num_heads, mlp_ratio,
                 qkv_bias, backend)
        mlp_hidden_dim = int(hidden_size * mlp_ratio)
        self.edit_mod = Modulation(hidden_size, double=True)
        self.edit_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
        self.edit_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias)

        self.edit_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
        self.edit_mlp = nn.Sequential(
            nn.Linear(hidden_size, mlp_hidden_dim, bias=True),
            nn.GELU(approximate="tanh"),
            nn.Linear(mlp_hidden_dim, hidden_size, bias=True),
        )

    def forward(self, x: Tensor, vec: Tensor,
                pe: Tensor, mask: Tensor = None,
                txt_length=None,
                edit_length=None):
        if edit_length is not None:
            txt, edit, img = x[:, :txt_length], x[:, txt_length:txt_length + edit_length], x[:, txt_length + edit_length:]
        else:
            txt, img = x[:, :txt_length], x[:, txt_length:]
        img_mod1, img_mod2 = self.img_mod(vec)
        txt_mod1, txt_mod2 = self.txt_mod(vec)
        # prepare image for attention
        img_modulated = self.img_norm1(img)
        img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift
        img_qkv = self.img_attn.qkv(img_modulated)
        img_q, img_k, img_v = rearrange(img_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
        img_q, img_k = self.img_attn.norm(img_q, img_k, img_v)
        # prepare txt for attention
        txt_modulated = self.txt_norm1(txt)
        txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift
        txt_qkv = self.txt_attn.qkv(txt_modulated)
        txt_q, txt_k, txt_v = rearrange(txt_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
        txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k, txt_v)

        if edit_length is not None:
            edit_mod1, edit_mod2 = self.edit_mod(vec)
            # prepare edit for attention
            edit_modulated = self.edit_norm1(edit)
            edit_modulated = (1 + edit_mod1.scale) * edit_modulated + edit_mod1.shift
            edit_qkv = self.edit_attn.qkv(edit_modulated)
            edit_q, edit_k, edit_v = rearrange(edit_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
            edit_q, edit_k = self.edit_attn.norm(edit_q, edit_k, edit_v)
        else:
            edit_q, edit_k, edit_v = None, None, None


        # run actual attention