initial commit

.gitignore

@@ -0,0 +1,2 @@
+venv
+.idea

app.py

@@ -0,0 +1,76 @@
+import os
+import gradio as gr
+
+# Set necessary environment variables for ACE++
+os.environ["FLUX_FILL_PATH"] = "hf://black-forest-labs/FLUX.1-Fill-dev"
+os.environ["PORTRAIT_MODEL_PATH"] = "ms://iic/ACE_Plus@portrait/comfyui_portrait_lora64.safetensors"
+os.environ["SUBJECT_MODEL_PATH"] = "ms://iic/ACE_Plus@subject/comfyui_subject_lora16.safetensors"
+os.environ["LOCAL_MODEL_PATH"] = "ms://iic/ACE_Plus@local_editing/comfyui_local_lora16.safetensors"
+
+# Import ACEInference and Config from the ACE_plus repo
+from inference.ace_plus_inference import ACEInference
+from scepter.modules.utils.config import Config
+
+# Define a minimal configuration dictionary.
+# Adjust the "MODEL" field as required by your ACE++ setup.
+config_dict = {
+    "MODEL": {
+        "type": "YourACEModelType",  # Replace with the actual model type string used in ACE_plus.
+        "pretrained_path": os.getenv("PORTRAIT_MODEL_PATH")
+    },
+    "MAX_SEQ_LEN": 77,
+    "SAMPLE_ARGS": {
+        "prompt": "Face swap"
+    },
+    "DTYPE": "bfloat16"
+}
+cfg = Config(config_dict)
+
+# Instantiate the ACEInference object.
+ace_infer = ACEInference(cfg)
+
+def face_swap_app(target_img, face_img):
+    """
+    Swaps the face in the target image using the provided face image via ACE++.
+
+    Parameters:
+      target_img: The image in which you want to swap a face.
+      face_img: The reference face image to insert.
+
+    Returns:
+      The output image after applying ACE++ face swapping.
+    """
+    # For ACEInference, we pass:
+    # - reference_image: the target image,
+    # - edit_image: the new face image,
+    # - edit_mask: set to None so the image processor will create it,
+    # - prompt: "Face swap" instructs the model to perform face swapping.
+    # Other parameters (output dimensions, sampler, etc.) are set here as desired.
+    output_img, edit_image, change_image, mask, seed = ace_infer(
+        reference_image=target_img,
+        edit_image=face_img,
+        edit_mask=None,          # No manual mask provided; let ACE++ handle it
+        prompt="Face swap",
+        output_height=1024,
+        output_width=1024,
+        sampler='flow_euler',
+        sample_steps=28,
+        guide_scale=50,
+        seed=-1                  # Use a random seed if not specified
+    )
+    return output_img
+
+# Create the Gradio interface.
+iface = gr.Interface(
+    fn=face_swap_app,
+    inputs=[
+        gr.Image(type="pil", label="Target Image"),
+        gr.Image(type="pil", label="Face Image")
+    ],
+    outputs=gr.Image(type="pil", label="Swapped Face Output"),
+    title="ACE++ Face Swap Demo",
+    description="Upload a target image and a face image to swap the face using the ACE++ model."
+)
+
+if __name__ == "__main__":
+    iface.launch()

inference/__init__.py

@@ -0,0 +1,2 @@
+from .ace_plus_diffusers import ACEPlusDiffuserInference
+from .ace_plus_inference import ACEInference

inference/ace_plus_diffusers.py

