From 7aef73a761cee1791c2b881ffd45de60a311395f Mon Sep 17 00:00:00 2001 From: "huizheng.hz" Date: Thu, 24 Aug 2023 21:20:47 +0800 Subject: [PATCH] [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 | 3 + modelscope/models/cv/__init__.py | 9 +- .../models/cv/image_editing/__init__.py | 23 ++ .../models/cv/image_editing/masactrl.py | 77 ++++ .../models/cv/image_editing/masactrl_utils.py | 124 ++++++ modelscope/outputs/outputs.py | 1 + modelscope/pipeline_inputs.py | 4 + modelscope/pipelines/cv/__init__.py | 2 + .../pipelines/cv/image_editing_pipeline.py | 365 ++++++++++++++++++ modelscope/utils/constant.py | 1 + .../pipelines/test_image_editing_masactrl.py | 48 +++ 11 files changed, 653 insertions(+), 4 deletions(-) create mode 100644 modelscope/models/cv/image_editing/__init__.py create mode 100644 modelscope/models/cv/image_editing/masactrl.py create mode 100644 modelscope/models/cv/image_editing/masactrl_utils.py create mode 100644 modelscope/pipelines/cv/image_editing_pipeline.py create mode 100644 tests/pipelines/test_image_editing_masactrl.py diff --git a/modelscope/metainfo.py b/modelscope/metainfo.py index 92702be8..7bb539c0 100644 --- a/modelscope/metainfo.py +++ b/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: diff --git a/modelscope/models/cv/__init__.py b/modelscope/models/cv/__init__.py index 39acec69..c0614a5c 100644 --- a/modelscope/models/cv/__init__.py +++ b/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, diff --git a/modelscope/models/cv/image_editing/__init__.py b/modelscope/models/cv/image_editing/__init__.py new file mode 100644 index 00000000..35341a18 --- /dev/null +++ b/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={}, + ) diff --git a/modelscope/models/cv/image_editing/masactrl.py b/modelscope/models/cv/image_editing/masactrl.py new file mode 100644 index 00000000..e1e505d1 --- /dev/null +++ b/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 diff --git a/modelscope/models/cv/image_editing/masactrl_utils.py b/modelscope/models/cv/image_editing/masactrl_utils.py new file mode 100644 index 00000000..a59e987f --- /dev/null +++ b/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 diff --git a/modelscope/outputs/outputs.py b/modelscope/outputs/outputs.py index f6ef0f6b..f9996b7a 100644 --- a/modelscope/outputs/outputs.py +++ b/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], diff --git a/modelscope/pipeline_inputs.py b/modelscope/pipeline_inputs.py index 712f6ff9..61d1043b 100644 --- a/modelscope/pipeline_inputs.py +++ b/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 + } } diff --git a/modelscope/pipelines/cv/__init__.py b/modelscope/pipelines/cv/__init__.py index 71d63d10..2cd7ede0 100644 --- a/modelscope/pipelines/cv/__init__.py +++ b/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': diff --git a/modelscope/pipelines/cv/image_editing_pipeline.py b/modelscope/pipelines/cv/image_editing_pipeline.py new file mode 100644 index 00000000..15e21eaf --- /dev/null +++ b/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 diff --git a/modelscope/utils/constant.py b/modelscope/utils/constant.py index b9b6f811..53006477 100644 --- a/modelscope/utils/constant.py +++ b/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' diff --git a/tests/pipelines/test_image_editing_masactrl.py b/tests/pipelines/test_image_editing_masactrl.py new file mode 100644 index 00000000..beaabfca --- /dev/null +++ b/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()