Class ZigMa

model_zigma.py:544–1009

A DiT-styled Mamba model with ZigZag scan.

class ZigMa(nn.Module):
    """
    A DiT-styled Mamba model with ZigZag scan.
    """

    def __init__(
        self,
        in_channels: int,
        embed_dim: int,
        depth: int,
        img_dim: int,
        patch_size: int = 1,
        has_text: bool = False,
        num_classes=-1,
        drop_path_rate=0.1,
        n_context_token: int = 0,
        d_context: int = 0,
        ssm_cfg=None,
        norm_epsilon: float = 1e-5,
        rms_norm: bool = True,
        fused_add_norm=True,
        residual_in_fp32=True,
        initializer_cfg=None,
        scan_type="v2",
        video_frames=0,
        tpe=False,  # apply temporal positional encoding for video-related task
        device="cuda",
        use_pe=0,
        use_jit=True,
        m_init=True,
        use_checkpoint=False,
        dtype=torch.float32,
    ):
        # assert num_classes == -1, "num_classes should be -1"
        # assert n_context_token == 0, "n_context_token should be 0"

        self.factory_kwargs = factory_kwargs = {"device": device, "dtype": dtype}
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = in_channels
        self.patch_size = patch_size
        self.embed_dim = embed_dim
        self.tpe = tpe

        self.residual_in_fp32 = residual_in_fp32
        self.fused_add_norm = fused_add_norm
        self.video_frames = video_frames
        self.use_pe = use_pe
        num_patches = (img_dim // patch_size) ** 2
        self.use_checkpoint = use_checkpoint
        print(
            "use_checkpoint",
            use_checkpoint,
            "use_pe",
            use_pe,
            "use tpe",
            tpe,
            "num_patches",