1036 lines
42 KiB
Python
1036 lines
42 KiB
Python
#taken from: https://github.com/lllyasviel/ControlNet
|
|
#and modified
|
|
#and then taken from comfy/cldm/cldm.py and modified again
|
|
|
|
import math
|
|
from typing import Iterable, Union
|
|
import torch
|
|
import torch as th
|
|
import torch.nn as nn
|
|
from torch import Tensor
|
|
from einops import rearrange, repeat
|
|
|
|
from comfy.ldm.modules.diffusionmodules.util import (
|
|
zero_module,
|
|
timestep_embedding,
|
|
)
|
|
|
|
from comfy.cldm.cldm import ControlNet as ControlNet_cldm
|
|
from comfy.ldm.modules.attention import SpatialTransformer
|
|
from comfy.ldm.modules.diffusionmodules.openaimodel import TimestepEmbedSequential, ResBlock, Downsample
|
|
from comfy.ldm.util import exists
|
|
from comfy.ldm.modules.attention import default, optimized_attention
|
|
from comfy.ldm.modules.attention import FeedForward, SpatialTransformer
|
|
from comfy.controlnet import broadcast_image_to
|
|
from comfy.utils import repeat_to_batch_size
|
|
import comfy.ops
|
|
|
|
from .utils import TimestepKeyframeGroup, disable_weight_init_clean_groupnorm, prepare_mask_batch
|
|
|
|
|
|
class SparseControlNet(ControlNet_cldm):
|
|
def __init__(self, *args,**kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
hint_channels = kwargs.get("hint_channels")
|
|
operations: disable_weight_init_clean_groupnorm = kwargs.get("operations", disable_weight_init_clean_groupnorm)
|
|
device = kwargs.get("device", None)
|
|
use_simplified_conditioning_embedding = kwargs.get("use_simplified_conditioning_embedding", False)
|
|
if use_simplified_conditioning_embedding:
|
|
self.input_hint_block = TimestepEmbedSequential(
|
|
operations.conv_nd(self.dims, hint_channels, self.model_channels, 3, padding=1, dtype=self.dtype, device=device),
|
|
)
|
|
|
|
def forward(self, x, hint, timesteps, context, y=None, **kwargs):
|
|
t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False).to(x.dtype)
|
|
emb = self.time_embed(t_emb)
|
|
|
|
x = torch.zeros_like(x)
|
|
|
|
conditioning_mask1 = torch.ones_like(hint[:, :1])
|
|
conditioning_mask2 = torch.zeros_like(hint[:, :1])
|
|
conditioning_mask = conditioning_mask2
|
|
conditioning_mask[0] = conditioning_mask1[0]
|
|
conditioning_mask[16] = conditioning_mask1[16]
|
|
#conditioning_mask[15] = conditioning_mask1[15]
|
|
#conditioning_mask[31] = conditioning_mask1[31]
|
|
modified_hint = torch.zeros_like(hint)
|
|
modified_hint[0] = hint[0]
|
|
modified_hint[16] = hint[16]
|
|
#modified_hint[15] = hint[15]
|
|
#modified_hint[31] = hint[31]
|
|
hint = torch.cat([modified_hint, conditioning_mask], dim=1)
|
|
guided_hint = self.input_hint_block(hint, emb, context)
|
|
|
|
outs = []
|
|
|
|
hs = []
|
|
if self.num_classes is not None:
|
|
assert y.shape[0] == x.shape[0]
|
|
emb = emb + self.label_emb(y)
|
|
|
|
h = x
|
|
for module, zero_conv in zip(self.input_blocks, self.zero_convs):
|
|
if guided_hint is not None:
|
|
h = module(h, emb, context)
|
|
h += guided_hint
|
|
guided_hint = None
|
|
else:
|
|
h = module(h, emb, context)
|
|
outs.append(zero_conv(h, emb, context))
|
|
|
|
h = self.middle_block(h, emb, context)
|
|
outs.append(self.middle_block_out(h, emb, context))
|
|
|
|
return outs
|
|
|
|
|
|
|
|
# main class for holding SparseControlNet
|
|
class SparseControlNetOld(nn.Module):
|
|
def __init__(
|
|
self,
|
|
image_size,
|
|
in_channels,
|
|
model_channels,
|
|
hint_channels,
|
|
num_res_blocks,
|
|
dropout=0,
|
|
channel_mult=(1, 2, 4, 8),
|
|
conv_resample=True,
|
|
dims=2,
|
|
num_classes=None,
|
|
use_checkpoint=False,
|
|
dtype=torch.float32,
|
|
num_heads=-1,
|
|
num_head_channels=-1,
|
|
num_heads_upsample=-1,
|
|
use_scale_shift_norm=False,
|
|
resblock_updown=False,
|
|
use_new_attention_order=False,
|
|
use_spatial_transformer=False, # custom transformer support
|
|
transformer_depth=1, # custom transformer support
|
|
context_dim=None, # custom transformer support
|
|
n_embed=None, # custom support for prediction of discrete ids into codebook of first stage vq model
|
|
legacy=True,
|
|
disable_self_attentions=None,
|
|
num_attention_blocks=None,
|
|
disable_middle_self_attn=False,
|
|
use_linear_in_transformer=False,
|
|
adm_in_channels=None,
|
|
transformer_depth_middle=None,
|
|
transformer_depth_output=None,
|
|
device=None,
|
|
operations=disable_weight_init_clean_groupnorm,
|
|
**kwargs,
|
|
):
|
|
super().__init__()
|
|
assert use_spatial_transformer == True, "use_spatial_transformer has to be true"
|
|
if use_spatial_transformer:
|
|
assert context_dim is not None, 'Fool!! You forgot to include the dimension of your cross-attention conditioning...'
|
|
|
|
if context_dim is not None:
|
|
assert use_spatial_transformer, 'Fool!! You forgot to use the spatial transformer for your cross-attention conditioning...'
|
|
# from omegaconf.listconfig import ListConfig
|
|
# if type(context_dim) == ListConfig:
|
|
# context_dim = list(context_dim)
|
|
if num_heads_upsample == -1:
|
|
num_heads_upsample = num_heads
|
|
|
|
if num_heads == -1:
|
|
assert num_head_channels != -1, 'Either num_heads or num_head_channels has to be set'
|
|
|
|
if num_head_channels == -1:
|
|
assert num_heads != -1, 'Either num_heads or num_head_channels has to be set'
|
|
|
|
self.dims = dims
|
|
self.image_size = image_size
|
|
self.in_channels = in_channels
|
|
self.model_channels = model_channels
|
|
|
|
if isinstance(num_res_blocks, int):
|
|
self.num_res_blocks = len(channel_mult) * [num_res_blocks]
|
|
else:
|
|
if len(num_res_blocks) != len(channel_mult):
|
|
raise ValueError("provide num_res_blocks either as an int (globally constant) or "
|
|
"as a list/tuple (per-level) with the same length as channel_mult")
|
|
self.num_res_blocks = num_res_blocks
|
|
|
|
if disable_self_attentions is not None:
|
|
# should be a list of booleans, indicating whether to disable self-attention in TransformerBlocks or not
|
|
assert len(disable_self_attentions) == len(channel_mult)
|
|
if num_attention_blocks is not None:
|
|
assert len(num_attention_blocks) == len(self.num_res_blocks)
|
|
assert all(map(lambda i: self.num_res_blocks[i] >= num_attention_blocks[i], range(len(num_attention_blocks))))
|
|
|
|
transformer_depth = transformer_depth[:]
|
|
|
|
self.dropout = dropout
|
|
self.channel_mult = channel_mult
|
|
self.conv_resample = conv_resample
|
|
self.num_classes = num_classes
|
|
self.use_checkpoint = use_checkpoint
|
|
self.dtype = dtype
|
|
self.num_heads = num_heads
|
|
self.num_head_channels = num_head_channels
|
|
self.num_heads_upsample = num_heads_upsample
|
|
self.predict_codebook_ids = n_embed is not None
|
|
|
|
time_embed_dim = model_channels * 4
|
|
self.time_embed = nn.Sequential(
|
|
operations.Linear(model_channels, time_embed_dim, dtype=self.dtype, device=device),
|
|
nn.SiLU(),
|
|
operations.Linear(time_embed_dim, time_embed_dim, dtype=self.dtype, device=device),
|
|
)
|
|
|
|
if self.num_classes is not None:
|
|
if isinstance(self.num_classes, int):
|
|
self.label_emb = nn.Embedding(num_classes, time_embed_dim)
|
|
elif self.num_classes == "continuous":
|
|
print("setting up linear c_adm embedding layer")
|
|
self.label_emb = nn.Linear(1, time_embed_dim)
|
|
elif self.num_classes == "sequential":
|
|
assert adm_in_channels is not None
|
|
self.label_emb = nn.Sequential(
|
|
nn.Sequential(
|
|
operations.Linear(adm_in_channels, time_embed_dim, dtype=self.dtype, device=device),
|
|
nn.SiLU(),
|
|
operations.Linear(time_embed_dim, time_embed_dim, dtype=self.dtype, device=device),
|
|
)
|
|
)
|
|
else:
|
|
raise ValueError()
|
|
|
|
self.input_blocks = nn.ModuleList(
|
|
[
|
|
TimestepEmbedSequential(
|
|
operations.conv_nd(dims, in_channels, model_channels, 3, padding=1, dtype=self.dtype, device=device)
|
|
)
|
|
]
|
|
)
|
|
self.zero_convs = nn.ModuleList([self.make_zero_conv(model_channels, operations=operations, dtype=self.dtype, device=device)])
|
|
|
|
self.input_hint_block = TimestepEmbedSequential(
|
|
operations.conv_nd(dims, hint_channels, 16, 3, padding=1, dtype=self.dtype, device=device),
|
|
nn.SiLU(),
|
|
operations.conv_nd(dims, 16, 16, 3, padding=1, dtype=self.dtype, device=device),
|
|
nn.SiLU(),
|
|
operations.conv_nd(dims, 16, 32, 3, padding=1, stride=2, dtype=self.dtype, device=device),
|
|
nn.SiLU(),
|
|
operations.conv_nd(dims, 32, 32, 3, padding=1, dtype=self.dtype, device=device),
|
|
nn.SiLU(),
|
|
operations.conv_nd(dims, 32, 96, 3, padding=1, stride=2, dtype=self.dtype, device=device),
|
|
nn.SiLU(),
|
|
operations.conv_nd(dims, 96, 96, 3, padding=1, dtype=self.dtype, device=device),
|
|
nn.SiLU(),
|
|
operations.conv_nd(dims, 96, 256, 3, padding=1, stride=2, dtype=self.dtype, device=device),
|
|
nn.SiLU(),
|
|
operations.conv_nd(dims, 256, model_channels, 3, padding=1, dtype=self.dtype, device=device)
|
|
)
|
|
|
|
self._feature_size = model_channels
|
|
input_block_chans = [model_channels]
|
|
ch = model_channels
|
|
ds = 1
|
|
for level, mult in enumerate(channel_mult):
|
|
for nr in range(self.num_res_blocks[level]):
|
|
layers = [
|
|
ResBlock(
|
|
ch,
|
|
time_embed_dim,
|
|
dropout,
|
|
out_channels=mult * model_channels,
|
|
dims=dims,
|
|
use_checkpoint=use_checkpoint,
|
|
use_scale_shift_norm=use_scale_shift_norm,
|
|
dtype=self.dtype,
|
|
device=device,
|
|
operations=operations,
|
|
)
|
|
]
|
|
ch = mult * model_channels
|
|
num_transformers = transformer_depth.pop(0)
|
|
if num_transformers > 0:
|
|
if num_head_channels == -1:
|
|
dim_head = ch // num_heads
|
|
else:
|
|
num_heads = ch // num_head_channels
|
|
dim_head = num_head_channels
|
|
if legacy:
|
|
#num_heads = 1
|
|
dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
|
|
if exists(disable_self_attentions):
|
|
disabled_sa = disable_self_attentions[level]
|
|
else:
|
|
disabled_sa = False
|
|
|
|
if not exists(num_attention_blocks) or nr < num_attention_blocks[level]:
|
|
layers.append(
|
|
SpatialTransformer(
|
|
ch, num_heads, dim_head, depth=num_transformers, context_dim=context_dim,
|
|
disable_self_attn=disabled_sa, use_linear=use_linear_in_transformer,
|
|
use_checkpoint=use_checkpoint, dtype=self.dtype, device=device, operations=operations
|
|
)
|
|
)
|
|
self.input_blocks.append(TimestepEmbedSequential(*layers))
|
|
self.zero_convs.append(self.make_zero_conv(ch, operations=operations, dtype=self.dtype, device=device))
|
|
self._feature_size += ch
|
|
input_block_chans.append(ch)
|
|
if level != len(channel_mult) - 1:
|
|
out_ch = ch
|
|
self.input_blocks.append(
|
|
TimestepEmbedSequential(
|
|
ResBlock(
|
|
ch,
|
|
time_embed_dim,
|
|
dropout,
|
|
out_channels=out_ch,
|
|
dims=dims,
|
|
use_checkpoint=use_checkpoint,
|
|
use_scale_shift_norm=use_scale_shift_norm,
|
|
down=True,
|
|
dtype=self.dtype,
|
|
device=device,
|
|
operations=operations
|
|
)
|
|
if resblock_updown
|
|
else Downsample(
|
|
ch, conv_resample, dims=dims, out_channels=out_ch, dtype=self.dtype, device=device, operations=operations
|
|
)
|
|
)
|
|
)
|
|
ch = out_ch
|
|
input_block_chans.append(ch)
|
|
self.zero_convs.append(self.make_zero_conv(ch, operations=operations, dtype=self.dtype, device=device))
|
|
ds *= 2
|
|
self._feature_size += ch
|
|
|
|
if num_head_channels == -1:
|
|
dim_head = ch // num_heads
|
|
else:
|
|
num_heads = ch // num_head_channels
|
|
dim_head = num_head_channels
|
|
if legacy:
|
|
#num_heads = 1
|
|
dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
|
|
mid_block = [
|
|
ResBlock(
|
|
ch,
|
|
time_embed_dim,
|
|
dropout,
|
|
dims=dims,
|
|
use_checkpoint=use_checkpoint,
|
|
use_scale_shift_norm=use_scale_shift_norm,
|
|
dtype=self.dtype,
|
|
device=device,
|
|
operations=operations
|
|
)]
|
|
if transformer_depth_middle >= 0:
|
|
mid_block += [SpatialTransformer( # always uses a self-attn
|
|
ch, num_heads, dim_head, depth=transformer_depth_middle, context_dim=context_dim,
|
|
disable_self_attn=disable_middle_self_attn, use_linear=use_linear_in_transformer,
|
|
use_checkpoint=use_checkpoint, dtype=self.dtype, device=device, operations=operations
|
|
),
|
|
ResBlock(
|
|
ch,
|
|
time_embed_dim,
|
|
dropout,
|
|
dims=dims,
|
|
use_checkpoint=use_checkpoint,
|
|
use_scale_shift_norm=use_scale_shift_norm,
|
|
dtype=self.dtype,
|
|
device=device,
|
|
operations=operations
|
|
)]
|
|
self.middle_block = TimestepEmbedSequential(*mid_block)
|
|
self.middle_block_out = self.make_zero_conv(ch, operations=operations, dtype=self.dtype, device=device)
|
|
self._feature_size += ch
|
|
|
|
#self._motion_wrapper: SparseCtrlMotionWrapper = None
|
|
|
|
def make_zero_conv(self, channels, operations=None, dtype=None, device=None):
|
|
return TimestepEmbedSequential(operations.conv_nd(self.dims, channels, channels, 1, padding=0, dtype=dtype, device=device))
|
|
|
|
def forward(self, x, hint, timesteps, context, y=None, **kwargs):
|
|
t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False).to(x.dtype)
|
|
emb = self.time_embed(t_emb)
|
|
|
|
x = torch.zeros_like(x)
|
|
|
|
conditioning_mask1 = torch.ones_like(hint[:, :1])
|
|
conditioning_mask2 = torch.zeros_like(hint[:, :1])
|
|
conditioning_mask = conditioning_mask2
|
|
conditioning_mask[0] = conditioning_mask1[0]
|
|
conditioning_mask[16] = conditioning_mask1[16]
|
|
#conditioning_mask[15] = conditioning_mask1[15]
|
|
#conditioning_mask[31] = conditioning_mask1[31]
|
|
modified_hint = torch.zeros_like(hint)
|
|
modified_hint[0] = hint[0]
|
|
modified_hint[16] = hint[16]
|
|
#modified_hint[15] = hint[15]
|
|
#modified_hint[31] = hint[31]
|
|
hint = torch.cat([modified_hint, conditioning_mask], dim=1)
|
|
guided_hint = self.input_hint_block(hint, emb, context)
|
|
|
|
outs = []
|
|
|
|
hs = []
|
|
if self.num_classes is not None:
|
|
assert y.shape[0] == x.shape[0]
|
|
emb = emb + self.label_emb(y)
|
|
|
|
h = x
|
|
for module, zero_conv in zip(self.input_blocks, self.zero_convs):
|
|
if guided_hint is not None:
|
|
h = module(h, emb, context)
|
|
h += guided_hint
|
|
guided_hint = None
|
|
else:
|
|
h = module(h, emb, context)
|
|
outs.append(zero_conv(h, emb, context))
|
|
|
|
h = self.middle_block(h, emb, context)
|
|
outs.append(self.middle_block_out(h, emb, context))
|
|
|
|
return outs
|
|
|
|
|
|
# motion-related portion of controlnet
|
|
class BlockType:
|
|
UP = "up"
|
|
DOWN = "down"
|
|
MID = "mid"
|
|
|
|
def get_down_block_max(mm_state_dict: dict[str, Tensor]) -> int:
|
|
return get_block_max(mm_state_dict, "down_blocks")
|
|
|
|
def get_up_block_max(mm_state_dict: dict[str, Tensor]) -> int:
|
|
return get_block_max(mm_state_dict, "up_blocks")
|
|
|
|
def get_block_max(mm_state_dict: dict[str, Tensor], block_name: str) -> int:
|
|
# keep track of biggest down_block count in module
|
|
biggest_block = -1
|
|
for key in mm_state_dict.keys():
|
|
if block_name in key:
|
|
try:
|
|
block_int = key.split(".")[1]
|
|
block_num = int(block_int)
|
|
if block_num > biggest_block:
|
|
biggest_block = block_num
|
|
except ValueError:
|
|
pass
|
|
return biggest_block
|
|
|
|
def has_mid_block(mm_state_dict: dict[str, Tensor]):
|
|
# check if keys contain mid_block
|
|
for key in mm_state_dict.keys():
|
|
if key.startswith("mid_block."):
|
|
return True
|
|
return False
|
|
|
|
def get_position_encoding_max_len(mm_state_dict: dict[str, Tensor], mm_name: str=None) -> int:
|
|
# use pos_encoder.pe entries to determine max length - [1, {max_length}, {320|640|1280}]
|
|
for key in mm_state_dict.keys():
|
|
if key.endswith("pos_encoder.pe"):
|
|
return mm_state_dict[key].size(1) # get middle dim
|
|
raise ValueError(f"No pos_encoder.pe found in SparseCtrl state_dict - {mm_name} is not a valid SparseCtrl model!")
|
|
|
|
|
|
class SparseCtrlMotionWrapper(nn.Module):
|
|
def __init__(self, mm_state_dict: dict[str, Tensor]):
|
|
super().__init__()
|
|
self.down_blocks: Iterable[MotionModule] = None
|
|
self.up_blocks: Iterable[MotionModule] = None
|
|
self.mid_block: MotionModule = None
|
|
self.encoding_max_len = get_position_encoding_max_len(mm_state_dict, "")
|
|
layer_channels = (320, 640, 1280, 1280)
|
|
if get_down_block_max(mm_state_dict) > -1:
|
|
self.down_blocks = nn.ModuleList([])
|
|
for c in layer_channels:
|
|
self.down_blocks.append(MotionModule(c, temporal_position_encoding_max_len=self.encoding_max_len, block_type=BlockType.DOWN))
|
|
if get_up_block_max(mm_state_dict) > -1:
|
|
self.up_blocks = nn.ModuleList([])
|
|
for c in reversed(layer_channels):
|
|
self.up_blocks.append(MotionModule(c, temporal_position_encoding_max_len=self.encoding_max_len, block_type=BlockType.UP))
|
|
if has_mid_block(mm_state_dict):
|
|
self.mid_block = MotionModule(1280, temporal_position_encoding_max_len=self.encoding_max_len, block_type=BlockType.MID)
|
|
|
|
def inject(self, unet: SparseControlNet):
|
|
# inject input (down) blocks
|
|
self._inject(unet.input_blocks, self.down_blocks)
|
|
# inject mid block, if present
|
|
if self.mid_block is not None:
|
|
self._inject([unet.middle_block], [self.mid_block])
|
|
#unet._motion_wrapper = self
|
|
|
|
def _inject(self, unet_blocks: nn.ModuleList, mm_blocks: nn.ModuleList):
|
|
# Rules for injection:
|
|
# For each component list in a unet block:
|
|
# if SpatialTransformer exists in list, place next block after last occurrence
|
|
# elif ResBlock exists in list, place next block after first occurrence
|
|
# else don't place block
|
|
injection_count = 0
|
|
unet_idx = 0
|
|
# details about blocks passed in
|
|
per_block = len(mm_blocks[0].motion_modules)
|
|
injection_goal = len(mm_blocks) * per_block
|
|
# only stop injecting when modules exhausted
|
|
while injection_count < injection_goal:
|
|
# figure out which VanillaTemporalModule from mm to inject
|
|
mm_blk_idx, mm_vtm_idx = injection_count // per_block, injection_count % per_block
|
|
# figure out layout of unet block components
|
|
st_idx = -1 # SpatialTransformer index
|
|
res_idx = -1 # first ResBlock index
|
|
# first, figure out indeces of relevant blocks
|
|
for idx, component in enumerate(unet_blocks[unet_idx]):
|
|
if type(component) == SpatialTransformer:
|
|
st_idx = idx
|
|
elif type(component).__name__ == "ResBlock" and res_idx < 0:
|
|
res_idx = idx
|
|
# if SpatialTransformer exists, inject right after
|
|
if st_idx >= 0:
|
|
unet_blocks[unet_idx].insert(st_idx+1, mm_blocks[mm_blk_idx].motion_modules[mm_vtm_idx])
|
|
injection_count += 1
|
|
# otherwise, if only ResBlock exists, inject right after
|
|
elif res_idx >= 0:
|
|
unet_blocks[unet_idx].insert(res_idx+1, mm_blocks[mm_blk_idx].motion_modules[mm_vtm_idx])
|
|
injection_count += 1
|
|
# increment unet_idx
|
|
unet_idx += 1
|
|
|
|
def eject(self, unet: SparseControlNet):
|
|
# remove from input blocks (downblocks)
|
|
self._eject(unet.input_blocks)
|
|
# remove from middle block (encapsulate in list to make compatible)
|
|
self._eject([unet.middle_block])
|
|
#del unet._motion_wrapper
|
|
|
|
def _eject(self, unet_blocks: nn.ModuleList):
|
|
# eject all VanillaTemporalModule objects from all blocks
|
|
for block in unet_blocks:
|
|
idx_to_pop = []
|
|
for idx, component in enumerate(block):
|
|
if type(component) == VanillaTemporalModule:
|
|
idx_to_pop.append(idx)
|
|
# pop in backwards order, as to not disturb what the indeces refer to
|
|
for idx in sorted(idx_to_pop, reverse=True):
|
|
block.pop(idx)
|
|
|
|
def set_video_length(self, video_length: int, full_length: int):
|
|
self.AD_video_length = video_length
|
|
for block in self.down_blocks:
|
|
block.set_video_length(video_length, full_length)
|
|
for block in self.up_blocks:
|
|
block.set_video_length(video_length, full_length)
|
|
if self.mid_block is not None:
|
|
self.mid_block.set_video_length(video_length, full_length)
|
|
|
|
def set_scale_multiplier(self, multiplier: Union[float, None]):
|
|
for block in self.down_blocks:
|
|
block.set_scale_multiplier(multiplier)
|
|
for block in self.up_blocks:
|
|
block.set_scale_multiplier(multiplier)
|
|
if self.mid_block is not None:
|
|
self.mid_block.set_scale_multiplier(multiplier)
|
|
|
|
def reset_temp_vars(self):
|
|
for block in self.down_blocks:
|
|
block.reset_temp_vars()
|
|
for block in self.up_blocks:
|
|
block.reset_temp_vars()
|
|
if self.mid_block is not None:
|
|
self.mid_block.reset_temp_vars()
|
|
|
|
def reset_scale_multiplier(self):
|
|
self.set_scale_multiplier(None)
|
|
|
|
def reset(self):
|
|
self.reset_scale_multiplier()
|
|
self.reset_temp_vars()
|
|
|
|
|
|
class MotionModule(nn.Module):
|
|
def __init__(self, in_channels, temporal_position_encoding_max_len=24, block_type: str=BlockType.DOWN):
|
|
super().__init__()
|
|
if block_type == BlockType.MID:
|
|
# mid blocks contain only a single VanillaTemporalModule
|
|
self.motion_modules: Iterable[VanillaTemporalModule] = nn.ModuleList([get_motion_module(in_channels, temporal_position_encoding_max_len)])
|
|
else:
|
|
# down blocks contain two VanillaTemporalModules
|
|
self.motion_modules: Iterable[VanillaTemporalModule] = nn.ModuleList(
|
|
[
|
|
get_motion_module(in_channels, temporal_position_encoding_max_len),
|
|
get_motion_module(in_channels, temporal_position_encoding_max_len)
|
|
]
|
|
)
|
|
# up blocks contain one additional VanillaTemporalModule
|
|
if block_type == BlockType.UP:
|
|
self.motion_modules.append(get_motion_module(in_channels, temporal_position_encoding_max_len))
|
|
|
|
def set_video_length(self, video_length: int, full_length: int):
|
|
for motion_module in self.motion_modules:
|
|
motion_module.set_video_length(video_length, full_length)
|
|
|
|
def set_scale_multiplier(self, multiplier: Union[float, None]):
|
|
for motion_module in self.motion_modules:
|
|
motion_module.set_scale_multiplier(multiplier)
|
|
|
|
def set_masks(self, masks: Tensor, min_val: float, max_val: float):
|
|
for motion_module in self.motion_modules:
|
|
motion_module.set_masks(masks, min_val, max_val)
|
|
|
|
def set_sub_idxs(self, sub_idxs: list[int]):
|
|
for motion_module in self.motion_modules:
|
|
motion_module.set_sub_idxs(sub_idxs)
|
|
|
|
def reset_temp_vars(self):
|
|
for motion_module in self.motion_modules:
|
|
motion_module.reset_temp_vars()
|
|
|
|
|
|
def get_motion_module(in_channels, temporal_position_encoding_max_len):
|
|
# unlike normal AD, there is only one attention block expected in SparseCtrl models
|
|
return VanillaTemporalModule(in_channels=in_channels, attention_block_types=("Temporal_Self",), temporal_position_encoding_max_len=temporal_position_encoding_max_len)
|
|
|
|
|
|
class VanillaTemporalModule(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_channels,
|
|
num_attention_heads=8,
|
|
num_transformer_block=1,
|
|
attention_block_types=("Temporal_Self", "Temporal_Self"),
|
|
cross_frame_attention_mode=None,
|
|
temporal_position_encoding=True,
|
|
temporal_position_encoding_max_len=24,
|
|
temporal_attention_dim_div=1,
|
|
zero_initialize=True,
|
|
):
|
|
super().__init__()
|
|
|
|
self.temporal_transformer = TemporalTransformer3DModel(
|
|
in_channels=in_channels,
|
|
num_attention_heads=num_attention_heads,
|
|
attention_head_dim=in_channels
|
|
// num_attention_heads
|
|
// temporal_attention_dim_div,
|
|
num_layers=num_transformer_block,
|
|
attention_block_types=attention_block_types,
|
|
cross_frame_attention_mode=cross_frame_attention_mode,
|
|
temporal_position_encoding=temporal_position_encoding,
|
|
temporal_position_encoding_max_len=temporal_position_encoding_max_len,
|
|
)
|
|
|
|
if zero_initialize:
|
|
self.temporal_transformer.proj_out = zero_module(
|
|
self.temporal_transformer.proj_out
|
|
)
|
|
|
|
def set_video_length(self, video_length: int, full_length: int):
|
|
self.temporal_transformer.set_video_length(video_length, full_length)
|
|
|
|
def set_scale_multiplier(self, multiplier: Union[float, None]):
|
|
self.temporal_transformer.set_scale_multiplier(multiplier)
|
|
|
|
def set_masks(self, masks: Tensor, min_val: float, max_val: float):
|
|
self.temporal_transformer.set_masks(masks, min_val, max_val)
|
|
|
|
def set_sub_idxs(self, sub_idxs: list[int]):
|
|
self.temporal_transformer.set_sub_idxs(sub_idxs)
|
|
|
|
def reset_temp_vars(self):
|
|
self.temporal_transformer.reset_temp_vars()
|
|
|
|
def forward(self, input_tensor, encoder_hidden_states=None, attention_mask=None):
|
|
return self.temporal_transformer(input_tensor, encoder_hidden_states, attention_mask)
|
|
|
|
|
|
class TemporalTransformer3DModel(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_channels,
|
|
num_attention_heads,
|
|
attention_head_dim,
|
|
num_layers,
|
|
attention_block_types=(
|
|
"Temporal_Self",
|
|
"Temporal_Self",
|
|
),
|
|
dropout=0.0,
|
|
norm_num_groups=32,
|
|
cross_attention_dim=768,
|
|
activation_fn="geglu",
|
|
attention_bias=False,
|
|
upcast_attention=False,
|
|
cross_frame_attention_mode=None,
|
|
temporal_position_encoding=False,
|
|
temporal_position_encoding_max_len=24,
|
|
):
|
|
super().__init__()
|
|
self.video_length = 16
|
|
self.full_length = 16
|
|
self.scale_min = 1.0
|
|
self.scale_max = 1.0
|
|
self.raw_scale_mask: Union[Tensor, None] = None
|
|
self.temp_scale_mask: Union[Tensor, None] = None
|
|
self.sub_idxs: Union[list[int], None] = None
|
|
self.prev_hidden_states_batch = 0
|
|
|
|
|
|
inner_dim = num_attention_heads * attention_head_dim
|
|
|
|
self.norm = disable_weight_init_clean_groupnorm.GroupNorm(
|
|
num_groups=norm_num_groups, num_channels=in_channels, eps=1e-6, affine=True
|
|
)
|
|
self.proj_in = nn.Linear(in_channels, inner_dim)
|
|
|
|
self.transformer_blocks: Iterable[TemporalTransformerBlock] = nn.ModuleList(
|
|
[
|
|
TemporalTransformerBlock(
|
|
dim=inner_dim,
|
|
num_attention_heads=num_attention_heads,
|
|
attention_head_dim=attention_head_dim,
|
|
attention_block_types=attention_block_types,
|
|
dropout=dropout,
|
|
norm_num_groups=norm_num_groups,
|
|
cross_attention_dim=cross_attention_dim,
|
|
activation_fn=activation_fn,
|
|
attention_bias=attention_bias,
|
|
upcast_attention=upcast_attention,
|
|
cross_frame_attention_mode=cross_frame_attention_mode,
|
|
temporal_position_encoding=temporal_position_encoding,
|
|
temporal_position_encoding_max_len=temporal_position_encoding_max_len,
|
|
)
|
|
for d in range(num_layers)
|
|
]
|
|
)
|
|
self.proj_out = nn.Linear(inner_dim, in_channels)
|
|
|
|
def set_video_length(self, video_length: int, full_length: int):
|
|
self.video_length = video_length
|
|
self.full_length = full_length
|
|
|
|
def set_scale_multiplier(self, multiplier: Union[float, None]):
|
|
for block in self.transformer_blocks:
|
|
block.set_scale_multiplier(multiplier)
|
|
|
|
def set_masks(self, masks: Tensor, min_val: float, max_val: float):
|
|
self.scale_min = min_val
|
|
self.scale_max = max_val
|
|
self.raw_scale_mask = masks
|
|
|
|
def set_sub_idxs(self, sub_idxs: list[int]):
|
|
self.sub_idxs = sub_idxs
|
|
for block in self.transformer_blocks:
|
|
block.set_sub_idxs(sub_idxs)
|
|
|
|
def reset_temp_vars(self):
|
|
del self.temp_scale_mask
|
|
self.temp_scale_mask = None
|
|
self.prev_hidden_states_batch = 0
|
|
|
|
def get_scale_mask(self, hidden_states: Tensor) -> Union[Tensor, None]:
|
|
# if no raw mask, return None
|
|
if self.raw_scale_mask is None:
|
|
return None
|
|
shape = hidden_states.shape
|
|
batch, channel, height, width = shape
|
|
# if temp mask already calculated, return it
|
|
if self.temp_scale_mask != None:
|
|
# check if hidden_states batch matches
|
|
if batch == self.prev_hidden_states_batch:
|
|
if self.sub_idxs is not None:
|
|
return self.temp_scale_mask[:, self.sub_idxs, :]
|
|
return self.temp_scale_mask
|
|
# if does not match, reset cached temp_scale_mask and recalculate it
|
|
del self.temp_scale_mask
|
|
self.temp_scale_mask = None
|
|
# otherwise, calculate temp mask
|
|
self.prev_hidden_states_batch = batch
|
|
mask = prepare_mask_batch(self.raw_scale_mask, shape=(self.full_length, 1, height, width))
|
|
mask = repeat_to_batch_size(mask, self.full_length)
|
|
# if mask not the same amount length as full length, make it match
|
|
if self.full_length != mask.shape[0]:
|
|
mask = broadcast_image_to(mask, self.full_length, 1)
|
|
# reshape mask to attention K shape (h*w, latent_count, 1)
|
|
batch, channel, height, width = mask.shape
|
|
# first, perform same operations as on hidden_states,
|
|
# turning (b, c, h, w) -> (b, h*w, c)
|
|
mask = mask.permute(0, 2, 3, 1).reshape(batch, height*width, channel)
|
|
# then, make it the same shape as attention's k, (h*w, b, c)
|
|
mask = mask.permute(1, 0, 2)
|
|
# make masks match the expected length of h*w
|
|
batched_number = shape[0] // self.video_length
|
|
if batched_number > 1:
|
|
mask = torch.cat([mask] * batched_number, dim=0)
|
|
# cache mask and set to proper device
|
|
self.temp_scale_mask = mask
|
|
# move temp_scale_mask to proper dtype + device
|
|
self.temp_scale_mask = self.temp_scale_mask.to(dtype=hidden_states.dtype, device=hidden_states.device)
|
|
# return subset of masks, if needed
|
|
if self.sub_idxs is not None:
|
|
return self.temp_scale_mask[:, self.sub_idxs, :]
|
|
return self.temp_scale_mask
|
|
|
|
def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None):
|
|
batch, channel, height, width = hidden_states.shape
|
|
residual = hidden_states
|
|
scale_mask = self.get_scale_mask(hidden_states)
|
|
# add some casts for fp8 purposes - does not affect speed otherwise
|
|
hidden_states = self.norm(hidden_states).to(hidden_states.dtype)
|
|
inner_dim = hidden_states.shape[1]
|
|
hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(
|
|
batch, height * width, inner_dim
|
|
)
|
|
hidden_states = self.proj_in(hidden_states).to(hidden_states.dtype)
|
|
|
|
# Transformer Blocks
|
|
for block in self.transformer_blocks:
|
|
hidden_states = block(
|
|
hidden_states,
|
|
encoder_hidden_states=encoder_hidden_states,
|
|
attention_mask=attention_mask,
|
|
video_length=self.video_length,
|
|
scale_mask=scale_mask
|
|
)
|
|
|
|
# output
|
|
hidden_states = self.proj_out(hidden_states)
|
|
hidden_states = (
|
|
hidden_states.reshape(batch, height, width, inner_dim)
|
|
.permute(0, 3, 1, 2)
|
|
.contiguous()
|
|
)
|
|
|
|
output = hidden_states + residual
|
|
|
|
return output
|
|
|
|
|
|
class TemporalTransformerBlock(nn.Module):
|
|
def __init__(
|
|
self,
|
|
dim,
|
|
num_attention_heads,
|
|
attention_head_dim,
|
|
attention_block_types=(
|
|
"Temporal_Self",
|
|
"Temporal_Self",
|
|
),
|
|
dropout=0.0,
|
|
norm_num_groups=32,
|
|
cross_attention_dim=768,
|
|
activation_fn="geglu",
|
|
attention_bias=False,
|
|
upcast_attention=False,
|
|
cross_frame_attention_mode=None,
|
|
temporal_position_encoding=False,
|
|
temporal_position_encoding_max_len=24,
|
|
):
|
|
super().__init__()
|
|
|
|
attention_blocks = []
|
|
norms = []
|
|
|
|
for block_name in attention_block_types:
|
|
attention_blocks.append(
|
|
VersatileAttention(
|
|
attention_mode=block_name.split("_")[0],
|
|
context_dim=cross_attention_dim # called context_dim for ComfyUI impl
|
|
if block_name.endswith("_Cross")
|
|
else None,
|
|
query_dim=dim,
|
|
heads=num_attention_heads,
|
|
dim_head=attention_head_dim,
|
|
dropout=dropout,
|
|
#bias=attention_bias, # remove for Comfy CrossAttention
|
|
#upcast_attention=upcast_attention, # remove for Comfy CrossAttention
|
|
cross_frame_attention_mode=cross_frame_attention_mode,
|
|
temporal_position_encoding=temporal_position_encoding,
|
|
temporal_position_encoding_max_len=temporal_position_encoding_max_len,
|
|
)
|
|
)
|
|
norms.append(nn.LayerNorm(dim))
|
|
|
|
self.attention_blocks: Iterable[VersatileAttention] = nn.ModuleList(attention_blocks)
|
|
self.norms = nn.ModuleList(norms)
|
|
|
|
self.ff = FeedForward(dim, dropout=dropout, glu=(activation_fn == "geglu"))
|
|
self.ff_norm = nn.LayerNorm(dim)
|
|
|
|
def set_scale_multiplier(self, multiplier: Union[float, None]):
|
|
for block in self.attention_blocks:
|
|
block.set_scale_multiplier(multiplier)
|
|
|
|
def set_sub_idxs(self, sub_idxs: list[int]):
|
|
for block in self.attention_blocks:
|
|
block.set_sub_idxs(sub_idxs)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states,
|
|
encoder_hidden_states=None,
|
|
attention_mask=None,
|
|
video_length=None,
|
|
scale_mask=None
|
|
):
|
|
for attention_block, norm in zip(self.attention_blocks, self.norms):
|
|
norm_hidden_states = norm(hidden_states).to(hidden_states.dtype)
|
|
hidden_states = (
|
|
attention_block(
|
|
norm_hidden_states,
|
|
encoder_hidden_states=encoder_hidden_states
|
|
if attention_block.is_cross_attention
|
|
else None,
|
|
attention_mask=attention_mask,
|
|
video_length=video_length,
|
|
scale_mask=scale_mask
|
|
)
|
|
+ hidden_states
|
|
)
|
|
|
|
hidden_states = self.ff(self.ff_norm(hidden_states)) + hidden_states
|
|
|
|
output = hidden_states
|
|
return output
|
|
|
|
|
|
class PositionalEncoding(nn.Module):
|
|
def __init__(self, d_model, dropout=0.0, max_len=24):
|
|
super().__init__()
|
|
self.dropout = nn.Dropout(p=dropout)
|
|
position = torch.arange(max_len).unsqueeze(1)
|
|
div_term = torch.exp(
|
|
torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)
|
|
)
|
|
pe = torch.zeros(1, max_len, d_model)
|
|
pe[0, :, 0::2] = torch.sin(position * div_term)
|
|
pe[0, :, 1::2] = torch.cos(position * div_term)
|
|
self.register_buffer("pe", pe)
|
|
self.sub_idxs = None
|
|
|
|
def set_sub_idxs(self, sub_idxs: list[int]):
|
|
self.sub_idxs = sub_idxs
|
|
|
|
def forward(self, x):
|
|
#if self.sub_idxs is not None:
|
|
# x = x + self.pe[:, self.sub_idxs]
|
|
#else:
|
|
x = x + self.pe[:, : x.size(1)]
|
|
return self.dropout(x)
|
|
|
|
|
|
class CrossAttentionMM(nn.Module):
|
|
def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0., dtype=None, device=None,
|
|
operations=comfy.ops.disable_weight_init):
|
|
super().__init__()
|
|
inner_dim = dim_head * heads
|
|
context_dim = default(context_dim, query_dim)
|
|
|
|
self.heads = heads
|
|
self.dim_head = dim_head
|
|
self.scale = None
|
|
|
|
self.to_q = operations.Linear(query_dim, inner_dim, bias=False, dtype=dtype, device=device)
|
|
self.to_k = operations.Linear(context_dim, inner_dim, bias=False, dtype=dtype, device=device)
|
|
self.to_v = operations.Linear(context_dim, inner_dim, bias=False, dtype=dtype, device=device)
|
|
|
|
self.to_out = nn.Sequential(operations.Linear(inner_dim, query_dim, dtype=dtype, device=device), nn.Dropout(dropout))
|
|
|
|
def forward(self, x, context=None, value=None, mask=None, scale_mask=None):
|
|
q = self.to_q(x)
|
|
context = default(context, x)
|
|
k: Tensor = self.to_k(context)
|
|
if value is not None:
|
|
v = self.to_v(value)
|
|
del value
|
|
else:
|
|
v = self.to_v(context)
|
|
|
|
# apply custom scale by multiplying k by scale factor
|
|
if self.scale is not None:
|
|
k *= self.scale
|
|
|
|
# apply scale mask, if present
|
|
if scale_mask is not None:
|
|
k *= scale_mask
|
|
|
|
out = optimized_attention(q, k, v, self.heads, mask)
|
|
return self.to_out(out)
|
|
|
|
|
|
class VersatileAttention(CrossAttentionMM):
|
|
def __init__(
|
|
self,
|
|
attention_mode=None,
|
|
cross_frame_attention_mode=None,
|
|
temporal_position_encoding=False,
|
|
temporal_position_encoding_max_len=24,
|
|
*args,
|
|
**kwargs,
|
|
):
|
|
super().__init__(*args, **kwargs)
|
|
assert attention_mode == "Temporal"
|
|
|
|
self.attention_mode = attention_mode
|
|
self.is_cross_attention = kwargs["context_dim"] is not None
|
|
|
|
self.pos_encoder = (
|
|
PositionalEncoding(
|
|
kwargs["query_dim"],
|
|
dropout=0.0,
|
|
max_len=temporal_position_encoding_max_len,
|
|
)
|
|
if (temporal_position_encoding and attention_mode == "Temporal")
|
|
else None
|
|
)
|
|
|
|
def extra_repr(self):
|
|
return f"(Module Info) Attention_Mode: {self.attention_mode}, Is_Cross_Attention: {self.is_cross_attention}"
|
|
|
|
def set_scale_multiplier(self, multiplier: Union[float, None]):
|
|
if multiplier is None or math.isclose(multiplier, 1.0):
|
|
self.scale = None
|
|
else:
|
|
self.scale = multiplier
|
|
|
|
def set_sub_idxs(self, sub_idxs: list[int]):
|
|
if self.pos_encoder != None:
|
|
self.pos_encoder.set_sub_idxs(sub_idxs)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: Tensor,
|
|
encoder_hidden_states=None,
|
|
attention_mask=None,
|
|
video_length=None,
|
|
scale_mask=None,
|
|
):
|
|
if self.attention_mode != "Temporal":
|
|
raise NotImplementedError
|
|
|
|
d = hidden_states.shape[1]
|
|
hidden_states = rearrange(
|
|
hidden_states, "(b f) d c -> (b d) f c", f=video_length
|
|
)
|
|
|
|
if self.pos_encoder is not None:
|
|
hidden_states = self.pos_encoder(hidden_states).to(hidden_states.dtype)
|
|
|
|
encoder_hidden_states = (
|
|
repeat(encoder_hidden_states, "b n c -> (b d) n c", d=d)
|
|
if encoder_hidden_states is not None
|
|
else encoder_hidden_states
|
|
)
|
|
|
|
hidden_states = super().forward(
|
|
hidden_states,
|
|
encoder_hidden_states,
|
|
value=None,
|
|
mask=attention_mask,
|
|
scale_mask=scale_mask,
|
|
)
|
|
|
|
hidden_states = rearrange(hidden_states, "(b d) f c -> (b f) d c", d=d)
|
|
|
|
return hidden_states
|