gamecraftvae_blocks ¶

def __init__(
    self,
    chan_in: int,
    chan_out: int,
    kernel_size: Union[int, Tuple[int, int, int]] = 3,
    stride: Union[int, Tuple[int, int, int]] = 1,
    dilation: Union[int, Tuple[int, int, int]] = 1,
    pad_mode: str = "replicate",
    disable_causal: bool = False,
    **kwargs,
):
    super().__init__()
    self.pad_mode = pad_mode
    k = kernel_size if isinstance(kernel_size, int) else kernel_size[0]
    if disable_causal:
        padding = (k // 2, k // 2, k // 2, k // 2, k // 2, k // 2)
    else:
        padding = (k // 2, k // 2, k // 2, k // 2, k - 1, 0)
    self.time_causal_padding = padding
    self.conv = nn.Conv3d(
        chan_in, chan_out, kernel_size, stride=stride, dilation=dilation, **kwargs
    )

def __init__(
    self,
    in_channels: int,
    out_channels: int,
    num_layers: int = 2,
    resnet_eps: float = 1e-6,
    resnet_act_fn: str = "swish",
    resnet_groups: int = 32,
    add_downsample: bool = True,
    downsample_stride: Union[int, Tuple[int, int, int]] = 2,
    downsample_padding: int = 0,
    disable_causal: bool = False,
):
    super().__init__()
    self.resnets = nn.ModuleList()
    for i in range(num_layers):
        inc = in_channels if i == 0 else out_channels
        self.resnets.append(
            ResnetBlockCausal3D(
                in_channels=inc,
                out_channels=out_channels,
                temb_channels=None,
                eps=resnet_eps,
                groups=resnet_groups,
                dropout=0.0,
                non_linearity=resnet_act_fn,
                disable_causal=disable_causal,
            )
        )

    self.downsamplers = None
    if add_downsample:
        self.downsamplers = nn.ModuleList([
            DownsampleCausal3D(
                out_channels,
                out_channels=out_channels,
                padding=downsample_padding,
                stride=downsample_stride,
                disable_causal=disable_causal,
            )
        ])

def __init__(
    self,
    channels: int,
    out_channels: Optional[int] = None,
    padding: int = 1,
    stride: Union[int, Tuple[int, int, int]] = 2,
    kernel_size: int = 3,
    bias: bool = True,
    disable_causal: bool = False,
):
    super().__init__()
    self.out_channels = out_channels or channels
    self.conv = CausalConv3d(
        channels,
        self.out_channels,
        kernel_size=kernel_size,
        stride=stride,
        padding=0,
        disable_causal=disable_causal,
        bias=bias,
    )

def __init__(
    self,
    in_channels: int,
    heads: int,
    dim_head: int,
    eps: float = 1e-6,
    norm_num_groups: Optional[int] = 32,
    bias: bool = True,
):
    super().__init__()
    self.heads = heads
    self.dim_head = dim_head
    inner_dim = heads * dim_head
    self.group_norm = nn.GroupNorm(
        norm_num_groups or in_channels, in_channels, eps=eps
    )
    self.to_q = nn.Linear(in_channels, inner_dim, bias=bias)
    self.to_k = nn.Linear(in_channels, inner_dim, bias=bias)
    self.to_v = nn.Linear(in_channels, inner_dim, bias=bias)
    self.to_out = nn.Sequential(nn.Linear(inner_dim, in_channels, bias=bias))

def __init__(
    self,
    in_channels: int,
    out_channels: Optional[int] = None,
    temb_channels: Optional[int] = None,
    eps: float = 1e-6,
    groups: int = 32,
    dropout: float = 0.0,
    non_linearity: str = "swish",
    disable_causal: bool = False,
):
    super().__init__()
    out_channels = out_channels or in_channels
    self.norm1 = nn.GroupNorm(groups, in_channels, eps=eps)
    self.conv1 = CausalConv3d(in_channels, out_channels, 3, 1, disable_causal=disable_causal)
    self.norm2 = nn.GroupNorm(groups, out_channels, eps=eps)
    self.conv2 = CausalConv3d(out_channels, out_channels, 3, 1, disable_causal=disable_causal)
    self.dropout = nn.Dropout(dropout)
    self.conv_shortcut = (
        CausalConv3d(in_channels, out_channels, 1, 1, disable_causal=disable_causal)
        if in_channels != out_channels
        else None
    )
    self.nonlinearity = getattr(F, non_linearity, F.silu)

def __init__(
    self,
    in_channels: int,
    temb_channels: Optional[int] = None,
    num_layers: int = 1,
    resnet_eps: float = 1e-6,
    resnet_act_fn: str = "swish",
    resnet_groups: int = 32,
    add_attention: bool = True,
    attention_head_dim: int = 1,
    disable_causal: bool = False,
    causal_attention: bool = False,
):
    super().__init__()
    self.add_attention = add_attention
    self.causal_attention = causal_attention

    self.resnets = nn.ModuleList()
    self.attentions = nn.ModuleList()

    self.resnets.append(
        ResnetBlockCausal3D(
            in_channels=in_channels,
            out_channels=in_channels,
            temb_channels=temb_channels,
            eps=resnet_eps,
            groups=resnet_groups,
            dropout=0.0,
            non_linearity=resnet_act_fn,
            disable_causal=disable_causal,
        )
    )

    for _ in range(num_layers):
        if add_attention:
            self.attentions.append(
                GameCraftVAEAttention(
                    in_channels=in_channels,
                    heads=in_channels // attention_head_dim,
                    dim_head=attention_head_dim,
                    eps=resnet_eps,
                    norm_num_groups=resnet_groups,
                    bias=True,
                )
            )
        else:
            self.attentions.append(None)
        self.resnets.append(
            ResnetBlockCausal3D(
                in_channels=in_channels,
                out_channels=in_channels,
                temb_channels=temb_channels,
                eps=resnet_eps,
                groups=resnet_groups,
                dropout=0.0,
                non_linearity=resnet_act_fn,
                disable_causal=disable_causal,
            )
        )

def __init__(
    self,
    in_channels: int,
    out_channels: int,
    num_layers: int = 3,
    resnet_eps: float = 1e-6,
    resnet_act_fn: str = "swish",
    resnet_groups: int = 32,
    add_upsample: bool = True,
    upsample_scale_factor: Tuple[int, int, int] = (2, 2, 2),
    temb_channels: Optional[int] = None,
    disable_causal: bool = False,
):
    super().__init__()
    self.resnets = nn.ModuleList()
    for i in range(num_layers):
        inc = in_channels if i == 0 else out_channels
        self.resnets.append(
            ResnetBlockCausal3D(
                in_channels=inc,
                out_channels=out_channels,
                temb_channels=temb_channels,
                eps=resnet_eps,
                groups=resnet_groups,
                dropout=0.0,
                non_linearity=resnet_act_fn,
                disable_causal=disable_causal,
            )
        )

    self.upsamplers = None
    if add_upsample:
        self.upsamplers = nn.ModuleList([
            UpsampleCausal3D(
                out_channels,
                out_channels=out_channels,
                upsample_factor=upsample_scale_factor,
                disable_causal=disable_causal,
            )
        ])

def __init__(
    self,
    channels: int,
    out_channels: Optional[int] = None,
    kernel_size: int = 3,
    upsample_factor: Tuple[int, int, int] = (2, 2, 2),
    disable_causal: bool = False,
    bias: bool = True,
):
    super().__init__()
    self.out_channels = out_channels or channels
    self.upsample_factor = upsample_factor
    self.disable_causal = disable_causal
    self.conv = CausalConv3d(
        channels,
        self.out_channels,
        kernel_size=kernel_size,
        stride=1,
        disable_causal=disable_causal,
        bias=bias,
    )