@@ -0,0 +1,121 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import random
+from collections import OrderedDict
+
+import torch, os
+from diffusers import FluxFillPipeline
+from scepter.modules.utils.config import Config
+from scepter.modules.utils.distribute import we
+from scepter.modules.utils.file_system import FS
+from scepter.modules.utils.logger import get_logger
+from transformers import T5TokenizerFast
+from .utils import ACEPlusImageProcessor
+
+class ACEPlusDiffuserInference():
+    def __init__(self, logger=None):
+        if logger is None:
+            logger = get_logger(name='ace_plus')
+        self.logger = logger
+        self.input = {}
+
+    def load_default(self, cfg):
+        if cfg is not None:
+            self.input_cfg = {k.lower(): v for k, v in cfg.INPUT.items()}
+            self.input = {k.lower(): dict(v).get('DEFAULT', None) if isinstance(v, (dict, OrderedDict, Config)) else v for k, v in cfg.INPUT.items()}
+            self.output = {k.lower(): v for k, v in cfg.OUTPUT.items()}
+
+    def init_from_cfg(self, cfg):
+        self.max_seq_len = cfg.get("MAX_SEQ_LEN", 4096)
+        self.image_processor = ACEPlusImageProcessor(max_seq_len=self.max_seq_len)
+
+        local_folder = FS.get_dir_to_local_dir(cfg.MODEL.PRETRAINED_MODEL)
+
+        self.pipe = FluxFillPipeline.from_pretrained(local_folder, torch_dtype=torch.bfloat16).to(we.device_id)
+
+        tokenizer_2 = T5TokenizerFast.from_pretrained(os.path.join(local_folder, "tokenizer_2"),
+                                                      additional_special_tokens=["{image}"])
+        self.pipe.tokenizer_2 = tokenizer_2
+        self.load_default(cfg.DEFAULT_PARAS)
+
+    def prepare_input(self,
+                      image,
+                      mask,
+                      batch_size=1,
+                      dtype = torch.bfloat16,
+                      num_images_per_prompt=1,
+                      height=512,
+                      width=512,
+                      generator=None):
+        num_channels_latents = self.pipe.vae.config.latent_channels
+        # import pdb;pdb.set_trace()
+        mask, masked_image_latents = self.pipe.prepare_mask_latents(
+            mask.unsqueeze(0),
+            image.unsqueeze(0).to(we.device_id, dtype = dtype),
+            batch_size,
+            num_channels_latents,
+            num_images_per_prompt,
+            height,
+            width,
+            dtype,
+            we.device_id,
+            generator,
+        )
+        # import pdb;pdb.set_trace()
+        masked_image_latents = torch.cat((masked_image_latents, mask), dim=-1)
+        return masked_image_latents
+
+    @torch.no_grad()
+    def __call__(self,
+                 reference_image=None,
+                 edit_image=None,
+                 edit_mask=None,
+                 prompt='',
+                 task=None,
+                 output_height=1024,
+                 output_width=1024,
+                 sampler='flow_euler',
+                 sample_steps=28,
+                 guide_scale=50,
+                 lora_path=None,
+                 seed=-1,
+                 tar_index=0,
+                 align=0,
+                 repainting_scale=0,
+                 **kwargs):
+        if isinstance(prompt, str):
+            prompt = [prompt]
+        seed = seed if seed >= 0 else random.randint(0, 2 ** 32 - 1)
+        # edit_image, edit_mask, change_image, content_image, out_h, out_w, slice_w
+        image, mask, _, _, out_h, out_w, slice_w = self.image_processor.preprocess(reference_image, edit_image, edit_mask,
+                                                                             width = output_width,
+                                                                             height = output_height,
+                                                                             repainting_scale = repainting_scale)
+        h, w = image.shape[1:]
+        generator = torch.Generator("cpu").manual_seed(seed)
+        masked_image_latents = self.prepare_input(image, mask,
+                                               batch_size=len(prompt) , height=h, width=w, generator = generator)
+
+        if lora_path is not None:
+            with FS.get_from(lora_path) as local_path:
+                self.pipe.load_lora_weights(local_path)
+
+
+
+        image = self.pipe(
+            prompt=prompt,
+            masked_image_latents=masked_image_latents,
+            height=h,
+            width=w,
+            guidance_scale=guide_scale,
+            num_inference_steps=sample_steps,
+            max_sequence_length=512,
+            generator=generator
+        ).images[0]
+        if lora_path is not None:
+            self.pipe.unload_lora_weights()
+        return self.image_processor.postprocess(image, slice_w, out_w, out_h), seed
+
+
+if __name__ == '__main__':
+    pass

inference/ace_plus_inference.py

