geometrycrafter/pmap_vae.py

from typing import Dict, Tuple, Union

import torch
import torch.nn as nn
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.utils.accelerate_utils import apply_forward_hook
from diffusers.models.attention_processor import CROSS_ATTENTION_PROCESSORS, AttentionProcessor, AttnProcessor
from diffusers.models.modeling_utils import ModelMixin
from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution, Encoder
from diffusers.utils import is_torch_version
from diffusers.models.unets.unet_3d_blocks import UpBlockTemporalDecoder, MidBlockTemporalDecoder
from diffusers.models.resnet import SpatioTemporalResBlock

def zero_module(module):
    """
    Zero out the parameters of a module and return it.
    """
    for p in module.parameters():
        p.detach().zero_()
    return module

class PMapTemporalDecoder(nn.Module):
    def __init__(
        self,
        in_channels: int = 4,
        out_channels: Tuple[int] = (1, 1, 1),
        block_out_channels: Tuple[int] = (128, 256, 512, 512),
        layers_per_block: int = 2,
    ):
        super().__init__()

        self.conv_in = nn.Conv2d(
            in_channels, 
            block_out_channels[-1], 
            kernel_size=3, 
            stride=1, 
            padding=1
        )
        self.mid_block = MidBlockTemporalDecoder(
            num_layers=layers_per_block,
            in_channels=block_out_channels[-1],
            out_channels=block_out_channels[-1],
            attention_head_dim=block_out_channels[-1],
        )

        # up
        self.up_blocks = nn.ModuleList([])
        reversed_block_out_channels = list(reversed(block_out_channels))
        output_channel = reversed_block_out_channels[0]
        for i in range(len(block_out_channels)):
            prev_output_channel = output_channel
            output_channel = reversed_block_out_channels[i]
            is_final_block = i == len(block_out_channels) - 1
            up_block = UpBlockTemporalDecoder(
                num_layers=layers_per_block + 1,
                in_channels=prev_output_channel,
                out_channels=output_channel,
                add_upsample=not is_final_block,
            )
            self.up_blocks.append(up_block)
            prev_output_channel = output_channel

        self.out_blocks = nn.ModuleList([])
        self.time_conv_outs = nn.ModuleList([])
        for out_channel in out_channels:
            self.out_blocks.append(
                nn.ModuleList([
                    nn.GroupNorm(num_channels=block_out_channels[0], num_groups=32, eps=1e-6),
                    nn.ReLU(inplace=True),
                    nn.Conv2d(
                        block_out_channels[0], 
                        block_out_channels[0] // 2, 
                        kernel_size=3, 
                        padding=1
                    ),
                    SpatioTemporalResBlock(
                        in_channels=block_out_channels[0] // 2,
                        out_channels=block_out_channels[0] // 2,
                        temb_channels=None,
                        eps=1e-6,
                        temporal_eps=1e-5,
                        merge_factor=0.0,
                        merge_strategy="learned",
                        switch_spatial_to_temporal_mix=True
                    ),
                    nn.ReLU(inplace=True),
                    nn.Conv2d(
                        block_out_channels[0] // 2, 
                        out_channel, 
                        kernel_size=1,
                    )
                ])
            )

            conv_out_kernel_size = (3, 1, 1)
            padding = [int(k // 2) for k in conv_out_kernel_size]
            self.time_conv_outs.append(nn.Conv3d(
                in_channels=out_channel,
                out_channels=out_channel,
                kernel_size=conv_out_kernel_size,
                padding=padding,
            ))

        self.gradient_checkpointing = False

    def forward(
        self,
        sample: torch.Tensor,
        image_only_indicator: torch.Tensor,
        num_frames: int = 1,
    ):
        sample = self.conv_in(sample)

        upscale_dtype = next(iter(self.up_blocks.parameters())).dtype

        if self.training and self.gradient_checkpointing:
            def create_custom_forward(module):
                def custom_forward(*inputs):
                    return module(*inputs)

                return custom_forward

            if is_torch_version(">=", "1.11.0"):
                # middle
                sample = torch.utils.checkpoint.checkpoint(
                    create_custom_forward(self.mid_block),
                    sample,
                    image_only_indicator,
                    use_reentrant=False,
                )
                sample = sample.to(upscale_dtype)

                # up
                for up_block in self.up_blocks:
                    sample = torch.utils.checkpoint.checkpoint(
                        create_custom_forward(up_block),
                        sample,
                        image_only_indicator,
                        use_reentrant=False,
                    )
            else:
                # middle
                sample = torch.utils.checkpoint.checkpoint(
                    create_custom_forward(self.mid_block),
                    sample,
                    image_only_indicator,
                )
                sample = sample.to(upscale_dtype)

                # up
                for up_block in self.up_blocks:
                    sample = torch.utils.checkpoint.checkpoint(
                        create_custom_forward(up_block),
                        sample,
                        image_only_indicator,
                    )
        else:
            # middle
            sample = self.mid_block(sample, image_only_indicator=image_only_indicator)
            sample = sample.to(upscale_dtype)

            # up
            for up_block in self.up_blocks:
                sample = up_block(sample, image_only_indicator=image_only_indicator)

        # post-process

        output = []
        
        for out_block, time_conv_out in zip(self.out_blocks, self.time_conv_outs):
            x = sample
            for layer in out_block:
                if isinstance(layer, SpatioTemporalResBlock):
                    x = layer(x, None, image_only_indicator)
                else:
                    x = layer(x)
            
            
            batch_frames, channels, height, width = x.shape
            batch_size = batch_frames // num_frames
            x = x[None, :].reshape(batch_size, num_frames, channels, height, width).permute(0, 2, 1, 3, 4)
            x = time_conv_out(x)
            x = x.permute(0, 2, 1, 3, 4).reshape(batch_frames, channels, height, width)
            output.append(x)

        return output

class PMapAutoencoderKLTemporalDecoder(ModelMixin, ConfigMixin):
    
    _supports_gradient_checkpointing = True

    @register_to_config
    def __init__(
        self,
        in_channels: int = 4,
        latent_channels: int = 4,
        enc_down_block_types: Tuple[str] = (
            "DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D"
        ),
        enc_block_out_channels: Tuple[int] = (128, 256, 512, 512),
        enc_layers_per_block: int = 2,
        dec_block_out_channels: Tuple[int] = (128, 256, 512, 512),
        dec_layers_per_block: int = 2,
        out_channels: Tuple[int] = (1, 1, 1),
        mid_block_add_attention: bool = True,
        offset_scale_factor: float = 0.1,
        **kwargs  
    ):
        super().__init__()

        self.encoder = Encoder(
            in_channels=in_channels,
            out_channels=latent_channels,
            down_block_types=enc_down_block_types,
            block_out_channels=enc_block_out_channels,
            layers_per_block=enc_layers_per_block,
            double_z=False,
            mid_block_add_attention=mid_block_add_attention
        )
        zero_module(self.encoder.conv_out)

        self.offset_scale_factor = offset_scale_factor

        self.decoder = PMapTemporalDecoder(
            in_channels=latent_channels,
            block_out_channels=dec_block_out_channels,
            layers_per_block=dec_layers_per_block,
            out_channels=out_channels
        )

    def _set_gradient_checkpointing(self, module, value=False):
        if isinstance(module, (Encoder, PMapTemporalDecoder)):
            module.gradient_checkpointing = value

    @property
    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors
    def attn_processors(self) -> Dict[str, AttentionProcessor]:
        r"""
        Returns:
            `dict` of attention processors: A dictionary containing all attention processors used in the model with
            indexed by its weight name.
        """
        # set recursively
        processors = {}

        def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
            if hasattr(module, "get_processor"):
                processors[f"{name}.processor"] = module.get_processor()

            for sub_name, child in module.named_children():
                fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)

            return processors

        for name, module in self.named_children():
            fn_recursive_add_processors(name, module, processors)

        return processors

    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor
    def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
        r"""
        Sets the attention processor to use to compute attention.

        Parameters:
            processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
                The instantiated processor class or a dictionary of processor classes that will be set as the processor
                for **all** `Attention` layers.

                If `processor` is a dict, the key needs to define the path to the corresponding cross attention
                processor. This is strongly recommended when setting trainable attention processors.

        """
        count = len(self.attn_processors.keys())

        if isinstance(processor, dict) and len(processor) != count:
            raise ValueError(
                f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
                f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
            )

        def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
            if hasattr(module, "set_processor"):
                if not isinstance(processor, dict):
                    module.set_processor(processor)
                else:
                    module.set_processor(processor.pop(f"{name}.processor"))

            for sub_name, child in module.named_children():
                fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)

        for name, module in self.named_children():
            fn_recursive_attn_processor(name, module, processor)

    def set_default_attn_processor(self):
        """
        Disables custom attention processors and sets the default attention implementation.
        """
        if all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
            processor = AttnProcessor()
        else:
            raise ValueError(
                f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}"
            )

        self.set_attn_processor(processor)

    @apply_forward_hook
    def encode(
        self, 
        x: torch.Tensor, 
        latent_dist: DiagonalGaussianDistribution
    ) -> DiagonalGaussianDistribution:
        h = self.encoder(x)
        offset = h * self.offset_scale_factor
        param = latent_dist.parameters.to(h.dtype)
        mean, logvar = torch.chunk(param, 2, dim=1)
        posterior = DiagonalGaussianDistribution(torch.cat([mean + offset, logvar], dim=1))
        return posterior

    @apply_forward_hook
    def decode(
        self,
        z: torch.Tensor,
        num_frames: int
    ) -> torch.Tensor:
        batch_size = z.shape[0] // num_frames
        image_only_indicator = torch.zeros(batch_size, num_frames, dtype=z.dtype, device=z.device)
        decoded = self.decoder(z, num_frames=num_frames, image_only_indicator=image_only_indicator)
        return decoded