DiffSynth-Studio/diffsynth/models/cog_vae.py

import torch
from einops import rearrange, repeat
from .tiler import TileWorker2Dto3D



class Downsample3D(torch.nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int = 3,
        stride: int = 2,
        padding: int = 0,
        compress_time: bool = False,
    ):
        super().__init__()

        self.conv = torch.nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding)
        self.compress_time = compress_time

    def forward(self, x: torch.Tensor, xq: torch.Tensor) -> torch.Tensor:
        if self.compress_time:
            batch_size, channels, frames, height, width = x.shape

            # (batch_size, channels, frames, height, width) -> (batch_size, height, width, channels, frames) -> (batch_size * height * width, channels, frames)
            x = x.permute(0, 3, 4, 1, 2).reshape(batch_size * height * width, channels, frames)

            if x.shape[-1] % 2 == 1:
                x_first, x_rest = x[..., 0], x[..., 1:]
                if x_rest.shape[-1] > 0:
                    # (batch_size * height * width, channels, frames - 1) -> (batch_size * height * width, channels, (frames - 1) // 2)
                    x_rest = torch.nn.functional.avg_pool1d(x_rest, kernel_size=2, stride=2)

                x = torch.cat([x_first[..., None], x_rest], dim=-1)
                # (batch_size * height * width, channels, (frames // 2) + 1) -> (batch_size, height, width, channels, (frames // 2) + 1) -> (batch_size, channels, (frames // 2) + 1, height, width)
                x = x.reshape(batch_size, height, width, channels, x.shape[-1]).permute(0, 3, 4, 1, 2)
            else:
                # (batch_size * height * width, channels, frames) -> (batch_size * height * width, channels, frames // 2)
                x = torch.nn.functional.avg_pool1d(x, kernel_size=2, stride=2)
                # (batch_size * height * width, channels, frames // 2) -> (batch_size, height, width, channels, frames // 2) -> (batch_size, channels, frames // 2, height, width)
                x = x.reshape(batch_size, height, width, channels, x.shape[-1]).permute(0, 3, 4, 1, 2)

        # Pad the tensor
        pad = (0, 1, 0, 1)
        x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
        batch_size, channels, frames, height, width = x.shape
        # (batch_size, channels, frames, height, width) -> (batch_size, frames, channels, height, width) -> (batch_size * frames, channels, height, width)
        x = x.permute(0, 2, 1, 3, 4).reshape(batch_size * frames, channels, height, width)
        x = self.conv(x)
        # (batch_size * frames, channels, height, width) -> (batch_size, frames, channels, height, width) -> (batch_size, channels, frames, height, width)
        x = x.reshape(batch_size, frames, x.shape[1], x.shape[2], x.shape[3]).permute(0, 2, 1, 3, 4)
        return x



class Upsample3D(torch.nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int = 3,
        stride: int = 1,
        padding: int = 1,
        compress_time: bool = False,
    ) -> None:
        super().__init__()
        self.conv = torch.nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding)
        self.compress_time = compress_time

    def forward(self, inputs: torch.Tensor, xq: torch.Tensor) -> torch.Tensor:
        if self.compress_time:
            if inputs.shape[2] > 1 and inputs.shape[2] % 2 == 1:
                # split first frame
                x_first, x_rest = inputs[:, :, 0], inputs[:, :, 1:]

                x_first = torch.nn.functional.interpolate(x_first, scale_factor=2.0)
                x_rest = torch.nn.functional.interpolate(x_rest, scale_factor=2.0)
                x_first = x_first[:, :, None, :, :]
                inputs = torch.cat([x_first, x_rest], dim=2)
            elif inputs.shape[2] > 1:
                inputs = torch.nn.functional.interpolate(inputs, scale_factor=2.0)
            else:
                inputs = inputs.squeeze(2)
                inputs = torch.nn.functional.interpolate(inputs, scale_factor=2.0)
                inputs = inputs[:, :, None, :, :]
        else:
            # only interpolate 2D
            b, c, t, h, w = inputs.shape
            inputs = inputs.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w)
            inputs = torch.nn.functional.interpolate(inputs, scale_factor=2.0)
            inputs = inputs.reshape(b, t, c, *inputs.shape[2:]).permute(0, 2, 1, 3, 4)

        b, c, t, h, w = inputs.shape
        inputs = inputs.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w)
        inputs = self.conv(inputs)
        inputs = inputs.reshape(b, t, *inputs.shape[1:]).permute(0, 2, 1, 3, 4)

        return inputs



