Merge pull request #1 from ExponentialML/feat/modelscope_comfy
Add files via upload
This commit is contained in:
ExponentialML
@@ -0,0 +1,3 @@
|
||||
from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
||||
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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",
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user