Spaces:
Running
on
Zero
Running
on
Zero
| # -*- coding: utf-8 -*- | |
| # Copyright (c) Alibaba, Inc. and its affiliates. | |
| import argparse | |
| import glob | |
| import io | |
| import os | |
| from PIL import Image | |
| from scepter.modules.transform.io import pillow_convert | |
| from scepter.modules.utils.config import Config | |
| from scepter.modules.utils.file_system import FS | |
| from examples.examples import all_examples | |
| from inference.ace_plus_diffusers import ACEPlusDiffuserInference | |
| inference_dict = { | |
| "ACE_DIFFUSER_PLUS": ACEPlusDiffuserInference | |
| } | |
| fs_list = [ | |
| Config(cfg_dict={"NAME": "HuggingfaceFs", "TEMP_DIR": "./cache"}, load=False), | |
| Config(cfg_dict={"NAME": "ModelscopeFs", "TEMP_DIR": "./cache"}, load=False), | |
| Config(cfg_dict={"NAME": "HttpFs", "TEMP_DIR": "./cache"}, load=False), | |
| Config(cfg_dict={"NAME": "LocalFs", "TEMP_DIR": "./cache"}, load=False), | |
| ] | |
| for one_fs in fs_list: | |
| FS.init_fs_client(one_fs) | |
| def run_one_case(pipe, | |
| input_image = None, | |
| input_mask = None, | |
| input_reference_image = None, | |
| save_path = "examples/output/example.png", | |
| instruction = "", | |
| output_h = 1024, | |
| output_w = 1024, | |
| seed = -1, | |
| sample_steps = None, | |
| guide_scale = None, | |
| repainting_scale = None, | |
| model_path = None, | |
| **kwargs): | |
| if input_image is not None: | |
| input_image = Image.open(io.BytesIO(FS.get_object(input_image))) | |
| input_image = pillow_convert(input_image, "RGB") | |
| if input_mask is not None: | |
| input_mask = Image.open(io.BytesIO(FS.get_object(input_mask))) | |
| input_mask = pillow_convert(input_mask, "L") | |
| if input_reference_image is not None: | |
| input_reference_image = Image.open(io.BytesIO(FS.get_object(input_reference_image))) | |
| input_reference_image = pillow_convert(input_reference_image, "RGB") | |
| image, seed = pipe( | |
| reference_image=input_reference_image, | |
| edit_image=input_image, | |
| edit_mask=input_mask, | |
| prompt=instruction, | |
| output_height=output_h, | |
| output_width=output_w, | |
| sampler='flow_euler', | |
| sample_steps=sample_steps or pipe.input.get("sample_steps", 28), | |
| guide_scale=guide_scale or pipe.input.get("guide_scale", 50), | |
| seed=seed, | |
| repainting_scale=repainting_scale or pipe.input.get("repainting_scale", 1.0), | |
| lora_path = model_path | |
| ) | |
| with FS.put_to(save_path) as local_path: | |
| image.save(local_path) | |
| return local_path, seed | |
| def run(): | |
| parser = argparse.ArgumentParser(description='Argparser for Scepter:\n') | |
| parser.add_argument('--instruction', | |
| dest='instruction', | |
| help='The instruction for editing or generating!', | |
| default="") | |
| parser.add_argument('--output_h', | |
| dest='output_h', | |
| help='The height of output image for generation tasks!', | |
| type=int, | |
| default=1024) | |
| parser.add_argument('--output_w', | |
| dest='output_w', | |
| help='The width of output image for generation tasks!', | |
| type=int, | |
| default=1024) | |
| parser.add_argument('--input_reference_image', | |
| dest='input_reference_image', | |
| help='The input reference image!', | |
| default=None | |
| ) | |
| parser.add_argument('--input_image', | |
| dest='input_image', | |
| help='The input image!', | |
| default=None | |
| ) | |
| parser.add_argument('--input_mask', | |
| dest='input_mask', | |
| help='The input mask!', | |
| default=None | |
| ) | |
| parser.add_argument('--save_path', | |
| dest='save_path', | |
| help='The save path for output image!', | |
| default='examples/output_images/output.png' | |
| ) | |
| parser.add_argument('--seed', | |
| dest='seed', | |
| help='The seed for generation!', | |
| type=int, | |
| default=-1) | |
| parser.add_argument('--step', | |
| dest='step', | |
| help='The sample step for generation!', | |
| type=int, | |
| default=None) | |
| parser.add_argument('--guide_scale', | |
| dest='guide_scale', | |
| help='The guide scale for generation!', | |
| type=int, | |
| default=None) | |
| parser.add_argument('--repainting_scale', | |
| dest='repainting_scale', | |
| help='The repainting scale for content filling generation!', | |
| type=int, | |
| default=None) | |
| parser.add_argument('--task_type', | |
| dest='task_type', | |
| choices=['portrait', 'subject', 'local_editing'], | |
| help="Choose the task type.", | |
| default='') | |
| parser.add_argument('--task_model', | |
| dest='task_model', | |
| help='The models list for different tasks!', | |
| default="./models/model_zoo.yaml") | |
| parser.add_argument('--infer_type', | |
| dest='infer_type', | |
| choices=['diffusers'], | |
| default='diffusers', | |
| help="Choose the inference scripts. 'native' refers to using the official implementation of ace++, " | |
| "while 'diffusers' refers to using the adaptation for diffusers") | |
| parser.add_argument('--cfg_folder', | |
| dest='cfg_folder', | |
| help='The inference config!', | |
| default="./config") | |
| cfg = Config(load=True, parser_ins=parser) | |
| model_yamls = glob.glob(os.path.join(cfg.args.cfg_folder, '*.yaml')) | |
| model_choices = dict() | |
| for i in model_yamls: | |
| model_cfg = Config(load=True, cfg_file=i) | |
| model_name = model_cfg.NAME | |
| model_choices[model_name] = model_cfg | |
| if cfg.args.infer_type == "native": | |
| infer_name = "ace_plus_native_infer" | |
| elif cfg.args.infer_type == "diffusers": | |
| infer_name = "ace_plus_diffuser_infer" | |
| else: | |
| raise ValueError("infer_type should be native or diffusers") | |
| assert infer_name in model_choices | |
| # choose different model | |
| task_model_cfg = Config(load=True, cfg_file=cfg.args.task_model) | |
| task_model_dict = {} | |
| for task_name, task_model in task_model_cfg.MODEL.items(): | |
| task_model_dict[task_name] = task_model | |
| # choose the inference scripts. | |
| pipe_cfg = model_choices[infer_name] | |
| infer_name = pipe_cfg.get("INFERENCE_TYPE", "ACE_PLUS") | |
| pipe = inference_dict[infer_name]() | |
| pipe.init_from_cfg(pipe_cfg) | |
| if cfg.args.instruction == "" and cfg.args.input_image is None and cfg.args.input_reference_image is None: | |
| params = { | |
| "output_h": cfg.args.output_h, | |
| "output_w": cfg.args.output_w, | |
| "sample_steps": cfg.args.step, | |
| "guide_scale": cfg.args.guide_scale | |
| } | |
| # run examples | |
| for example in all_examples: | |
| example["model_path"] = FS.get_from(task_model_dict[example["task_type"].upper()]["MODEL_PATH"]) | |
| example.update(params) | |
| if example["edit_type"] == "repainting": | |
| example["repainting_scale"] = 1.0 | |
| else: | |
| example["repainting_scale"] = task_model_dict[example["task_type"].upper()].get("REPAINTING_SCALE", 1.0) | |
| print(example) | |
| local_path, seed = run_one_case(pipe, **example) | |
| else: | |
| assert cfg.args.task_type.upper() in task_model_cfg | |
| params = { | |
| "input_image": cfg.args.input_image, | |
| "input_mask": cfg.args.input_mask, | |
| "input_reference_image": cfg.args.input_reference_image, | |
| "save_path": cfg.args.save_path, | |
| "instruction": cfg.args.instruction, | |
| "output_h": cfg.args.output_h, | |
| "output_w": cfg.args.output_w, | |
| "sample_steps": cfg.args.step, | |
| "guide_scale": cfg.args.guide_scale, | |
| "repainting_scale": cfg.args.repainting_scale, | |
| "model_path": FS.get_from(task_model_dict[cfg.args.task_type.upper()]["MODEL_PATH"]) | |
| } | |
| local_path, seed = run_one_case(pipe, **params) | |
| print(local_path, seed) | |
| if __name__ == '__main__': | |
| run() | |