148 lines
6.1 KiB
Python
148 lines
6.1 KiB
Python
import torch
|
|
from torch import Tensor
|
|
import numpy as np
|
|
import os
|
|
import json
|
|
|
|
import folder_paths
|
|
|
|
from .ad_settings import AnimateDiffSettings
|
|
from .adapter_motionctrl import (ObjectControlModelPatcher, inject_motionctrl_cmcm, load_motionctrl_omcm,
|
|
convert_cameractrl_poses_to_RT)
|
|
|
|
from .model_injection import MotionModelPatcher, MotionModelGroup, load_motion_module_gen2, get_mm_attachment
|
|
from .motion_lora import MotionLoraList
|
|
|
|
from .nodes_gen2 import ApplyAnimateDiffModelNode
|
|
from .utils_model import get_available_motion_models
|
|
from .utils_motion import ADKeyframeGroup, AllPerBlocks
|
|
|
|
|
|
class LoadMotionCtrlCMCM:
|
|
NodeID = "ADE_LoadMotionCtrl_CMCMMOdel"
|
|
NodeName = "Load AnimateDiff+MotionCtrl Camera Model 🎭🅐🅓②"
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model_name": (get_available_motion_models(),),
|
|
"motionctrl_cmcm": (get_available_motion_models(),),
|
|
},
|
|
"optional": {
|
|
"override_ad_weights": ("BOOLEAN", {"default": True}),
|
|
"ad_settings": ("AD_SETTINGS",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MOTION_MODEL_ADE",)
|
|
RETURN_NAMES = ("MOTION_MODEL",)
|
|
CATEGORY = "Animate Diff 🎭🅐🅓/② Gen2 nodes ②/MotionCtrl"
|
|
FUNCTION = "load_motionctrl_cmcm"
|
|
|
|
def load_motionctrl_cmcm(self, model_name: str, motionctrl_cmcm: str,
|
|
override_ad_weights=True, ad_settings: AnimateDiffSettings=None,):
|
|
motion_model = load_motion_module_gen2(model_name=model_name, motion_model_settings=ad_settings)
|
|
inject_motionctrl_cmcm(motion_model.model, cmcm_name=motionctrl_cmcm, apply_non_ccs=override_ad_weights)
|
|
return (motion_model,)
|
|
|
|
|
|
class LoadMotionCtrlOMCM:
|
|
NodeID = "ADE_LoadMotionCtrl_OMCMMOdel"
|
|
NodeName = "Load MotionCtrl Object Model 🎭🅐🅓②"
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"motionctrl_omcm": (get_available_motion_models(),),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("OMCM_MOTIONCTRL",)
|
|
CATEGORY = "Animate Diff 🎭🅐🅓/② Gen2 nodes ②/MotionCtrl"
|
|
FUNCTION = "load_motionctrl_omcm"
|
|
|
|
def load_motionctrl_omcm(self, motionctrl_omcm: str):
|
|
omcm_modelpatcher = load_motionctrl_omcm(motionctrl_omcm)
|
|
return (omcm_modelpatcher,)
|
|
|
|
|
|
class LoadMotionCtrlCameraPosesFromFile:
|
|
NodeID = "ADE_LoadMotionCtrlCameraPosesFromFile"
|
|
NodeName = "Load MotionCtrl Camera Poses 🎭🅐🅓"
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
input_dir = folder_paths.get_input_directory()
|
|
files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))]
|
|
files = [f for f in files if f.endswith(".json")]
|
|
return {
|
|
"required": {
|
|
"pose_filename": (sorted(files),),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("CAMERA_MOTIONCTRL",)
|
|
CATEGORY = "Animate Diff 🎭🅐🅓/② Gen2 nodes ②/MotionCtrl"
|
|
FUNCTION = "load_camera_poses"
|
|
|
|
def load_camera_poses(self, pose_filename):
|
|
file_path = folder_paths.get_annotated_filepath(pose_filename)
|
|
with open(file_path, 'r') as f:
|
|
RT = json.load(f)
|
|
RT = np.array(RT)
|
|
RT = torch.tensor(RT).float() # [t, 12]
|
|
return (RT,)
|
|
|
|
|
|
class ApplyAnimateDiffMotionCtrlModel:
|
|
NodeID = "ADE_ApplyAnimateDiffModelWithMotionCtrl"
|
|
NodeName = "Apply AnimateDiff+MotionCtrl Model 🎭🅐🅓②"
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"motion_model": ("MOTION_MODEL_ADE",),
|
|
"start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
|
|
"end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001}),
|
|
},
|
|
"optional": {
|
|
"omcm_motionctrl": ("OMCM_MOTIONCTRL",),
|
|
"cameractrl_poses": ("CAMERACTRL_POSES",),
|
|
"motion_lora": ("MOTION_LORA",),
|
|
"scale_multival": ("MULTIVAL",),
|
|
"effect_multival": ("MULTIVAL",),
|
|
"ad_keyframes": ("AD_KEYFRAMES",),
|
|
"prev_m_models": ("M_MODELS",),
|
|
"per_block": ("PER_BLOCK",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("M_MODELS",)
|
|
CATEGORY = "Animate Diff 🎭🅐🅓/② Gen2 nodes ②/MotionCtrl"
|
|
FUNCTION = "apply_motion_model"
|
|
|
|
def apply_motion_model(self, motion_model: MotionModelPatcher, start_percent: float=0.0, end_percent: float=1.0,
|
|
omcm_motionctrl: ObjectControlModelPatcher=None,
|
|
cameractrl_poses: list[list[float]]=None,
|
|
motion_lora: MotionLoraList=None, ad_keyframes: ADKeyframeGroup=None,
|
|
scale_multival=None, effect_multival=None, per_block: AllPerBlocks=None,
|
|
prev_m_models: MotionModelGroup=None,):
|
|
(new_m_models,) = ApplyAnimateDiffModelNode.apply_motion_model(self, motion_model, start_percent=start_percent, end_percent=end_percent,
|
|
motion_lora=motion_lora, ad_keyframes=ad_keyframes, per_block=per_block,
|
|
scale_multival=scale_multival, effect_multival=effect_multival, prev_m_models=prev_m_models)
|
|
# most recent added model will always be first in list
|
|
curr_model = new_m_models.models[0]
|
|
# check if model has CMCM; if so, make sure something is provided for it
|
|
if curr_model.model.is_motionctrl_cc_enabled():
|
|
attachment = get_mm_attachment(curr_model)
|
|
if cameractrl_poses is not None:
|
|
RT = convert_cameractrl_poses_to_RT(cameractrl_poses)
|
|
attachment.orig_RT = RT
|
|
else:
|
|
attachment.orig_RT = torch.zeros((1, 12))
|
|
# attachment.orig_RT = cameractrl_poses
|
|
# else:
|
|
# attachment.orig_RT = torch.zeros([])
|
|
|
|
# check if OMCM is provided; if so, make sure something is provided for it
|
|
return (new_m_models,)
|