diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..3d998ea --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] diff --git a/configs/modelscope_config.py b/configs/modelscope_config.py new file mode 100644 index 0000000..699f0a1 --- /dev/null +++ b/configs/modelscope_config.py @@ -0,0 +1,83 @@ +import torch +from comfy import model_base +from comfy import utils + +from comfy import sd1_clip +from comfy import sd2_clip + +from comfy import supported_models_base +from comfy import latent_formats + +SD15_CONFIG = { + 'use_checkpoint': False, + 'image_size': 32, + 'out_channels': 4, + 'use_spatial_transformer': True, + 'legacy': False, + 'adm_in_channels': None, + 'in_channels': 4, + 'model_channels': 320, + 'num_res_blocks': [2, 2, 2, 2], + 'transformer_depth': [1, 1, 1, 1, 1, 1, 0, 0], + 'channel_mult': [1, 2, 4, 4], + 'transformer_depth_middle': 1, + 'use_linear_in_transformer': False, + 'context_dim': 768, + 'num_heads': 8, + 'transformer_depth_output': [1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0], + 'use_temporal_attention': False, + 'use_temporal_resblock': False +} + +MODELSCOPE_UNET_CONFIG = { + "use_spatial_transformer": True, + "model_channels": 320, + "in_channels": 4, + "out_channels": 4, + "context_dim": 1024, + 'num_res_blocks': [2, 2, 2, 2], + "channel_mult": (1, 2, 4, 4), + "num_heads": 8, + "num_head_channels": 64, + "image_size": 32, + "use_linear_in_transformer": False, + "adm_in_channels": None, + "use_temporal_attention": False, + "merge_strategy": "fixed", + 'transformer_depth': [1, 1, 1, 1, 1, 1, 0, 0], + 'transformer_depth_middle': 1, + 'transformer_depth_output': [1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0] +} + +class SD15ModelScope(supported_models_base.BASE): + unet_extra_config = { + "num_heads": MODELSCOPE_UNET_CONFIG['num_heads'], + "num_head_channels": MODELSCOPE_UNET_CONFIG['num_head_channels'], + } + latent_format = latent_formats.SD15 + + def model_type(self, state_dict, prefix=""): + if self.unet_config["in_channels"] == 4: #SD2.0 inpainting models are not v prediction + k = "{}output_blocks.11.1.transformer_blocks.0.norm1.bias".format(prefix) + out = state_dict[k] + if torch.std(out, unbiased=False) > 0.09: # not sure how well this will actually work. I guess we will find out. + return model_base.ModelType.V_PREDICTION + return model_base.ModelType.EPS + + def process_clip_state_dict(self, state_dict): + replace_prefix = {} + replace_prefix["conditioner.embedders.0.model."] = "clip_h." #SD2 in sgm format + replace_prefix["cond_stage_model.model."] = "clip_h." + state_dict = utils.state_dict_prefix_replace(state_dict, replace_prefix, filter_keys=True) + state_dict = utils.clip_text_transformers_convert(state_dict, "clip_h.", "clip_h.transformer.") + return state_dict + + def process_clip_state_dict_for_saving(self, state_dict): + replace_prefix = {} + replace_prefix["clip_h"] = "cond_stage_model.model" + state_dict = utils.state_dict_prefix_replace(state_dict, replace_prefix) + state_dict = diffusers_convert.convert_text_enc_state_dict_v20(state_dict) + return state_dict + + def clip_target(self): + return supported_models_base.ClipTarget(sd2_clip.SD2Tokenizer, sd2_clip.SD2ClipModel) \ No newline at end of file diff --git a/modules/modelscope_modules.py b/modules/modelscope_modules.py new file mode 100644 index 0000000..3e99394 --- /dev/null +++ b/modules/modelscope_modules.py @@ -0,0 +1,357 @@ +import torch +import torch.nn as nn +import comfy.ops +from einops import rearrange, repeat +from comfy.ldm.modules.attention import SpatialTransformer, BasicTransformerBlock +from comfy.ldm.modules.diffusionmodules.openaimodel import TimestepBlock, Downsample, Upsample + +operations = comfy.ops.disable_weight_init +ops = operations + +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 get_batch_size(b: int, use_cfg: True): + return 2 if use_cfg else 1 + +def use_attn_strength(strength): + return 1.0 - strength > 0.0 + +class TemporalConvBlock_v2(nn.Module): + + def __init__(self, + in_dim, + out_dim=None, + dropout=0.0, + device='cuda', + dtype=torch.float16, + ops=operations, + temporal_conv_strength=1.0 + ): + super(TemporalConvBlock_v2, self).__init__() + if out_dim is None: + out_dim = in_dim # int(1.5*in_dim) + self.in_dim = in_dim + self.out_dim = out_dim + self.temporal_conv_strength = temporal_conv_strength + + # conv layers + self.conv1 = nn.Sequential( + ops.GroupNorm(32, in_dim, dtype=dtype, device=device), nn.SiLU(), + ops.Conv3d(in_dim, out_dim, (3, 1, 1), padding=(1, 0, 0), dtype=dtype, device=device)) + self.conv2 = nn.Sequential( + ops.GroupNorm(32, out_dim, dtype=dtype, device=device), nn.SiLU(), nn.Dropout(dropout), + ops.Conv3d(out_dim, in_dim, (3, 1, 1), padding=(1, 0, 0), dtype=dtype, device=device)) + self.conv3 = nn.Sequential( + ops.GroupNorm(32, out_dim, dtype=dtype, device=device), nn.SiLU(), nn.Dropout(dropout), + ops.Conv3d(out_dim, in_dim, (3, 1, 1), padding=(1, 0, 0), dtype=dtype, device=device)) + self.conv4 = nn.Sequential( + ops.GroupNorm(32, out_dim, dtype=dtype, device=device), nn.SiLU(), nn.Dropout(dropout), + ops.Conv3d(out_dim, in_dim, (3, 1, 1), padding=(1, 0, 0), dtype=dtype, device=device)) + + # zero out the last layer params,so the conv block is identity + nn.init.zeros_(self.conv4[-1].weight) + nn.init.zeros_(self.conv4[-1].bias) + + def forward(self, x): + identity = x + x = self.conv1(x) + x = self.conv2(x) + x = self.conv3(x) + x = self.conv4(x) + + if use_attn_strength(self.temporal_conv_strength): + x = identity + self.temporal_conv_strength * x + else: + x = identity + x + return x + +class ResBlock(TimestepBlock): + """ + A residual block that can optionally change the number of channels. + :param channels: the number of input channels. + :param emb_channels: the number of timestep embedding channels. + :param dropout: the rate of dropout. + :param out_channels: if specified, the number of out channels. + :param use_conv: if True and out_channels is specified, use a spatial + convolution instead of a smaller 1x1 convolution to change the + channels in the skip connection. + :param dims: determines if the signal is 1D, 2D, or 3D. + :param use_checkpoint: if True, use gradient checkpointing on this module. + :param up: if True, use this block for upsampling. + :param down: if True, use this block for downsampling. + """ + + def __init__( + self, + channels, + emb_channels, + dropout, + out_channels=None, + use_conv=False, + use_scale_shift_norm=False, + dims=2, + use_checkpoint=False, + up=False, + down=False, + kernel_size=3, + exchange_temb_dims=False, + skip_t_emb=False, + dtype=None, + device=None, + operations=ops, + use_temporal_conv=True, + temporal_conv_strength=1.0, + ): + super().__init__() + self.channels = channels + self.emb_channels = emb_channels + self.dropout = dropout + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.use_checkpoint = use_checkpoint + self.use_scale_shift_norm = use_scale_shift_norm + self.exchange_temb_dims = exchange_temb_dims + self.use_temporal_conv = use_temporal_conv + self.use_cfg = True + + if isinstance(kernel_size, list): + padding = [k // 2 for k in kernel_size] + else: + padding = kernel_size // 2 + + self.in_layers = nn.Sequential( + operations.GroupNorm(32, channels, dtype=dtype, device=device), + nn.SiLU(), + operations.conv_nd(dims, channels, self.out_channels, 3, padding=1, dtype=dtype, device=device), + ) + + self.updown = up or down + + if up: + self.h_upd = Upsample(channels, False, dims, dtype=dtype, device=device) + self.x_upd = Upsample(channels, False, dims, dtype=dtype, device=device) + elif down: + self.h_upd = Downsample(channels, False, dims, dtype=dtype, device=device) + self.x_upd = Downsample(channels, False, dims, dtype=dtype, device=device) + else: + self.h_upd = self.x_upd = nn.Identity() + + self.skip_t_emb = skip_t_emb + if self.skip_t_emb: + self.emb_layers = None + self.exchange_temb_dims = False + else: + self.emb_layers = nn.Sequential( + nn.SiLU(), + operations.Linear( + emb_channels, + 2 * self.out_channels if use_scale_shift_norm else self.out_channels, dtype=dtype, device=device + ), + ) + self.out_layers = nn.Sequential( + operations.GroupNorm(32, self.out_channels, dtype=dtype, device=device), + nn.SiLU(), + nn.Dropout(p=dropout), + zero_module(operations.conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1, dtype=dtype, device=device)) + ) + + if self.out_channels == channels: + self.skip_connection = nn.Identity() + elif use_conv: + self.skip_connection = operations.conv_nd( + dims, channels, self.out_channels, 3, padding=1, dtype=dtype, device=device + ) + else: + self.skip_connection = operations.conv_nd(dims, channels, self.out_channels, 1, dtype=dtype, device=device) + + if use_temporal_conv: + self.temopral_conv = TemporalConvBlock_v2( + self.out_channels, + self.out_channels, + dropout=0.1, + temporal_conv_strength=temporal_conv_strength + ) + + def forward(self, x, emb): + """ + Apply the block to a Tensor, conditioned on a timestep embedding. + :param x: an [N x C x ...] Tensor of features. + :param emb: an [N x emb_channels] Tensor of timestep embeddings. + :return: an [N x C x ...] Tensor of outputs. + """ + return self._forward(x, emb) + + def _forward(self, x, emb): + b = x.shape[0] + if self.updown: + in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1] + h = in_rest(x) + h = self.h_upd(h) + x = self.x_upd(x) + h = in_conv(h) + else: + h = self.in_layers(x) + + emb_out = None + if not self.skip_t_emb: + emb_out = self.emb_layers(emb).type(h.dtype) + while len(emb_out.shape) < len(h.shape): + emb_out = emb_out[..., None] + if self.use_scale_shift_norm: + out_norm, out_rest = self.out_layers[0], self.out_layers[1:] + h = out_norm(h) + if emb_out is not None: + scale, shift = th.chunk(emb_out, 2, dim=1) + h *= (1 + scale) + h += shift + h = out_rest(h) + else: + if emb_out is not None: + if self.exchange_temb_dims: + emb_out = rearrange(emb_out, "b t c ... -> b c t ...") + h = h + emb_out + h = self.out_layers(h) + + h = self.skip_connection(x) + h + + if self.use_temporal_conv: + batch_size = get_batch_size(b, self.use_cfg) + h = rearrange(h, '(b f) c h w -> b c f h w', b=batch_size) + h = self.temopral_conv(h) + h = rearrange(h, 'b c f h w -> (b f) c h w') + + return h + +class TemporalTransformer(nn.Module): # nn.Module -> SpatialTransformer + """ + Transformer block for image-like data in temporal axis. + First, reshape to b, t, d. + Then apply standard transformer action. + Finally, reshape to image + """ + + def __init__(self, + in_channels, + n_heads, + d_head, + depth=1, + dropout=0.0, + context_dim=None, + disable_self_attn=False, + use_linear=False, + use_checkpoint=True, + only_self_att=True, + multiply_zero=False, + device='cuda', + dtype=torch.float16, + ops=operations, + use_temporal_attention=True, + temporal_attn_strength=1.0 + ): + super().__init__() + self.multiply_zero = multiply_zero + self.only_self_att = only_self_att + self.use_temporal_attention = use_temporal_attention + self.use_cfg = True + self.temporal_attn_strength = temporal_attn_strength + + if self.only_self_att: + context_dim = None + if not isinstance(context_dim, list): + context_dim = [context_dim] + self.in_channels = in_channels + inner_dim = n_heads * d_head + self.norm = ops.GroupNorm( + num_groups=32, num_channels=in_channels, eps=1e-6, affine=True, device=device, dtype=dtype) + if not use_linear: + self.proj_in = nn.Conv1d( + in_channels, inner_dim, kernel_size=1, stride=1, padding=0, device=device, dtype=dtype) + else: + self.proj_in = ops.Linear(in_channels, inner_dim, device=device, dtype=dtype) + + self.transformer_blocks = nn.ModuleList([ + BasicTransformerBlock( + inner_dim, + n_heads, + d_head, + dropout=dropout, + context_dim=context_dim[d], + checkpoint=use_checkpoint, + device=device, + dtype=dtype + ) for d in range(depth) + ]) + if not use_linear: + self.proj_out = zero_module( + nn.Conv1d( + inner_dim, in_channels, kernel_size=1, stride=1, + padding=0)).to(dtype=dtype) + else: + self.proj_out = zero_module(ops.Linear(in_channels, inner_dim, device=device, dtype=dtype)) + self.use_linear = use_linear + + def forward(self, x, context=None): + if not self.use_temporal_attention: + return x + # note: if no context is given, cross-attention defaults to self-attention + if self.only_self_att: + context = None + if not isinstance(context, list): + context = [context] + x_in = x + b = get_batch_size(x_in.shape[0], self.use_cfg) + x = rearrange(x, '(b f) c h w -> b c f h w', b=b) + b, c, f, h, w = x.shape + + x = self.norm(x) + + if not self.use_linear: + x = rearrange(x, 'b c f h w -> (b h w) c f').contiguous() + x = self.proj_in(x) + if self.use_linear: + x = rearrange( + x, '(b f) c h w -> b (h w) f c', f=self.frames).contiguous() + x = self.proj_in(x) + + if self.only_self_att: + x = rearrange(x, 'bhw c f -> bhw f c').contiguous() + for i, block in enumerate(self.transformer_blocks): + x = block(x) + x = rearrange(x, '(b hw) f c -> b hw f c', b=b).contiguous() + else: + x = rearrange(x, '(b hw) c f -> b hw f c', b=b).contiguous() + for i, block in enumerate(self.transformer_blocks): + context[i] = rearrange( + context[i], '(b f) l con -> b f l con', + f=self.frames).contiguous() + # calculate each batch one by one (since number in shape could not greater then 65,535 for some package) + for j in range(b): + context_i_j = repeat( + context[i][j], + 'f l con -> (f r) l con', + r=(h * w) // self.frames, + f=self.frames).contiguous() + x[j] = block(x[j], context=context_i_j) + + if self.use_linear: + x = self.proj_out(x) + x = rearrange(x, 'b (h w) f c -> b f c h w', h=h, w=w).contiguous() + if not self.use_linear: + x = rearrange(x, 'b hw f c -> (b hw) c f').contiguous() + x = self.proj_out(x) + x = rearrange( + x, '(b h w) c f -> b c f h w', b=b, h=h, w=w).contiguous() + + x = rearrange(x, 'b c f h w -> (b f) c h w') + + if use_attn_strength(self.temporal_attn_strength): + x = self.temporal_attn_strength * x + x_in + else: + x = x + x_in + return x \ No newline at end of file diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..8009876 --- /dev/null +++ b/nodes.py @@ -0,0 +1,234 @@ +import os +import torch +import comfy +import yaml +import folder_paths + +from comfy import model_base, model_management, model_detection, latent_formats, model_sampling +from comfy.supported_models import SD15 +from comfy.ldm.modules.diffusionmodules.openaimodel import UNetModel +from .configs.modelscope_config import MODELSCOPE_UNET_CONFIG, SD15_CONFIG +from .utils.module_utils import inject_temporal_modules, TEMPORAL_TRANSFORMER_KEY_MAP + +MODEL_DIR= "text2video" +model_dir_path = os.path.join(folder_paths.models_dir, MODEL_DIR) + +def load_modelscope_checkpoint( + sd_dict, + modelscope_dict, + ckpt_path, + output_model=True, + enable_attn=True, + enable_conv = True, + temporal_attn_strength = 1.0, + temporal_conv_strength = 1.0, + sd_15_model=None + ): + sd = sd_dict + modelscope_dict = modelscope_dict + sd_keys = sd.keys() + model_patcher = None + model_scope_config = MODELSCOPE_UNET_CONFIG + modelscope_dict_process = None + + if sd_15_model is not None: + model_scope_config = SD15_CONFIG + if enable_attn: + print("Modelscope Temporal Attention cannot be enabled when loading a model. Disabling...") + enable_attn = False + + modelscope_dict_process = modelscope_dict.copy() + + for k in modelscope_dict_process.keys(): + if any([key_map in k for key_map in TEMPORAL_TRANSFORMER_KEY_MAP]): + modelscope_dict.pop(k) + + if all([not enable_conv, not enable_attn]): + modelscope_dict = {} + + parameters = comfy.utils.calculate_parameters(sd, "model.diffusion_model.") + load_device = model_management.get_torch_device() + + num_head_channels = MODELSCOPE_UNET_CONFIG['num_head_channels'] + model_config = SD15(model_scope_config) + + if enable_attn: + model_config.unet_config['num_head_channels'] = num_head_channels + + unet_dtype = model_management.unet_dtype(model_params=parameters, supported_dtypes=model_config.supported_inference_dtypes) + manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device, model_config.supported_inference_dtypes) + model_config.set_inference_dtype(unet_dtype, manual_cast_dtype) + + device = model_management.intermediate_device() + + if model_config is None: + raise RuntimeError("ERROR: Could not detect model type of: {}".format(ckpt_path)) + + if output_model: + inital_load_device = model_management.unet_inital_load_device(parameters, unet_dtype) + offload_device = model_management.unet_offload_device() + + model = model_base.BaseModel(model_config, model_type=model_base.ModelType.EPS, device=inital_load_device) + inject_temporal_modules( + model, + model_scope_config, + modelscope_dict, + unet_dtype, + model_management.intermediate_device(), + enable_attn=enable_attn, + enable_conv=enable_conv, + temporal_attn_strength=temporal_attn_strength, + temporal_conv_strength=temporal_conv_strength + ) + + # We need to remove the leftover keys if we're using conv processing. + if all([modelscope_dict_process is not None, enable_conv]): + for k, v in modelscope_dict_process.items(): + if 'middle_block.3.' in k: + del modelscope_dict[k] + k = k.replace("middle_block.3.", "middle_block.2.") + modelscope_dict[k] = v + + if 'output_blocks.5.3' in k or 'output_blocks.8.3' in k: + del modelscope_dict[k] + + sd.update(modelscope_dict) + model.load_model_weights(sd, "model.diffusion_model.") + + del modelscope_dict_process + + left_over = sd.keys() + if len(left_over) > 0: + print("left over keys:", left_over) + + model_patcher = comfy.model_patcher.ModelPatcher( + model, + load_device=load_device, + offload_device=model_management.unet_offload_device(), + current_device=inital_load_device + ) + + # TODO + def update_cfg(*args, **kwargs): + model = kwargs.get("model") + forward = args[0] + model_input = args[1] + x, t, c, c_or_uc = model_input['input'], model_input['timestep'], model_input['c'], model_input['cond_or_uncond'] + is_cfg = len(c_or_uc) > 1 + return forward(x, t, **c) + + if inital_load_device != torch.device("cpu"): + print("loaded straight to GPU") + model_management.load_model_gpu(model_patcher) + + return model_patcher + +def get_sd_15_keys(): + sd_15_base = UNetModel(**SD15_CONFIG).cpu() + sd15_keys = sd_15_base.state_dict().keys() + del sd_15_base + + return sd15_keys + +def set_prefix(k): + if "model.diffusion_model." not in k: + k = f"model.diffusion_model.{k}" + return k + +def maybe_existing_weight(input_dict: dict, k, v): + if input_dict is None: + return v + try: + return input_dict[k] + except: + print("Failed", k) + return v + +def check_and_apply_keys(state_dict: dict, sd_15_model=None): + out_dict = {} + temporal_dict = {} + + sd_15_keys = get_sd_15_keys() + + loaded_dict = ( + sd_15_model.model.diffusion_model.state_dict() \ + if sd_15_model is not None else None + ) + + for k, v in state_dict.items(): + # Handle SD 1.5 Base Model + if k in sd_15_keys or "op." in k: + if any(['proj_in' in k, 'proj_out' in k]) and len(v.shape) == 2 and loaded_dict is None: + v = maybe_existing_weight(loaded_dict, k, v)[:, :, None, None] + + if ".op." in k: + k = k.replace(".op", ".0.op") + out_dict[set_prefix(k)] = maybe_existing_weight(loaded_dict, k, v) + + # Create separate dict for temporal modules and parameters. + if k not in sd_15_keys: + temporal_dict[set_prefix(k)] = v + + if loaded_dict is not None: + out_dict = {set_prefix(k): v for k, v in loaded_dict.items()} + + print(f"Spatial Keys (SD): {len(out_dict.keys())}") + print(f"Total Possible Spatial Keys (SD): {len(sd_15_keys)}") + print(f"Temporal Keys (ModelScope): {len(temporal_dict.keys())}") + + return out_dict, temporal_dict + +class ModelScopeT2VLoader: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model_path": (os.listdir(model_dir_path), ), + "enable_attn": ("BOOLEAN", {"default": True}, ), + "enable_conv": ("BOOLEAN", {"default": True}, ), + "temporal_attn_strength": ("FLOAT", {"default": 1.0, "min": 0., "max": 1., "step": 0.1}, ), + "temporal_conv_strength": ("FLOAT", {"default": 1.0, "min": 0., "max": 1., "step": 0.1}, ), + }, + "optional": { + "sd_15_model": ("MODEL", ) + } + } + + RETURN_TYPES = ("MODEL",) + FUNCTION = "load_modelscopet2v" + + def load_modelscopet2v( + self, + model_path, + enable_attn, + enable_conv, + temporal_attn_strength=1.0, + temporal_conv_strength=1.0, + sd_15_model=None + ): + model_path = os.path.join(model_dir_path, model_path) + + if os.path.exists(model_path): + load_modelscope_dict = comfy.utils.load_torch_file(model_path) + sd_dict, modelscope_state_dict = check_and_apply_keys(load_modelscope_dict, sd_15_model) + model = load_modelscope_checkpoint( + sd_dict, + modelscope_state_dict, + ckpt_path=None, + output_model=True, + enable_attn=enable_attn, + enable_conv=enable_conv, + temporal_attn_strength=temporal_attn_strength, + temporal_conv_strength=temporal_conv_strength, + sd_15_model=sd_15_model + ) + + return (model, ) + +NODE_CLASS_MAPPINGS = { + "ModelScopeT2VLoader": ModelScopeT2VLoader, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ModelScopeT2VLoader": "ModelScopeT2VLoader", +} \ No newline at end of file diff --git a/utils/module_utils.py b/utils/module_utils.py new file mode 100644 index 0000000..6fa7692 --- /dev/null +++ b/utils/module_utils.py @@ -0,0 +1,155 @@ +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) + \ No newline at end of file