192 lines
6.8 KiB
Python
192 lines
6.8 KiB
Python
# 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)
|