129 lines
7.0 KiB
Markdown
129 lines
7.0 KiB
Markdown
# Documentation
|
|
- Class name: MagicAnimateModelLoader
|
|
- Category: ComfyUI Magic Animate
|
|
- Output node: False
|
|
- Repo Ref: https://github.com/thecooltechguy/ComfyUI-MagicAnimate
|
|
|
|
MagicAnimateModelLoader类负责管理和加载魔法动画模型的各个组件。它作为控制网络、外观编码器、运动模块和其他相关模块的集中存储库。load_model方法旨在清除任何现有的模型，加载新配置，并初始化用于动画目的的模型组件。
|
|
|
|
# Input types
|
|
## Optional
|
|
- controlnet
|
|
- controlnet参数指定控制网络的检查点或路径，这是动画模型的关键部分。它用于指导生成过程。
|
|
- Comfy dtype: str
|
|
- Python dtype: str
|
|
- appearance_encoder
|
|
- appearance_encoder参数表示外观编码器的检查点或文件路径，它负责编码输入的视觉特征。
|
|
- Comfy dtype: str
|
|
- Python dtype: str
|
|
- motion_module
|
|
- motion_module参数指向运动模块的检查点或文件路径，这对于根据给定的运动对输入进行动画制作至关重要。
|
|
- Comfy dtype: str
|
|
- Python dtype: str
|
|
- device
|
|
- device参数确定用于模型执行的计算设备，通常为'cuda'以进行GPU加速或'cpu'。
|
|
- Comfy dtype: str
|
|
- Python dtype: str
|
|
|
|
# Output types
|
|
- MAGIC_ANIMATE_MODEL
|
|
- 输出是一个包含已加载模型和配置的字典。它包括VAE、文本编码器、外观编码器、分词器、UNet、控制网络和动画流水线。
|
|
- Comfy dtype: dict
|
|
- Python dtype: Dict[str, torch.nn.Module]
|
|
|
|
# Usage tips
|
|
- Infra type: GPU
|
|
|
|
# Source code
|
|
```
|
|
class MagicAnimateModelLoader:
|
|
|
|
def __init__(self):
|
|
self.models = {}
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
magic_animate_checkpoints = folder_paths.get_filename_list('magic_animate')
|
|
devices = []
|
|
if True:
|
|
devices.append('cuda')
|
|
devices.append('cpu')
|
|
return {'required': {'controlnet': (magic_animate_checkpoints, {'default': magic_animate_checkpoints[0]}), 'appearance_encoder': (magic_animate_checkpoints, {'default': magic_animate_checkpoints[0]}), 'motion_module': (magic_animate_checkpoints, {'default': magic_animate_checkpoints[0]}), 'device': (devices,)}}
|
|
RETURN_TYPES = ('MAGIC_ANIMATE_MODEL',)
|
|
FUNCTION = 'load_model'
|
|
CATEGORY = 'ComfyUI Magic Animate'
|
|
|
|
def load_model(self, controlnet, appearance_encoder, motion_module, device):
|
|
if self.models:
|
|
all_keys = list(self.models.keys())
|
|
for key in all_keys:
|
|
del self.models[key]
|
|
self.models = {}
|
|
gc.collect()
|
|
current_dir = os.path.dirname(os.path.realpath(__file__))
|
|
config = OmegaConf.load(os.path.join(current_dir, 'configs', 'prompts', 'animation.yaml'))
|
|
inference_config = OmegaConf.load(os.path.join(current_dir, 'configs', 'inference', 'inference.yaml'))
|
|
magic_animate_models_dir = folder_paths.get_folder_paths('magic_animate')[0]
|
|
config.pretrained_model_path = os.path.join(magic_animate_models_dir, 'stable-diffusion-v1-5')
|
|
config.pretrained_vae_path = os.path.join(magic_animate_models_dir, 'sd-vae-ft-mse')
|
|
config.pretrained_appearance_encoder_path = os.path.join(magic_animate_models_dir, os.path.dirname(appearance_encoder))
|
|
config.pretrained_controlnet_path = os.path.join(magic_animate_models_dir, os.path.dirname(controlnet))
|
|
motion_module = os.path.join(magic_animate_models_dir, motion_module)
|
|
config.motion_module = motion_module
|
|
tokenizer = CLIPTokenizer.from_pretrained(config.pretrained_model_path, subfolder='tokenizer')
|
|
text_encoder = CLIPTextModel.from_pretrained(config.pretrained_model_path, subfolder='text_encoder')
|
|
if config.pretrained_unet_path:
|
|
unet = UNet3DConditionModel.from_pretrained_2d(config.pretrained_unet_path, unet_additional_kwargs=OmegaConf.to_container(inference_config.unet_additional_kwargs))
|
|
else:
|
|
unet = UNet3DConditionModel.from_pretrained_2d(config.pretrained_model_path, subfolder='unet', unet_additional_kwargs=OmegaConf.to_container(inference_config.unet_additional_kwargs))
|
|
appearance_encoder = AppearanceEncoderModel.from_pretrained(config.pretrained_appearance_encoder_path).to(device)
|
|
reference_control_writer = ReferenceAttentionControl(appearance_encoder, do_classifier_free_guidance=True, mode='write', fusion_blocks=config.fusion_blocks)
|
|
reference_control_reader = ReferenceAttentionControl(unet, do_classifier_free_guidance=True, mode='read', fusion_blocks=config.fusion_blocks)
|
|
vae = AutoencoderKL.from_pretrained(config.pretrained_vae_path)
|
|
controlnet = ControlNetModel.from_pretrained(config.pretrained_controlnet_path)
|
|
vae.to(torch.float16)
|
|
unet.to(torch.float16)
|
|
text_encoder.to(torch.float16)
|
|
appearance_encoder.to(torch.float16)
|
|
controlnet.to(torch.float16)
|
|
pipeline = AnimationPipeline(vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, unet=unet, controlnet=controlnet, scheduler=DDIMScheduler(**OmegaConf.to_container(inference_config.noise_scheduler_kwargs)))
|
|
motion_module_state_dict = torch.load(motion_module, map_location='cpu')
|
|
motion_module_state_dict = motion_module_state_dict['state_dict'] if 'state_dict' in motion_module_state_dict else motion_module_state_dict
|
|
try:
|
|
state_dict = OrderedDict()
|
|
for key in motion_module_state_dict.keys():
|
|
if key.startswith('module.'):
|
|
_key = key.split('module.')[-1]
|
|
state_dict[_key] = motion_module_state_dict[key]
|
|
else:
|
|
state_dict[key] = motion_module_state_dict[key]
|
|
motion_module_state_dict = state_dict
|
|
del state_dict
|
|
(missing, unexpected) = pipeline.unet.load_state_dict(motion_module_state_dict, strict=False)
|
|
assert len(unexpected) == 0
|
|
except:
|
|
_tmp_ = OrderedDict()
|
|
for key in motion_module_state_dict.keys():
|
|
if 'motion_modules' in key:
|
|
if key.startswith('unet.'):
|
|
_key = key.split('unet.')[-1]
|
|
_tmp_[_key] = motion_module_state_dict[key]
|
|
else:
|
|
_tmp_[key] = motion_module_state_dict[key]
|
|
(missing, unexpected) = unet.load_state_dict(_tmp_, strict=False)
|
|
assert len(unexpected) == 0
|
|
del _tmp_
|
|
del motion_module_state_dict
|
|
pipeline.to(device)
|
|
self.models['vae'] = vae
|
|
self.models['text_encoder'] = text_encoder
|
|
self.models['appearance_encoder'] = appearance_encoder
|
|
self.models['tokenizer'] = tokenizer
|
|
self.models['unet'] = unet
|
|
self.models['controlnet'] = controlnet
|
|
self.models['pipeline'] = pipeline
|
|
self.models['config'] = config
|
|
self.models['reference_control_writer'] = reference_control_writer
|
|
self.models['reference_control_reader'] = reference_control_reader
|
|
return (self.models,)
|
|
``` |