class CogVideoXSpatialNorm3D(torch.nn.Module):
    def __init__(self, f_channels, zq_channels, groups):
        super().__init__()
        self.norm_layer = torch.nn.GroupNorm(num_channels=f_channels, num_groups=groups, eps=1e-6, affine=True)
        self.conv_y = torch.nn.Conv3d(zq_channels, f_channels, kernel_size=1, stride=1)
        self.conv_b = torch.nn.Conv3d(zq_channels, f_channels, kernel_size=1, stride=1)


    def forward(self, f: torch.Tensor, zq: torch.Tensor) -> torch.Tensor:
        if f.shape[2] > 1 and f.shape[2] % 2 == 1:
            f_first, f_rest = f[:, :, :1], f[:, :, 1:]
            f_first_size, f_rest_size = f_first.shape[-3:], f_rest.shape[-3:]
            z_first, z_rest = zq[:, :, :1], zq[:, :, 1:]
            z_first = torch.nn.functional.interpolate(z_first, size=f_first_size)
            z_rest = torch.nn.functional.interpolate(z_rest, size=f_rest_size)
            zq = torch.cat([z_first, z_rest], dim=2)
        else:
            zq = torch.nn.functional.interpolate(zq, size=f.shape[-3:])

        norm_f = self.norm_layer(f)
        new_f = norm_f * self.conv_y(zq) + self.conv_b(zq)
        return new_f



class Resnet3DBlock(torch.nn.Module):
    def __init__(self, in_channels, out_channels, spatial_norm_dim, groups, eps=1e-6, use_conv_shortcut=False):
        super().__init__()
        self.nonlinearity = torch.nn.SiLU()
        if spatial_norm_dim is None:
            self.norm1 = torch.nn.GroupNorm(num_channels=in_channels, num_groups=groups, eps=eps)
            self.norm2 = torch.nn.GroupNorm(num_channels=out_channels, num_groups=groups, eps=eps)
        else:
            self.norm1 = CogVideoXSpatialNorm3D(in_channels, spatial_norm_dim, groups)
            self.norm2 = CogVideoXSpatialNorm3D(out_channels, spatial_norm_dim, groups)

        self.conv1 = CachedConv3d(in_channels, out_channels, kernel_size=3, padding=(0, 1, 1))

        self.conv2 = CachedConv3d(out_channels, out_channels, kernel_size=3, padding=(0, 1, 1))

        if in_channels != out_channels:
            if use_conv_shortcut:
                self.conv_shortcut = CachedConv3d(in_channels, out_channels, kernel_size=3, padding=(0, 1, 1))
            else:
                self.conv_shortcut = torch.nn.Conv3d(in_channels, out_channels, kernel_size=1)
        else:
            self.conv_shortcut = lambda x: x


    def forward(self, hidden_states, zq):
        residual = hidden_states

        hidden_states = self.norm1(hidden_states, zq) if isinstance(self.norm1, CogVideoXSpatialNorm3D) else self.norm1(hidden_states)
        hidden_states = self.nonlinearity(hidden_states)
        hidden_states = self.conv1(hidden_states)

        hidden_states = self.norm2(hidden_states, zq) if isinstance(self.norm2, CogVideoXSpatialNorm3D) else self.norm2(hidden_states)
        hidden_states = self.nonlinearity(hidden_states)
        hidden_states = self.conv2(hidden_states)

        hidden_states = hidden_states + self.conv_shortcut(residual)

        return hidden_states
    


