diff --git a/modelscope/metainfo.py b/modelscope/metainfo.py index c907f482..24b274a0 100644 --- a/modelscope/metainfo.py +++ b/modelscope/metainfo.py @@ -304,6 +304,7 @@ class Pipelines(object): ddcolor_image_colorization = 'ddcolor-image-colorization' image_fewshot_detection = 'image-fewshot-detection' image_face_fusion = 'image-face-fusion' + motion_generattion = 'mdm-motion-generation' # nlp tasks automatic_post_editing = 'automatic-post-editing' diff --git a/modelscope/models/cv/motion_generation/__init__.py b/modelscope/models/cv/motion_generation/__init__.py new file mode 100644 index 00000000..0f8cbad7 --- /dev/null +++ b/modelscope/models/cv/motion_generation/__init__.py @@ -0,0 +1,24 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +from typing import TYPE_CHECKING + +from modelscope.utils.import_utils import LazyImportModule + +if TYPE_CHECKING: + + from .model import create_model, load_model_wo_clip + from .modules.cfg_sampler import ClassifierFreeSampleModel +else: + _import_structure = { + 'model': ['create_model', 'load_model_wo_clip'], + 'modules.cfg_sampler': ['ClassifierFreeSampleModel'] + } + + import sys + + sys.modules[__name__] = LazyImportModule( + __name__, + globals()['__file__'], + _import_structure, + module_spec=__spec__, + extra_objects={}, + ) diff --git a/modelscope/models/cv/motion_generation/model.py b/modelscope/models/cv/motion_generation/model.py new file mode 100644 index 00000000..aa944ada --- /dev/null +++ b/modelscope/models/cv/motion_generation/model.py @@ -0,0 +1,65 @@ +# This code is borrowed and modified from Human Motion Diffusion Model, +# made publicly available under MIT license at https://github.com/GuyTevet/motion-diffusion-model + +from .modules import gaussian_diffusion as gd +from .modules.mdm import MDM +from .modules.respace import SpacedDiffusion, space_timesteps + + +def load_model_wo_clip(model, state_dict): + missing_keys, unexpected_keys = model.load_state_dict( + state_dict, strict=False) + assert len(unexpected_keys) == 0 + assert all([k.startswith('clip_model.') for k in missing_keys]) + + +def create_model(cfg): + model = MDM( + '', + njoints=263, + nfeats=1, + num_actions=1, + translation=True, + pose_rep='rot6d', + glob=True, + glob_rot=True, + latent_dim=512, + ff_size=1024, + smpl_data_path=cfg.smpl_data_path, + data_rep='hml_vec', + dataset='humanml', + clip_version='ViT-B/32', + **{ + 'cond_mode': 'text', + 'cond_mask_prob': 0.1, + 'action_emb': 'tensor' + }) + + predict_xstart = True # we always predict x_start (a.k.a. x0), that's our deal! + steps = cfg.sample_steps + scale_beta = 1. # no scaling + timestep_respacing = '' # can be used for ddim sampling, we don't use it. + learn_sigma = False + rescale_timesteps = False + + betas = gd.get_named_beta_schedule('cosine', steps, scale_beta) + loss_type = gd.LossType.MSE + + if not timestep_respacing: + timestep_respacing = [steps] + + diffusion = SpacedDiffusion( + use_timesteps=space_timesteps(steps, timestep_respacing), + betas=betas, + model_mean_type=(gd.ModelMeanType.EPSILON + if not predict_xstart else gd.ModelMeanType.START_X), + model_var_type=((gd.ModelVarType.FIXED_LARGE + if not True else gd.ModelVarType.FIXED_SMALL) + if not learn_sigma else gd.ModelVarType.LEARNED_RANGE), + loss_type=loss_type, + rescale_timesteps=rescale_timesteps, + lambda_vel=0.0, + lambda_rcxyz=0.0, + lambda_fc=0.0, + ) + return model, diffusion diff --git a/modelscope/models/cv/motion_generation/modules/__init__.py b/modelscope/models/cv/motion_generation/modules/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/modelscope/models/cv/motion_generation/modules/cfg_sampler.py b/modelscope/models/cv/motion_generation/modules/cfg_sampler.py new file mode 100644 index 00000000..cd07362d --- /dev/null +++ b/modelscope/models/cv/motion_generation/modules/cfg_sampler.py @@ -0,0 +1,33 @@ +# This code is borrowed and modified from Human Motion Diffusion Model, +# made publicly available under MIT license at https://github.com/GuyTevet/motion-diffusion-model +from copy import deepcopy + +import torch.nn as nn + + +# A wrapper model for Classifier-free guidance **SAMPLING** only +# https://arxiv.org/abs/2207.12598 +class ClassifierFreeSampleModel(nn.Module): + + def __init__(self, model): + super().__init__() + self.model = model # model is the actual model to run + + assert self.model.cond_mask_prob > 0 + + # pointers to inner model + self.rot2xyz = self.model.rot2xyz + self.translation = self.model.translation + self.njoints = self.model.njoints + self.nfeats = self.model.nfeats + self.data_rep = self.model.data_rep + self.cond_mode = self.model.cond_mode + + def forward(self, x, timesteps, y=None): + cond_mode = self.model.cond_mode + assert cond_mode in ['text', 'action'] + y_uncond = deepcopy(y) + y_uncond['uncond'] = True + out = self.model(x, timesteps, y) + out_uncond = self.model(x, timesteps, y_uncond) + return out_uncond + (y['scale'].view(-1, 1, 1, 1) * (out - out_uncond)) diff --git a/modelscope/models/cv/motion_generation/modules/gaussian_diffusion.py b/modelscope/models/cv/motion_generation/modules/gaussian_diffusion.py new file mode 100644 index 00000000..2d283642 --- /dev/null +++ b/modelscope/models/cv/motion_generation/modules/gaussian_diffusion.py @@ -0,0 +1,666 @@ +# This code is borrowed and modified from Human Motion Diffusion Model, +# made publicly available under MIT license at https://github.com/GuyTevet/motion-diffusion-model + +import enum +import math +from copy import deepcopy + +import numpy as np +import torch as th + + +def get_named_beta_schedule(schedule_name, + num_diffusion_timesteps, + scale_betas=1.): + """ + Get a pre-defined beta schedule for the given name. + + The beta schedule library consists of beta schedules which remain similar + in the limit of num_diffusion_timesteps. + Beta schedules may be added, but should not be removed or changed once + they are committed to maintain backwards compatibility. + """ + if schedule_name == 'linear': + # Linear schedule from Ho et al, extended to work for any number of + # diffusion steps. + scale = scale_betas * 1000 / num_diffusion_timesteps + beta_start = scale * 0.0001 + beta_end = scale * 0.02 + return np.linspace( + beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64) + elif schedule_name == 'cosine': + return betas_for_alpha_bar( + num_diffusion_timesteps, + lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2)**2, + ) + else: + raise NotImplementedError(f'unknown beta schedule: {schedule_name}') + + +def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999): + """ + Create a beta schedule that discretizes the given alpha_t_bar function, + which defines the cumulative product of (1-beta) over time from t = [0,1]. + + :param num_diffusion_timesteps: the number of betas to produce. + :param alpha_bar: a lambda that takes an argument t from 0 to 1 and + produces the cumulative product of (1-beta) up to that + part of the diffusion process. + :param max_beta: the maximum beta to use; use values lower than 1 to + prevent singularities. + """ + betas = [] + for i in range(num_diffusion_timesteps): + t1 = i / num_diffusion_timesteps + t2 = (i + 1) / num_diffusion_timesteps + betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta)) + return np.array(betas) + + +class ModelMeanType(enum.Enum): + """ + Which type of output the model predicts. + """ + + PREVIOUS_X = enum.auto() # the model predicts x_{t-1} + START_X = enum.auto() # the model predicts x_0 + EPSILON = enum.auto() # the model predicts epsilon + + +class ModelVarType(enum.Enum): + """ + What is used as the model's output variance. + + The LEARNED_RANGE option has been added to allow the model to predict + values between FIXED_SMALL and FIXED_LARGE, making its job easier. + """ + + LEARNED = enum.auto() + FIXED_SMALL = enum.auto() + FIXED_LARGE = enum.auto() + LEARNED_RANGE = enum.auto() + + +class LossType(enum.Enum): + MSE = enum.auto() # use raw MSE loss (and KL when learning variances) + RESCALED_MSE = ( + enum.auto() + ) # use raw MSE loss (with RESCALED_KL when learning variances) + KL = enum.auto() # use the variational lower-bound + RESCALED_KL = enum.auto() # like KL, but rescale to estimate the full VLB + + def is_vb(self): + return self == LossType.KL or self == LossType.RESCALED_KL + + +class GaussianDiffusion: + """ + Utilities for training and sampling diffusion models. + + Ported directly from here, and then adapted over time to further experimentation. + https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py#L42 + + :param betas: a 1-D numpy array of betas for each diffusion timestep, + starting at T and going to 1. + :param model_mean_type: a ModelMeanType determining what the model outputs. + :param model_var_type: a ModelVarType determining how variance is output. + :param loss_type: a LossType determining the loss function to use. + :param rescale_timesteps: if True, pass floating point timesteps into the + model so that they are always scaled like in the + original paper (0 to 1000). + """ + + def __init__( + self, + *, + betas, + model_mean_type, + model_var_type, + loss_type, + rescale_timesteps=False, + lambda_rcxyz=0., + lambda_vel=0., + lambda_pose=1., + lambda_orient=1., + lambda_loc=1., + data_rep='rot6d', + lambda_root_vel=0., + lambda_vel_rcxyz=0., + lambda_fc=0., + ): + self.model_mean_type = model_mean_type + self.model_var_type = model_var_type + self.loss_type = loss_type + self.rescale_timesteps = rescale_timesteps + self.data_rep = data_rep + + if data_rep != 'rot_vel' and lambda_pose != 1.: + raise ValueError( + 'lambda_pose is relevant only when training on velocities!') + self.lambda_pose = lambda_pose + self.lambda_orient = lambda_orient + self.lambda_loc = lambda_loc + + self.lambda_rcxyz = lambda_rcxyz + self.lambda_vel = lambda_vel + self.lambda_root_vel = lambda_root_vel + self.lambda_vel_rcxyz = lambda_vel_rcxyz + self.lambda_fc = lambda_fc + + if self.lambda_rcxyz > 0. or self.lambda_vel > 0. or self.lambda_root_vel > 0. or \ + self.lambda_vel_rcxyz > 0. or self.lambda_fc > 0.: + assert self.loss_type == LossType.MSE, 'Geometric losses are supported by MSE loss type only!' + + # Use float64 for accuracy. + betas = np.array(betas, dtype=np.float64) + self.betas = betas + assert len(betas.shape) == 1, 'betas must be 1-D' + assert (betas > 0).all() and (betas <= 1).all() + + self.num_timesteps = int(betas.shape[0]) + + alphas = 1.0 - betas + self.alphas_cumprod = np.cumprod(alphas, axis=0) + self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1]) + self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0) + assert self.alphas_cumprod_prev.shape == (self.num_timesteps, ) + + # calculations for diffusion q(x_t | x_{t-1}) and others + self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod) + self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod) + self.log_one_minus_alphas_cumprod = np.log(1.0 - self.alphas_cumprod) + self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod) + self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod + - 1) + + # calculations for posterior q(x_{t-1} | x_t, x_0) + self.posterior_variance = betas * (1.0 - self.alphas_cumprod_prev) / ( + 1.0 - self.alphas_cumprod) + # log calculation clipped because the posterior variance is 0 at the + # beginning of the diffusion chain. + self.posterior_log_variance_clipped = np.log( + np.append(self.posterior_variance[1], self.posterior_variance[1:])) + self.posterior_mean_coef1 = betas * np.sqrt( + self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod) + self.posterior_mean_coef2 = (1.0 - self.alphas_cumprod_prev) * np.sqrt( + alphas) / (1.0 - self.alphas_cumprod) + + self.l2_loss = lambda a, b: ( + a - b + )**2 # th.nn.MSELoss(reduction='none') # must be None for handling mask later on. + + def q_mean_variance(self, x_start, t): + """ + Get the distribution q(x_t | x_0). + + :param x_start: the [N x C x ...] tensor of noiseless inputs. + :param t: the number of diffusion steps (minus 1). Here, 0 means one step. + :return: A tuple (mean, variance, log_variance), all of x_start's shape. + """ + mean = ( + _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) + * x_start) + variance = _extract_into_tensor(1.0 - self.alphas_cumprod, t, + x_start.shape) + log_variance = _extract_into_tensor(self.log_one_minus_alphas_cumprod, + t, x_start.shape) + return mean, variance, log_variance + + def q_sample(self, x_start, t, noise=None): + """ + Diffuse the dataset for a given number of diffusion steps. + + In other words, sample from q(x_t | x_0). + + :param x_start: the initial dataset batch. + :param t: the number of diffusion steps (minus 1). Here, 0 means one step. + :param noise: if specified, the split-out normal noise. + :return: A noisy version of x_start. + """ + if noise is None: + noise = th.randn_like(x_start) + assert noise.shape == x_start.shape + return ( + _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) + * x_start + _extract_into_tensor( + self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise) + + def q_posterior_mean_variance(self, x_start, x_t, t): + """ + Compute the mean and variance of the diffusion posterior: + + q(x_{t-1} | x_t, x_0) + + """ + assert x_start.shape == x_t.shape + posterior_mean = ( + _extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) + * x_start + + _extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) + * x_t) + posterior_variance = _extract_into_tensor(self.posterior_variance, t, + x_t.shape) + posterior_log_variance_clipped = _extract_into_tensor( + self.posterior_log_variance_clipped, t, x_t.shape) + return posterior_mean, posterior_variance, posterior_log_variance_clipped + + def p_mean_variance(self, + model, + x, + t, + clip_denoised=True, + denoised_fn=None, + model_kwargs=None): + """ + Apply the model to get p(x_{t-1} | x_t), as well as a prediction of + the initial x, x_0. + + :param model: the model, which takes a signal and a batch of timesteps + as input. + :param x: the [N x C x ...] tensor at time t. + :param t: a 1-D Tensor of timesteps. + :param clip_denoised: if True, clip the denoised signal into [-1, 1]. + :param denoised_fn: if not None, a function which applies to the + x_start prediction before it is used to sample. Applies before + clip_denoised. + :param model_kwargs: if not None, a dict of extra keyword arguments to + pass to the model. This can be used for conditioning. + :return: a dict with the following keys: + - 'mean': the model mean output. + - 'variance': the model variance output. + - 'log_variance': the log of 'variance'. + - 'pred_xstart': the prediction for x_0. + """ + if model_kwargs is None: + model_kwargs = {} + + B, C = x.shape[:2] + assert t.shape == (B, ) + model_output = model(x, self._scale_timesteps(t), **model_kwargs) + + if 'inpainting_mask' in model_kwargs['y'].keys( + ) and 'inpainted_motion' in model_kwargs['y'].keys(): + inpainting_mask, inpainted_motion = model_kwargs['y'][ + 'inpainting_mask'], model_kwargs['y']['inpainted_motion'] + assert self.model_mean_type == ModelMeanType.START_X, 'This feature supports only X_start pred for mow!' + assert model_output.shape == inpainting_mask.shape == inpainted_motion.shape + model_output = (model_output * ~inpainting_mask) + ( + inpainted_motion * inpainting_mask) + + if self.model_var_type in [ + ModelVarType.LEARNED, ModelVarType.LEARNED_RANGE + ]: + assert model_output.shape == (B, C * 2, *x.shape[2:]) + model_output, model_var_values = th.split(model_output, C, dim=1) + if self.model_var_type == ModelVarType.LEARNED: + model_log_variance = model_var_values + model_variance = th.exp(model_log_variance) + else: + min_log = _extract_into_tensor( + self.posterior_log_variance_clipped, t, x.shape) + max_log = _extract_into_tensor(np.log(self.betas), t, x.shape) + # The model_var_values is [-1, 1] for [min_var, max_var]. + frac = (model_var_values + 1) / 2 + model_log_variance = frac * max_log + (1 - frac) * min_log + model_variance = th.exp(model_log_variance) + else: + model_variance, model_log_variance = { + ModelVarType.FIXED_LARGE: ( + np.append(self.posterior_variance[1], self.betas[1:]), + np.log( + np.append(self.posterior_variance[1], self.betas[1:])), + ), + ModelVarType.FIXED_SMALL: ( + self.posterior_variance, + self.posterior_log_variance_clipped, + ), + }[self.model_var_type] + + model_variance = _extract_into_tensor(model_variance, t, x.shape) + model_log_variance = _extract_into_tensor(model_log_variance, t, + x.shape) + + def process_xstart(x): + if denoised_fn is not None: + x = denoised_fn(x) + if clip_denoised: + return x.clamp(-1, 1) + return x + + if self.model_mean_type == ModelMeanType.PREVIOUS_X: + pred_xstart = process_xstart( + self._predict_xstart_from_xprev( + x_t=x, t=t, xprev=model_output)) + model_mean = model_output + elif self.model_mean_type in [ + ModelMeanType.START_X, ModelMeanType.EPSILON + ]: # THIS IS US! + if self.model_mean_type == ModelMeanType.START_X: + pred_xstart = process_xstart(model_output) + else: + pred_xstart = process_xstart( + self._predict_xstart_from_eps( + x_t=x, t=t, eps=model_output)) + model_mean, _, _ = self.q_posterior_mean_variance( + x_start=pred_xstart, x_t=x, t=t) + else: + raise NotImplementedError(self.model_mean_type) + + return { + 'mean': model_mean, + 'variance': model_variance, + 'log_variance': model_log_variance, + 'pred_xstart': pred_xstart, + } + + def _predict_xstart_from_eps(self, x_t, t, eps): + assert x_t.shape == eps.shape + return ( + _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) + * x_t - _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, + x_t.shape) * eps) + + def _predict_xstart_from_xprev(self, x_t, t, xprev): + assert x_t.shape == xprev.shape + return ( # (xprev - coef2*x_t) / coef1 + _extract_into_tensor(1.0 / self.posterior_mean_coef1, t, x_t.shape) + * xprev - _extract_into_tensor( + self.posterior_mean_coef2 / self.posterior_mean_coef1, t, + x_t.shape) * x_t) + + def _predict_eps_from_xstart(self, x_t, t, pred_xstart): + return ( + _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) + * x_t - pred_xstart) / _extract_into_tensor( + self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) + + def _scale_timesteps(self, t): + if self.rescale_timesteps: + return t.float() * (1000.0 / self.num_timesteps) + return t + + def condition_mean(self, cond_fn, p_mean_var, x, t, model_kwargs=None): + """ + Compute the mean for the previous step, given a function cond_fn that + computes the gradient of a conditional log probability with respect to + x. In particular, cond_fn computes grad(log(p(y|x))), and we want to + condition on y. + + This uses the conditioning strategy from Sohl-Dickstein et al. (2015). + """ + gradient = cond_fn(x, self._scale_timesteps(t), **model_kwargs) + new_mean = ( + p_mean_var['mean'].float() + + p_mean_var['variance'] * gradient.float()) + return new_mean + + def condition_mean_with_grad(self, + cond_fn, + p_mean_var, + x, + t, + model_kwargs=None): + """ + Compute the mean for the previous step, given a function cond_fn that + computes the gradient of a conditional log probability with respect to + x. In particular, cond_fn computes grad(log(p(y|x))), and we want to + condition on y. + + This uses the conditioning strategy from Sohl-Dickstein et al. (2015). + """ + gradient = cond_fn(x, t, p_mean_var, **model_kwargs) + new_mean = ( + p_mean_var['mean'].float() + + p_mean_var['variance'] * gradient.float()) + return new_mean + + def p_sample( + self, + model, + x, + t, + clip_denoised=True, + denoised_fn=None, + cond_fn=None, + model_kwargs=None, + const_noise=False, + ): + """ + Sample x_{t-1} from the model at the given timestep. + + :param model: the model to sample from. + :param x: the current tensor at x_{t-1}. + :param t: the value of t, starting at 0 for the first diffusion step. + :param clip_denoised: if True, clip the x_start prediction to [-1, 1]. + :param denoised_fn: if not None, a function which applies to the + x_start prediction before it is used to sample. + :param cond_fn: if not None, this is a gradient function that acts + similarly to the model. + :param model_kwargs: if not None, a dict of extra keyword arguments to + pass to the model. This can be used for conditioning. + :return: a dict containing the following keys: + - 'sample': a random sample from the model. + - 'pred_xstart': a prediction of x_0. + """ + out = self.p_mean_variance( + model, + x, + t, + clip_denoised=clip_denoised, + denoised_fn=denoised_fn, + model_kwargs=model_kwargs, + ) + noise = th.randn_like(x) + if const_noise: + noise = noise[[0]].repeat(x.shape[0], 1, 1, 1) + + nonzero_mask = ((t != 0).float().view(-1, *([1] * (len(x.shape) - 1))) + ) # no noise when t == 0 + if cond_fn is not None: + out['mean'] = self.condition_mean( + cond_fn, out, x, t, model_kwargs=model_kwargs) + sample = out['mean'] + nonzero_mask * th.exp( + 0.5 * out['log_variance']) * noise + return {'sample': sample, 'pred_xstart': out['pred_xstart']} + + def p_sample_with_grad( + self, + model, + x, + t, + clip_denoised=True, + denoised_fn=None, + cond_fn=None, + model_kwargs=None, + ): + """ + Sample x_{t-1} from the model at the given timestep. + + :param model: the model to sample from. + :param x: the current tensor at x_{t-1}. + :param t: the value of t, starting at 0 for the first diffusion step. + :param clip_denoised: if True, clip the x_start prediction to [-1, 1]. + :param denoised_fn: if not None, a function which applies to the + x_start prediction before it is used to sample. + :param cond_fn: if not None, this is a gradient function that acts + similarly to the model. + :param model_kwargs: if not None, a dict of extra keyword arguments to + pass to the model. This can be used for conditioning. + :return: a dict containing the following keys: + - 'sample': a random sample from the model. + - 'pred_xstart': a prediction of x_0. + """ + with th.enable_grad(): + x = x.detach().requires_grad_() + out = self.p_mean_variance( + model, + x, + t, + clip_denoised=clip_denoised, + denoised_fn=denoised_fn, + model_kwargs=model_kwargs, + ) + noise = th.randn_like(x) + nonzero_mask = ((t != 0).float().view(-1, + *([1] * (len(x.shape) - 1))) + ) # no noise when t == 0 + if cond_fn is not None: + out['mean'] = self.condition_mean_with_grad( + cond_fn, out, x, t, model_kwargs=model_kwargs) + sample = out['mean'] + nonzero_mask * th.exp( + 0.5 * out['log_variance']) * noise + return {'sample': sample, 'pred_xstart': out['pred_xstart'].detach()} + + def p_sample_loop( + self, + model, + shape, + noise=None, + clip_denoised=True, + denoised_fn=None, + cond_fn=None, + model_kwargs=None, + device=None, + progress=False, + skip_timesteps=0, + init_image=None, + randomize_class=False, + cond_fn_with_grad=False, + dump_steps=None, + const_noise=False, + ): + """ + Generate samples from the model. + + :param model: the model module. + :param shape: the shape of the samples, (N, C, H, W). + :param noise: if specified, the noise from the encoder to sample. + Should be of the same shape as `shape`. + :param clip_denoised: if True, clip x_start predictions to [-1, 1]. + :param denoised_fn: if not None, a function which applies to the + x_start prediction before it is used to sample. + :param cond_fn: if not None, this is a gradient function that acts + similarly to the model. + :param model_kwargs: if not None, a dict of extra keyword arguments to + pass to the model. This can be used for conditioning. + :param device: if specified, the device to create the samples on. + If not specified, use a model parameter's device. + :param progress: if True, show a tqdm progress bar. + :param const_noise: If True, will noise all samples with the same noise throughout sampling + :return: a non-differentiable batch of samples. + """ + final = None + if dump_steps is not None: + dump = [] + + for i, sample in enumerate( + self.p_sample_loop_progressive( + model, + shape, + noise=noise, + clip_denoised=clip_denoised, + denoised_fn=denoised_fn, + cond_fn=cond_fn, + model_kwargs=model_kwargs, + device=device, + progress=progress, + skip_timesteps=skip_timesteps, + init_image=init_image, + randomize_class=randomize_class, + cond_fn_with_grad=cond_fn_with_grad, + const_noise=const_noise, + )): + if dump_steps is not None and i in dump_steps: + dump.append(deepcopy(sample['sample'])) + final = sample + if dump_steps is not None: + return dump + return final['sample'] + + def p_sample_loop_progressive( + self, + model, + shape, + noise=None, + clip_denoised=True, + denoised_fn=None, + cond_fn=None, + model_kwargs=None, + device=None, + progress=False, + skip_timesteps=0, + init_image=None, + randomize_class=False, + cond_fn_with_grad=False, + const_noise=False, + ): + """ + Generate samples from the model and yield intermediate samples from + each timestep of diffusion. + + Arguments are the same as p_sample_loop(). + Returns a generator over dicts, where each dict is the return value of + p_sample(). + """ + if device is None: + device = next(model.parameters()).device + assert isinstance(shape, (tuple, list)) + if noise is not None: + img = noise + else: + img = th.randn(*shape, device=device) + + if skip_timesteps and init_image is None: + init_image = th.zeros_like(img) + + indices = list(range(self.num_timesteps - skip_timesteps))[::-1] + + if init_image is not None: + my_t = th.ones([shape[0]], device=device, + dtype=th.long) * indices[0] + img = self.q_sample(init_image, my_t, img) + + if progress: + # Lazy import so that we don't depend on tqdm. + from tqdm.auto import tqdm + + indices = tqdm(indices) + + for i in indices: + t = th.tensor([i] * shape[0], device=device) + if randomize_class and 'y' in model_kwargs: + model_kwargs['y'] = th.randint( + low=0, + high=model.num_classes, + size=model_kwargs['y'].shape, + device=model_kwargs['y'].device) + with th.no_grad(): + sample_fn = self.p_sample_with_grad if cond_fn_with_grad else self.p_sample + out = sample_fn( + model, + img, + t, + clip_denoised=clip_denoised, + denoised_fn=denoised_fn, + cond_fn=cond_fn, + model_kwargs=model_kwargs, + const_noise=const_noise, + ) + yield out + img = out['sample'] + + +def _extract_into_tensor(arr, timesteps, broadcast_shape): + """ + Extract values from a 1-D numpy array for a batch of indices. + + :param arr: the 1-D numpy array. + :param timesteps: a tensor of indices into the array to extract. + :param broadcast_shape: a larger shape of K dimensions with the batch + dimension equal to the length of timesteps. + :return: a tensor of shape [batch_size, 1, ...] where the shape has K dims. + """ + res = th.from_numpy(arr).to(device=timesteps.device)[timesteps].float() + while len(res.shape) < len(broadcast_shape): + res = res[..., None] + return res.expand(broadcast_shape) diff --git a/modelscope/models/cv/motion_generation/modules/mdm.py b/modelscope/models/cv/motion_generation/modules/mdm.py new file mode 100644 index 00000000..716acd83 --- /dev/null +++ b/modelscope/models/cv/motion_generation/modules/mdm.py @@ -0,0 +1,364 @@ +# This code is borrowed and modified from Human Motion Diffusion Model, +# made publicly available under MIT license at https://github.com/GuyTevet/motion-diffusion-model + +import clip +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .rotation2xyz import Rotation2xyz + + +class MDM(nn.Module): + + def __init__(self, + modeltype, + njoints, + nfeats, + num_actions, + translation, + pose_rep, + glob, + glob_rot, + latent_dim=256, + ff_size=1024, + num_layers=8, + num_heads=4, + dropout=0.1, + smpl_data_path=None, + ablation=None, + activation='gelu', + legacy=False, + data_rep='rot6d', + dataset='amass', + clip_dim=512, + arch='trans_enc', + emb_trans_dec=False, + clip_version=None, + **kargs): + super().__init__() + + self.legacy = legacy + self.modeltype = modeltype + self.njoints = njoints + self.nfeats = nfeats + self.num_actions = num_actions + self.data_rep = data_rep + self.dataset = dataset + + self.pose_rep = pose_rep + self.glob = glob + self.glob_rot = glob_rot + self.translation = translation + + self.latent_dim = latent_dim + + self.ff_size = ff_size + self.num_layers = num_layers + self.num_heads = num_heads + self.dropout = dropout + + self.ablation = ablation + self.activation = activation + self.clip_dim = clip_dim + self.action_emb = kargs.get('action_emb', None) + + self.input_feats = self.njoints * self.nfeats + + self.normalize_output = kargs.get('normalize_encoder_output', False) + + self.cond_mode = kargs.get('cond_mode', 'no_cond') + self.cond_mask_prob = kargs.get('cond_mask_prob', 0.) + self.arch = arch + self.gru_emb_dim = self.latent_dim if self.arch == 'gru' else 0 + self.input_process = InputProcess(self.data_rep, + self.input_feats + self.gru_emb_dim, + self.latent_dim) + + self.sequence_pos_encoder = PositionalEncoding(self.latent_dim, + self.dropout) + self.emb_trans_dec = emb_trans_dec + + if self.arch == 'trans_enc': + print('TRANS_ENC init') + seqTransEncoderLayer = nn.TransformerEncoderLayer( + d_model=self.latent_dim, + nhead=self.num_heads, + dim_feedforward=self.ff_size, + dropout=self.dropout, + activation=self.activation) + + self.seqTransEncoder = nn.TransformerEncoder( + seqTransEncoderLayer, num_layers=self.num_layers) + elif self.arch == 'trans_dec': + print('TRANS_DEC init') + seqTransDecoderLayer = nn.TransformerDecoderLayer( + d_model=self.latent_dim, + nhead=self.num_heads, + dim_feedforward=self.ff_size, + dropout=self.dropout, + activation=activation) + self.seqTransDecoder = nn.TransformerDecoder( + seqTransDecoderLayer, num_layers=self.num_layers) + elif self.arch == 'gru': + print('GRU init') + self.gru = nn.GRU( + self.latent_dim, + self.latent_dim, + num_layers=self.num_layers, + batch_first=True) + else: + raise ValueError( + 'Please choose correct architecture [trans_enc, trans_dec, gru]' + ) + + self.embed_timestep = TimestepEmbedder(self.latent_dim, + self.sequence_pos_encoder) + + if self.cond_mode != 'no_cond': + if 'text' in self.cond_mode: + self.embed_text = nn.Linear(self.clip_dim, self.latent_dim) + print('EMBED TEXT') + print('Loading CLIP...') + self.clip_version = clip_version + self.clip_model = self.load_and_freeze_clip(clip_version) + if 'action' in self.cond_mode: + self.embed_action = EmbedAction(self.num_actions, + self.latent_dim) + print('EMBED ACTION') + + self.output_process = OutputProcess(self.data_rep, self.input_feats, + self.latent_dim, self.njoints, + self.nfeats) + + self.rot2xyz = Rotation2xyz( + device='cpu', smpl_data_path=smpl_data_path, dataset=self.dataset) + + def parameters_wo_clip(self): + return [ + p for name, p in self.named_parameters() + if not name.startswith('clip_model.') + ] + + def load_and_freeze_clip(self, clip_version): + clip_model, clip_preprocess = clip.load( + clip_version, device='cpu', + jit=False) # Must set jit=False for training + # clip.model.convert_weights( + # clip_model) # Actually this line is unnecessary since clip by default already on float16 + + # Freeze CLIP weights + clip_model.eval() + for p in clip_model.parameters(): + p.requires_grad = False + + return clip_model + + def mask_cond(self, cond, force_mask=False): + bs, d = cond.shape + if force_mask: + return torch.zeros_like(cond) + elif self.training and self.cond_mask_prob > 0.: + mask = torch.bernoulli( + torch.ones(bs, device=cond.device) * self.cond_mask_prob).view( + bs, 1) # 1-> use null_cond, 0-> use real cond + return cond * (1. - mask) + else: + return cond + + def encode_text(self, raw_text): + device = next(self.parameters()).device + max_text_len = 20 if self.dataset in [ + 'humanml', 'kit' + ] else None # Specific hardcoding for humanml dataset + if max_text_len is not None: + default_context_length = 77 + context_length = max_text_len + 2 # start_token + 20 + end_token + assert context_length < default_context_length + texts = clip.tokenize( + raw_text, context_length=context_length, + truncate=True).to(device) + zero_pad = torch.zeros( + [texts.shape[0], default_context_length - context_length], + dtype=texts.dtype, + device=texts.device) + texts = torch.cat([texts, zero_pad], dim=1) + else: + texts = clip.tokenize(raw_text, truncate=True).to(device) + return self.clip_model.encode_text(texts).float() + + def forward(self, x, timesteps, y=None): + """ + x: [batch_size, njoints, nfeats, max_frames], denoted x_t in the paper + timesteps: [batch_size] (int) + """ + bs, njoints, nfeats, nframes = x.shape + emb = self.embed_timestep(timesteps) # [1, bs, d] + + force_mask = y.get('uncond', False) + if 'text' in self.cond_mode: + enc_text = self.encode_text(y['text']) + emb += self.embed_text( + self.mask_cond(enc_text, force_mask=force_mask)) + if 'action' in self.cond_mode: + action_emb = self.embed_action(y['action']) + emb += self.mask_cond(action_emb, force_mask=force_mask) + + if self.arch == 'gru': + x_reshaped = x.reshape(bs, njoints * nfeats, 1, nframes) + emb_gru = emb.repeat(nframes, 1, 1) # [#frames, bs, d] + emb_gru = emb_gru.permute(1, 2, 0) # [bs, d, #frames] + emb_gru = emb_gru.reshape(bs, self.latent_dim, 1, + nframes) # [bs, d, 1, #frames] + x = torch.cat((x_reshaped, emb_gru), + axis=1) # [bs, d+joints*feat, 1, #frames] + + x = self.input_process(x) + + if self.arch == 'trans_enc': + # adding the timestep embed + xseq = torch.cat((emb, x), axis=0) # [seqlen+1, bs, d] + xseq = self.sequence_pos_encoder(xseq) # [seqlen+1, bs, d] + output = self.seqTransEncoder(xseq)[ + 1:] # , src_key_padding_mask=~maskseq) # [seqlen, bs, d] + + elif self.arch == 'trans_dec': + if self.emb_trans_dec: + xseq = torch.cat((emb, x), axis=0) + else: + xseq = x + xseq = self.sequence_pos_encoder(xseq) # [seqlen+1, bs, d] + if self.emb_trans_dec: + output = self.seqTransDecoder( + tgt=xseq, memory=emb + )[1:] # [seqlen, bs, d] # FIXME - maybe add a causal mask + else: + output = self.seqTransDecoder(tgt=xseq, memory=emb) + elif self.arch == 'gru': + xseq = x + xseq = self.sequence_pos_encoder(xseq) # [seqlen, bs, d] + output, _ = self.gru(xseq) + + output = self.output_process(output) # [bs, njoints, nfeats, nframes] + return output + + def _apply(self, fn): + super()._apply(fn) + self.rot2xyz.smpl_model._apply(fn) + + def train(self, *args, **kwargs): + super().train(*args, **kwargs) + self.rot2xyz.smpl_model.train(*args, **kwargs) + + +class PositionalEncoding(nn.Module): + + def __init__(self, d_model, dropout=0.1, max_len=5000): + super(PositionalEncoding, self).__init__() + self.dropout = nn.Dropout(p=dropout) + + pe = torch.zeros(max_len, d_model) + position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) + div_term = torch.exp( + torch.arange(0, d_model, 2).float() * (-np.log(10000.0) / d_model)) + pe[:, 0::2] = torch.sin(position * div_term) + pe[:, 1::2] = torch.cos(position * div_term) + pe = pe.unsqueeze(0).transpose(0, 1) + + self.register_buffer('pe', pe) + + def forward(self, x): + # not used in the final model + x = x + self.pe[:x.shape[0], :] + return self.dropout(x) + + +class TimestepEmbedder(nn.Module): + + def __init__(self, latent_dim, sequence_pos_encoder): + super().__init__() + self.latent_dim = latent_dim + self.sequence_pos_encoder = sequence_pos_encoder + + time_embed_dim = self.latent_dim + self.time_embed = nn.Sequential( + nn.Linear(self.latent_dim, time_embed_dim), + nn.SiLU(), + nn.Linear(time_embed_dim, time_embed_dim), + ) + + def forward(self, timesteps): + return self.time_embed( + self.sequence_pos_encoder.pe[timesteps]).permute(1, 0, 2) + + +class InputProcess(nn.Module): + + def __init__(self, data_rep, input_feats, latent_dim): + super().__init__() + self.data_rep = data_rep + self.input_feats = input_feats + self.latent_dim = latent_dim + self.poseEmbedding = nn.Linear(self.input_feats, self.latent_dim) + if self.data_rep == 'rot_vel': + self.velEmbedding = nn.Linear(self.input_feats, self.latent_dim) + + def forward(self, x): + bs, njoints, nfeats, nframes = x.shape + x = x.permute((3, 0, 1, 2)).reshape(nframes, bs, njoints * nfeats) + + if self.data_rep in ['rot6d', 'xyz', 'hml_vec']: + x = self.poseEmbedding(x) # [seqlen, bs, d] + return x + elif self.data_rep == 'rot_vel': + first_pose = x[[0]] # [1, bs, 150] + first_pose = self.poseEmbedding(first_pose) # [1, bs, d] + vel = x[1:] # [seqlen-1, bs, 150] + vel = self.velEmbedding(vel) # [seqlen-1, bs, d] + return torch.cat((first_pose, vel), axis=0) # [seqlen, bs, d] + else: + raise ValueError + + +class OutputProcess(nn.Module): + + def __init__(self, data_rep, input_feats, latent_dim, njoints, nfeats): + super().__init__() + self.data_rep = data_rep + self.input_feats = input_feats + self.latent_dim = latent_dim + self.njoints = njoints + self.nfeats = nfeats + self.poseFinal = nn.Linear(self.latent_dim, self.input_feats) + if self.data_rep == 'rot_vel': + self.velFinal = nn.Linear(self.latent_dim, self.input_feats) + + def forward(self, output): + nframes, bs, d = output.shape + if self.data_rep in ['rot6d', 'xyz', 'hml_vec']: + output = self.poseFinal(output) # [seqlen, bs, 150] + elif self.data_rep == 'rot_vel': + first_pose = output[[0]] # [1, bs, d] + first_pose = self.poseFinal(first_pose) # [1, bs, 150] + vel = output[1:] # [seqlen-1, bs, d] + vel = self.velFinal(vel) # [seqlen-1, bs, 150] + output = torch.cat((first_pose, vel), axis=0) # [seqlen, bs, 150] + else: + raise ValueError + output = output.reshape(nframes, bs, self.njoints, self.nfeats) + output = output.permute(1, 2, 3, 0) # [bs, njoints, nfeats, nframes] + return output + + +class EmbedAction(nn.Module): + + def __init__(self, num_actions, latent_dim): + super().__init__() + self.action_embedding = nn.Parameter( + torch.randn(num_actions, latent_dim)) + + def forward(self, input): + idx = input[:, 0].to(torch.long) # an index array must be long + output = self.action_embedding[idx] + return output diff --git a/modelscope/models/cv/motion_generation/modules/respace.py b/modelscope/models/cv/motion_generation/modules/respace.py new file mode 100644 index 00000000..45c37ee6 --- /dev/null +++ b/modelscope/models/cv/motion_generation/modules/respace.py @@ -0,0 +1,132 @@ +# This code is borrowed and modified from Human Motion Diffusion Model, +# made publicly available under MIT license at https://github.com/GuyTevet/motion-diffusion-model + +import numpy as np +import torch as th + +from .gaussian_diffusion import GaussianDiffusion + + +def space_timesteps(num_timesteps, section_counts): + """ + Create a list of timesteps to use from an original diffusion process, + given the number of timesteps we want to take from equally-sized portions + of the original process. + + For example, if there's 300 timesteps and the section counts are [10,15,20] + then the first 100 timesteps are strided to be 10 timesteps, the second 100 + are strided to be 15 timesteps, and the final 100 are strided to be 20. + + If the stride is a string starting with "ddim", then the fixed striding + from the DDIM paper is used, and only one section is allowed. + + :param num_timesteps: the number of diffusion steps in the original + process to divide up. + :param section_counts: either a list of numbers, or a string containing + comma-separated numbers, indicating the step count + per section. As a special case, use "ddimN" where N + is a number of steps to use the striding from the + DDIM paper. + :return: a set of diffusion steps from the original process to use. + """ + if isinstance(section_counts, str): + if section_counts.startswith('ddim'): + desired_count = int(section_counts[len('ddim'):]) + for i in range(1, num_timesteps): + if len(range(0, num_timesteps, i)) == desired_count: + return set(range(0, num_timesteps, i)) + raise ValueError( + f'cannot create exactly {num_timesteps} steps with an integer stride' + ) + section_counts = [int(x) for x in section_counts.split(',')] + size_per = num_timesteps // len(section_counts) + extra = num_timesteps % len(section_counts) + start_idx = 0 + all_steps = [] + for i, section_count in enumerate(section_counts): + size = size_per + (1 if i < extra else 0) + if size < section_count: + raise ValueError( + f'cannot divide section of {size} steps into {section_count}') + if section_count <= 1: + frac_stride = 1 + else: + frac_stride = (size - 1) / (section_count - 1) + cur_idx = 0.0 + taken_steps = [] + for _ in range(section_count): + taken_steps.append(start_idx + round(cur_idx)) + cur_idx += frac_stride + all_steps += taken_steps + start_idx += size + return set(all_steps) + + +class SpacedDiffusion(GaussianDiffusion): + """ + A diffusion process which can skip steps in a base diffusion process. + + :param use_timesteps: a collection (sequence or set) of timesteps from the + original diffusion process to retain. + :param kwargs: the kwargs to create the base diffusion process. + """ + + def __init__(self, use_timesteps, **kwargs): + self.use_timesteps = set(use_timesteps) + self.timestep_map = [] + self.original_num_steps = len(kwargs['betas']) + + base_diffusion = GaussianDiffusion(**kwargs) # pylint: disable=missing-kwoa + last_alpha_cumprod = 1.0 + new_betas = [] + for i, alpha_cumprod in enumerate(base_diffusion.alphas_cumprod): + if i in self.use_timesteps: + new_betas.append(1 - alpha_cumprod / last_alpha_cumprod) + last_alpha_cumprod = alpha_cumprod + self.timestep_map.append(i) + kwargs['betas'] = np.array(new_betas) + super().__init__(**kwargs) + + def p_mean_variance(self, model, *args, **kwargs): # pylint: disable=signature-differs + return super().p_mean_variance( + self._wrap_model(model), *args, **kwargs) + + def training_losses(self, model, *args, **kwargs): # pylint: disable=signature-differs + return super().training_losses( + self._wrap_model(model), *args, **kwargs) + + def condition_mean(self, cond_fn, *args, **kwargs): + return super().condition_mean( + self._wrap_model(cond_fn), *args, **kwargs) + + def condition_score(self, cond_fn, *args, **kwargs): + return super().condition_score( + self._wrap_model(cond_fn), *args, **kwargs) + + def _wrap_model(self, model): + if isinstance(model, _WrappedModel): + return model + return _WrappedModel(model, self.timestep_map, self.rescale_timesteps, + self.original_num_steps) + + def _scale_timesteps(self, t): + # Scaling is done by the wrapped model. + return t + + +class _WrappedModel: + + def __init__(self, model, timestep_map, rescale_timesteps, + original_num_steps): + self.model = model + self.timestep_map = timestep_map + self.rescale_timesteps = rescale_timesteps + self.original_num_steps = original_num_steps + + def __call__(self, x, ts, **kwargs): + map_tensor = th.tensor( + self.timestep_map, device=ts.device, dtype=ts.dtype) + new_ts = map_tensor[ts] + if self.rescale_timesteps: + new_ts = new_ts.float() * (1000.0 / self.original_num_steps) + return self.model(x, new_ts, **kwargs) diff --git a/modelscope/models/cv/motion_generation/modules/rotation2xyz.py b/modelscope/models/cv/motion_generation/modules/rotation2xyz.py new file mode 100644 index 00000000..0b00015e --- /dev/null +++ b/modelscope/models/cv/motion_generation/modules/rotation2xyz.py @@ -0,0 +1,112 @@ +# This code is borrowed and modified from Human Motion Diffusion Model, +# made publicly available under MIT license at https://github.com/GuyTevet/motion-diffusion-model + +import torch + +from modelscope.utils.cv.motion_utils import rotation_conversions as geometry +from .smpl import JOINTSTYPE_ROOT, SMPL + +JOINTSTYPES = ['a2m', 'a2mpl', 'smpl', 'vibe', 'vertices'] + + +class Rotation2xyz: + + def __init__(self, device, smpl_data_path, dataset='amass'): + self.device = device + self.dataset = dataset + self.smpl_model = SMPL(smpl_data_path).eval().to(device) + + def __call__(self, + x, + mask, + pose_rep, + translation, + glob, + jointstype, + vertstrans, + betas=None, + beta=0, + glob_rot=None, + get_rotations_back=False, + **kwargs): + if pose_rep == 'xyz': + return x + + if mask is None: + mask = torch.ones((x.shape[0], x.shape[-1]), + dtype=bool, + device=x.device) + + if not glob and glob_rot is None: + raise TypeError( + 'You must specify global rotation if glob is False') + + if jointstype not in JOINTSTYPES: + raise NotImplementedError('This jointstype is not implemented.') + + if translation: + x_translations = x[:, -1, :3] + x_rotations = x[:, :-1] + else: + x_rotations = x + + x_rotations = x_rotations.permute(0, 3, 1, 2) + nsamples, time, njoints, feats = x_rotations.shape + + # Compute rotations (convert only masked sequences output) + if pose_rep == 'rotvec': + rotations = geometry.axis_angle_to_matrix(x_rotations[mask]) + elif pose_rep == 'rotmat': + rotations = x_rotations[mask].view(-1, njoints, 3, 3) + elif pose_rep == 'rotquat': + rotations = geometry.quaternion_to_matrix(x_rotations[mask]) + elif pose_rep == 'rot6d': + rotations = geometry.rotation_6d_to_matrix(x_rotations[mask]) + else: + raise NotImplementedError('No geometry for this one.') + + if not glob: + global_orient = torch.tensor(glob_rot, device=x.device) + global_orient = geometry.axis_angle_to_matrix(global_orient).view( + 1, 1, 3, 3) + global_orient = global_orient.repeat(len(rotations), 1, 1, 1) + else: + global_orient = rotations[:, 0] + rotations = rotations[:, 1:] + + if betas is None: + betas = torch.zeros( + [rotations.shape[0], self.smpl_model.num_betas], + dtype=rotations.dtype, + device=rotations.device) + betas[:, 1] = beta + # import ipdb; ipdb.set_trace() + out = self.smpl_model( + body_pose=rotations, global_orient=global_orient, betas=betas) + + # get the desirable joints + joints = out[jointstype] + + x_xyz = torch.empty( + nsamples, time, joints.shape[1], 3, device=x.device, dtype=x.dtype) + x_xyz[~mask] = 0 + x_xyz[mask] = joints + + x_xyz = x_xyz.permute(0, 2, 3, 1).contiguous() + + # the first translation root at the origin on the prediction + if jointstype != 'vertices': + rootindex = JOINTSTYPE_ROOT[jointstype] + x_xyz = x_xyz - x_xyz[:, [rootindex], :, :] + + if translation and vertstrans: + # the first translation root at the origin + x_translations = x_translations - x_translations[:, :, [0]] + + # add the translation to all the joints + x_xyz = x_xyz + x_translations[:, None, :, :] + + if get_rotations_back: + return x_xyz, rotations, global_orient + else: + return x_xyz diff --git a/modelscope/models/cv/motion_generation/modules/smpl.py b/modelscope/models/cv/motion_generation/modules/smpl.py new file mode 100644 index 00000000..60b027de --- /dev/null +++ b/modelscope/models/cv/motion_generation/modules/smpl.py @@ -0,0 +1,117 @@ +# This code is borrowed and modified from Human Motion Diffusion Model, +# made publicly available under MIT license at https://github.com/GuyTevet/motion-diffusion-model + +import contextlib +import os.path as osp + +import numpy as np +import torch +from smplx import SMPLLayer as _SMPLLayer +from smplx.lbs import vertices2joints + +action2motion_joints = [ + 8, 1, 2, 3, 4, 5, 6, 7, 0, 9, 10, 11, 12, 13, 14, 21, 24, 38 +] + +JOINTSTYPE_ROOT = { + 'a2m': 0, # action2motion + 'smpl': 0, + 'a2mpl': 0, # set(smpl, a2m) + 'vibe': 8 +} # 0 is the 8 position: OP MidHip below + +JOINT_MAP = { + 'OP Nose': 24, + 'OP Neck': 12, + 'OP RShoulder': 17, + 'OP RElbow': 19, + 'OP RWrist': 21, + 'OP LShoulder': 16, + 'OP LElbow': 18, + 'OP LWrist': 20, + 'OP MidHip': 0, + 'OP RHip': 2, + 'OP RKnee': 5, + 'OP RAnkle': 8, + 'OP LHip': 1, + 'OP LKnee': 4, + 'OP LAnkle': 7, + 'OP REye': 25, + 'OP LEye': 26, + 'OP REar': 27, + 'OP LEar': 28, + 'OP LBigToe': 29, + 'OP LSmallToe': 30, + 'OP LHeel': 31, + 'OP RBigToe': 32, + 'OP RSmallToe': 33, + 'OP RHeel': 34, + 'Right Ankle': 8, + 'Right Knee': 5, + 'Right Hip': 45, + 'Left Hip': 46, + 'Left Knee': 4, + 'Left Ankle': 7, + 'Right Wrist': 21, + 'Right Elbow': 19, + 'Right Shoulder': 17, + 'Left Shoulder': 16, + 'Left Elbow': 18, + 'Left Wrist': 20, + 'Neck (LSP)': 47, + 'Top of Head (LSP)': 48, + 'Pelvis (MPII)': 49, + 'Thorax (MPII)': 50, + 'Spine (H36M)': 51, + 'Jaw (H36M)': 52, + 'Head (H36M)': 53, + 'Nose': 24, + 'Left Eye': 26, + 'Right Eye': 25, + 'Left Ear': 28, + 'Right Ear': 27 +} + +JOINT_NAMES = list(JOINT_MAP.keys()) + + +class SMPL(_SMPLLayer): + """ Extension of the official SMPL implementation to support more joints """ + + def __init__(self, smpl_data_path, **kwargs): + kwargs['model_path'] = osp.join(smpl_data_path, 'SMPL_NEUTRAL.pkl') + + # remove the verbosity for the 10-shapes beta parameters + with contextlib.redirect_stdout(None): + super(SMPL, self).__init__(**kwargs) + + J_regressor_extra = np.load( + osp.join(smpl_data_path, 'J_regressor_extra.npy')) + self.register_buffer( + 'J_regressor_extra', + torch.tensor(J_regressor_extra, dtype=torch.float32)) + vibe_indexes = np.array([JOINT_MAP[i] for i in JOINT_NAMES]) + a2m_indexes = vibe_indexes[action2motion_joints] + smpl_indexes = np.arange(24) + a2mpl_indexes = np.unique(np.r_[smpl_indexes, a2m_indexes]) + + self.maps = { + 'vibe': vibe_indexes, + 'a2m': a2m_indexes, + 'smpl': smpl_indexes, + 'a2mpl': a2mpl_indexes + } + + def forward(self, *args, **kwargs): + smpl_output = super(SMPL, self).forward(*args, **kwargs) + + extra_joints = vertices2joints(self.J_regressor_extra, + smpl_output.vertices) + all_joints = torch.cat([smpl_output.joints, extra_joints], dim=1) + + output = {'vertices': smpl_output.vertices} + + for joinstype, indexes in self.maps.items(): + output[joinstype] = all_joints[:, indexes] + + return output diff --git a/modelscope/outputs/outputs.py b/modelscope/outputs/outputs.py index 76a5c779..73e9401d 100644 --- a/modelscope/outputs/outputs.py +++ b/modelscope/outputs/outputs.py @@ -935,6 +935,13 @@ TASK_OUTPUTS = { # "masks": [np.array # 3D array with shape [frame_num, height, width]] # } Tasks.video_object_segmentation: [OutputKeys.MASKS], + + # motion generation result for a single input + # { + # "keypoints": [np.array # 3D array with shape [frame_num, joint_num, 3]] + # "output_video": "path_to_rendered_video" + # } + Tasks.motion_generation: [OutputKeys.KEYPOINTS, OutputKeys.OUTPUT_VIDEO], } diff --git a/modelscope/pipelines/cv/motion_generation_pipeline.py b/modelscope/pipelines/cv/motion_generation_pipeline.py new file mode 100644 index 00000000..0d8a21c9 --- /dev/null +++ b/modelscope/pipelines/cv/motion_generation_pipeline.py @@ -0,0 +1,128 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. + +import os.path as osp +import tempfile +from typing import Any, Dict + +import numpy as np +import torch + +from modelscope.metainfo import Pipelines +from modelscope.models.cv.motion_generation import (ClassifierFreeSampleModel, + create_model, + load_model_wo_clip) +from modelscope.outputs import OutputKeys +from modelscope.pipelines.base import Input, Pipeline +from modelscope.pipelines.builder import PIPELINES +from modelscope.utils.config import Config +from modelscope.utils.constant import ModelFile, Tasks +from modelscope.utils.cv.motion_utils.motion_process import recover_from_ric +from modelscope.utils.cv.motion_utils.plot_script import plot_3d_motion +from modelscope.utils.logger import get_logger + +logger = get_logger() + + +@PIPELINES.register_module( + Tasks.motion_generation, module_name=Pipelines.motion_generattion) +class MDMMotionGeneration(Pipeline): + + def __init__(self, model: str, **kwargs): + """ + use `model` to create motion generation pipeline for prediction + Args: + model: model id on modelscope hub. + """ + super().__init__(model=model, **kwargs) + model_path = osp.join(self.model, ModelFile.TORCH_MODEL_FILE) + logger.info(f'loading model from {model_path}') + config_path = osp.join(self.model, ModelFile.CONFIGURATION) + logger.info(f'loading config from {config_path}') + self.mean = np.load(osp.join(self.model, 'Mean.npy')) + self.std = np.load(osp.join(self.model, 'Std.npy')) + self.cfg = Config.from_file(config_path) + self.cfg.update({'smpl_data_path': osp.join(self.model, 'smpl')}) + self.cfg.update(kwargs) + self.n_joints = 22 + self.fps = 20 + self.n_frames = 120 + self.mdm, self.diffusion = create_model(self.cfg) + state_dict = torch.load(model_path, map_location='cpu') + load_model_wo_clip(self.mdm, state_dict) + self.mdm = ClassifierFreeSampleModel(self.mdm) + self.mdm.to(self.device) + self.mdm.eval() + logger.info('load model done') + + def preprocess(self, input: Input) -> Dict[str, Any]: + if isinstance(input, str): + input_text = input + else: + raise TypeError(f'input should be a str,' + f' but got {type(input)}') + result = {'input_text': input_text} + return result + + def forward(self, input: Dict[str, Any]) -> Dict[str, Any]: + texts = [input['input_text']] + model_kwargs = { + 'y': { + 'mask': torch.ones(1, 1, 1, self.n_frames) > 0, + 'lengths': torch.tensor([self.n_frames]), + 'tokens': None, + 'text': texts, + 'scale': torch.ones(1, device=self.device) * 2.5 + } + } + sample_fn = self.diffusion.p_sample_loop + sample = sample_fn( + self.mdm, + (1, self.mdm.njoints, self.mdm.nfeats, self.n_frames), + clip_denoised=False, + model_kwargs=model_kwargs, + skip_timesteps=0, + init_image=None, + progress=True, + dump_steps=None, + noise=None, + const_noise=False, + ) + sample = (sample.cpu().permute(0, 2, 3, 1) * self.std + + self.mean).float() + sample = recover_from_ric(sample, self.n_joints) + sample = sample.view(-1, *sample.shape[2:]).permute(0, 2, 3, 1) + + sample = self.mdm.rot2xyz( + x=sample, + mask=None, + pose_rep='xyz', + glob=True, + translation=True, + jointstype='smpl', + vertstrans=True, + betas=None, + beta=0, + glob_rot=None, + get_rotations_back=False) + motion = sample.cpu().numpy() + motion = motion[0].transpose(2, 0, 1) + out = {OutputKeys.KEYPOINTS: motion, 'text': input['input_text']} + return out + + def postprocess(self, inputs: Dict[str, Any], **kwargs) -> Dict[str, Any]: + output_video_path = kwargs.get( + 'output_video', + tempfile.NamedTemporaryFile(suffix='.mp4').name) + kinematic_chain = [[0, 2, 5, 8, 11], [0, 1, 4, 7, 10], + [0, 3, 6, 9, 12, 15], [9, 14, 17, 19, 21], + [9, 13, 16, 18, 20]] + if output_video_path is not None: + plot_3d_motion( + output_video_path, + kinematic_chain, + inputs[OutputKeys.KEYPOINTS], + inputs.pop('text'), + dataset='humanml', + fps=20) + inputs.update({OutputKeys.OUTPUT_VIDEO: output_video_path}) + return inputs diff --git a/modelscope/utils/constant.py b/modelscope/utils/constant.py index 4a83b278..1710d4c9 100644 --- a/modelscope/utils/constant.py +++ b/modelscope/utils/constant.py @@ -123,6 +123,9 @@ class CVTasks(object): # domain specific object detection domain_specific_object_detection = 'domain-specific-object-detection' + # motion generation + motion_generation = 'motion-generation' + class NLPTasks(object): # nlp tasks diff --git a/modelscope/utils/cv/motion_utils/__init__.py b/modelscope/utils/cv/motion_utils/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/modelscope/utils/cv/motion_utils/motion_process.py b/modelscope/utils/cv/motion_utils/motion_process.py new file mode 100644 index 00000000..30c37ac1 --- /dev/null +++ b/modelscope/utils/cv/motion_utils/motion_process.py @@ -0,0 +1,72 @@ +# This code is borrowed and modified from Actor, +# made publicly available under MIT license at https://github.com/Mathux/ACTOR + +import torch + + +def qinv(q): + assert q.shape[-1] == 4, 'q must be a tensor of shape (*, 4)' + mask = torch.ones_like(q) + mask[..., 1:] = -mask[..., 1:] + return q * mask + + +def qrot(q, v): + """ + Rotate vector(s) v about the rotation described by quaternion(s) q. + Expects a tensor of shape (*, 4) for q and a tensor of shape (*, 3) for v, + where * denotes any number of dimensions. + Returns a tensor of shape (*, 3). + """ + assert q.shape[-1] == 4 + assert v.shape[-1] == 3 + assert q.shape[:-1] == v.shape[:-1] + + original_shape = list(v.shape) + # print(q.shape) + q = q.contiguous().view(-1, 4) + v = v.contiguous().view(-1, 3) + + qvec = q[:, 1:] + uv = torch.cross(qvec, v, dim=1) + uuv = torch.cross(qvec, uv, dim=1) + return (v + 2 * (q[:, :1] * uv + uuv)).view(original_shape) + + +def recover_root_rot_pos(data): + rot_vel = data[..., 0] + r_rot_ang = torch.zeros_like(rot_vel).to(data.device) + '''Get Y-axis rotation from rotation velocity''' + r_rot_ang[..., 1:] = rot_vel[..., :-1] + r_rot_ang = torch.cumsum(r_rot_ang, dim=-1) + + r_rot_quat = torch.zeros(data.shape[:-1] + (4, )).to(data.device) + r_rot_quat[..., 0] = torch.cos(r_rot_ang) + r_rot_quat[..., 2] = torch.sin(r_rot_ang) + + r_pos = torch.zeros(data.shape[:-1] + (3, )).to(data.device) + r_pos[..., 1:, [0, 2]] = data[..., :-1, 1:3] + '''Add Y-axis rotation to root position''' + r_pos = qrot(qinv(r_rot_quat), r_pos) + + r_pos = torch.cumsum(r_pos, dim=-2) + + r_pos[..., 1] = data[..., 3] + return r_rot_quat, r_pos + + +def recover_from_ric(data, joints_num): + r_rot_quat, r_pos = recover_root_rot_pos(data) + positions = data[..., 4:(joints_num - 1) * 3 + 4] + positions = positions.view(positions.shape[:-1] + (-1, 3)) + '''Add Y-axis rotation to local joints''' + positions = qrot( + qinv(r_rot_quat[..., None, :]).expand(positions.shape[:-1] + (4, )), + positions) + '''Add root XZ to joints''' + positions[..., 0] += r_pos[..., 0:1] + positions[..., 2] += r_pos[..., 2:3] + '''Concate root and joints''' + positions = torch.cat([r_pos.unsqueeze(-2), positions], dim=-2) + + return positions diff --git a/modelscope/utils/cv/motion_utils/plot_script.py b/modelscope/utils/cv/motion_utils/plot_script.py new file mode 100644 index 00000000..94aab9f6 --- /dev/null +++ b/modelscope/utils/cv/motion_utils/plot_script.py @@ -0,0 +1,122 @@ +# This code is borrowed and modified from Actor, +# made publicly available under MIT license at https://github.com/Mathux/ACTOR + +import math +from textwrap import wrap + +import matplotlib +import matplotlib.pyplot as plt +import mpl_toolkits.mplot3d.axes3d as p3 +import numpy as np +from matplotlib.animation import FuncAnimation +from mpl_toolkits.mplot3d.art3d import Poly3DCollection + + +def list_cut_average(ll, intervals): + if intervals == 1: + return ll + + bins = math.ceil(len(ll) * 1.0 / intervals) + ll_new = [] + for i in range(bins): + l_low = intervals * i + l_high = l_low + intervals + l_high = l_high if l_high < len(ll) else len(ll) + ll_new.append(np.mean(ll[l_low:l_high])) + return ll_new + + +def plot_3d_motion(save_path, + kinematic_tree, + joints, + title, + dataset, + figsize=(3, 3), + fps=120, + radius=3, + vis_mode='default', + gt_frames=[]): + matplotlib.use('Agg') + + title = '\n'.join(wrap(title, 30)) + + def init(): + ax.set_xlim3d([-radius / 2, radius / 2]) + ax.set_ylim3d([0, radius]) + ax.set_zlim3d([-radius / 3., radius * 2 / 3.]) + fig.suptitle(title, fontsize=10) + ax.grid(b=False) + + def plot_xzPlane(minx, maxx, miny, minz, maxz): + verts = [[minx, miny, minz], [minx, miny, maxz], [maxx, miny, maxz], + [maxx, miny, minz]] + xz_plane = Poly3DCollection([verts]) + xz_plane.set_facecolor((0.5, 0.5, 0.5, 0.5)) + ax.add_collection3d(xz_plane) + + data = joints.copy().reshape(len(joints), -1, 3) + + if dataset == 'kit': + data *= 0.003 # scale for visualization + elif dataset == 'humanml': + data *= 1.3 # scale for visualization + elif dataset in ['humanact12', 'uestc']: + data *= -1.5 # reverse axes, scale for visualization + + fig = plt.figure(figsize=figsize) + plt.tight_layout() + ax = p3.Axes3D(fig) + init() + MINS = data.min(axis=0).min(axis=0) + MAXS = data.max(axis=0).max(axis=0) + colors_blue = ['#4D84AA', '#5B9965', '#61CEB9', '#34C1E2', + '#80B79A'] # GT color + colors_orange = ['#DD5A37', '#D69E00', '#B75A39', '#FF6D00', + '#DDB50E'] # Generation color + colors = colors_orange + if vis_mode == 'upper_body': # lower body taken fixed to input motion + colors[0] = colors_blue[0] + colors[1] = colors_blue[1] + elif vis_mode == 'gt': + colors = colors_blue + + frame_number = data.shape[0] + # print(dataset.shape) + + height_offset = MINS[1] + data[:, :, 1] -= height_offset + trajec = data[:, 0, [0, 2]] + + data[..., 0] -= data[:, 0:1, 0] + data[..., 2] -= data[:, 0:1, 2] + + def update(index): + ax.lines.clear() + ax.collections.clear() + ax.view_init(elev=120, azim=-90) + ax.dist = 7.5 + plot_xzPlane(MINS[0] - trajec[index, 0], MAXS[0] - trajec[index, 0], 0, + MINS[2] - trajec[index, 1], MAXS[2] - trajec[index, 1]) + + used_colors = colors_blue if index in gt_frames else colors + for i, (chain, color) in enumerate(zip(kinematic_tree, used_colors)): + if i < 5: + linewidth = 4.0 + else: + linewidth = 2.0 + ax.plot3D( + data[index, chain, 0], + data[index, chain, 1], + data[index, chain, 2], + linewidth=linewidth, + color=color) + plt.axis('off') + ax.set_xticklabels([]) + ax.set_yticklabels([]) + ax.set_zticklabels([]) + + ani = FuncAnimation( + fig, update, frames=frame_number, interval=1000 / fps, repeat=False) + ani.save(save_path, fps=fps) + + plt.close() diff --git a/modelscope/utils/cv/motion_utils/rotation_conversions.py b/modelscope/utils/cv/motion_utils/rotation_conversions.py new file mode 100644 index 00000000..5f0ee947 --- /dev/null +++ b/modelscope/utils/cv/motion_utils/rotation_conversions.py @@ -0,0 +1,132 @@ +# This code is borrowed and modified from Actor, +# made publicly available under MIT license at https://github.com/Mathux/ACTOR + +import functools + +import torch +import torch.nn.functional as F + + +def quaternion_to_matrix(quaternions): + """ + Convert rotations given as quaternions to rotation matrices. + + Args: + quaternions: quaternions with real part first, + as tensor of shape (..., 4). + + Returns: + Rotation matrices as tensor of shape (..., 3, 3). + """ + r, i, j, k = torch.unbind(quaternions, -1) + two_s = 2.0 / (quaternions * quaternions).sum(-1) + + o = torch.stack( + ( + 1 - two_s * (j * j + k * k), + two_s * (i * j - k * r), + two_s * (i * k + j * r), + two_s * (i * j + k * r), + 1 - two_s * (i * i + k * k), + two_s * (j * k - i * r), + two_s * (i * k - j * r), + two_s * (j * k + i * r), + 1 - two_s * (i * i + j * j), + ), + -1, + ) + return o.reshape(quaternions.shape[:-1] + (3, 3)) + + +def _axis_angle_rotation(axis: str, angle): + """ + Return the rotation matrices for one of the rotations about an axis + of which Euler angles describe, for each value of the angle given. + + Args: + axis: Axis label "X" or "Y or "Z". + angle: any shape tensor of Euler angles in radians + + Returns: + Rotation matrices as tensor of shape (..., 3, 3). + """ + + cos = torch.cos(angle) + sin = torch.sin(angle) + one = torch.ones_like(angle) + zero = torch.zeros_like(angle) + + if axis == 'X': + R_flat = (one, zero, zero, zero, cos, -sin, zero, sin, cos) + if axis == 'Y': + R_flat = (cos, zero, sin, zero, one, zero, -sin, zero, cos) + if axis == 'Z': + R_flat = (cos, -sin, zero, sin, cos, zero, zero, zero, one) + + return torch.stack(R_flat, -1).reshape(angle.shape + (3, 3)) + + +def euler_angles_to_matrix(euler_angles, convention: str): + """ + Convert rotations given as Euler angles in radians to rotation matrices. + + Args: + euler_angles: Euler angles in radians as tensor of shape (..., 3). + convention: Convention string of three uppercase letters from + {"X", "Y", and "Z"}. + + Returns: + Rotation matrices as tensor of shape (..., 3, 3). + """ + if euler_angles.dim() == 0 or euler_angles.shape[-1] != 3: + raise ValueError('Invalid input euler angles.') + if len(convention) != 3: + raise ValueError('Convention must have 3 letters.') + if convention[1] in (convention[0], convention[2]): + raise ValueError(f'Invalid convention {convention}.') + for letter in convention: + if letter not in ('X', 'Y', 'Z'): + raise ValueError(f'Invalid letter {letter} in convention string.') + matrices = map(_axis_angle_rotation, convention, + torch.unbind(euler_angles, -1)) + return functools.reduce(torch.matmul, matrices) + + +def axis_angle_to_matrix(axis_angle): + """ + Convert rotations given as axis/angle to rotation matrices. + + Args: + axis_angle: Rotations given as a vector in axis angle form, + as a tensor of shape (..., 3), where the magnitude is + the angle turned anticlockwise in radians around the + vector's direction. + + Returns: + Rotation matrices as tensor of shape (..., 3, 3). + """ + return quaternion_to_matrix(axis_angle_to_quaternion(axis_angle)) + + +def rotation_6d_to_matrix(d6: torch.Tensor) -> torch.Tensor: + """ + Converts 6D rotation representation by Zhou et al. [1] to rotation matrix + using Gram--Schmidt orthogonalisation per Section B of [1]. + Args: + d6: 6D rotation representation, of size (*, 6) + + Returns: + batch of rotation matrices of size (*, 3, 3) + + [1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H. + On the Continuity of Rotation Representations in Neural Networks. + IEEE Conference on Computer Vision and Pattern Recognition, 2019. + Retrieved from http://arxiv.org/abs/1812.07035 + """ + + a1, a2 = d6[..., :3], d6[..., 3:] + b1 = F.normalize(a1, dim=-1) + b2 = a2 - (b1 * a2).sum(-1, keepdim=True) * b1 + b2 = F.normalize(b2, dim=-1) + b3 = torch.cross(b1, b2, dim=-1) + return torch.stack((b1, b2, b3), dim=-2) diff --git a/requirements/cv.txt b/requirements/cv.txt index 17bd6ecf..a68a9a35 100644 --- a/requirements/cv.txt +++ b/requirements/cv.txt @@ -1,6 +1,7 @@ albumentations>=1.0.3 av>=9.2.0 bmt_clipit>=1.0 +chumpy clip>=1.0 easydict easyrobust @@ -34,6 +35,7 @@ scikit-image>=0.19.3 scikit-learn>=0.20.1 shapely shotdetect_scenedetect_lgss +smplx tensorflow-estimator>=1.15.1 tf_slim timm>=0.4.9 diff --git a/tests/pipelines/test_motion_generation.py b/tests/pipelines/test_motion_generation.py new file mode 100644 index 00000000..7938611c --- /dev/null +++ b/tests/pipelines/test_motion_generation.py @@ -0,0 +1,32 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import unittest + +from modelscope.outputs import OutputKeys +from modelscope.pipelines import pipeline +from modelscope.utils.constant import Tasks +from modelscope.utils.demo_utils import DemoCompatibilityCheck +from modelscope.utils.test_utils import test_level + + +class MDMMotionGenerationTest(unittest.TestCase, DemoCompatibilityCheck): + + def setUp(self) -> None: + self.task = Tasks.motion_generation + self.model_id = 'damo/cv_mdm_motion-generation' + + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + def test_run(self): + motion_generation_pipline = pipeline(self.task, model=self.model_id) + result = motion_generation_pipline( + 'the person walked forward and is picking up his toolbox') + print('motion generation data shape:', + result[OutputKeys.KEYPOINTS].shape) + print('motion generation video file:', result[OutputKeys.OUTPUT_VIDEO]) + + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + def test_demo_compatibility(self): + self.compatibility_check() + + +if __name__ == '__main__': + unittest.main()