upload disco guided diffusion

把https://huggingface.co/IDEA-CCNL/Taiyi-Diffusion-532M-Cyberpunk-Chinese和https://huggingface.co/IDEA-CCNL/Taiyi-Diffusion-532M-Nature-Chinese迁移到MaaS-lib上。该project基于disco diffusion+guided diffusion。

Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/11818412

* upload disco guided diffusion

modelscope/metainfo.py

@@ -486,6 +486,7 @@ class Pipelines(object):
     video_captioning = 'video-captioning'
     video_question_answering = 'video-question-answering'
     diffusers_stable_diffusion = 'diffusers-stable-diffusion'
+    disco_guided_diffusion = 'disco_guided_diffusion'
     document_vl_embedding = 'document-vl-embedding'
     chinese_stable_diffusion = 'chinese-stable-diffusion'
     text_to_video_synthesis = 'latent-text-to-video-synthesis'  # latent-text-to-video-synthesis

modelscope/models/multi_modal/guided_diffusion/__init__.py

@@ -0,0 +1,23 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+from typing import TYPE_CHECKING
+
+from modelscope.utils.import_utils import LazyImportModule
+
+if TYPE_CHECKING:
+    from .unet import HFUNetModel
+    from .script import create_diffusion
+else:
+    _import_structure = {
+        'unet': ['HFUNetModel'],
+        'script': ['create_diffusion']
+    }
+
+    import sys
+
+    sys.modules[__name__] = LazyImportModule(
+        __name__,
+        globals()['__file__'],
+        _import_structure,
+        module_spec=__spec__,
+        extra_objects={},
+    )

modelscope/models/multi_modal/guided_diffusion/gaussian_diffusion.py

