import torch import os import sys from pathlib import Path EXTENSION_PATH = Path(__file__).parent sys.path.insert(0, str(EXTENSION_PATH.resolve())) import custom_mmpkg.custom_mmcv as mmcv import numpy as np import torch from mogen.models import build_architecture from custom_mmpkg.custom_mmcv.runner import load_checkpoint from mogen.utils.plot_utils import ( plot_3d_motion, t2m_kinematic_chain ) from comfy.model_management import get_torch_device, soft_empty_cache from .config import get_model_dataset_dict from .utils import * from mogen.smpl.simplify_loc2rot import joints2smpl from mogen.smpl.rotation2xyz import Rotation2xyz #from custom_mmpkg.custom_mmhuman3d.core.conventions.keypoints_mapping import convert_kps #from custom_mmpkg.custom_mmhuman3d.core.visualization.visualize_keypoints3d import visualize_kp3d from mogen.smpl.render_mesh import render_from_smpl import gc from pathlib import Path rot2xyz = Rotation2xyz(device=get_torch_device()) def create_mdm_model(model_config): cfg = mmcv.Config.fromstring(model_config.config_code, '.py') mdm = build_architecture(cfg.model) load_checkpoint(mdm, str(model_config.ckpt_path), map_location='cpu') mdm.eval().cpu() return mdm class MotionDiffModelWrapper(torch.nn.Module): #Anything beside CLIP (mdm.model) def __init__(self, mdm, dataset) -> None: super(MotionDiffModelWrapper, self).__init__() self.loss_recon=mdm.loss_recon self.diffusion_train=mdm.diffusion_train self.diffusion_test=mdm.diffusion_test self.sampler=mdm.sampler self.dataset = dataset def forward(self, clip, cond_dict, **kwargs): motion, motion_mask = kwargs['motion'].float(), kwargs['motion_mask'].float() sample_idx = kwargs.get('sample_idx', None) clip_feat = kwargs.get('clip_feat', None) sampler = kwargs.get('sampler', 'ddpm') B, T = motion.shape[:2] dim_pose = kwargs['motion'].shape[-1] model_kwargs = cond_dict model_kwargs['motion_mask'] = motion_mask model_kwargs['sample_idx'] = sample_idx inference_kwargs = kwargs.get('inference_kwargs', {}) if sampler == 'ddpm': output = self.diffusion_test.p_sample_loop( clip, (B, T, dim_pose), clip_denoised=False, progress=False, model_kwargs=model_kwargs, **inference_kwargs ) else: output = self.diffusion_test.ddim_sample_loop( clip, (B, T, dim_pose), clip_denoised=False, progress=False, model_kwargs=model_kwargs, eta=0, **inference_kwargs ) if getattr(clip, "post_process") is not None: output = clip.post_process(output) results = kwargs results['pred_motion'] = output results = self.split_results(results) return results def split_results(self, results): B = results['motion'].shape[0] output = [] for i in range(B): batch_output = dict() batch_output['motion'] = to_cpu(results['motion'][i]) batch_output['pred_motion'] = to_cpu(results['pred_motion'][i]) batch_output['motion_length'] = to_cpu(results['motion_length'][i]) batch_output['motion_mask'] = to_cpu(results['motion_mask'][i]) if 'pred_motion_length' in results.keys(): batch_output['pred_motion_length'] = to_cpu(results['pred_motion_length'][i]) else: batch_output['pred_motion_length'] = to_cpu(results['motion_length'][i]) if 'pred_motion_mask' in results: batch_output['pred_motion_mask'] = to_cpu(results['pred_motion_mask'][i]) else: batch_output['pred_motion_mask'] = to_cpu(results['motion_mask'][i]) if 'motion_metas' in results.keys(): motion_metas = results['motion_metas'][i] if 'text' in motion_metas.keys(): batch_output['text'] = motion_metas['text'] if 'token' in motion_metas.keys(): batch_output['token'] = motion_metas['token'] output.append(batch_output) return output class MotionDiffCLIPWrapper(torch.nn.Module): def __init__(self, mdm): super(MotionDiffCLIPWrapper, self).__init__() self.model = mdm.model def forward(self, text, motion_data): self.model.to(get_torch_device()) B, T = motion_data["motion"].shape[:2] texts = [] for _ in range(B): texts.append(text) out = self.model.get_precompute_condition(device=get_torch_device(), text=texts, **motion_data) self.model.cpu() return out model_dataset_dict = None class MotionDiffLoader: @classmethod def INPUT_TYPES(s): global model_dataset_dict model_dataset_dict = get_model_dataset_dict() return { "required": { "model_dataset": ( list(model_dataset_dict.keys()), { "default": "remodiffuse-human_ml3d" } ) }, } RETURN_TYPES = ("MD_MODEL", "MD_CLIP") CATEGORY = "MotionDiff" FUNCTION = "load_mdm" def load_mdm(self, model_dataset): global model_dataset_dict model_config = model_dataset_dict[model_dataset]() mdm = create_mdm_model(model_config) return (MotionDiffModelWrapper(mdm, dataset=model_config.dataset), MotionDiffCLIPWrapper(mdm)) class MotionCLIPTextEncode: @classmethod def INPUT_TYPES(s): return { "required": { "md_clip": ("MD_CLIP", ), "motion_data": ("MOTION_DATA", ), "text": ("STRING", {"default": "a person performs a cartwheel" ,"multiline": True}) }, } RETURN_TYPES = ("MD_CONDITIONING",) CATEGORY = "MotionDiff" FUNCTION = "encode_text" def encode_text(self, md_clip, motion_data, text): return (md_clip(text, motion_data), ) class EmptyMotionData: @classmethod def INPUT_TYPES(s): return { "required": { "frames": ("INT", {"default": 196, "min": 1, "max": 196}) } } RETURN_TYPES = ("MOTION_DATA", ) CATEGORY = "MotionDiff" FUNCTION = "encode_text" def encode_text(self, frames): return ({ 'motion': torch.zeros(1, frames, 263), 'motion_mask': torch.ones(1, frames), 'motion_length': torch.Tensor([frames]).long(), }, ) class MotionDiffSimpleSampler: @classmethod def INPUT_TYPES(s): return { "required": { "sampler_name": (["ddpm", "ddim"], ), "md_model": ("MD_MODEL", ), "md_clip": ("MD_CLIP", ), "md_cond": ("MD_CONDITIONING", ), "motion_data": ("MOTION_DATA",) } } RETURN_TYPES = ("MOTION_DATA",) CATEGORY = "MotionDiff" FUNCTION = "sample" def sample(self, sampler_name, md_model: MotionDiffModelWrapper, md_clip, md_cond, motion_data): md_model.to(get_torch_device()) md_clip.to(get_torch_device()) for key in motion_data: motion_data[key] = to_gpu(motion_data[key]) kwargs = { **motion_data, 'inference_kwargs': {}, 'sampler': sampler_name, } with torch.no_grad(): output = md_model(md_clip.model, cond_dict=md_cond, **kwargs)[0]['pred_motion'] pred_motion = output * md_model.dataset.std + md_model.dataset.mean pred_motion = pred_motion.cpu().detach() md_model.cpu(), md_clip.cpu() for key in motion_data: motion_data[key] = to_cpu(motion_data[key]) return ({ 'motion': pred_motion, 'motion_mask': motion_data['motion_mask'], 'motion_length': motion_data['motion_length'], }, ) class MotionDataVisualizer: @classmethod def INPUT_TYPES(s): return { "required": { "motion_data": ("MOTION_DATA", ), "visualization": (["original", "pseudo-openpose"], {"default": "pseudo-openpose"}), "distance": ("FLOAT", {"default": 7.0, "min": 0.0, "max": 10.0, "step": 0.1}), "elevation": ("FLOAT", {"default": 120, "min": 0.0, "max": 300.0, "step": 0.1}), "rotation": ("FLOAT", {"default": -90, "min": -180, "max": 180, "step": 1}), "poselinewidth": ("FLOAT", {"default": 2, "min": 0, "max": 50, "step": 0.1}), }, "optional": { "opt_title": ("STRING", {"default": '' ,"multiline": False}), } } RETURN_TYPES = ("IMAGE",) CATEGORY = "MotionDiff" FUNCTION = "visualize" def visualize(self, motion_data, visualization, distance, elevation, rotation, poselinewidth, opt_title=None): joint = motion_data_to_joints(motion_data["motion"]) pil_frames = plot_3d_motion( None, t2m_kinematic_chain, joint, distance, elevation, rotation, poselinewidth, title=opt_title if opt_title is not None else '', fps=1, save_as_pil_lists=True, visualization=visualization ) tensor_frames = [] for pil_image in pil_frames: np_image = np.array(pil_image.convert("RGB")).astype(np.float32) / 255.0 tensor_frames.append(torch.from_numpy(np_image)) return (torch.stack(tensor_frames, dim=0), ) class SmplifyMotionData: @classmethod def INPUT_TYPES(s): return { "required": { "motion_data": ("MOTION_DATA", ), "num_smplify_iters": ("INT", {"min": 10, "max": 1000, "default": 150}), "smplify_step_size": ("FLOAT", {"min": 1e-4, "max": 5e-1, "step": 1e-4, "default": 1e-2}) } } RETURN_TYPES = ("SMPL",) CATEGORY = "MotionDiff/smpl" FUNCTION = "convent" def convent(self, motion_data, num_smplify_iters, smplify_step_size): joints = motion_data_to_joints(motion_data["motion"]) with torch.inference_mode(False): convention = joints2smpl( num_frames=joints.shape[0], device=get_torch_device(), num_smplify_iters=num_smplify_iters, smplify_step_size=smplify_step_size ) motion_tensor, meta = convention.joint2smpl(joints) motion_tensor = motion_tensor.cpu().detach() for key in meta: meta[key] = meta[key].cpu().detach() gc.collect() soft_empty_cache() return ((motion_tensor, meta), ) #Caching class RenderOpenPoseFromSMPL: @classmethod def INPUT_TYPES(s): return { "required": { "smpl": ("SMPL", ) } } RETURN_TYPES = ("IMAGE",) CATEGORY = "MotionDiff/smpl" FUNCTION = "convent" def convent(self, smpl): kps = smpl[1]["pose"] kp3d_openpose, _ = convert_kps(kps, src='smpl_45', dst='openpose_25') cv2_frames = visualize_kp3d(kp3d_openpose.cpu().numpy(), data_source='openpose_25', return_array=True, resolution=(1024, 1024)) return (torch.from_numpy(cv2_frames), ) class RenderSMPLMesh: @classmethod def INPUT_TYPES(s): return { "required": { "smpl": ("SMPL", ), "draw_platform": ("BOOLEAN", {"default": False}), "depth_only": ("BOOLEAN", {"default": False}), "yfov": ("FLOAT", {"default": 3,"min": 0.1, "max": 10, "step": 0.1}), "move_x": ("FLOAT", {"default": 0,"min": -500, "max": 500, "step": 0.1}), "move_y": ("FLOAT", {"default": 0,"min": -500, "max": 500, "step": 0.1}), "move_z": ("FLOAT", {"default": 0,"min": -500, "max": 500, "step": 0.1}), } } RETURN_TYPES = ("IMAGE", "IMAGE") RETURN_NAMES = ("IMAGE", "DEPTH_MAP") CATEGORY = "MotionDiff/smpl" FUNCTION = "render" def render(self, smpl, yfov, move_x, move_y, move_z, draw_platform, depth_only): motion_tensor, _ = smpl color_frames, depth_frames = render_from_smpl(motion_tensor.to(get_torch_device()),yfov, move_x, move_y, move_z, draw_platform,depth_only) color_frames = torch.from_numpy(color_frames[..., :3].astype(np.float32) / 255.) #Normalize to [0, 1] normalized_depth = (depth_frames - depth_frames.min()) / (depth_frames.max() - depth_frames.min()) #Pyrender's depths are the distance in meters to the camera, which is the inverse of depths in normal context #Ref: https://github.com/mmatl/pyrender/issues/10#issuecomment-468995891 normalized_depth[normalized_depth != 0] = 1 - normalized_depth[normalized_depth != 0] #https://github.com/Fannovel16/comfyui_controlnet_aux/blob/main/src/controlnet_aux/util.py#L24 depth_frames = [torch.from_numpy(np.concatenate([x, x, x], axis=2)) for x in normalized_depth[..., None]] depth_frames = torch.stack(depth_frames, dim=0) return (color_frames, depth_frames,) NODE_CLASS_MAPPINGS = { "MotionDiffLoader": MotionDiffLoader, "MotionCLIPTextEncode": MotionCLIPTextEncode, "MotionDiffSimpleSampler": MotionDiffSimpleSampler, "EmptyMotionData": EmptyMotionData, "MotionDataVisualizer": MotionDataVisualizer, "SmplifyMotionData": SmplifyMotionData, "RenderSMPLMesh": RenderSMPLMesh #"RenderOpenPoseFromSMPL": RenderOpenPoseFromSMPL }