[to #42322933] add image editing model masactrl

Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/13671142
* add image editing model MasaCtrl

* add image editing model MasaCtrl

* Merge remote-tracking branch 'origin/master' into cv/image-editing-masactrl

modelscope/metainfo.py

@@ -279,6 +279,7 @@ class Pipelines(object):
     universal_matting = 'unet-universal-matting'
     image_denoise = 'nafnet-image-denoise'
     image_deblur = 'nafnet-image-deblur'
+    image_editing = 'masactrl-image-editing'
     person_image_cartoon = 'unet-person-image-cartoon'
     ocr_detection = 'resnet18-ocr-detection'
     table_recognition = 'dla34-table-recognition'
@@ -603,6 +604,8 @@ DEFAULT_MODEL_FOR_PIPELINE = {
                             'damo/cv_nafnet_image-denoise_sidd'),
     Tasks.image_deblurring: (Pipelines.image_deblur,
                              'damo/cv_nafnet_image-deblur_gopro'),
+    Tasks.image_editing: (Pipelines.image_editing,
+                          'damo/cv_masactrl_image-editing'),
     Tasks.video_stabilization: (Pipelines.video_stabilization,
                                 'damo/cv_dut-raft_video-stabilization_base'),
     Tasks.video_super_resolution:

modelscope/models/cv/__init__.py

@@ -7,10 +7,11 @@ from . import (action_recognition, animal_recognition, bad_image_detecting,
                crowd_counting, face_detection, face_generation,
                face_reconstruction, human_reconstruction, image_classification,
                image_color_enhance, image_colorization, image_defrcn_fewshot,
-               image_denoise, image_inpainting, image_instance_segmentation,
-               image_matching, image_mvs_depth_estimation,
-               image_panoptic_segmentation, image_portrait_enhancement,
-               image_probing_model, image_quality_assessment_degradation,
+               image_denoise, image_editing, image_inpainting,
+               image_instance_segmentation, image_matching,
+               image_mvs_depth_estimation, image_panoptic_segmentation,
+               image_portrait_enhancement, image_probing_model,
+               image_quality_assessment_degradation,
                image_quality_assessment_man, image_quality_assessment_mos,
                image_reid_person, image_restoration,
                image_semantic_segmentation, image_to_image_generation,

modelscope/models/cv/image_editing/__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 .masactrl import MutualSelfAttentionControl
+    from .masactrl_utils import regiter_attention_editor_diffusers
+else:
+    _import_structure = {
+        'masactrl': ['MutualSelfAttentionControl'],
+        'masactrl_utils': ['regiter_attention_editor_diffusers']
+    }
+
+    import sys
+
+    sys.modules[__name__] = LazyImportModule(
+        __name__,
+        globals()['__file__'],
+        _import_structure,
+        module_spec=__spec__,
+        extra_objects={},
+    )

modelscope/models/cv/image_editing/masactrl.py

@@ -0,0 +1,77 @@
+# ------------------------------------------------------------------------
+# Modified from https://github.com/TencentARC/MasaCtrl/blob/main/masactrl/masactrl.py
+# Copyright (c) 2023 TencentARC. All Rights Reserved.
+# ------------------------------------------------------------------------
+
+import torch
+from einops import rearrange
+
+from .masactrl_utils import AttentionBase
+
+
+class MutualSelfAttentionControl(AttentionBase):
+
+    def __init__(self,
+                 start_step=4,
+                 start_layer=10,
+                 layer_idx=None,
+                 step_idx=None,
+                 total_steps=50):
+        """
+        Mutual self-attention control for Stable-Diffusion model
+        Args:
+            start_step: the step to start mutual self-attention control
+            start_layer: the layer to start mutual self-attention control
+            layer_idx: list of the layers to apply mutual self-attention control
+            step_idx: list the steps to apply mutual self-attention control
+            total_steps: the total number of steps
+        """
+        super().__init__()
+        self.total_steps = total_steps
+        self.start_step = start_step
+        self.start_layer = start_layer
+        self.layer_idx = layer_idx if layer_idx is not None else list(
+            range(start_layer, 16))
+        self.step_idx = step_idx if step_idx is not None else list(
+            range(start_step, total_steps))  # denoise index
+        print('step_idx: ', self.step_idx)
+        print('layer_idx: ', self.layer_idx)
+
+    def attn_batch(self, q, k, v, sim, attn, is_cross, place_in_unet,
+                   num_heads, **kwargs):
+        b = q.shape[0] // num_heads
+        q = rearrange(q, '(b h) n d -> h (b n) d', h=num_heads)
+        k = rearrange(k, '(b h) n d -> h (b n) d', h=num_heads)
+        v = rearrange(v, '(b h) n d -> h (b n) d', h=num_heads)
+
+        sim = torch.einsum('h i d, h j d -> h i j', q, k) * kwargs.get('scale')
+        attn = sim.softmax(-1)
+        out = torch.einsum('h i j, h j d -> h i d', attn, v)
+        out = rearrange(out, 'h (b n) d -> b n (h d)', b=b)
+        return out
+
+    def forward(self, q, k, v, sim, attn, is_cross, place_in_unet, num_heads,
+                **kwargs):
+        """
+        Attention forward function
+        """
+        if is_cross or self.cur_step not in self.step_idx or self.cur_att_layer // 2 not in self.layer_idx:
+            return super().forward(q, k, v, sim, attn, is_cross, place_in_unet,
+                                   num_heads, **kwargs)
+
+        qu, qc = q.chunk(2)  # uncond, cond
+        ku, kc = k.chunk(2)
+        vu, vc = v.chunk(2)
+        attnu, attnc = attn.chunk(2)
+        # uncond
+        # ku[:num_heads], vu[:num_heads] -> source
+        # qu -> [source, target]
+        out_u = self.attn_batch(qu, ku[:num_heads], vu[:num_heads],
+                                sim[:num_heads], attnu, is_cross,
+                                place_in_unet, num_heads, **kwargs)
+        out_c = self.attn_batch(qc, kc[:num_heads], vc[:num_heads],
+                                sim[:num_heads], attnc, is_cross,
+                                place_in_unet, num_heads, **kwargs)
+        out = torch.cat([out_u, out_c], dim=0)
+
+        return out

modelscope/models/cv/image_editing/masactrl_utils.py

@@ -0,0 +1,124 @@
+# ------------------------------------------------------------------------
+# Modified from https://github.com/TencentARC/MasaCtrl/blob/main/masactrl/masactrl_utils.py
+# Copyright (c) 2023 TencentARC. All Rights Reserved.
+# ------------------------------------------------------------------------
+
+import torch
+import torch.nn as nn
+from einops import rearrange, repeat
+
+
+class AttentionBase:
+
+    def __init__(self):
+        self.cur_step = 0
+        self.num_att_layers = -1
+        self.cur_att_layer = 0
+
+    def after_step(self):
+        pass
+
+    def __call__(self, q, k, v, sim, attn, is_cross, place_in_unet, num_heads,
+                 **kwargs):
+        out = self.forward(q, k, v, sim, attn, is_cross, place_in_unet,
+                           num_heads, **kwargs)
+        self.cur_att_layer += 1
+        if self.cur_att_layer == self.num_att_layers:
+            self.cur_att_layer = 0
+            self.cur_step += 1
+            # after step
+            self.after_step()
+        return out
+
+    def forward(self, q, k, v, sim, attn, is_cross, place_in_unet, num_heads,
+                **kwargs):
+        out = torch.einsum('b i j, b j d -> b i d', attn, v)
+        out = rearrange(out, '(b h) n d -> b n (h d)', h=num_heads)
+        return out
+
+    def reset(self):
+        self.cur_step = 0
+        self.cur_att_layer = 0
+
+
+def regiter_attention_editor_diffusers(model, editor: AttentionBase):
+    """
+    Register a attention editor to Diffuser Pipeline, refer from [Prompt-to-Prompt]
+    """
+
+    def ca_forward(self, place_in_unet):
+
+        def forward(x,
+                    encoder_hidden_states=None,
+                    attention_mask=None,
+                    context=None,
+                    mask=None):
+            """
+            The attention is similar to the original implementation of LDM CrossAttention class
+            except adding some modifications on the attention
+            """
+            if encoder_hidden_states is not None:
+                context = encoder_hidden_states
+            if attention_mask is not None:
+                mask = attention_mask
+
+            to_out = self.to_out
+            if isinstance(to_out, nn.modules.container.ModuleList):
+                to_out = self.to_out[0]
+            else:
+                to_out = self.to_out
+
+            h = self.heads
+            q = self.to_q(x)
+            is_cross = context is not None
+            context = context if is_cross else x
+            k = self.to_k(context)
+            v = self.to_v(context)
+            q, k, v = map(
+                lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h),
+                (q, k, v))
+
+            sim = torch.einsum('b i d, b j d -> b i j', q, k) * self.scale
+
+            if mask is not None:
+                mask = rearrange(mask, 'b ... -> b (...)')
+                max_neg_value = -torch.finfo(sim.dtype).max
+                mask = repeat(mask, 'b j -> (b h) () j', h=h)
+                mask = mask[:, None, :].repeat(h, 1, 1)
+                sim.masked_fill_(~mask, max_neg_value)
+
+            attn = sim.softmax(dim=-1)
+            # the only difference
+            out = editor(
+                q,
+                k,
+                v,
+                sim,
+                attn,
+                is_cross,
+                place_in_unet,
+                self.heads,
+                scale=self.scale)
+
+            return to_out(out)
+
+        return forward
+
+    def register_editor(net, count, place_in_unet):
+        for name, subnet in net.named_children():
+            if net.__class__.__name__ == 'Attention':  # spatial Transformer layer
+                net.forward = ca_forward(net, place_in_unet)
+                return count + 1
+            elif hasattr(net, 'children'):
+                count = register_editor(subnet, count, place_in_unet)
+        return count
+
+    cross_att_count = 0
+    for net_name, net in model.unet.named_children():
+        if 'down' in net_name:
+            cross_att_count += register_editor(net, 0, 'down')
+        elif 'mid' in net_name:
+            cross_att_count += register_editor(net, 0, 'mid')
+        elif 'up' in net_name:
+            cross_att_count += register_editor(net, 0, 'up')
+    editor.num_att_layers = cross_att_count

modelscope/outputs/outputs.py

@@ -730,6 +730,7 @@ TASK_OUTPUTS = {
     Tasks.image_colorization: [OutputKeys.OUTPUT_IMG],
     Tasks.image_color_enhancement: [OutputKeys.OUTPUT_IMG],
     Tasks.image_denoising: [OutputKeys.OUTPUT_IMG],
+    Tasks.image_editing: [OutputKeys.OUTPUT_IMG],
     Tasks.image_portrait_enhancement: [OutputKeys.OUTPUT_IMG],
     Tasks.crowd_counting: [OutputKeys.SCORES, OutputKeys.OUTPUT_IMG],
     Tasks.image_inpainting: [OutputKeys.OUTPUT_IMG],

modelscope/pipeline_inputs.py

@@ -453,4 +453,8 @@ TASK_INPUTS = {
     Tasks.text_to_360panorama_image: {
         'prompt': InputType.TEXT,
     },
+    Tasks.image_editing: {
+        'img': InputType.IMAGE,
+        'prompts': InputType.LIST
+    }
 }

modelscope/pipelines/cv/__init__.py

@@ -30,6 +30,7 @@ if TYPE_CHECKING:
     from .image_colorization_pipeline import ImageColorizationPipeline
     from .image_denoise_pipeline import ImageDenoisePipeline
     from .image_deblur_pipeline import ImageDeblurPipeline
+    from .image_editing_pipeline import ImageEditingPipeline
     from .image_instance_segmentation_pipeline import ImageInstanceSegmentationPipeline
     from .image_matting_pipeline import ImageMattingPipeline
     from .image_portrait_enhancement_pipeline import ImagePortraitEnhancementPipeline
@@ -136,6 +137,7 @@ else:
         'image_cartoon_pipeline': ['ImageCartoonPipeline'],
         'image_denoise_pipeline': ['ImageDenoisePipeline'],
         'image_deblur_pipeline': ['ImageDeblurPipeline'],
+        'image_editing_pipeline': ['ImageEditingPipeline'],
         'image_color_enhance_pipeline': ['ImageColorEnhancePipeline'],
         'image_colorization_pipeline': ['ImageColorizationPipeline'],
         'image_instance_segmentation_pipeline':

modelscope/pipelines/cv/image_editing_pipeline.py

@@ -0,0 +1,365 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import os.path
+from typing import Any, Dict, Optional, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from diffusers import DDIMScheduler, StableDiffusionPipeline
+from PIL import Image
+from torchvision import transforms
+from tqdm import tqdm
+
+from modelscope.metainfo import Pipelines
+from modelscope.models.cv.image_editing import (
+    MutualSelfAttentionControl, regiter_attention_editor_diffusers)
+from modelscope.outputs import OutputKeys
+from modelscope.pipelines.builder import PIPELINES
+from modelscope.pipelines.multi_modal.diffusers_wrapped.diffusers_pipeline import \
+    DiffusersPipeline
+from modelscope.preprocessors import LoadImage
+from modelscope.utils.constant import Tasks
+from modelscope.utils.logger import get_logger
+
+logger = get_logger()
+
+__all__ = ['ImageEditingPipeline']
+
+
+@PIPELINES.register_module(
+    Tasks.image_editing, module_name=Pipelines.image_editing)
+class ImageEditingPipeline(DiffusersPipeline):
+
+    def __init__(self, model=str, preprocessor=None, **kwargs):
+        """  MasaCtrl Image Editing Pipeline.
+
+        Examples:
+
+        >>> import cv2
+        >>> from modelscope.pipelines import pipeline
+        >>> from modelscope.utils.constant import Tasks
+
+        >>> prompts = [
+        >>>     "",                           # source prompt
+        >>>     "a photo of a running corgi"  # target prompt
+        >>> ]
+        >>> output_image_path = './result.png'
+        >>> img = 'https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/public/ModelScope/test/images/corgi.jpg'
+        >>> input = {'img': img, 'prompts': prompts}
+        >>>
+        >>> pipe = pipeline(
+        >>>     Tasks.image_editing,
+        >>>     model='damo/cv_masactrl_image-editing')
+        >>>
+        >>> output = pipe(input)['output_img']
+        >>> cv2.imwrite(output_image_path, output)
+        >>> print('pipeline: the output image path is {}'.format(output_image_path))
+        """
+        super().__init__(model=model, preprocessor=preprocessor, **kwargs)
+
+        torch_dtype = kwargs.get('torch_dtype', torch.float32)
+        self._device = getattr(
+            kwargs, 'device',
+            torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
+        logger.info('load image editing pipeline done')
+        scheduler = DDIMScheduler.from_pretrained(
+            os.path.join(model, 'stable-diffusion-v1-4'),
+            subfolder='scheduler')
+        self.pipeline = _MasaCtrlPipeline.from_pretrained(
+            os.path.join(model, 'stable-diffusion-v1-4'),
+            scheduler=scheduler,
+            torch_dtype=torch_dtype,
+            use_safetensors=True).to(self._device)
+
+    def preprocess(self, input: Dict[str, Any]) -> Dict[str, Any]:
+        img = LoadImage.convert_to_img(input.get('img'))
+        test_transforms = transforms.Compose(
+            [transforms.ToTensor(),
+             transforms.Normalize([0.5], [0.5])])  # [-1, 1]
+        img = test_transforms(img).unsqueeze(0)
+        img = F.interpolate(img, (512, 512))
+        input['img'] = img.to(self._device)
+        return input
+
+    def forward(self, input: Dict[str, Any]) -> Dict[str, Any]:
+
+        if not isinstance(input, dict):
+            raise ValueError(
+                f'Expected the input to be a dictionary, but got {type(input)}'
+            )
+        prompts = input.get('prompts')
+        start_code, latents_list = self.pipeline.invert(
+            input.get('img'),
+            prompts[0],
+            guidance_scale=7.5,
+            num_inference_steps=50,
+            return_intermediates=True)
+        start_code = start_code.expand(len(prompts), -1, -1, -1)
+        STEP, LAYER = 4, 10
+        editor = MutualSelfAttentionControl(STEP, LAYER)
+        regiter_attention_editor_diffusers(self.pipeline, editor)
+
+        # inference the synthesized image
+        output = self.pipeline(
+            prompts,
+            latents=start_code,
+            guidance_scale=input.get('guidance_scale', 7.5),
+        )[-1:]
+
+        return {'output_tensor': output}
+
+    def postprocess(self, input: Dict[str, Any]) -> Dict[str, Any]:
+        output_img = (input['output_tensor'].squeeze(0) * 255).cpu().permute(
+            1, 2, 0).numpy().astype('uint8')
+        return {OutputKeys.OUTPUT_IMG: output_img[:, :, ::-1]}
+
+
+class _MasaCtrlPipeline(StableDiffusionPipeline):
+
+    def next_step(
+        self,
+        model_output: torch.FloatTensor,
+        timestep: int,
+        x: torch.FloatTensor,
+        eta=0,
+        verbose=False,
+    ):
+        """
+        Inverse sampling for DDIM Inversion
+        x_t -> x_(t+1)
+        """
+        if verbose:
+            print('timestep: ', timestep)
+        next_step = timestep
+        timestep = min(
+            timestep - self.scheduler.config.num_train_timesteps
+            // self.scheduler.num_inference_steps, 999)
+        alpha_prod_t = self.scheduler.alphas_cumprod[
+            timestep] if timestep >= 0 else self.scheduler.final_alpha_cumprod
+        alpha_prod_t_next = self.scheduler.alphas_cumprod[next_step]
+        beta_prod_t = 1 - alpha_prod_t
+        pred_x0 = (x - beta_prod_t**0.5 * model_output) / alpha_prod_t**0.5
+        pred_dir = (1 - alpha_prod_t_next)**0.5 * model_output
+        x_next = alpha_prod_t_next**0.5 * pred_x0 + pred_dir
+        return x_next, pred_x0
+
+    def step(
+        self,
+        model_output: torch.FloatTensor,
+        timestep: int,
+        x: torch.FloatTensor,
+        eta: float = 0.0,
+        verbose=False,
+    ):
+        """
+        predict the sample the next step in the denoise process.
+        x_t -> x_(t-1)
+        """
+        prev_timestep = timestep - self.scheduler.config.num_train_timesteps // self.scheduler.num_inference_steps
+        alpha_prod_t = self.scheduler.alphas_cumprod[timestep]
+        alpha_prod_t_prev = self.scheduler.alphas_cumprod[
+            prev_timestep] if prev_timestep > 0 else self.scheduler.final_alpha_cumprod
+        beta_prod_t = 1 - alpha_prod_t
+        pred_x0 = (x - beta_prod_t**0.5 * model_output) / alpha_prod_t**0.5
+        pred_dir = (1 - alpha_prod_t_prev)**0.5 * model_output
+        x_prev = alpha_prod_t_prev**0.5 * pred_x0 + pred_dir
+        return x_prev, pred_x0
+
+    @torch.no_grad()
+    def image2latent(self, image):
+        DEVICE = self._execution_device
+        if type(image) is Image:
+            image = np.array(image)
+            image = torch.from_numpy(image).float() / 127.5 - 1
+            image = image.permute(2, 0, 1).unsqueeze(0).to(DEVICE)
+        # input image density range [-1, 1]
+        latents = self.vae.encode(image)['latent_dist'].mean
+        latents = latents * 0.18215
+        return latents
+
+    @torch.no_grad()
+    def latent2image(self, latents, return_type='pt'):
+        latents = 1 / 0.18215 * latents.detach()
+        image = self.vae.decode(latents)['sample']
+        if return_type == 'np':
+            image = (image / 2 + 0.5).clamp(0, 1)
+            image = image.cpu().permute(0, 2, 3, 1).numpy()[0]
+            image = (image * 255).astype(np.uint8)
+        elif return_type == 'pt':
+            image = (image / 2 + 0.5).clamp(0, 1)
+
+        return image
+
+    @torch.no_grad()
+    def __call__(self,
+                 prompt,
+                 batch_size=1,
+                 height=512,
+                 width=512,
+                 num_inference_steps=50,
+                 guidance_scale=7.5,
+                 eta=0.0,
+                 latents=None,
+                 unconditioning=None,
+                 neg_prompt=None,
+                 ref_intermediate_latents=None,
+                 return_intermediates=False,
+                 **kwds):
+        DEVICE = self._execution_device
+        if isinstance(prompt, list):
+            batch_size = len(prompt)
+        elif isinstance(prompt, str):
+            if batch_size > 1:
+                prompt = [prompt] * batch_size
+
+        # text embeddings
+        text_input = self.tokenizer(
+            prompt, padding='max_length', max_length=77, return_tensors='pt')
+
+        text_embeddings = self.text_encoder(text_input.input_ids.to(DEVICE))[0]
+        print('input text embeddings :', text_embeddings.shape)
+
+        # define initial latents
+        latents_shape = (batch_size, self.unet.in_channels, height // 8,
+                         width // 8)
+        if latents is None:
+            latents = torch.randn(latents_shape, device=DEVICE)
+        else:
+            assert latents.shape == latents_shape, f'The shape of input latent tensor {latents.shape} should equal ' \
+                                                   f'to predefined one.'
+
+        # unconditional embedding for classifier free guidance
+        if guidance_scale > 1.:
+            if neg_prompt:
+                uc_text = neg_prompt
+            else:
+                uc_text = ''
+            unconditional_input = self.tokenizer(
+                [uc_text] * batch_size,
+                padding='max_length',
+                max_length=77,
+                return_tensors='pt')
+            unconditional_embeddings = self.text_encoder(
+                unconditional_input.input_ids.to(DEVICE))[0]
+            text_embeddings = torch.cat(
+                [unconditional_embeddings, text_embeddings], dim=0)
+
+        print('latents shape: ', latents.shape)
+        # iterative sampling
+        self.scheduler.set_timesteps(num_inference_steps)
+        latents_list = [latents]
+        pred_x0_list = [latents]
+        for i, t in enumerate(
+                tqdm(self.scheduler.timesteps, desc='DDIM Sampler')):
+            if ref_intermediate_latents is not None:
+                # note that the batch_size >= 2
+                latents_ref = ref_intermediate_latents[-1 - i]
+                _, latents_cur = latents.chunk(2)
+                latents = torch.cat([latents_ref, latents_cur])
+
+            if guidance_scale > 1.:
+                model_inputs = torch.cat([latents] * 2)
+            else:
+                model_inputs = latents
+            if unconditioning is not None and isinstance(unconditioning, list):
+                _, text_embeddings = text_embeddings.chunk(2)
+                text_embeddings = torch.cat([
+                    unconditioning[i].expand(*text_embeddings.shape),
+                    text_embeddings
+                ])
+            # predict the noise
+            noise_pred = self.unet(
+                model_inputs, t, encoder_hidden_states=text_embeddings).sample
+            if guidance_scale > 1.:
+                noise_pred_uncon, noise_pred_con = noise_pred.chunk(2, dim=0)
+                noise_pred = noise_pred_uncon + guidance_scale * (
+                    noise_pred_con - noise_pred_uncon)
+            # compute the previous noise sample x_t -> x_t-1
+            latents, pred_x0 = self.step(noise_pred, t, latents)
+            latents_list.append(latents)
+            pred_x0_list.append(pred_x0)
+
+        image = self.latent2image(latents, return_type='pt')
+        if return_intermediates:
+            pred_x0_list = [
+                self.latent2image(img, return_type='pt')
+                for img in pred_x0_list
+            ]
+            latents_list = [
+                self.latent2image(img, return_type='pt')
+                for img in latents_list
+            ]
+            return image, pred_x0_list, latents_list
+        return image
+
+    @torch.no_grad()
+    def invert(self,
+               image: torch.Tensor,
+               prompt,
+               num_inference_steps=50,
+               guidance_scale=7.5,
+               eta=0.0,
+               return_intermediates=False,
+               **kwds):
+        """
+        invert a real image into noise map with determinisc DDIM inversion
+        """
+        DEVICE = self._execution_device
+        batch_size = image.shape[0]
+        if isinstance(prompt, list):
+            if batch_size == 1:
+                image = image.expand(len(prompt), -1, -1, -1)
+        elif isinstance(prompt, str):
+            if batch_size > 1:
+                prompt = [prompt] * batch_size
+
+        # text embeddings
+        text_input = self.tokenizer(
+            prompt, padding='max_length', max_length=77, return_tensors='pt')
+        text_embeddings = self.text_encoder(text_input.input_ids.to(DEVICE))[0]
+        print('input text embeddings :', text_embeddings.shape)
+        # define initial latents
+        latents = self.image2latent(image)
+        start_latents = latents
+
+        # unconditional embedding for classifier free guidance
+        if guidance_scale > 1.:
+            unconditional_input = self.tokenizer(
+                [''] * batch_size,
+                padding='max_length',
+                max_length=77,
+                return_tensors='pt')
+            unconditional_embeddings = self.text_encoder(
+                unconditional_input.input_ids.to(DEVICE))[0]
+            text_embeddings = torch.cat(
+                [unconditional_embeddings, text_embeddings], dim=0)
+
+        print('latents shape: ', latents.shape)
+        self.scheduler.set_timesteps(num_inference_steps)
+        print('Valid timesteps: ', reversed(self.scheduler.timesteps))
+        latents_list = [latents]
+        pred_x0_list = [latents]
+        for i, t in enumerate(
+                tqdm(
+                    reversed(self.scheduler.timesteps),
+                    desc='DDIM Inversion')):
+            if guidance_scale > 1.:
+                model_inputs = torch.cat([latents] * 2)
+            else:
+                model_inputs = latents
+
+            # predict the noise
+            noise_pred = self.unet(
+                model_inputs, t, encoder_hidden_states=text_embeddings).sample
+            if guidance_scale > 1.:
+                noise_pred_uncon, noise_pred_con = noise_pred.chunk(2, dim=0)
+                noise_pred = noise_pred_uncon + guidance_scale * (
+                    noise_pred_con - noise_pred_uncon)
+            # compute the previous noise sample x_t-1 -> x_t
+            latents, pred_x0 = self.next_step(noise_pred, t, latents)
+            latents_list.append(latents)
+            pred_x0_list.append(pred_x0)
+
+        if return_intermediates:
+            return latents, latents_list
+        return latents, start_latents

modelscope/utils/constant.py

@@ -85,6 +85,7 @@ class CVTasks(object):
     image_paintbyexample = 'image-paintbyexample'
     image_skychange = 'image-skychange'
     image_demoireing = 'image-demoireing'
+    image_editing = 'image-editing'
 
     # image generation
     image_to_image_translation = 'image-to-image-translation'

tests/pipelines/test_image_editing_masactrl.py

@@ -0,0 +1,48 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+
+import unittest
+
+import cv2
+
+from modelscope.hub.snapshot_download import snapshot_download
+from modelscope.outputs import OutputKeys
+from modelscope.pipelines import pipeline
+from modelscope.pipelines.cv import ImageEditingPipeline
+from modelscope.utils.constant import Tasks
+from modelscope.utils.test_utils import test_level
+
+
+class ImageEditingTest(unittest.TestCase):
+
+    def setUp(self) -> None:
+        self.task = Tasks.image_editing
+        self.model_id = 'damo/cv_masactrl_image-editing'
+        prompts = [
+            '',  # source prompt
+            'a photo of a running corgi'  # target prompt
+        ]
+        img = 'https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/public/ModelScope/test/images/corgi.jpg'
+        self.input = {'img': img, 'prompts': prompts}
+        self.output_image_path = './result.png'
+
+    @unittest.skipUnless(test_level() >= 2, 'skip test in current test level')
+    def test_run_by_direct_model_download(self):
+        cache_path = snapshot_download(self.model_id)
+        pipeline = ImageEditingPipeline(cache_path)
+        pipeline.group_key = self.task
+        edited_img = pipeline(input=self.input)[OutputKeys.OUTPUT_IMG]  # BGR
+        cv2.imwrite(self.output_image_path, edited_img)
+        print('MasaCtrl pipeline: the edited image path is {}'.format(
+            self.output_image_path))
+
+    @unittest.skipUnless(test_level() >= 1, 'skip test in current test level')
+    def test_run_with_model_name(self):
+        pipeline_ins = pipeline(task=Tasks.image_editing, model=self.model_id)
+        edited_img = pipeline_ins(self.input)[OutputKeys.OUTPUT_IMG]  # BGR
+        cv2.imwrite(self.output_image_path, edited_img)
+        print('MasaCtrl pipeline: the edited image path is {}'.format(
+            self.output_image_path))
+
+
+if __name__ == '__main__':
+    unittest.main()