vllm_omni.diffusion.models.seedvr2.vae ¶

temporal_padding = 2 * padding[0]

conv_in = CausalConv3d(latent_channels, channels[-1])

conv_norm_out = nn.GroupNorm(
    groups, channels[0], eps=1e-06
)

conv_out = CausalConv3d(channels[0], 3)

mid_block = MidBlock3d(channels[-1], groups)

up_blocks = nn.ModuleList(
    [
        DecoderBlock3d(
            widths[max(0, i - 1)], width, groups, i
        )
        for i, width in enumerate(widths)
    ]
)

resnets = nn.ModuleList(
    [ResnetBlock3d(in_channels, out_channels, groups)]
    + [
        ResnetBlock3d(out_channels, out_channels, groups)
        for _ in range(2)
    ]
)

upsamplers = nn.ModuleList(
    [Upsample3d(out_channels, temporal=index < 2)]
    if index < 3
    else []
)

conv = CausalConv3d(
    channels,
    channels,
    (3 if temporal else 1, 3, 3),
    (2 if temporal else 1, 2, 2),
    (1 if temporal else 0, 0, 0),
)

conv_in = CausalConv3d(3, channels[0])

conv_norm_out = nn.GroupNorm(
    groups, channels[-1], eps=1e-06
)

conv_out = CausalConv3d(channels[-1], 2 * latent_channels)

down_blocks = nn.ModuleList(
    [
        EncoderBlock3d(
            channels[max(0, i - 1)], width, groups, i
        )
        for i, width in enumerate(channels)
    ]
)

mid_block = MidBlock3d(channels[-1], groups)

downsamplers = nn.ModuleList(
    [Downsample3d(out_channels, temporal=index >= 1)]
    if index < 3
    else []
)

resnets = nn.ModuleList(
    [
        ResnetBlock3d(in_channels, out_channels, groups),
        ResnetBlock3d(out_channels, out_channels, groups),
    ]
)

attentions = nn.ModuleList(
    [
        Attention(
            channels,
            heads=1,
            dim_head=channels,
            eps=1e-06,
            norm_num_groups=groups,
            residual_connection=True,
            bias=True,
            upcast_softmax=True,
            _from_deprecated_attn_block=True,
        )
    ]
)

resnets = nn.ModuleList(
    [
        ResnetBlock3d(channels, channels, groups)
        for _ in range(2)
    ]
)

conv1 = CausalConv3d(in_channels, out_channels)

conv2 = CausalConv3d(out_channels, out_channels)

conv_shortcut = (
    CausalConv3d(
        in_channels,
        out_channels,
        (1, 1, 1),
        padding=(0, 0, 0),
    )
    if in_channels != out_channels
    else None
)

norm1 = nn.GroupNorm(groups, in_channels, eps=1e-06)

norm2 = nn.GroupNorm(groups, out_channels, eps=1e-06)

decoder = Decoder3d(channels, groups, latent_channels)

encoder = Encoder3d(channels, groups, latent_channels)

spatial_scale_factor = 8

temporal_scale_factor = 4

use_tiling = False

group: ProcessGroup

rank: int

units: tuple[int, ...]

cache_temporal: bool = False

first: bool = True

history: dict[Module, Tensor] = field(default_factory=dict)

spatial: SpatialContext | None = None

tile_spatial: bool = False

conv = CausalConv3d(channels, channels)

temporal_ratio = 2 if temporal else 1

upscale_conv = nn.Conv3d(
    channels,
    channels * 4 * self.temporal_ratio,
    kernel_size=1,
)