class CachedConv3d(torch.nn.Conv3d):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0):
        super().__init__(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding)
        self.cached_tensor = None


    def clear_cache(self):
        self.cached_tensor = None
    

    def forward(self, input: torch.Tensor, use_cache = True) -> torch.Tensor:
        if use_cache:
            if self.cached_tensor is None:
                self.cached_tensor = torch.concat([input[:, :, :1]] * 2, dim=2)
            input = torch.concat([self.cached_tensor, input], dim=2)
            self.cached_tensor = input[:, :, -2:]
        return super().forward(input)



class CogVAEDecoder(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.scaling_factor = 0.7
        self.conv_in = CachedConv3d(16, 512, kernel_size=3, stride=1, padding=(0, 1, 1))

        self.blocks = torch.nn.ModuleList([
            Resnet3DBlock(512, 512, 16, 32),
            Resnet3DBlock(512, 512, 16, 32),
            Resnet3DBlock(512, 512, 16, 32),
            Resnet3DBlock(512, 512, 16, 32),
            Resnet3DBlock(512, 512, 16, 32),
            Resnet3DBlock(512, 512, 16, 32),
            Upsample3D(512, 512, compress_time=True),
            Resnet3DBlock(512, 256, 16, 32),
            Resnet3DBlock(256, 256, 16, 32),
            Resnet3DBlock(256, 256, 16, 32),
            Resnet3DBlock(256, 256, 16, 32),
            Upsample3D(256, 256, compress_time=True),
            Resnet3DBlock(256, 256, 16, 32),
            Resnet3DBlock(256, 256, 16, 32),
            Resnet3DBlock(256, 256, 16, 32),
            Resnet3DBlock(256, 256, 16, 32),
            Upsample3D(256, 256, compress_time=False),
            Resnet3DBlock(256, 128, 16, 32),
            Resnet3DBlock(128, 128, 16, 32),
            Resnet3DBlock(128, 128, 16, 32),
            Resnet3DBlock(128, 128, 16, 32),
        ])

        self.norm_out = CogVideoXSpatialNorm3D(128, 16, 32)
        self.conv_act = torch.nn.SiLU()
        self.conv_out = CachedConv3d(128, 3, kernel_size=3, stride=1, padding=(0, 1, 1))


    def forward(self, sample):
        sample = sample / self.scaling_factor
        hidden_states = self.conv_in(sample)

        for block in self.blocks:
            hidden_states = block(hidden_states, sample)
        
        hidden_states = self.norm_out(hidden_states, sample)
        hidden_states = self.conv_act(hidden_states)
        hidden_states = self.conv_out(hidden_states)

        return hidden_states
    

    def decode_video(self, sample, tiled=True, tile_size=(60, 90), tile_stride=(30, 45), progress_bar=lambda x:x):
        if tiled:
            B, C, T, H, W = sample.shape
            return TileWorker2Dto3D().tiled_forward(
                forward_fn=lambda x: self.decode_small_video(x),
                model_input=sample,
                tile_size=tile_size, tile_stride=tile_stride,
                tile_device=sample.device, tile_dtype=sample.dtype,
                computation_device=sample.device, computation_dtype=sample.dtype,
                scales=(3/16, (T//2*8+T%2)/T, 8, 8),
                progress_bar=progress_bar
            )
        else:
            return self.decode_small_video(sample)
    

    def decode_small_video(self, sample):
        B, C, T, H, W = sample.shape
        computation_device = self.conv_in.weight.device
        computation_dtype = self.conv_in.weight.dtype
        value = []
        for i in range(T//2):
            tl = i*2 + T%2 - (T%2 and i==0)
            tr = i*2 + 2 + T%2
            model_input = sample[:, :, tl: tr, :, :].to(dtype=computation_dtype, device=computation_device)
            model_output = self.forward(model_input).to(dtype=sample.dtype, device=sample.device)
            value.append(model_output)
        value = torch.concat(value, dim=2)
        for name, module in self.named_modules():
            if isinstance(module, CachedConv3d):
                module.clear_cache()
        return value
    

    @staticmethod
    def state_dict_converter():
        return CogVAEDecoderStateDictConverter()
    


class CogVAEEncoder(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.scaling_factor = 0.7
        self.conv_in = CachedConv3d(3, 128, kernel_size=3, stride=1, padding=(0, 1, 1))

        self.blocks = torch.nn.ModuleList([
            Resnet3DBlock(128, 128, None, 32),
            Resnet3DBlock(128, 128, None, 32),
            Resnet3DBlock(128, 128, None, 32),
            Downsample3D(128, 128, compress_time=True),
            Resnet3DBlock(128, 256, None, 32),
            Resnet3DBlock(256, 256, None, 32),
            Resnet3DBlock(256, 256, None, 32),
            Downsample3D(256, 256, compress_time=True),
            Resnet3DBlock(256, 256, None, 32),
            Resnet3DBlock(256, 256, None, 32),
            Resnet3DBlock(256, 256, None, 32),
            Downsample3D(256, 256, compress_time=False),
            Resnet3DBlock(256, 512, None, 32),
            Resnet3DBlock(512, 512, None, 32),
            Resnet3DBlock(512, 512, None, 32),
            Resnet3DBlock(512, 512, None, 32),
            Resnet3DBlock(512, 512, None, 32),
        ])

        self.norm_out = torch.nn.GroupNorm(32, 512, eps=1e-06, affine=True)
        self.conv_act = torch.nn.SiLU()
        self.conv_out = CachedConv3d(512, 32, kernel_size=3, stride=1, padding=(0, 1, 1))


    def forward(self, sample):
        hidden_states = self.conv_in(sample)

        for block in self.blocks:
            hidden_states = block(hidden_states, sample)
        
        hidden_states = self.norm_out(hidden_states)
        hidden_states = self.conv_act(hidden_states)
        hidden_states = self.conv_out(hidden_states)[:, :16]
        hidden_states = hidden_states * self.scaling_factor

        return hidden_states
    

    def encode_video(self, sample, tiled=True, tile_size=(60, 90), tile_stride=(30, 45), progress_bar=lambda x:x):
        if tiled:
            B, C, T, H, W = sample.shape
            return TileWorker2Dto3D().tiled_forward(
                forward_fn=lambda x: self.encode_small_video(x),
                model_input=sample,
                tile_size=(i * 8 for i in tile_size), tile_stride=(i * 8 for i in tile_stride),
                tile_device=sample.device, tile_dtype=sample.dtype,
                computation_device=sample.device, computation_dtype=sample.dtype,
                scales=(16/3, (T//4+T%2)/T, 1/8, 1/8),
                progress_bar=progress_bar
            )
        else:
            return self.encode_small_video(sample)
    

    def encode_small_video(self, sample):
        B, C, T, H, W = sample.shape
        computation_device = self.conv_in.weight.device
        computation_dtype = self.conv_in.weight.dtype
        value = []
        for i in range(T//8):
            t = i*8 + T%2 - (T%2 and i==0)
            t_ = i*8 + 8 + T%2
            model_input = sample[:, :, t: t_, :, :].to(dtype=computation_dtype, device=computation_device)
            model_output = self.forward(model_input).to(dtype=sample.dtype, device=sample.device)
            value.append(model_output)
        value = torch.concat(value, dim=2)
        for name, module in self.named_modules():
            if isinstance(module, CachedConv3d):
                module.clear_cache()
        return value
    

    @staticmethod
    def state_dict_converter():
        return CogVAEEncoderStateDictConverter()



class CogVAEEncoderStateDictConverter:
    def __init__(self):
        pass


    def from_diffusers(self, state_dict):
        rename_dict = {
            "encoder.conv_in.conv.weight": "conv_in.weight",
            "encoder.conv_in.conv.bias": "conv_in.bias",
            "encoder.down_blocks.0.downsamplers.0.conv.weight": "blocks.3.conv.weight",
            "encoder.down_blocks.0.downsamplers.0.conv.bias": "blocks.3.conv.bias",
            "encoder.down_blocks.1.downsamplers.0.conv.weight": "blocks.7.conv.weight",
            "encoder.down_blocks.1.downsamplers.0.conv.bias": "blocks.7.conv.bias",
            "encoder.down_blocks.2.downsamplers.0.conv.weight": "blocks.11.conv.weight",
            "encoder.down_blocks.2.downsamplers.0.conv.bias": "blocks.11.conv.bias",
            "encoder.norm_out.weight": "norm_out.weight",
            "encoder.norm_out.bias": "norm_out.bias",
            "encoder.conv_out.conv.weight": "conv_out.weight",
            "encoder.conv_out.conv.bias": "conv_out.bias",
        }
        prefix_dict = {
            "encoder.down_blocks.0.resnets.0.": "blocks.0.",
            "encoder.down_blocks.0.resnets.1.": "blocks.1.",
            "encoder.down_blocks.0.resnets.2.": "blocks.2.",
            "encoder.down_blocks.1.resnets.0.": "blocks.4.",
            "encoder.down_blocks.1.resnets.1.": "blocks.5.",
            "encoder.down_blocks.1.resnets.2.": "blocks.6.",
            "encoder.down_blocks.2.resnets.0.": "blocks.8.",
            "encoder.down_blocks.2.resnets.1.": "blocks.9.",
            "encoder.down_blocks.2.resnets.2.": "blocks.10.",
            "encoder.down_blocks.3.resnets.0.": "blocks.12.",
            "encoder.down_blocks.3.resnets.1.": "blocks.13.",
            "encoder.down_blocks.3.resnets.2.": "blocks.14.",
            "encoder.mid_block.resnets.0.": "blocks.15.",
            "encoder.mid_block.resnets.1.": "blocks.16.",
        }
        suffix_dict = {
            "norm1.norm_layer.weight": "norm1.norm_layer.weight",
            "norm1.norm_layer.bias": "norm1.norm_layer.bias",
            "norm1.conv_y.conv.weight": "norm1.conv_y.weight",
            "norm1.conv_y.conv.bias": "norm1.conv_y.bias",
            "norm1.conv_b.conv.weight": "norm1.conv_b.weight",
            "norm1.conv_b.conv.bias": "norm1.conv_b.bias",
            "norm2.norm_layer.weight": "norm2.norm_layer.weight",
            "norm2.norm_layer.bias": "norm2.norm_layer.bias",
            "norm2.conv_y.conv.weight": "norm2.conv_y.weight",
            "norm2.conv_y.conv.bias": "norm2.conv_y.bias",
            "norm2.conv_b.conv.weight": "norm2.conv_b.weight",
            "norm2.conv_b.conv.bias": "norm2.conv_b.bias",
            "conv1.conv.weight": "conv1.weight",
            "conv1.conv.bias": "conv1.bias",
            "conv2.conv.weight": "conv2.weight",
            "conv2.conv.bias": "conv2.bias",
            "conv_shortcut.weight": "conv_shortcut.weight",
            "conv_shortcut.bias": "conv_shortcut.bias",
            "norm1.weight": "norm1.weight",
            "norm1.bias": "norm1.bias",
            "norm2.weight": "norm2.weight",
            "norm2.bias": "norm2.bias",
        }
        state_dict_ = {}
        for name, param in state_dict.items():
            if name in rename_dict:
                state_dict_[rename_dict[name]] = param
            else:
                for prefix in prefix_dict:
                    if name.startswith(prefix):
                        suffix = name[len(prefix):]
                        state_dict_[prefix_dict[prefix] + suffix_dict[suffix]] = param
        return state_dict_
    

    def from_civitai(self, state_dict):
        return self.from_diffusers(state_dict)



class CogVAEDecoderStateDictConverter:
    def __init__(self):
        pass


    def from_diffusers(self, state_dict):
        rename_dict = {
            "decoder.conv_in.conv.weight": "conv_in.weight",
            "decoder.conv_in.conv.bias": "conv_in.bias",
            "decoder.up_blocks.0.upsamplers.0.conv.weight": "blocks.6.conv.weight",
            "decoder.up_blocks.0.upsamplers.0.conv.bias": "blocks.6.conv.bias",
            "decoder.up_blocks.1.upsamplers.0.conv.weight": "blocks.11.conv.weight",
            "decoder.up_blocks.1.upsamplers.0.conv.bias": "blocks.11.conv.bias",
            "decoder.up_blocks.2.upsamplers.0.conv.weight": "blocks.16.conv.weight",
            "decoder.up_blocks.2.upsamplers.0.conv.bias": "blocks.16.conv.bias",
            "decoder.norm_out.norm_layer.weight": "norm_out.norm_layer.weight",
            "decoder.norm_out.norm_layer.bias": "norm_out.norm_layer.bias",
            "decoder.norm_out.conv_y.conv.weight": "norm_out.conv_y.weight",
            "decoder.norm_out.conv_y.conv.bias": "norm_out.conv_y.bias",
            "decoder.norm_out.conv_b.conv.weight": "norm_out.conv_b.weight",
            "decoder.norm_out.conv_b.conv.bias": "norm_out.conv_b.bias",
            "decoder.conv_out.conv.weight": "conv_out.weight",
            "decoder.conv_out.conv.bias": "conv_out.bias"
        }
        prefix_dict = {
            "decoder.mid_block.resnets.0.": "blocks.0.",
            "decoder.mid_block.resnets.1.": "blocks.1.",
            "decoder.up_blocks.0.resnets.0.": "blocks.2.",
            "decoder.up_blocks.0.resnets.1.": "blocks.3.",
            "decoder.up_blocks.0.resnets.2.": "blocks.4.",
            "decoder.up_blocks.0.resnets.3.": "blocks.5.",
            "decoder.up_blocks.1.resnets.0.": "blocks.7.",
            "decoder.up_blocks.1.resnets.1.": "blocks.8.",
            "decoder.up_blocks.1.resnets.2.": "blocks.9.",
            "decoder.up_blocks.1.resnets.3.": "blocks.10.",
            "decoder.up_blocks.2.resnets.0.": "blocks.12.",
            "decoder.up_blocks.2.resnets.1.": "blocks.13.",
            "decoder.up_blocks.2.resnets.2.": "blocks.14.",
            "decoder.up_blocks.2.resnets.3.": "blocks.15.",
            "decoder.up_blocks.3.resnets.0.": "blocks.17.",
            "decoder.up_blocks.3.resnets.1.": "blocks.18.",
            "decoder.up_blocks.3.resnets.2.": "blocks.19.",
            "decoder.up_blocks.3.resnets.3.": "blocks.20.",
        }
        suffix_dict = {
            "norm1.norm_layer.weight": "norm1.norm_layer.weight",
            "norm1.norm_layer.bias": "norm1.norm_layer.bias",
            "norm1.conv_y.conv.weight": "norm1.conv_y.weight",
            "norm1.conv_y.conv.bias": "norm1.conv_y.bias",
            "norm1.conv_b.conv.weight": "norm1.conv_b.weight",
            "norm1.conv_b.conv.bias": "norm1.conv_b.bias",
            "norm2.norm_layer.weight": "norm2.norm_layer.weight",
            "norm2.norm_layer.bias": "norm2.norm_layer.bias",
            "norm2.conv_y.conv.weight": "norm2.conv_y.weight",
            "norm2.conv_y.conv.bias": "norm2.conv_y.bias",
            "norm2.conv_b.conv.weight": "norm2.conv_b.weight",
            "norm2.conv_b.conv.bias": "norm2.conv_b.bias",
            "conv1.conv.weight": "conv1.weight",
            "conv1.conv.bias": "conv1.bias",
            "conv2.conv.weight": "conv2.weight",
            "conv2.conv.bias": "conv2.bias",
            "conv_shortcut.weight": "conv_shortcut.weight",
            "conv_shortcut.bias": "conv_shortcut.bias",
        }
        state_dict_ = {}
        for name, param in state_dict.items():
            if name in rename_dict:
                state_dict_[rename_dict[name]] = param
            else:
                for prefix in prefix_dict:
                    if name.startswith(prefix):
                        suffix = name[len(prefix):]
                        state_dict_[prefix_dict[prefix] + suffix_dict[suffix]] = param
        return state_dict_
    

    def from_civitai(self, state_dict):
        return self.from_diffusers(state_dict)