Class SingleStreamBlockC

modules/layers.py:379–405

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

        return x

class SingleStreamBlockC(SingleStreamBlock):
    """
    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 = 4.0,
                    qk_scale: float | None = None,
                    backend='pytorch',
                    abondon_cond = False):
        super().__init__(hidden_size, num_heads, mlp_ratio,
                 qk_scale, backend)
        self.abondon_cond = abondon_cond

    def forward(self, x: Tensor, vec: Tensor, pe: Tensor, mask: Tensor = None,
                uncondi_length = None, uncondi_pe = None, mask_uncond = None) -> Tensor:
        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("single stream block", x.shape, pe.shape)
        x = super().forward(x, vec, pe, mask)
        return x


class DoubleStreamBlockD(DoubleStreamBlock):