@@ -0,0 +1,83 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import random
+from collections import OrderedDict
+
+import torch, numpy as np
+from PIL import Image
+from scepter.modules.model.registry import MODELS
+from scepter.modules.utils.config import Config
+from scepter.modules.utils.distribute import we
+from .registry import BaseInference, INFERENCES
+from .utils import ACEPlusImageProcessor
+
+@INFERENCES.register_class()
+class ACEInference(BaseInference):
+    '''
+        reuse the ldm code
+    '''
+    def __init__(self, cfg, logger=None):
+        super().__init__(cfg, logger)
+        self.pipe = MODELS.build(cfg.MODEL, logger=self.logger).eval().to(we.device_id)
+        self.image_processor = ACEPlusImageProcessor(max_seq_len=cfg.MAX_SEQ_LEN)
+        self.input = {k.lower(): dict(v).get('DEFAULT', None) if isinstance(v, (dict, OrderedDict, Config)) else v for
+                      k, v in cfg.SAMPLE_ARGS.items()}
+        self.dtype = getattr(torch, cfg.get("DTYPE", "bfloat16"))
+    @torch.no_grad()
+    def __call__(self,
+                 reference_image=None,
+                 edit_image=None,
+                 edit_mask=None,
+                 prompt='',
+                 edit_type=None,
+                 output_height=1024,
+                 output_width=1024,
+                 sampler='flow_euler',
+                 sample_steps=28,
+                 guide_scale=50,
+                 lora_path=None,
+                 seed=-1,
+                 repainting_scale=0,
+                 use_change=False,
+                 keep_pixels=False,
+                 keep_pixels_rate=0.8,
+                 **kwargs):
+        # convert the input info to the input of ldm.
+        if isinstance(prompt, str):
+            prompt = [prompt]
+        seed = seed if seed >= 0 else random.randint(0, 2 ** 24 - 1)
+        image, mask, change_image, content_image, out_h, out_w, slice_w = self.image_processor.preprocess(reference_image, edit_image, edit_mask,
+                                                                             height=output_height, width=output_width,
+                                                                             repainting_scale=repainting_scale,
+                                                                             keep_pixels=keep_pixels,
+                                                                             keep_pixels_rate=keep_pixels_rate,
+                                                                             use_change = use_change)
+        change_image = [None] if change_image is None else [change_image.to(we.device_id)]
+        image, mask = [image.to(we.device_id)], [mask.to(we.device_id)]
+
+        (src_image_list, src_mask_list, modify_image_list,
+         edit_id, prompt) = [image], [mask], [change_image], [[0]], [prompt]
+
+        with torch.amp.autocast(enabled=True, dtype=self.dtype, device_type='cuda'):
+            out_image = self.pipe(
+                src_image_list=src_image_list,
+                modify_image_list= modify_image_list,
+                src_mask_list=src_mask_list,
+                edit_id=edit_id,
+                image=image,
+                image_mask=mask,
+                prompt=prompt,
+                sampler='flow_euler',
+                sample_steps=sample_steps,
+                seed=seed,
+                guide_scale=guide_scale,
+                show_process=True,
+            )
+        imgs = [x_i['reconstruct_image'].float().permute(1, 2, 0).cpu().numpy()
+            for x_i in out_image
+        ]
+        imgs = [Image.fromarray((img * 255).astype(np.uint8)) for img in imgs]
+        edit_image = Image.fromarray((torch.clamp(image[0] / 2 + 0.5, min=0.0, max=1.0)*255).float().permute(1, 2, 0).cpu().numpy().astype(np.uint8))
+        change_image = Image.fromarray((torch.clamp(change_image[0] / 2 + 0.5, min=0.0, max=1.0)*255).float().permute(1, 2, 0).cpu().numpy().astype(np.uint8))
+        mask = Image.fromarray((mask[0] * 255).squeeze(0).cpu().numpy().astype(np.uint8))
+        return self.image_processor.postprocess(imgs[0], slice_w, out_w, out_h), edit_image, change_image, mask, seed

inference/registry.py

