[to #42322933] add new motion-generation model

新增运动生成模型，根据文本描述，生成人体的运动数据

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'

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={},
+    )

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

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))

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)

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

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)

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

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

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],
 }
 
 

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

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

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

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()

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)

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

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()