155 lines
5.6 KiB
Python
155 lines
5.6 KiB
Python
import torch.nn as nn
|
|
import comfy.ops
|
|
import inspect
|
|
import copy
|
|
from ..modules.modelscope_modules import ResBlock, TemporalTransformer
|
|
from comfy.ldm.modules.diffusionmodules.openaimodel import Upsample, Downsample, SpatialTransformer, TimestepEmbedSequential
|
|
|
|
operations = comfy.ops.disable_weight_init
|
|
|
|
TEMPORAL_TRANSFORMER_KEY_MAP = ["proj_in", "proj_out", "transformer_blocks", "norm"]
|
|
TEMPORAL_RESNET_KEY_MAP = ["temopral_conv"]
|
|
|
|
def remove_param_str(k):
|
|
return k.replace(".weight", "").replace(".bias", "")
|
|
|
|
def get_keys_from_dict(input_dict: dict, key_map: list):
|
|
out = []
|
|
for name in input_dict.keys():
|
|
if any([k in name for k in key_map]) and "weight" in name:
|
|
out.append(remove_param_str(name))
|
|
return out
|
|
|
|
def strip_module_name(keys: list, key_map: list):
|
|
processed_full_keys = []
|
|
out_module_names = []
|
|
|
|
for full_key in keys:
|
|
module_name = None
|
|
for key_name in key_map:
|
|
if key_name in full_key:
|
|
module_name = key_name
|
|
break
|
|
|
|
if module_name is not None:
|
|
if full_key not in processed_full_keys:
|
|
processed_full_keys.append(full_key)
|
|
full_key_split = full_key.split(".")
|
|
stop_idx = full_key_split.index(module_name)
|
|
|
|
out_full_key = ".".join(full_key_split[:stop_idx])
|
|
out_module_names.append(out_full_key)
|
|
|
|
# Remove duplicates
|
|
processed_module_names = []
|
|
for module_name in out_module_names:
|
|
if module_name not in processed_module_names:
|
|
processed_module_names.append(module_name)
|
|
|
|
return processed_module_names
|
|
|
|
def get_submodule_paths(key:str, prefix: str = "model.diffusion_model."):
|
|
# Remove the full path if we have it set.
|
|
if prefix:
|
|
key = key.replace(prefix, "")
|
|
|
|
key = key.split(".")
|
|
parent, k = key[:len(key) - 1], key[-1]
|
|
|
|
module_path = ".".join(parent)
|
|
module_index = int(k)
|
|
|
|
return module_path, module_index
|
|
|
|
def add_init_temporal_attention(model, unet_config: dict, temporal_attn_strength=1.0):
|
|
model.input_blocks[0].insert(1,
|
|
TemporalTransformer(
|
|
unet_config['model_channels'],
|
|
unet_config['num_heads'],
|
|
unet_config['num_head_channels'],
|
|
context_dim=unet_config['context_dim'],
|
|
depth=1,
|
|
temporal_attn_strength=temporal_attn_strength
|
|
)
|
|
)
|
|
|
|
|
|
def add_temporal_attention_blocks(model, modelscope_dict: dict, temporal_attn_strength=1.0):
|
|
processed_temporal_keys = get_keys_from_dict(modelscope_dict, TEMPORAL_TRANSFORMER_KEY_MAP)
|
|
temporal_transformer_keys = strip_module_name(processed_temporal_keys, TEMPORAL_TRANSFORMER_KEY_MAP)
|
|
|
|
for key in temporal_transformer_keys:
|
|
|
|
module_path, module_index = get_submodule_paths(key)
|
|
|
|
# Get previous module, which is a SpatialTransformer (idx - 1).
|
|
# The current module_index is where the new Temporal Transformer should be.
|
|
module = model.get_submodule(module_path)[module_index - 1]
|
|
|
|
if not isinstance(module, SpatialTransformer):
|
|
continue
|
|
|
|
n_heads = module.transformer_blocks[0].n_heads
|
|
d_head = module.transformer_blocks[0].d_head
|
|
|
|
temporal_transformer_block = TemporalTransformer(
|
|
module.in_channels,
|
|
n_heads,
|
|
d_head,
|
|
depth=1,
|
|
temporal_attn_strength=temporal_attn_strength
|
|
)
|
|
|
|
model.get_submodule(module_path).insert(module_index, temporal_transformer_block)
|
|
#print("Verify", model.get_submodule(module_path)[module_index])
|
|
|
|
|
|
def replace_proj_resblock(model, modelscope_dict: dict, dtype, device, temporal_conv_strength=1.0):
|
|
processed_resnet_keys = get_keys_from_dict(modelscope_dict, TEMPORAL_RESNET_KEY_MAP)
|
|
temporal_resnet_keys = strip_module_name(processed_resnet_keys, TEMPORAL_RESNET_KEY_MAP)
|
|
|
|
for key in temporal_resnet_keys:
|
|
|
|
module_path, module_index = get_submodule_paths(key)
|
|
|
|
# Since the temporal module is within the ResBlock, we can just access / assign it directly.
|
|
if module_index == 3:
|
|
module_index -= 1
|
|
|
|
submodule = model.get_submodule(module_path)[module_index]
|
|
|
|
res_block = ResBlock(
|
|
channels=submodule.channels,
|
|
emb_channels=submodule.emb_channels,
|
|
dropout=0.0,
|
|
out_channels=submodule.out_channels,
|
|
use_scale_shift_norm=submodule.use_scale_shift_norm,
|
|
down=isinstance(submodule.h_upd, Downsample),
|
|
up=isinstance(submodule.h_upd, Upsample),
|
|
dtype=dtype,
|
|
device=device,
|
|
operations=operations,
|
|
use_temporal_conv=True,
|
|
temporal_conv_strength=temporal_conv_strength
|
|
)
|
|
|
|
model.get_submodule(module_path)[module_index] = res_block
|
|
#print("Verify", model.get_submodule(module_path)[module_index])
|
|
|
|
def inject_temporal_modules(
|
|
model,
|
|
unet_config,
|
|
modelscope_dict,
|
|
dtype,
|
|
device,
|
|
enable_attn=True,
|
|
enable_conv=True,
|
|
temporal_attn_strength=1.0,
|
|
temporal_conv_strength=1.0
|
|
):
|
|
if enable_conv:
|
|
replace_proj_resblock(model.diffusion_model, modelscope_dict, dtype, device, temporal_conv_strength)
|
|
if enable_attn:
|
|
add_init_temporal_attention(model.diffusion_model, unet_config, temporal_attn_strength)
|
|
add_temporal_attention_blocks(model.diffusion_model, modelscope_dict, temporal_attn_strength)
|
|
|