custom_nodes/ComfyUI-AnimateDiff-Evolved/animatediff/adapter_animatelcm_i2v.py

# Modified from code provided by Fu-Yun Wang (G-U-N on github)
import math
import numpy as np
import torch
import torch.nn.functional as F
from torch import Tensor, nn

import comfy.ops
import comfy.model_management


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


def avg_pool_nd(dims, *args, **kwargs):
    """
    Create a 1D, 2D, or 3D average pooling module.
    """
    if dims == 1:
        return nn.AvgPool1d(*args, **kwargs)
    elif dims == 2:
        return nn.AvgPool2d(*args, **kwargs)
    elif dims == 3:
        return nn.AvgPool3d(*args, **kwargs)
    raise ValueError(f"unsupported dimensions: {dims}")


# based on PositionalEncoding of AnimateDiff
def fixed_positional_embedding(t, d_model):
    position = torch.arange(0, t, dtype=torch.float).unsqueeze(1)
    div_term = torch.exp(torch.arange(0, d_model, 2).float()
                         * (-np.log(10000.0) / d_model))
    pos_embedding = torch.zeros(t, d_model)
    pos_embedding[:, 0::2] = torch.sin(position * div_term)
    pos_embedding[:, 1::2] = torch.cos(position * div_term)
    return pos_embedding


class AdapterEmbed(nn.Module):
    def __init__(self, channels=[320, 640, 1280, 1280],
                 nums_rb=3, cin=64, ksize=3, sk=False, use_conv=True,
                 ops=comfy.ops.disable_weight_init):
        super(AdapterEmbed, self).__init__()
        self.channels = channels
        self.nums_rb = nums_rb
        self.body = []
        for i in range(len(channels)):
            for j in range(nums_rb):
                if (i != 0) and (j == 0):
                    self.body.append(ResnetBlockEmbed(
                        channels[i-1], channels[i], down=True, ksize=ksize, sk=sk, use_conv=use_conv, ops=ops
                    ))
                else:
                    self.body.append(ResnetBlockEmbed(
                        channels[i], channels[i], down=False, ksize=ksize, sk=sk, use_conv=use_conv, ops=ops
                    ))
        self.body = nn.ModuleList(self.body)
        self.conv_in = zero_module(ops.Conv2d(in_channels=cin, out_channels=channels[0],
                                              kernel_size=3, stride=1, padding=1))
        self.d_model = channels[0]
        # settable
        self.ref_drift = 0.5
        self.insertion_weights = [1.0, 1.0, 1.0, 1.0]

    def set_ref_drift(self, ref_drift: float):
        if ref_drift is None:
            ref_drift = 0.5
        self.ref_drift = ref_drift
    
    def set_insertion_weights(self, insertion_weights: list[float]):
        if insertion_weights is None:
            insertion_weights = [1.0, 1.0, 1.0, 1.0]
        assert len(insertion_weights) == 4
        self.insertion_weights = insertion_weights
    
    def cleanup(self):
        self.set_ref_drift(None)
        self.set_insertion_weights(None)

    def forward(self, x: Tensor, video_length: int, batched_number: int):
        b, c, h, w = x.shape

        features = []

        use_dtype = comfy.model_management.unet_dtype()
        # allow fp8 to work
        if comfy.model_management.dtype_size(use_dtype) == 1:
            use_dtype = x.dtype

        x = self.conv_in(x.to(use_dtype))

        pos_embedding = fixed_positional_embedding(
            video_length, self.d_model).to(use_dtype).to(x.device)
        pos_embedding = pos_embedding.unsqueeze(-1).unsqueeze(-1)
        pos_embedding = pos_embedding.expand(-1, -1, h, w)
        # add x_pos with influence amount
        x = x + (pos_embedding * self.ref_drift)

        for i in range(len(self.channels)):
            for j in range(self.nums_rb):
                # get real index in self.body that corresponds to current channel/resnetblock
                idx = i*self.nums_rb + j
                x = self.body[idx](x)
            # match real_x to batched_number
            real_x = x.repeat(batched_number, 1, 1, 1)
            features.append(real_x)
        features = [weight * feature for weight, feature in zip(features, self.insertion_weights)]
        return features


class ResnetBlockEmbed(nn.Module):
    def __init__(self, in_c, out_c, down: bool, ksize=3, sk=False, use_conv=True,
                 ops=comfy.ops.disable_weight_init):
        super().__init__()
        ps = ksize // 2 # padding size
        if in_c != out_c or sk == False:
            self.in_conv = zero_module(ops.Conv2d(in_c, out_c, ksize, 1, ps))
        else:
            self.in_conv = None
        self.block1 = ops.Conv2d(out_c, out_c, 3, 1, 1)
        self.act = nn.ReLU()
        self.block2 = zero_module(ops.Conv2d(out_c, out_c, ksize, 1, ps))
        if sk == False:
            self.skep = ops.Conv2d(in_c, out_c, ksize, 1, ps)
        else:
            self.skep = None
        
        self.down = down
        if self.down == True:
            self.down_opt = DownsampleEmbed(in_c, use_conv=use_conv, ops=ops)
    
    def forward(self, x: Tensor):
        if self.down == True:
            x = self.down_opt(x)
        
        if self.in_conv is not None:
            x = self.in_conv(x)
        
        h = self.block1(x)
        h = self.act(h)
        h = self.block2(h)

        if self.skep is not None:
            return h + self.skep(x)
        else:
            return h + x


class DownsampleEmbed(nn.Module):
    """
    A downsampling layer with an optional convolution.
    :param channels: channels in the inputs and outputs.
    :param use_conv: a bool determining if a convolution is applied.
    :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
                 downsampling occurs in the inner-two dimensions.
    """

    def __init__(self, channels, use_conv: bool, dims=2, out_channels=None, padding=1,
                 ops=comfy.ops.disable_weight_init):
        super().__init__()
        self.channels = channels
        self.out_channels = out_channels or channels  # use channels if out_channels is None
        self.use_conv = use_conv
        self.dims = dims
        stride = 2 if dims != 3 else (1, 2, 2)
        if use_conv:
            self.operation = ops.conv_nd(dims, in_channels=self.channels, out_channels=self.out_channels,
                                  kernel_size=3, stride=stride, padding=padding)
        else:
            assert self.channels == self.out_channels
            self.operation = avg_pool_nd(dims, kernel_size=stride, stride=stride)  # both are stride value on purpose
    
    def forward(self, x: Tensor):
        assert x.shape[1] == self.channels

        kernel_size = (2, 2)

        input_height, input_width = x.size(-2), x.size(-1)

        padding_height = (
            math.ceil(input_height / kernel_size[0]) * kernel_size[0]) - input_height
        padding_width = (
            math.ceil(input_width / kernel_size[1]) * kernel_size[1]) - input_width

        x = F.pad(x, (0, padding_width, 0, padding_height), mode='replicate')

        return self.operation(x)