|
|
|
|
|
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): |
|
|
|
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 |
|
|
