234 lines
7.9 KiB
Python
234 lines
7.9 KiB
Python
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",
|
|
} |