@@ -0,0 +1,228 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+
+import torch
+from PIL.Image import Image
+from collections import OrderedDict
+from scepter.modules.utils.distribute import we
+from scepter.modules.utils.config import Config
+from scepter.modules.utils.logger import get_logger
+from scepter.studio.utils.env import get_available_memory
+from scepter.modules.model.registry import MODELS, BACKBONES, EMBEDDERS
+from scepter.modules.utils.registry import Registry, build_from_config
+def get_model(model_tuple):
+    assert 'model' in model_tuple
+    return model_tuple['model']
+
+class BaseInference():
+    '''
+        support to load the components dynamicly.
+        create and load model when run this model at the first time.
+    '''
+    def __init__(self, cfg, logger=None):
+        if logger is None:
+            logger = get_logger(name='scepter')
+        self.logger = logger
+        self.name = cfg.NAME
+
+    def init_from_modules(self, modules):
+        for k, v in modules.items():
+            self.__setattr__(k, v)
+
+    def infer_model(self, cfg, module_paras=None):
+        module = {
+            'model': None,
+            'cfg': cfg,
+            'device': 'offline',
+            'name': cfg.NAME,
+            'function_info': {},
+            'paras': {}
+        }
+        if module_paras is None:
+            return module
+        function_info = {}
+        paras = {
+            k.lower(): v
+            for k, v in module_paras.get('PARAS', {}).items()
+        }
+        for function in module_paras.get('FUNCTION', []):
+            input_dict = {}
+            for inp in function.get('INPUT', []):
+                if inp.lower() in self.input:
+                    input_dict[inp.lower()] = self.input[inp.lower()]
+            function_info[function.NAME] = {
+                'dtype': function.get('DTYPE', 'float32'),
+                'input': input_dict
+            }
+        module['paras'] = paras
+        module['function_info'] = function_info
+        return module
+
+    def init_from_ckpt(self, path, model, ignore_keys=list()):
+        if path.endswith('safetensors'):
+            from safetensors.torch import load_file as load_safetensors
+            sd = load_safetensors(path)
+        else:
+            sd = torch.load(path, map_location='cpu', weights_only=True)
+
+        new_sd = OrderedDict()
+        for k, v in sd.items():
+            ignored = False
+            for ik in ignore_keys:
+                if ik in k:
+                    if we.rank == 0:
+                        self.logger.info(
+                            'Ignore key {} from state_dict.'.format(k))
+                    ignored = True
+                    break
+            if not ignored:
+                new_sd[k] = v
+
+        missing, unexpected = model.load_state_dict(new_sd, strict=False)
+        if we.rank == 0:
+            self.logger.info(
+                f'Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys'
+            )
+            if len(missing) > 0:
+                self.logger.info(f'Missing Keys:\n {missing}')
+            if len(unexpected) > 0:
+                self.logger.info(f'\nUnexpected Keys:\n {unexpected}')
+
+    def load(self, module):
+        if module['device'] == 'offline':
+            from scepter.modules.utils.import_utils import LazyImportModule
+            if (LazyImportModule.get_module_type(('MODELS', module['cfg'].NAME)) or
+                    module['cfg'].NAME in MODELS.class_map):
+                model = MODELS.build(module['cfg'], logger=self.logger).eval()
+            elif (LazyImportModule.get_module_type(('BACKBONES', module['cfg'].NAME)) or
+                    module['cfg'].NAME in BACKBONES.class_map):
+                model = BACKBONES.build(module['cfg'],
+                                        logger=self.logger).eval()
+            elif (LazyImportModule.get_module_type(('EMBEDDERS', module['cfg'].NAME)) or
+                    module['cfg'].NAME in EMBEDDERS.class_map):
+                model = EMBEDDERS.build(module['cfg'],
+                                        logger=self.logger).eval()
+            else:
+                raise NotImplementedError
+            if 'DTYPE' in module['cfg'] and module['cfg']['DTYPE'] is not None:
+                model = model.to(getattr(torch, module['cfg'].DTYPE))
+            if module['cfg'].get('RELOAD_MODEL', None):
+                self.init_from_ckpt(module['cfg'].RELOAD_MODEL, model)
+            module['model'] = model
+            module['device'] = 'cpu'
+        if module['device'] == 'cpu':
+            module['device'] = we.device_id
+            module['model'] = module['model'].to(we.device_id)
+        return module
+
+    def unload(self, module):
+        if module is None:
+            return module
+        mem = get_available_memory()
+        free_mem = int(mem['available'] / (1024**2))
+        total_mem = int(mem['total'] / (1024**2))
+        if free_mem < 0.5 * total_mem:
+            if module['model'] is not None:
+                module['model'] = module['model'].to('cpu')
+                del module['model']
+            module['model'] = None
+            module['device'] = 'offline'
+            print('delete module')
+        else:
+            if module['model'] is not None:
+                module['model'] = module['model'].to('cpu')
+                module['device'] = 'cpu'
+            else:
+                module['device'] = 'offline'
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            torch.cuda.ipc_collect()
+        return module
+
+    def dynamic_load(self, module=None, name=''):
+        self.logger.info('Loading {} model'.format(name))
+        if name == 'all':
+            for subname in self.loaded_model_name:
+                self.loaded_model[subname] = self.dynamic_load(
+                    getattr(self, subname), subname)
+        elif name in self.loaded_model_name:
+            if name in self.loaded_model:
+                if module['cfg'] != self.loaded_model[name]['cfg']:
+                    self.unload(self.loaded_model[name])
+                    module = self.load(module)
+                    self.loaded_model[name] = module
+                    return module
+                elif module['device'] == 'cpu' or module['device'] == 'offline':
+                    module = self.load(module)
+                    return module
+                else:
+                    return module
+            else:
+                module = self.load(module)
+                self.loaded_model[name] = module
+                return module
+        else:
+            return self.load(module)
+
+    def dynamic_unload(self, module=None, name='', skip_loaded=False):
+        self.logger.info('Unloading {} model'.format(name))
+        if name == 'all':
+            for name, module in self.loaded_model.items():
+                module = self.unload(self.loaded_model[name])
+                self.loaded_model[name] = module
+        elif name in self.loaded_model_name:
+            if name in self.loaded_model:
+                if not skip_loaded:
+                    module = self.unload(self.loaded_model[name])
+                    self.loaded_model[name] = module
+            else:
+                self.unload(module)
+        else:
+            self.unload(module)
+
+    def load_default(self, cfg):
+        module_paras = {}
+        if cfg is not None:
+            self.paras = cfg.PARAS
+            self.input_cfg = {k.lower(): v for k, v in cfg.INPUT.items()}
+            self.input = {k.lower(): dict(v).get('DEFAULT', None) if isinstance(v, (dict, OrderedDict, Config)) else v for k, v in cfg.INPUT.items()}
+            self.output = {k.lower(): v for k, v in cfg.OUTPUT.items()}
+            module_paras = cfg.MODULES_PARAS
+        return module_paras
+
+    def load_image(self, image, num_samples=1):
+        if isinstance(image, torch.Tensor):
+            pass
+        elif isinstance(image, Image):
+            pass
+        elif isinstance(image, Image):
+            pass
+
+    def get_function_info(self, module, function_name=None):
+        all_function = module['function_info']
+        if function_name in all_function:
+            return function_name, all_function[function_name]['dtype']
+        if function_name is None and len(all_function) == 1:
+            for k, v in all_function.items():
+                return k, v['dtype']
+
+    @torch.no_grad()
+    def __call__(self,
+                 input,
+                 **kwargs):
+        return
+
+def build_inference(cfg, registry, logger=None, *args, **kwargs):
+    """ After build model, load pretrained model if exists key `pretrain`.
+
+    pretrain (str, dict): Describes how to load pretrained model.
+        str, treat pretrain as model path;
+        dict: should contains key `path`, and other parameters token by function load_pretrained();
+    """
+    if not isinstance(cfg, Config):
+        raise TypeError(f'Config must be type dict, got {type(cfg)}')
+    model = build_from_config(cfg, registry, logger=logger, *args, **kwargs)
+    return model
+
+# reigister cls for diffusion.
+INFERENCES = Registry('INFERENCE', build_func=build_inference)

