139 lines
5.4 KiB
Python
139 lines
5.4 KiB
Python
# main code adapted from https://github.com/TencentARC/MotionCtrl/tree/animatediff
|
|
from __future__ import annotations
|
|
from torch import nn, Tensor
|
|
import torch
|
|
|
|
from comfy.model_patcher import ModelPatcher
|
|
import comfy.model_management
|
|
import comfy.ops
|
|
import comfy.utils
|
|
|
|
from .adapter_cameractrl import ResnetBlockCameraCtrl
|
|
from .ad_settings import AnimateDiffSettings
|
|
from .motion_module_ad import AnimateDiffModel
|
|
from .model_injection import apply_mm_settings
|
|
from .utils_model import get_motion_model_path
|
|
|
|
|
|
# cmcm (Camera Control)
|
|
def inject_motionctrl_cmcm(motion_model: AnimateDiffModel, cmcm_name: str, ad_settings: AnimateDiffSettings=None,
|
|
apply_non_ccs=True):
|
|
cmcm_path = get_motion_model_path(cmcm_name)
|
|
state_dict = comfy.utils.load_torch_file(cmcm_path, safe_load=True)
|
|
_remove_module_prefix(state_dict)
|
|
# if applicable, apply ad_settings to cmcm to match expected behavior
|
|
if ad_settings is not None:
|
|
state_dict = apply_mm_settings(model_dict=state_dict, mm_settings=ad_settings)
|
|
motion_model.init_motionctrl_cc_projections(state_dict=state_dict)
|
|
# seperate out PE keys so can be applied separately in case dims don't match
|
|
apply_dict = {}
|
|
for key in list(state_dict.keys()):
|
|
if "cc_projection" in key:
|
|
apply_dict[key] = state_dict[key]
|
|
state_dict.pop(key)
|
|
pe_dict = {}
|
|
for key in list(state_dict.keys()):
|
|
if "pos_encoder" in key:
|
|
pe_dict[key] = state_dict[key]
|
|
state_dict.pop(key)
|
|
if apply_non_ccs:
|
|
apply_dict.update(state_dict)
|
|
for key, value in pe_dict.items():
|
|
comfy.utils.set_attr(motion_model, key, value)
|
|
_, unexpected = motion_model.load_state_dict(apply_dict, strict=False)
|
|
if len(unexpected) > 0:
|
|
raise Exception(f"MotionCtrl CMCM model had unexpected keys: {unexpected}")
|
|
# make sure model is still has proper dtype and offload device
|
|
motion_model.to(comfy.model_management.unet_dtype())
|
|
motion_model.to(comfy.model_management.unet_offload_device())
|
|
|
|
|
|
# omcm (Object Control)
|
|
def load_motionctrl_omcm(omcm_name: str):
|
|
omcm_path = get_motion_model_path(omcm_name)
|
|
state_dict = comfy.utils.load_torch_file(omcm_path, safe_load=True)
|
|
_remove_module_prefix(state_dict)
|
|
|
|
if comfy.model_management.unet_manual_cast(comfy.model_management.unet_dtype(), comfy.model_management.get_torch_device()) is None:
|
|
ops = comfy.ops.disable_weight_init
|
|
else:
|
|
ops = comfy.ops.manual_cast
|
|
adapter = MotionCtrlAdapter(ops=ops)
|
|
adapter.load_state_dict(state_dict=state_dict, strict=True)
|
|
adapter.to(
|
|
device = comfy.model_management.unet_offload_device(),
|
|
dtype = comfy.model_management.unet_dtype()
|
|
)
|
|
omcm_modelpatcher = _create_OMCMModelPatcher(model=adapter,
|
|
load_device=comfy.model_management.get_torch_device(),
|
|
offload_device=comfy.model_management.unet_offload_device())
|
|
return omcm_modelpatcher
|
|
|
|
|
|
def _create_OMCMModelPatcher(model, load_device, offload_device) -> ObjectControlModelPatcher:
|
|
patcher = ModelPatcher(model, load_device=load_device, offload_device=offload_device)
|
|
return patcher
|
|
|
|
|
|
def _remove_module_prefix(state_dict: dict[str, Tensor]):
|
|
for key in list(state_dict.keys()):
|
|
# remove 'module.' prefix
|
|
if key.startswith('module.'):
|
|
new_key = key.replace('module.', '')
|
|
state_dict[new_key] = state_dict[key]
|
|
state_dict.pop(key)
|
|
|
|
|
|
def convert_cameractrl_poses_to_RT(poses: list[list[float]]):
|
|
tensors = []
|
|
for pose in poses:
|
|
new_tensor = torch.tensor(pose[7:])
|
|
new_tensor = new_tensor.unsqueeze(0)
|
|
tensors.append(new_tensor)
|
|
RT = torch.cat(tensors, dim=0)
|
|
return RT
|
|
|
|
|
|
class ObjectControlModelPatcher(ModelPatcher):
|
|
'''Class only used for type hints.'''
|
|
def __init__(self):
|
|
self.model: MotionCtrlAdapter
|
|
|
|
|
|
class MotionCtrlAdapter(nn.Module):
|
|
def __init__(self,
|
|
downscale_factor=8,
|
|
channels=[320, 640, 1280, 1280],
|
|
nums_rb=2, cin=128, # 2*8*8
|
|
ksize=3, sk=True,
|
|
use_conv=False,
|
|
ops=comfy.ops.disable_weight_init):
|
|
super(MotionCtrlAdapter, self).__init__()
|
|
self.downscale_factor = downscale_factor
|
|
self.unshuffle = nn.PixelUnshuffle(downscale_factor)
|
|
self.channels = channels
|
|
self.nums_rb = nums_rb
|
|
self.body = []
|
|
for i in range(len(channels)):
|
|
for j in range(nums_rb):
|
|
if (i != 0) and (j == 0):
|
|
self.body.append(
|
|
ResnetBlockCameraCtrl(channels[i - 1], channels[i], down=True, ksize=ksize, sk=sk, use_conv=use_conv, ops=ops))
|
|
else:
|
|
self.body.append(
|
|
ResnetBlockCameraCtrl(channels[i], channels[i], down=False, ksize=ksize, sk=sk, use_conv=use_conv, ops=ops))
|
|
self.body = nn.ModuleList(self.body)
|
|
self.conv_in = ops.Conv2d(cin, channels[0], 3, 1, 1)
|
|
|
|
def forward(self, x: Tensor):
|
|
x = self.unshuffle(x)
|
|
# extract features
|
|
features = []
|
|
x = self.conv_in(x)
|
|
for i in range(len(self.channels)):
|
|
for j in range(self.nums_rb):
|
|
idx = i * self.nums_rb + j
|
|
x = self.body[idx](x)
|
|
features.append(x)
|
|
return features
|