diff --git a/animatediff/model_injection.py b/animatediff/model_injection.py index dc1f58c..aecc2c0 100644 --- a/animatediff/model_injection.py +++ b/animatediff/model_injection.py @@ -1262,6 +1262,14 @@ def inject_img_encoder_into_model(motion_model: MotionModelPatcher, w_encoder: M motion_model.model.img_encoder.load_state_dict(w_encoder.model.img_encoder.state_dict()) +def inject_pia_conv_in_into_model(motion_model: MotionModelPatcher, w_pia: MotionModelPatcher): + motion_model.model.init_conv_in(w_pia.model.state_dict()) + motion_model.model.conv_in.to(comfy.model_management.unet_dtype()) + motion_model.model.conv_in.to(comfy.model_management.unet_offload_device()) + motion_model.model.conv_in.load_state_dict(w_pia.model.conv_in.state_dict()) + motion_model.model.mm_info.mm_format = AnimateDiffFormat.PIA + + def inject_camera_encoder_into_model(motion_model: MotionModelPatcher, camera_ctrl_name: str): camera_ctrl_path = get_motion_model_path(camera_ctrl_name) full_state_dict = comfy.utils.load_torch_file(camera_ctrl_path, safe_load=True) diff --git a/animatediff/nodes.py b/animatediff/nodes.py index 61aa8c0..6a8e7cf 100644 --- a/animatediff/nodes.py +++ b/animatediff/nodes.py @@ -9,7 +9,7 @@ from .nodes_animatelcmi2v import (ApplyAnimateLCMI2VModel, LoadAnimateLCMI2VMode from .nodes_cameractrl import (LoadAnimateDiffModelWithCameraCtrl, ApplyAnimateDiffWithCameraCtrl, CameraCtrlADKeyframeNode, LoadCameraPoses, CameraCtrlPoseBasic, CameraCtrlPoseCombo, CameraCtrlPoseAdvanced, CameraCtrlManualAppendPose, CameraCtrlReplaceCameraParameters, CameraCtrlSetOriginalAspectRatio) -from .nodes_pia import (ApplyAnimateDiffPIAModel, InputPIA_MultivalNode, InputPIA_PaperPresetsNode, PIA_ADKeyframeNode) +from .nodes_pia import (ApplyAnimateDiffPIAModel, LoadAnimateDiffAndInjectPIANode, InputPIA_MultivalNode, InputPIA_PaperPresetsNode, PIA_ADKeyframeNode) from .nodes_multival import MultivalDynamicNode, MultivalScaledMaskNode, MultivalDynamicFloatInputNode, MultivalConvertToMaskNode from .nodes_conditioning import (MaskableLoraLoader, MaskableLoraLoaderModelOnly, MaskableSDModelLoader, MaskableSDModelLoaderModelOnly, SetModelLoraHook, SetClipLoraHook, @@ -142,6 +142,7 @@ NODE_CLASS_MAPPINGS = { "ADE_InputPIA_Multival": InputPIA_MultivalNode, "ADE_InputPIA_PaperPresets": InputPIA_PaperPresetsNode, "ADE_PIA_AnimateDiffKeyframe": PIA_ADKeyframeNode, + "ADE_InjectPIAIntoAnimateDiffModel": LoadAnimateDiffAndInjectPIANode, # Deprecated Nodes "AnimateDiffLoaderV1": AnimateDiffLoader_Deprecated, "ADE_AnimateDiffLoaderV1Advanced": AnimateDiffLoaderAdvanced_Deprecated, @@ -254,6 +255,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "ADE_InputPIA_Multival": "PIA Input [Multival] ๐ŸŽญ๐Ÿ…๐Ÿ…“โ‘ก", "ADE_InputPIA_PaperPresets": "PIA Input [Paper Presets] ๐ŸŽญ๐Ÿ…๐Ÿ…“โ‘ก", "ADE_PIA_AnimateDiffKeyframe": "AnimateDiff-PIA Keyframe ๐ŸŽญ๐Ÿ…๐Ÿ…“", + "ADE_InjectPIAIntoAnimateDiffModel": "๐ŸงชInject PIA into AnimateDiff Model ๐ŸŽญ๐Ÿ…๐Ÿ…“โ‘ก", # Deprecated Nodes "AnimateDiffLoaderV1": "๐ŸšซAnimateDiff Loader [DEPRECATED] ๐ŸŽญ๐Ÿ…๐Ÿ…“", "ADE_AnimateDiffLoaderV1Advanced": "๐ŸšซAnimateDiff Loader (Advanced) [DEPRECATED] ๐ŸŽญ๐Ÿ…๐Ÿ…“", diff --git a/animatediff/nodes_pia.py b/animatediff/nodes_pia.py index c72837e..fc8024d 100644 --- a/animatediff/nodes_pia.py +++ b/animatediff/nodes_pia.py @@ -5,11 +5,12 @@ import math from comfy.sd import VAE +from .ad_settings import AnimateDiffSettings from .logger import logger -from .utils_model import BIGMIN, BIGMAX +from .utils_model import BIGMIN, BIGMAX, get_available_motion_models from .utils_motion import ADKeyframeGroup, InputPIA, InputPIA_Multival, extend_list_to_batch_size, extend_to_batch_size, prepare_mask_batch from .motion_lora import MotionLoraList -from .model_injection import MotionModelGroup, MotionModelPatcher +from .model_injection import MotionModelGroup, MotionModelPatcher, load_motion_module_gen2, inject_pia_conv_in_into_model from .motion_module_ad import AnimateDiffFormat from .nodes_gen2 import ApplyAnimateDiffModelNode, ADKeyframeNode @@ -153,6 +154,35 @@ class ApplyAnimateDiffPIAModel: return new_m_models +class LoadAnimateDiffAndInjectPIANode: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model_name": (get_available_motion_models(),), + "motion_model": ("MOTION_MODEL_ADE",), + }, + "optional": { + "ad_settings": ("AD_SETTINGS",), + } + } + + RETURN_TYPES = ("MOTION_MODEL_ADE",) + RETURN_NAMES = ("MOTION_MODEL",) + + CATEGORY = "Animate Diff ๐ŸŽญ๐Ÿ…๐Ÿ…“/โ‘ก Gen2 nodes โ‘ก/PIA/๐Ÿงชexperimental" + FUNCTION = "load_motion_model" + + def load_motion_model(self, model_name: str, motion_model: MotionModelPatcher, ad_settings: AnimateDiffSettings=None): + # make sure model actually has PIA conv_in + if motion_model.model.conv_in is None: + raise Exception("Passed-in motion model was expected to be PIA (contain conv_in), but did not.") + # load motion module and motion settings, if included + loaded_motion_model = load_motion_module_gen2(model_name=model_name, motion_model_settings=ad_settings) + inject_pia_conv_in_into_model(motion_model=loaded_motion_model, w_pia=motion_model) + return (loaded_motion_model,) + + class PIA_ADKeyframeNode: @classmethod def INPUT_TYPES(s):