inference/utils.py

@@ -0,0 +1,132 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import math
+
+import torch
+import torchvision.transforms as T
+import numpy as np
+from scepter.modules.annotator.registry import ANNOTATORS
+from scepter.modules.utils.config import Config
+from PIL import Image
+
+
+def edit_preprocess(processor, device, edit_image, edit_mask):
+    if edit_image is None or processor is None:
+        return edit_image
+    processor = Config(cfg_dict=processor, load=False)
+    processor = ANNOTATORS.build(processor).to(device)
+    new_edit_image = processor(np.asarray(edit_image))
+    processor = processor.to("cpu")
+    del processor
+    new_edit_image = Image.fromarray(new_edit_image)
+    return Image.composite(new_edit_image, edit_image, edit_mask)
+
+class ACEPlusImageProcessor():
+    def __init__(self, max_aspect_ratio=4, d=16, max_seq_len=1024):
+        self.max_aspect_ratio = max_aspect_ratio
+        self.d = d
+        self.max_seq_len = max_seq_len
+        self.transforms = T.Compose([
+            T.ToTensor(),
+            T.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+        ])
+
+    def image_check(self, image):
+        if image is None:
+            return image
+        # preprocess
+        W, H = image.size
+        if H / W > self.max_aspect_ratio:
+            image = T.CenterCrop([int(self.max_aspect_ratio * W), W])(image)
+        elif W / H > self.max_aspect_ratio:
+            image = T.CenterCrop([H, int(self.max_aspect_ratio * H)])(image)
+        return self.transforms(image)
+
+
+    def preprocess(self,
+                   reference_image=None,
+                   edit_image=None,
+                   edit_mask=None,
+                   height=1024,
+                   width=1024,
+                   repainting_scale = 1.0,
+                   keep_pixels = False,
+                   keep_pixels_rate = 0.8,
+                   use_change = False):
+        reference_image = self.image_check(reference_image)
+        edit_image = self.image_check(edit_image)
+        # for reference generation
+        if edit_image is None:
+            edit_image = torch.zeros([3, height, width])
+            edit_mask = torch.ones([1, height, width])
+        else:
+            if edit_mask is None:
+                _, eH, eW = edit_image.shape
+                edit_mask = np.ones((eH, eW))
+            else:
+                edit_mask = np.asarray(edit_mask)
+                edit_mask = np.where(edit_mask > 128, 1, 0)
+            edit_mask = edit_mask.astype(
+                np.float32) if np.any(edit_mask) else np.ones_like(edit_mask).astype(
+                np.float32)
+            edit_mask = torch.tensor(edit_mask).unsqueeze(0)
+
+        edit_image = edit_image * (1 - edit_mask * repainting_scale)
+
+
+        out_h, out_w = edit_image.shape[-2:]
+
+        assert edit_mask is not None
+        if reference_image is not None:
+            _, H, W = reference_image.shape
+            _, eH, eW = edit_image.shape
+            if not keep_pixels:
+                # align height with edit_image
+                scale = eH / H
+                tH, tW = eH, int(W * scale)
+                reference_image = T.Resize((tH, tW), interpolation=T.InterpolationMode.BILINEAR, antialias=True)(
+                    reference_image)
+            else:
+                # padding
+                if H >= keep_pixels_rate * eH:
+                    tH = int(eH * keep_pixels_rate)
+                    scale = tH/H
+                    tW = int(W * scale)
+                    reference_image = T.Resize((tH, tW), interpolation=T.InterpolationMode.BILINEAR, antialias=True)(
+                        reference_image)
+                rH, rW = reference_image.shape[-2:]
+                delta_w = 0
+                delta_h = eH - rH
+                padding = (delta_w // 2, delta_h // 2, delta_w - (delta_w // 2), delta_h - (delta_h // 2))
+                reference_image = T.Pad(padding, fill=0, padding_mode="constant")(reference_image)
+            edit_image = torch.cat([reference_image, edit_image], dim=-1)
+            edit_mask = torch.cat([torch.zeros([1, reference_image.shape[1], reference_image.shape[2]]), edit_mask], dim=-1)
+            slice_w = reference_image.shape[-1]
+        else:
+            slice_w = 0
+
+        H, W = edit_image.shape[-2:]
+        scale = min(1.0, math.sqrt(self.max_seq_len * 2 / ((H / self.d) * (W / self.d))))
+        rH = int(H * scale) // self.d * self.d  # ensure divisible by self.d
+        rW = int(W * scale) // self.d * self.d
+        slice_w = int(slice_w * scale) // self.d * self.d
+
+        edit_image = T.Resize((rH, rW), interpolation=T.InterpolationMode.NEAREST_EXACT, antialias=True)(edit_image)
+        edit_mask = T.Resize((rH, rW), interpolation=T.InterpolationMode.NEAREST_EXACT, antialias=True)(edit_mask)
+        content_image = edit_image
+        if use_change:
+            change_image = edit_image * edit_mask
+            edit_image = edit_image * (1 - edit_mask)
+        else:
+            change_image = None
+        return edit_image, edit_mask, change_image, content_image, out_h, out_w, slice_w
+
+
+    def postprocess(self, image, slice_w, out_w, out_h):
+        w, h = image.size
+        if slice_w > 0:
+            output_image = image.crop((slice_w + 30, 0, w, h))
+            output_image = output_image.resize((out_w, out_h))
+        else:
+            output_image = image
+        return output_image

requirements.txt

@@ -0,0 +1,5 @@
+gradio
+scepter
+torch
+torchvision
+transformers