@@ -0,0 +1,930 @@
+# This code is borrowed and modified from Guided Diffusion Model,
+# made publicly available under MIT license
+# at https://github.com/IDEA-CCNL/Fengshenbang-LM/tree/main/fengshen/examples/disco_project
+
+import enum
+import math
+
+import numpy as np
+import torch as th
+
+
+def get_named_beta_schedule(schedule_name, num_diffusion_timesteps):
+    """
+    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 = 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,
+    ):
+        self.model_mean_type = model_mean_type
+        self.model_var_type = model_var_type
+        self.loss_type = loss_type
+        self.rescale_timesteps = rescale_timesteps
+
+        # 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)
+        v1 = betas * (1.0 - self.alphas_cumprod_prev)
+        v2 = 1.0 - self.alphas_cumprod
+        self.posterior_variance = v1 / v2
+        # 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:]))
+
+        v1 = betas * np.sqrt(self.alphas_cumprod_prev)
+        v2 = 1.0 - self.alphas_cumprod
+        self.posterior_mean_coef1 = v1 / v2
+
+        v1 = (1.0 - self.alphas_cumprod_prev) * np.sqrt(alphas)
+        v2 = 1.0 - self.alphas_cumprod
+        self.posterior_mean_coef2 = v1 / v2
+
+    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 data for a given number of diffusion steps.
+
+        In other words, sample from q(x_t | x_0).
+
+        :param x_start: the initial data 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)
+        assert posterior_mean.shape[0] == posterior_variance.shape[0]
+        assert posterior_mean.shape[0] == posterior_log_variance_clipped.shape[
+            0]
+        assert posterior_mean.shape[0] == x_start.shape[0]
+
+        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 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 = {
+                # for fixedlarge, we set the initial (log-)variance like so
+                # to get a better decoder log likelihood.
+                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
+        ]:
+            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 condition_score(self, cond_fn, p_mean_var, x, t, model_kwargs=None):
+        """
+        Compute what the p_mean_variance output would have been, should the
+        model's score function be conditioned by cond_fn.
+
+        See condition_mean() for details on cond_fn.
+
+        Unlike condition_mean(), this instead uses the conditioning strategy
+        from Song et al (2020).
+        """
+        alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
+
+        eps = self._predict_eps_from_xstart(x, t, p_mean_var['pred_xstart'])
+        eps = eps - (1 - alpha_bar).sqrt() * cond_fn(
+            x, self._scale_timesteps(t), **model_kwargs)
+
+        out = p_mean_var.copy()
+        out['pred_xstart'] = self._predict_xstart_from_eps(x, t, eps)
+        out['mean'], _, _ = self.q_posterior_mean_variance(
+            x_start=out['pred_xstart'], x_t=x, t=t)
+        return out
+
+    def condition_score_with_grad(self,
+                                  cond_fn,
+                                  p_mean_var,
+                                  x,
+                                  t,
+                                  model_kwargs=None):
+        """
+        Compute what the p_mean_variance output would have been, should the
+        model's score function be conditioned by cond_fn.
+
+        See condition_mean() for details on cond_fn.
+
+        Unlike condition_mean(), this instead uses the conditioning strategy
+        from Song et al (2020).
+        """
+        alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
+
+        eps = self._predict_eps_from_xstart(x, t, p_mean_var['pred_xstart'])
+
+        grad = cond_fn(x, t, p_mean_var, **model_kwargs)
+        eps = eps - (1 - alpha_bar).sqrt() * grad
+
+        out = p_mean_var.copy()
+        out['pred_xstart'] = self._predict_xstart_from_eps(x, t, eps)
+        out['mean'], _, _ = self.q_posterior_mean_variance(
+            x_start=out['pred_xstart'], x_t=x, t=t)
+        return out
+
+    def p_sample(
+        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.
+        """
+        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(
+                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,
+    ):
+        """
+        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.
+        :return: a non-differentiable batch of samples.
+        """
+        final = None
+        for sample in 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,
+        ):
+            final = sample
+        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,
+    ):
+        """
+        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, desc='Steps')
+
+        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,
+                )
+                yield out
+                img = out['sample']
+
+    def ddim_sample(
+        self,
+        model,
+        x,
+        t,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        eta=0.0,
+        inpainting_mode=False,
+        orig_img=None,
+        mask_inpaint=None,
+    ):
+        """
+        Sample x_{t-1} from the model using DDIM.
+
+        Same usage as p_sample().
+        """
+        alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
+        if inpainting_mode:
+            noised_orig_img = th.sqrt(alpha_bar) * orig_img + \
+                th.sqrt(1 - alpha_bar) * th.randn_like(x)
+            # noised_orig_img_pil = TF.to_pil_image(noised_orig_img[0].add(1).div(2).clamp(0, 1))
+            # noised_orig_img_pil.save(f'/content/drive/MyDrive/AI/Disco_Diffusion/images_out/InpaintingTest/inpainting_dump/noised_orig_{t[0].item()}.png')
+            x = (1 - mask_inpaint) * noised_orig_img + mask_inpaint * x
+            # mixed_x = TF.to_pil_image(x[0].add(1).div(2).clamp(0, 1))
+            # mixed_x.save(f'/content/drive/MyDrive/AI/Disco_Diffusion/images_out/InpaintingTest/inpainting_dump/mixed_x_{t[0].item()}.png')
+
+        out_orig = self.p_mean_variance(
+            model,
+            x,
+            t,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            model_kwargs=model_kwargs,
+        )
+        if cond_fn is not None:
+            out = self.condition_score(
+                cond_fn, out_orig, x, t, model_kwargs=model_kwargs)
+        else:
+            out = out_orig
+
+        # Usually our model outputs epsilon, but we re-derive it
+        # in case we used x_start or x_prev prediction.
+        eps = self._predict_eps_from_xstart(x, t, out['pred_xstart'])
+
+        alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod_prev, t,
+                                              x.shape)
+
+        v1 = eta * th.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar))
+        v2 = th.sqrt(1 - alpha_bar / alpha_bar_prev)
+        sigma = v1 * v2
+
+        # Equation 12.
+        noise = th.randn_like(x)
+        mean_pred = (
+            out['pred_xstart'] * th.sqrt(alpha_bar_prev)
+            + th.sqrt(1 - alpha_bar_prev - sigma**2) * eps)
+        nonzero_mask = ((t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
+                        )  # no noise when t == 0
+        sample = mean_pred + nonzero_mask * sigma * noise
+        return {'sample': sample, 'pred_xstart': out_orig['pred_xstart']}
+
+    def ddim_sample_with_grad(
+        self,
+        model,
+        x,
+        t,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        eta=0.0,
+    ):
+        """
+        Sample x_{t-1} from the model using DDIM.
+
+        Same usage as p_sample().
+        """
+        with th.enable_grad():
+            x = x.detach().requires_grad_()
+            out_orig = self.p_mean_variance(
+                model,
+                x,
+                t,
+                clip_denoised=clip_denoised,
+                denoised_fn=denoised_fn,
+                model_kwargs=model_kwargs,
+            )
+            if cond_fn is not None:
+                out = self.condition_score_with_grad(
+                    cond_fn, out_orig, x, t, model_kwargs=model_kwargs)
+            else:
+                out = out_orig
+
+        out['pred_xstart'] = out['pred_xstart'].detach()
+
+        # Usually our model outputs epsilon, but we re-derive it
+        # in case we used x_start or x_prev prediction.
+        eps = self._predict_eps_from_xstart(x, t, out['pred_xstart'])
+
+        alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
+        alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod_prev, t,
+                                              x.shape)
+
+        v1 = eta * th.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar))
+        v2 = th.sqrt(1 - alpha_bar / alpha_bar_prev)
+        sigma = v1 * v2
+
+        # Equation 12.
+        noise = th.randn_like(x)
+        mean_pred = (
+            out['pred_xstart'] * th.sqrt(alpha_bar_prev)
+            + th.sqrt(1 - alpha_bar_prev - sigma**2) * eps)
+        nonzero_mask = ((t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
+                        )  # no noise when t == 0
+        sample = mean_pred + nonzero_mask * sigma * noise
+        return {
+            'sample': sample,
+            'pred_xstart': out_orig['pred_xstart'].detach()
+        }
+
+    def ddim_sample_loop(
+        self,
+        model,
+        shape,
+        noise=None,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        device=None,
+        progress=False,
+        eta=0.0,
+        skip_timesteps=0,
+        init_image=None,
+        randomize_class=False,
+        cond_fn_with_grad=False,
+    ):
+        """
+        Generate samples from the model using DDIM.
+
+        Same usage as p_sample_loop().
+        """
+        final = None
+        for sample in self.ddim_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,
+                eta=eta,
+                skip_timesteps=skip_timesteps,
+                init_image=init_image,
+                randomize_class=randomize_class,
+                cond_fn_with_grad=cond_fn_with_grad,
+        ):
+            final = sample
+        return final['sample']
+
+    def ddim_sample_loop_progressive(self,
+                                     model,
+                                     shape,
+                                     noise=None,
+                                     clip_denoised=True,
+                                     denoised_fn=None,
+                                     cond_fn=None,
+                                     model_kwargs=None,
+                                     device=None,
+                                     progress=False,
+                                     eta=0.0,
+                                     skip_timesteps=0,
+                                     init_image=None,
+                                     randomize_class=False,
+                                     cond_fn_with_grad=False,
+                                     transformation_fn=None,
+                                     transformation_percent=[],
+                                     inpainting_mode=False,
+                                     mask_inpaint=None,
+                                     skip_timesteps_orig=None):
+        """
+        Use DDIM to sample from the model and yield intermediate samples from
+        each timestep of DDIM.
+
+        Same usage as p_sample_loop_progressive().
+        """
+
+        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]
+        transformation_steps = [
+            int(len(indices) * (1 - i)) for i in transformation_percent
+        ]
+
+        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, desc='Steps')
+
+        if inpainting_mode and skip_timesteps_orig is None:
+            skip_timesteps_orig = self.num_timesteps
+
+        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():
+                if i in transformation_steps and transformation_fn is not None:
+                    img = transformation_fn(img)
+                sample_fn = self.ddim_sample_with_grad if cond_fn_with_grad else self.ddim_sample
+                if inpainting_mode \
+                        and i >= self.num_timesteps - skip_timesteps_orig \
+                        and not cond_fn_with_grad:
+                    out = sample_fn(
+                        model,
+                        img,
+                        t,
+                        clip_denoised=clip_denoised,
+                        denoised_fn=denoised_fn,
+                        cond_fn=cond_fn,
+                        model_kwargs=model_kwargs,
+                        eta=eta,
+                        inpainting_mode=inpainting_mode,
+                        orig_img=init_image,
+                        mask_inpaint=mask_inpaint,
+                    )
+                else:
+                    out = sample_fn(
+                        model,
+                        img,
+                        t,
+                        clip_denoised=clip_denoised,
+                        denoised_fn=denoised_fn,
+                        cond_fn=cond_fn,
+                        model_kwargs=model_kwargs,
+                        eta=eta,
+                    )
+                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/multi_modal/guided_diffusion/respace.py

@@ -0,0 +1,78 @@
+# This code is borrowed and modified from Guided Diffusion Model,
+# made publicly available under MIT license
+# at https://github.com/IDEA-CCNL/Fengshenbang-LM/tree/main/fengshen/examples/disco_project
+
+import numpy as np
+import torch as th
+
+from .gaussian_diffusion import GaussianDiffusion
+
+
+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/multi_modal/guided_diffusion/script.py

@@ -0,0 +1,39 @@
+# This code is borrowed and modified from Guided Diffusion Model,
+# made publicly available under MIT license
+# at https://github.com/IDEA-CCNL/Fengshenbang-LM/tree/main/fengshen/examples/disco_project
+
+from modelscope.models.cv.motion_generation.modules.respace import \
+    space_timesteps
+from . import gaussian_diffusion as gd
+from .respace import SpacedDiffusion
+
+
+def create_diffusion(diffusion_config):
+    predict_xstart = False
+    sigma_small = False
+    learn_sigma = True
+
+    steps = diffusion_config['steps']
+    timestep_respacing = f'ddim{steps}'
+    diffusion_steps = 1000
+
+    rescale_timesteps = True
+
+    betas = gd.get_named_beta_schedule('linear', diffusion_steps)
+    loss_type = gd.LossType.MSE
+
+    if not timestep_respacing:
+        timestep_respacing = [diffusion_steps]
+
+    diffusion = SpacedDiffusion(
+        use_timesteps=space_timesteps(diffusion_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 sigma_small else gd.ModelVarType.FIXED_SMALL)
+                        if not learn_sigma else gd.ModelVarType.LEARNED_RANGE),
+        loss_type=loss_type,
+        rescale_timesteps=rescale_timesteps)
+
+    return diffusion

modelscope/pipelines/multi_modal/disco_guided_diffusion_pipeline/__init__.py

@@ -0,0 +1,23 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+from typing import TYPE_CHECKING
+
+from modelscope.utils.import_utils import LazyImportModule
+
+if TYPE_CHECKING:
+    from .disco_guided_diffusion import DiscoDiffusionPipeline
+    from .utils import resize
+else:
+    _import_structure = {
+        'disco_guided_diffusion': ['DiscoDiffusionPipeline'],
+        'utils': ['resize'],
+    }
+
+    import sys
+
+    sys.modules[__name__] = LazyImportModule(
+        __name__,
+        globals()['__file__'],
+        _import_structure,
+        module_spec=__spec__,
+        extra_objects={},
+    )

modelscope/pipelines/multi_modal/disco_guided_diffusion_pipeline/disco_guided_diffusion.py

@@ -0,0 +1,430 @@
+# This code is borrowed and modified from Guided Diffusion Model,
+# made publicly available under MIT license at
+# https://github.com/IDEA-CCNL/Fengshenbang-LM/tree/main/fengshen/examples/disco_project
+
+import gc
+import importlib
+import math
+import os
+
+import clip
+import cv2
+import json
+import numpy as np
+import torch
+import torch.nn as nn
+import torchvision.transforms as T
+import torchvision.transforms.functional as TF
+from PIL import Image
+from torch.nn import functional as F
+
+from modelscope.metainfo import Pipelines
+from modelscope.models.multi_modal.guided_diffusion.script import \
+    create_diffusion
+from modelscope.models.multi_modal.guided_diffusion.unet import HFUNetModel
+from modelscope.outputs import OutputKeys
+from modelscope.pipelines.builder import PIPELINES
+from modelscope.pipelines.multi_modal.diffusers_wrapped.diffusers_pipeline import \
+    DiffusersPipeline
+from modelscope.utils.constant import Tasks
+from .utils import resize
+
+
+def parse_prompt(prompt):
+    if prompt.startswith('http://') or prompt.startswith('https://'):
+        vals = prompt.rsplit(':', 2)
+        vals = [vals[0] + ':' + vals[1], *vals[2:]]
+    else:
+        vals = prompt.rsplit(':', 1)
+    vals = vals + ['', '1'][len(vals):]
+    return vals[0], float(vals[1])
+
+
+def sinc(x):
+    return torch.where(x != 0,
+                       torch.sin(math.pi * x) / (math.pi * x), x.new_ones([]))
+
+
+def lanczos(x, a):
+    cond = torch.logical_and(-a < x, x < a)
+    out = torch.where(cond, sinc(x) * sinc(x / a), x.new_zeros([]))
+    return out / out.sum()
+
+
+class MakeCutoutsDango(nn.Module):
+
+    def __init__(
+        self,
+        cut_size,
+        Overview=4,
+        InnerCrop=0,
+        IC_Size_Pow=0.5,
+        IC_Grey_P=0.2,
+    ):
+        super().__init__()
+        self.padargs = {}
+        self.cutout_debug = False
+        self.cut_size = cut_size
+        self.Overview = Overview
+        self.InnerCrop = InnerCrop
+        self.IC_Size_Pow = IC_Size_Pow
+        self.IC_Grey_P = IC_Grey_P
+        self.augs = T.Compose([
+            T.RandomHorizontalFlip(p=0.5),
+            T.Lambda(lambda x: x + torch.randn_like(x) * 0.01),
+            T.RandomAffine(
+                degrees=10,
+                translate=(0.05, 0.05),
+                interpolation=T.InterpolationMode.BILINEAR),
+            T.Lambda(lambda x: x + torch.randn_like(x) * 0.01),
+            T.RandomGrayscale(p=0.1),
+            T.Lambda(lambda x: x + torch.randn_like(x) * 0.01),
+            T.ColorJitter(
+                brightness=0.1, contrast=0.1, saturation=0.1, hue=0.1),
+        ])
+
+    def forward(self, input):
+        cutouts = []
+        gray = T.Grayscale(3)
+        sideY, sideX = input.shape[2:4]
+        max_size = min(sideX, sideY)
+        min_size = min(sideX, sideY, self.cut_size)
+        output_shape = [1, 3, self.cut_size, self.cut_size]
+        pad_input = F.pad(input,
+                          ((sideY - max_size) // 2, (sideY - max_size) // 2,
+                           (sideX - max_size) // 2, (sideX - max_size) // 2),
+                          **self.padargs)
+        cutout = resize(pad_input, out_shape=output_shape)
+
+        if self.Overview > 0:
+            if self.Overview <= 4:
+                if self.Overview >= 1:
+                    cutouts.append(cutout)
+                if self.Overview >= 2:
+                    cutouts.append(gray(cutout))
+                if self.Overview >= 3:
+                    cutouts.append(TF.hflip(cutout))
+                if self.Overview == 4:
+                    cutouts.append(gray(TF.hflip(cutout)))
+            else:
+                cutout = resize(pad_input, out_shape=output_shape)
+                for _ in range(self.Overview):
+                    cutouts.append(cutout)
+
+            if self.cutout_debug:
+                TF.to_pil_image(cutouts[0].clamp(0, 1).squeeze(0)).save(
+                    'cutout_overview0.jpg', quality=99)
+
+        if self.InnerCrop > 0:
+            for i in range(self.InnerCrop):
+                size = int(
+                    torch.rand([])**self.IC_Size_Pow * (max_size - min_size)
+                    + min_size)
+                offsetx = torch.randint(0, sideX - size + 1, ())
+                offsety = torch.randint(0, sideY - size + 1, ())
+                cutout = input[:, :, offsety:offsety + size,
+                               offsetx:offsetx + size]
+                if i <= int(self.IC_Grey_P * self.InnerCrop):
+                    cutout = gray(cutout)
+                cutout = resize(cutout, out_shape=output_shape)
+                cutouts.append(cutout)
+            if self.cutout_debug:
+                TF.to_pil_image(cutouts[-1].clamp(0, 1).squeeze(0)).save(
+                    'cutout_InnerCrop.jpg', quality=99)
+        cutouts = torch.cat(cutouts)
+
+        cutouts = self.augs(cutouts)
+        return cutouts
+
+
+def spherical_dist_loss(x, y):
+    x = F.normalize(x, dim=-1)
+    y = F.normalize(y, dim=-1)
+    return (x - y).norm(dim=-1).div(2).arcsin().pow(2).mul(2)
+
+
+def tv_loss(input):
+    """L2 total variation loss, as in Mahendran et al."""
+    input = F.pad(input, (0, 1, 0, 1), 'replicate')
+    x_diff = input[..., :-1, 1:] - input[..., :-1, :-1]
+    y_diff = input[..., 1:, :-1] - input[..., :-1, :-1]
+    return (x_diff**2 + y_diff**2).mean([1, 2, 3])
+
+
+def range_loss(input):
+    return (input - input.clamp(-1, 1)).pow(2).mean([1, 2, 3])
+
+
+normalize = T.Normalize(
+    mean=[0.48145466, 0.4578275, 0.40821073],
+    std=[0.26862954, 0.26130258, 0.27577711])
+
+
+@PIPELINES.register_module(
+    Tasks.text_to_image_synthesis,
+    module_name=Pipelines.disco_guided_diffusion)
+class DiscoDiffusionPipeline(DiffusersPipeline):
+
+    def __init__(self, model: str, device: str = 'gpu', **kwargs):
+        """  Chinese Disco Diffusion Pipeline.
+
+        Examples:
+
+        >>> import cv2
+        >>> from modelscope.pipelines import pipeline
+        >>> from modelscope.utils.constant import Tasks
+
+        >>> prompt = '赛博朋克，城市'
+        >>> output_image_path = './result.png'
+        >>> input = {
+        >>>     'text': prompt
+        >>> }
+        >>> pipe = pipeline(
+        >>>     Tasks.text_to_image_synthesis,
+        >>>     model='yyqoni/yinyueqin_cyberpunk',
+        >>>     model_revision='v1.0')
+        >>> output = pipe(input)['output_imgs'][0]
+        >>> cv2.imwrite(output_image_path, output)
+        >>> print('pipeline: the output image path is {}'.format(output_image_path))
+        """
+
+        super().__init__(model, device, **kwargs)
+
+        model_path = model
+
+        model_config = {'steps': 100, 'use_fp16': True}
+        self.diffusion = create_diffusion(model_config)
+
+        self.unet = HFUNetModel.from_pretrained(f'{model_path}/unet')
+
+        self.unet.requires_grad_(False).eval().to(self.device)
+        for name, param in self.unet.named_parameters():
+            if 'qkv' in name or 'norm' in name or 'proj' in name:
+                param.requires_grad_()
+        if model_config['use_fp16']:
+            self.unet.convert_to_fp16()
+
+        with open(
+                os.path.join(model_path, 'model_index.json'),
+                'r',
+                encoding='utf-8') as reader:
+            text = reader.read()
+        config_dict = json.loads(text)
+
+        library = importlib.import_module(config_dict['tokenizer'][0])
+        class_name = config_dict['tokenizer'][1]
+
+        self.taiyi_tokenizer = getattr(
+            library, class_name).from_pretrained(f'{model_path}/tokenizer')
+
+        library = importlib.import_module(config_dict['text_encoder'][0])
+        class_name = config_dict['text_encoder'][1]
+
+        self.taiyi_transformer = getattr(library, class_name).from_pretrained(
+            f'{model_path}/text_encoder').eval().to(self.device)
+
+        self.clip_models = []
+        self.clip_models.append(
+            clip.load('ViT-L/14',
+                      jit=False)[0].eval().requires_grad_(False).to(
+                          self.device))
+
+    def forward(self,
+                inputs,
+                init=None,
+                init_scale=2000,
+                skip_steps=10,
+                randomize_class=True,
+                eta=0.8,
+                output_type='pil',
+                return_dict=True,
+                clip_guidance_scale=7500):
+        if not isinstance(inputs, dict):
+            raise ValueError(
+                f'Expected the input to be a dictionary, but got {type(input)}'
+            )
+        if 'text' not in inputs:
+            raise ValueError('input should contain "text", but not found')
+
+        batch_size = 1
+        cutn_batches = 1
+
+        tv_scale = 0
+        range_scale = 150
+        sat_scale = 0
+
+        cut_overview = eval('[12]*400+[4]*600')
+        cut_innercut = eval('[4]*400+[12]*600')
+        cut_ic_pow = eval('[1]*1000')
+        cut_icgray_p = eval('[0.2]*400+[0]*600')
+
+        side_x = 512
+        side_y = 512
+
+        if 'width' in inputs:
+            side_x = inputs['width']
+        if 'height' in inputs:
+            side_y = inputs['height']
+        frame_prompt = [inputs.get('text')]
+        loss_values = []
+
+        model_stats = []
+        for clip_model in self.clip_models:
+            # cutn = 16
+            model_stat = {
+                'clip_model': None,
+                'target_embeds': [],
+                'make_cutouts': None,
+                'weights': []
+            }
+            model_stat['clip_model'] = clip_model
+
+            for prompt in frame_prompt:
+                txt, weight = parse_prompt(prompt)
+                # NOTE use chinese CLIP
+                txt = self.taiyi_transformer(
+                    self.taiyi_tokenizer(txt,
+                                         return_tensors='pt')['input_ids'].to(
+                                             self.device)).logits
+
+                model_stat['target_embeds'].append(txt)
+                model_stat['weights'].append(weight)
+
+            model_stat['target_embeds'] = torch.cat(
+                model_stat['target_embeds'])
+            model_stat['weights'] = torch.tensor(
+                model_stat['weights'], device=self.device)
+            if model_stat['weights'].sum().abs() < 1e-3:
+                raise RuntimeError('The weights must not sum to 0.')
+            model_stat['weights'] /= model_stat['weights'].sum().abs()
+            model_stats.append(model_stat)
+
+        init = None
+        cur_t = None
+
+        def cond_fn(x, t, y=None):
+            with torch.enable_grad():
+                x_is_NaN = False
+                x = x.detach().requires_grad_()
+                n = x.shape[0]
+
+                my_t = torch.ones([n], device=self.device,
+                                  dtype=torch.long) * cur_t
+                out = self.diffusion.p_mean_variance(
+                    self.unet,
+                    x,
+                    my_t,
+                    clip_denoised=False,
+                    model_kwargs={'y': y})
+                fac = self.diffusion.sqrt_one_minus_alphas_cumprod[cur_t]
+                x_in = out['pred_xstart'] * fac + x * (1 - fac)
+                x_in_grad = torch.zeros_like(x_in)
+
+                for model_stat in model_stats:
+                    for i in range(cutn_batches):
+                        t_int = int(t.item()) + 1
+                        input_resolution = model_stat[
+                            'clip_model'].visual.input_resolution
+
+                        cuts = MakeCutoutsDango(
+                            input_resolution,
+                            Overview=cut_overview[1000 - t_int],
+                            InnerCrop=cut_innercut[1000 - t_int],
+                            IC_Size_Pow=cut_ic_pow[1000 - t_int],
+                            IC_Grey_P=cut_icgray_p[1000 - t_int],
+                        )
+                        clip_in = normalize(cuts(x_in.add(1).div(2)))
+                        image_embeds = model_stat['clip_model'].encode_image(
+                            clip_in).float()
+                        dists = spherical_dist_loss(
+                            image_embeds.unsqueeze(1),
+                            model_stat['target_embeds'].unsqueeze(0))
+                        dists = dists.view([
+                            cut_overview[1000 - t_int]
+                            + cut_innercut[1000 - t_int], n, -1
+                        ])
+                        losses = dists.mul(
+                            model_stat['weights']).sum(2).mean(0)
+                        loss_values.append(losses.sum().item(
+                        ))  # log loss, probably shouldn't do per cutn_batch
+                        x_in_grad += torch.autograd.grad(
+                            losses.sum() * clip_guidance_scale,
+                            x_in)[0] / cutn_batches
+                tv_losses = tv_loss(x_in)
+                range_losses = range_loss(out['pred_xstart'])
+                sat_losses = torch.abs(x_in - x_in.clamp(min=-1, max=1)).mean()
+                loss = tv_losses.sum() * tv_scale + range_losses.sum(
+                ) * range_scale + sat_losses.sum() * sat_scale
+                if init is not None and init_scale:
+                    init_losses = self.lpips_model(x_in, init)
+                    loss = loss + init_losses.sum() * init_scale
+                x_in_grad += torch.autograd.grad(loss, x_in)[0]
+                if not torch.isnan(x_in_grad).any():
+                    grad = -torch.autograd.grad(x_in, x, x_in_grad)[0]
+                else:
+                    x_is_NaN = True
+                    grad = torch.zeros_like(x)
+            if not x_is_NaN:
+                magnitude = grad.square().mean().sqrt()
+                return grad * magnitude.clamp(max=0.05) / magnitude
+            return grad
+
+        sample_fn = self.diffusion.ddim_sample_loop_progressive
+
+        n_batches = 1
+
+        for i in range(n_batches):
+            gc.collect()
+            torch.cuda.empty_cache()
+            cur_t = self.diffusion.num_timesteps - skip_steps - 1
+
+            samples = sample_fn(
+                self.unet,
+                (batch_size, 3, side_y, side_x),
+                clip_denoised=False,
+                model_kwargs={},
+                cond_fn=cond_fn,
+                progress=True,
+                skip_timesteps=skip_steps,
+                init_image=init,
+                randomize_class=randomize_class,
+                eta=eta,
+            )
+
+            for j, sample in enumerate(samples):
+                image = sample['pred_xstart']
+                image = (image / 2 + 0.5).clamp(0, 1)
+                image = image.cpu().permute(0, 2, 3, 1).numpy()
+
+            if output_type == 'pil':
+                image = self.numpy_to_pil(image)
+                return image
+
+            if not return_dict:
+                return (image, None)
+
+    @staticmethod
+    def numpy_to_pil(images):
+        """
+        Convert a numpy image or a batch of images to a PIL image.
+        """
+        if images.ndim == 3:
+            images = images[None, ...]
+        images = (images * 255).round().astype('uint8')
+        if images.shape[-1] == 1:
+            # special case for grayscale (single channel) images
+            pil_images = [
+                Image.fromarray(image.squeeze(), mode='L') for image in images
+            ]
+        else:
+            pil_images = [Image.fromarray(image) for image in images]
+
+        return pil_images
+
+    def postprocess(self, inputs):
+        images = []
+        for img in inputs:
+            if isinstance(img, Image.Image):
+                img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
+                images.append(img)
+        return {OutputKeys.OUTPUT_IMGS: images}

modelscope/pipelines/multi_modal/disco_guided_diffusion_pipeline/utils.py

@@ -0,0 +1,468 @@
+# The implementation is adopted from https://github.com/assafshocher/ResizeRight
+import warnings
+from fractions import Fraction
+from math import ceil
+
+
+class NoneClass:
+    pass
+
+
+try:
+    import torch
+    from torch import nn
+    nnModuleWrapped = nn.Module
+except ImportError:
+    warnings.warn('No PyTorch found, will work only with Numpy')
+    torch = None
+    nnModuleWrapped = NoneClass
+
+try:
+    import numpy
+except ImportError:
+    warnings.warn('No Numpy found, will work only with PyTorch')
+    numpy = None
+
+if numpy is None and torch is None:
+    raise ImportError('Must have either Numpy or PyTorch but both not found')
+
+
+def set_framework_dependencies(x):
+    if type(x) is numpy.ndarray:
+
+        def to_dtype(a):
+            return a
+
+        fw = numpy
+    else:
+
+        def to_dtype(a):
+            return a.to(x.dtype)
+
+        fw = torch
+    eps = fw.finfo(fw.float32).eps
+    return fw, to_dtype, eps
+
+
+def support_sz(sz):
+
+    def wrapper(f):
+        f.support_sz = sz
+        return f
+
+    return wrapper
+
+
+@support_sz(4)
+def cubic(x):
+    fw, to_dtype, eps = set_framework_dependencies(x)
+    absx = fw.abs(x)
+    absx2 = absx**2
+    absx3 = absx**3
+    v1 = (1.5 * absx3 - 2.5 * absx2 + 1.) * to_dtype(absx <= 1.)
+    v2 = (-0.5 * absx3 + 2.5 * absx2 - 4. * absx
+          + 2.) * to_dtype((1. < absx) & (absx <= 2.))
+    return v1 + v2
+
+
+def resize(input,
+           scale_factors=None,
+           out_shape=None,
+           interp_method=cubic,
+           support_sz=None,
+           antialiasing=True,
+           by_convs=False,
+           scale_tolerance=None,
+           max_numerator=10,
+           pad_mode='constant'):
+    # get properties of the input tensor
+    in_shape, n_dims = input.shape, input.ndim
+
+    # fw stands for framework that can be either numpy or torch,
+    # determined by the input type
+    fw = numpy if type(input) is numpy.ndarray else torch
+    eps = fw.finfo(fw.float32).eps
+    device = input.device if fw is torch else None
+
+    # set missing scale factors or output shapem one according to another,
+    # scream if both missing. this is also where all the defults policies
+    # take place. also handling the by_convs attribute carefully.
+    scale_factors, out_shape, by_convs = set_scale_and_out_sz(
+        in_shape, out_shape, scale_factors, by_convs, scale_tolerance,
+        max_numerator, eps, fw)
+
+    # sort indices of dimensions according to scale of each dimension.
+    # since we are going dim by dim this is efficient
+    sorted_filtered_dims_and_scales = [
+        (dim, scale_factors[dim], by_convs[dim], in_shape[dim], out_shape[dim])
+        for dim in sorted(range(n_dims), key=lambda ind: scale_factors[ind])
+        if scale_factors[dim] != 1.
+    ]
+
+    # unless support size is specified by the user, it is an attribute
+    # of the interpolation method
+    if support_sz is None:
+        support_sz = interp_method.support_sz
+
+    # output begins identical to input and changes with each iteration
+    output = input
+
+    # iterate over dims
+    for (dim, scale_factor, dim_by_convs, in_sz,
+         out_sz) in sorted_filtered_dims_and_scales:
+        # STEP 1- PROJECTED GRID: The non-integer locations of the projection
+        # of output pixel locations to the input tensor
+        projected_grid = get_projected_grid(in_sz, out_sz, scale_factor, fw,
+                                            dim_by_convs, device)
+
+        # STEP 1.5: ANTIALIASING- If antialiasing is taking place, we modify
+        # the window size and the interpolation method (see inside function)
+        cur_interp_method, cur_support_sz = apply_antialiasing_if_needed(
+            interp_method, support_sz, scale_factor, antialiasing)
+
+        # STEP 2- FIELDS OF VIEW: for each output pixels, map the input pixels
+        # that influence it. Also calculate needed padding and update grid
+        # accoedingly
+        field_of_view = get_field_of_view(projected_grid, cur_support_sz, fw,
+                                          eps, device)
+
+        # STEP 2.5- CALCULATE PAD AND UPDATE: according to the field of view,
+        # the input should be padded to handle the boundaries, coordinates
+        # should be updated. actual padding only occurs when weights are
+        # aplied (step 4). if using by_convs for this dim, then we need to
+        # calc right and left boundaries for each filter instead.
+        pad_sz, projected_grid, field_of_view = calc_pad_sz(
+            in_sz, out_sz, field_of_view, projected_grid, scale_factor,
+            dim_by_convs, fw, device)
+
+        # STEP 3- CALCULATE WEIGHTS: Match a set of weights to the pixels in
+        # the field of view for each output pixel
+        weights = get_weights(cur_interp_method, projected_grid, field_of_view)
+
+        # STEP 4- APPLY WEIGHTS: Each output pixel is calculated by multiplying
+        # its set of weights with the pixel values in its field of view.
+        # We now multiply the fields of view with their matching weights.
+        # We do this by tensor multiplication and broadcasting.
+        # if by_convs is true for this dim, then we do this action by
+        # convolutions. this is equivalent but faster.
+        if not dim_by_convs:
+            output = apply_weights(output, field_of_view, weights, dim, n_dims,
+                                   pad_sz, pad_mode, fw)
+        else:
+            output = apply_convs(output, scale_factor, in_sz, out_sz, weights,
+                                 dim, pad_sz, pad_mode, fw)
+    return output
+
+
+def get_projected_grid(in_sz, out_sz, scale_factor, fw, by_convs, device=None):
+    # we start by having the ouput coordinates which are just integer locations
+    # in the special case when usin by_convs, we only need two cycles of grid
+    # points. the first and last.
+    grid_sz = out_sz if not by_convs else scale_factor.numerator
+    out_coordinates = fw_arange(grid_sz, fw, device)
+
+    # This is projecting the ouput pixel locations in 1d to the input tensor,
+    # as non-integer locations.
+    # the following fomrula is derived in the paper
+    # "From Discrete to Continuous Convolutions" by Shocher et al.
+    v1 = out_coordinates / float(scale_factor) + (in_sz - 1) / 2
+    v2 = (out_sz - 1) / (2 * float(scale_factor))
+    return v1 - v2
+
+
+def get_field_of_view(projected_grid, cur_support_sz, fw, eps, device):
+    # for each output pixel, map which input pixels influence it, in 1d.
+    # we start by calculating the leftmost neighbor, using half of the window
+    # size (eps is for when boundary is exact int)
+    left_boundaries = fw_ceil(projected_grid - cur_support_sz / 2 - eps, fw)
+
+    # then we simply take all the pixel centers in the field by counting
+    # window size pixels from the left boundary
+    ordinal_numbers = fw_arange(ceil(cur_support_sz - eps), fw, device)
+    return left_boundaries[:, None] + ordinal_numbers
+
+
+def calc_pad_sz(in_sz, out_sz, field_of_view, projected_grid, scale_factor,
+                dim_by_convs, fw, device):
+    if not dim_by_convs:
+        # determine padding according to neighbor coords out of bound.
+        # this is a generalized notion of padding, when pad<0 it means crop
+        pad_sz = [
+            -field_of_view[0, 0].item(),
+            field_of_view[-1, -1].item() - in_sz + 1
+        ]
+
+        # since input image will be changed by padding, coordinates of both
+        # field_of_view and projected_grid need to be updated
+        field_of_view += pad_sz[0]
+        projected_grid += pad_sz[0]
+
+    else:
+        # only used for by_convs, to calc the boundaries of each filter the
+        # number of distinct convolutions is the numerator of the scale factor
+        num_convs, stride = scale_factor.numerator, scale_factor.denominator
+
+        # calculate left and right boundaries for each conv. left can also be
+        # negative right can be bigger than in_sz. such cases imply padding if
+        # needed. however if# both are in-bounds, it means we need to crop,
+        # practically apply the conv only on part of the image.
+        left_pads = -field_of_view[:, 0]
+
+        # next calc is tricky, explanation by rows:
+        # 1) counting output pixels between the first position of each filter
+        #    to the right boundary of the input
+        # 2) dividing it by number of filters to count how many 'jumps'
+        #    each filter does
+        # 3) multiplying by the stride gives us the distance over the input
+        #    coords done by all these jumps for each filter
+        # 4) to this distance we add the right boundary of the filter when
+        #    placed in its leftmost position. so now we get the right boundary
+        #    of that filter in input coord.
+        # 5) the padding size needed is obtained by subtracting the rightmost
+        #    input coordinate. if the result is positive padding is needed. if
+        #    negative then negative padding means shaving off pixel columns.
+        right_pads = (((out_sz - fw_arange(num_convs, fw, device) - 1)  # (1)
+                       // num_convs)  # (2)
+                      * stride  # (3)
+                      + field_of_view[:, -1]  # (4)
+                      - in_sz + 1)  # (5)
+
+        # in the by_convs case pad_sz is a list of left-right pairs. one per
+        # each filter
+
+        pad_sz = list(zip(left_pads, right_pads))
+
+    return pad_sz, projected_grid, field_of_view
+
+
+def get_weights(interp_method, projected_grid, field_of_view):
+    # the set of weights per each output pixels is the result of the chosen
+    # interpolation method applied to the distances between projected grid
+    # locations and the pixel-centers in the field of view (distances are
+    # directed, can be positive or negative)
+    weights = interp_method(projected_grid[:, None] - field_of_view)
+
+    # we now carefully normalize the weights to sum to 1 per each output pixel
+    sum_weights = weights.sum(1, keepdims=True)
+    sum_weights[sum_weights == 0] = 1
+    return weights / sum_weights
+
+
+def apply_weights(input, field_of_view, weights, dim, n_dims, pad_sz, pad_mode,
+                  fw):
+    # for this operation we assume the resized dim is the first one.
+    # so we transpose and will transpose back after multiplying
+    tmp_input = fw_swapaxes(input, dim, 0, fw)
+
+    # apply padding
+    tmp_input = fw_pad(tmp_input, fw, pad_sz, pad_mode)
+
+    # field_of_view is a tensor of order 2: for each output (1d location
+    # along cur dim)- a list of 1d neighbors locations.
+    # note that this whole operations is applied to each dim separately,
+    # this is why it is all in 1d.
+    # neighbors = tmp_input[field_of_view] is a tensor of order image_dims+1:
+    # for each output pixel (this time indicated in all dims), these are the
+    # values of the neighbors in the 1d field of view. note that we only
+    # consider neighbors along the current dim, but such set exists for every
+    # multi-dim location, hence the final tensor order is image_dims+1.
+    neighbors = tmp_input[field_of_view]
+
+    # weights is an order 2 tensor: for each output location along 1d- a list
+    # of weights matching the field of view. we augment it with ones, for
+    # broadcasting, so that when multiplies some tensor the weights affect
+    # only its first dim.
+    tmp_weights = fw.reshape(weights, (*weights.shape, *[1] * (n_dims - 1)))
+
+    # now we simply multiply the weights with the neighbors, and then sum
+    # along the field of view, to get a single value per out pixel
+    tmp_output = (neighbors * tmp_weights).sum(1)
+
+    # we transpose back the resized dim to its original position
+    return fw_swapaxes(tmp_output, 0, dim, fw)
+
+
+def apply_convs(input, scale_factor, in_sz, out_sz, weights, dim, pad_sz,
+                pad_mode, fw):
+    # for this operations we assume the resized dim is the last one.
+    # so we transpose and will transpose back after multiplying
+    input = fw_swapaxes(input, dim, -1, fw)
+
+    # the stride for all convs is the denominator of the scale factor
+    stride, num_convs = scale_factor.denominator, scale_factor.numerator
+
+    # prepare an empty tensor for the output
+    tmp_out_shape = list(input.shape)
+    tmp_out_shape[-1] = out_sz
+    tmp_output = fw_empty(tuple(tmp_out_shape), fw, input.device)
+
+    # iterate over the conv operations. we have as many as the numerator
+    # of the scale-factor. for each we need boundaries and a filter.
+    for conv_ind, (pad_sz, filt) in enumerate(zip(pad_sz, weights)):
+        # apply padding (we pad last dim, padding can be negative)
+        pad_dim = input.ndim - 1
+        tmp_input = fw_pad(input, fw, pad_sz, pad_mode, dim=pad_dim)
+
+        # apply convolution over last dim. store in the output tensor with
+        # positional strides so that when the loop is comlete conv results are
+        # interwind
+        tmp_output[..., conv_ind::num_convs] = fw_conv(tmp_input, filt, stride)
+
+    return fw_swapaxes(tmp_output, -1, dim, fw)
+
+
+def set_scale_and_out_sz(in_shape, out_shape, scale_factors, by_convs,
+                         scale_tolerance, max_numerator, eps, fw):
+    # eventually we must have both scale-factors and out-sizes for all in/out
+    # dims. however, we support many possible partial arguments
+    if scale_factors is None and out_shape is None:
+        raise ValueError('either scale_factors or out_shape should be '
+                         'provided')
+    if out_shape is not None:
+        # if out_shape has less dims than in_shape, we defaultly resize the
+        # first dims for numpy and last dims for torch
+        out_shape = (
+            list(out_shape) + list(in_shape[len(out_shape):]) if fw is numpy
+            else list(in_shape[:-len(out_shape)]) + list(out_shape))
+        if scale_factors is None:
+            # if no scale given, we calculate it as the out to in ratio
+            # (not recomended)
+            scale_factors = [
+                out_sz / in_sz for out_sz, in_sz in zip(out_shape, in_shape)
+            ]
+
+    if scale_factors is not None:
+        # by default, if a single number is given as scale, we assume resizing
+        # two dims (most common are images with 2 spatial dims)
+        scale_factors = (
+            scale_factors if isinstance(scale_factors, (list, tuple)) else
+            [scale_factors, scale_factors])
+        # if less scale_factors than in_shape dims, we defaultly resize the
+        # first dims for numpy and last dims for torch
+        if fw is numpy:
+            scale_factors = list(scale_factors) + [1] * (
+                len(in_shape) - len(scale_factors))
+        else:
+            scale_factors = [1] * (len(in_shape)
+                                   - len(scale_factors)) + list(scale_factors)
+        if out_shape is None:
+            # when no out_shape given, it is calculated by multiplying the
+            # scale by the in_shape (not recomended)
+            out_shape = [
+                ceil(scale_factor * in_sz)
+                for scale_factor, in_sz in zip(scale_factors, in_shape)
+            ]
+        # next part intentionally after out_shape determined for stability
+        # we fix by_convs to be a list of truth values in case it is not
+        if not isinstance(by_convs, (list, tuple)):
+            by_convs = [by_convs] * len(out_shape)
+
+        # next loop fixes the scale for each dim to be either frac or float.
+        # this is determined by by_convs and by tolerance for scale accuracy.
+        for ind, (sf, dim_by_convs) in enumerate(zip(scale_factors, by_convs)):
+            # first we fractionaize
+            if dim_by_convs:
+                frac = Fraction(1 / sf).limit_denominator(max_numerator)
+                frac = Fraction(
+                    numerator=frac.denominator, denominator=frac.numerator)
+
+            # if accuracy is within tolerance scale will be frac. if not, then
+            # it will be float and the by_convs attr will be set false for
+            # this dim
+            if scale_tolerance is None:
+                scale_tolerance = eps
+            if dim_by_convs and abs(frac - sf) < scale_tolerance:
+                scale_factors[ind] = frac
+            else:
+                scale_factors[ind] = float(sf)
+                by_convs[ind] = False
+
+        return scale_factors, out_shape, by_convs
+
+
+def apply_antialiasing_if_needed(interp_method, support_sz, scale_factor,
+                                 antialiasing):
+    # antialiasing is "stretching" the field of view according to the scale
+    # factor (only for downscaling). this is low-pass filtering. this
+    # requires modifying both the interpolation (stretching the 1d
+    # function and multiplying by the scale-factor) and the window size.
+    scale_factor = float(scale_factor)
+    if scale_factor >= 1.0 or not antialiasing:
+        return interp_method, support_sz
+    cur_interp_method = (
+        lambda arg: scale_factor * interp_method(scale_factor * arg))
+    cur_support_sz = support_sz / scale_factor
+    return cur_interp_method, cur_support_sz
+
+
+def fw_ceil(x, fw):
+    if fw is numpy:
+        return fw.int_(fw.ceil(x))
+    else:
+        return x.ceil().long()
+
+
+def fw_floor(x, fw):
+    if fw is numpy:
+        return fw.int_(fw.floor(x))
+    else:
+        return x.floor().long()
+
+
+def fw_cat(x, fw):
+    if fw is numpy:
+        return fw.concatenate(x)
+    else:
+        return fw.cat(x)
+
+
+def fw_swapaxes(x, ax_1, ax_2, fw):
+    if fw is numpy:
+        return fw.swapaxes(x, ax_1, ax_2)
+    else:
+        return x.transpose(ax_1, ax_2)
+
+
+def fw_pad(x, fw, pad_sz, pad_mode, dim=0):
+    if pad_sz == (0, 0):
+        return x
+    if fw is numpy:
+        pad_vec = [(0, 0)] * x.ndim
+        pad_vec[dim] = pad_sz
+        return fw.pad(x, pad_width=pad_vec, mode=pad_mode)
+    else:
+        if x.ndim < 3:
+            x = x[None, None, ...]
+
+        pad_vec = [0] * ((x.ndim - 2) * 2)
+        pad_vec[0:2] = pad_sz
+        return fw.nn.functional.pad(
+            x.transpose(dim, -1), pad=pad_vec,
+            mode=pad_mode).transpose(dim, -1)
+
+
+def fw_conv(input, filter, stride):
+    # we want to apply 1d conv to any nd array. the way to do it is to reshape
+    # the input to a 4D tensor. first two dims are singeletons, 3rd dim stores
+    # all the spatial dims that we are not convolving along now. then we can
+    # apply conv2d with a 1xK filter. This convolves the same way all the other
+    # dims stored in the 3d dim. like depthwise conv over these.
+    # TODO: numpy support
+    reshaped_input = input.reshape(1, 1, -1, input.shape[-1])
+    reshaped_output = torch.nn.functional.conv2d(
+        reshaped_input, filter.view(1, 1, 1, -1), stride=(1, stride))
+    return reshaped_output.reshape(*input.shape[:-1], -1)
+
+
+def fw_arange(upper_bound, fw, device):
+    if fw is numpy:
+        return fw.arange(upper_bound)
+    else:
+        return fw.arange(upper_bound, device=device)
+
+
+def fw_empty(shape, fw, device):
+    if fw is numpy:
+        return fw.empty(shape)
+    else:
+        return fw.empty(size=(*shape, ), device=device)

tests/pipelines/test_disco_guided_diffusion.py

@@ -0,0 +1,46 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import unittest
+
+import cv2
+
+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 DiscoGuidedDiffusionTest(unittest.TestCase, DemoCompatibilityCheck):
+
+    def setUp(self) -> None:
+        self.task = Tasks.text_to_image_synthesis
+        self.model_id1 = 'yyqoni/yinyueqin_test'
+        self.model_id2 = 'yyqoni/yinyueqin_cyberpunk'
+
+    test_input1 = '夕阳西下'
+    test_input2 = '城市，赛博朋克'
+
+    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
+    def test_run(self):
+        diffusers_pipeline = pipeline(
+            task=self.task, model=self.model_id1, model_revision='v1.0')
+        output = diffusers_pipeline({
+            'text': self.test_input1,
+            'height': 256,
+            'width': 256
+        })
+        cv2.imwrite('output1.png', output['output_imgs'][0])
+        print('Image saved to output1.png')
+
+        diffusers_pipeline = pipeline(
+            task=self.task, model=self.model_id2, model_revision='v1.0')
+        output = diffusers_pipeline({
+            'text': self.test_input2,
+            'height': 256,
+            'width': 256
+        })
+        cv2.imwrite('output2.png', output['output_imgs'][0])
+        print('Image saved to output2.png')
+
+
+if __name__ == '__main__':
+    unittest.main()