Class DoubleStreamBlockC

modules/layers.py:349–377

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

class DoubleStreamBlockC(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',
                 abondon_cond = False):
        super().__init__(hidden_size, num_heads, mlp_ratio,
                 qkv_bias, backend)
        self.abondon_cond = abondon_cond

    def forward(self, x: Tensor, vec: Tensor,
                pe: Tensor, mask: Tensor = None,
                txt_length=None,
                uncondi_length=None,
                uncondi_pe = None,
                mask_uncond = None):
        # pad_sequence(tuple(x_list), batch_first=True)
        if self.abondon_cond:
            x = [ix[:u_l, :] for ix, u_l in zip(x, uncondi_length)]
            x = pad_sequence(x, batch_first=True)
        if not x.shape[1] == pe.shape[2]:
            pe = uncondi_pe
            mask = mask_uncond
        # print("double stream block", x.shape, pe.shape)
        x = super().forward(x, vec, pe, mask, txt_length)
        return x

class SingleStreamBlockC(SingleStreamBlock):
    """