229 lines
8.6 KiB
Python
229 lines
8.6 KiB
Python
# -*- 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()
|
|
|