70af4845af
new file: inpaint/__main__.py new file: inpaint/api.py new file: inpaint/batch_processing.py new file: inpaint/benchmark.py new file: inpaint/cli.py new file: inpaint/const.py new file: inpaint/download.py new file: inpaint/file_manager/__init__.py new file: inpaint/file_manager/file_manager.py new file: inpaint/file_manager/storage_backends.py new file: inpaint/file_manager/utils.py new file: inpaint/helper.py new file: inpaint/installer.py new file: inpaint/model/__init__.py new file: inpaint/model/anytext/__init__.py new file: inpaint/model/anytext/anytext_model.py new file: inpaint/model/anytext/anytext_pipeline.py new file: inpaint/model/anytext/anytext_sd15.yaml new file: inpaint/model/anytext/cldm/__init__.py new file: inpaint/model/anytext/cldm/cldm.py new file: inpaint/model/anytext/cldm/ddim_hacked.py new file: inpaint/model/anytext/cldm/embedding_manager.py new file: inpaint/model/anytext/cldm/hack.py new file: inpaint/model/anytext/cldm/model.py new file: inpaint/model/anytext/cldm/recognizer.py new file: inpaint/model/anytext/ldm/__init__.py new file: inpaint/model/anytext/ldm/models/__init__.py new file: inpaint/model/anytext/ldm/models/autoencoder.py new file: inpaint/model/anytext/ldm/models/diffusion/__init__.py new file: inpaint/model/anytext/ldm/models/diffusion/ddim.py new file: inpaint/model/anytext/ldm/models/diffusion/ddpm.py new file: inpaint/model/anytext/ldm/models/diffusion/dpm_solver/__init__.py new file: inpaint/model/anytext/ldm/models/diffusion/dpm_solver/dpm_solver.py new file: inpaint/model/anytext/ldm/models/diffusion/dpm_solver/sampler.py new file: inpaint/model/anytext/ldm/models/diffusion/plms.py new file: inpaint/model/anytext/ldm/models/diffusion/sampling_util.py new file: inpaint/model/anytext/ldm/modules/__init__.py new file: inpaint/model/anytext/ldm/modules/attention.py new file: inpaint/model/anytext/ldm/modules/diffusionmodules/__init__.py new file: inpaint/model/anytext/ldm/modules/diffusionmodules/model.py new file: inpaint/model/anytext/ldm/modules/diffusionmodules/openaimodel.py new file: inpaint/model/anytext/ldm/modules/diffusionmodules/upscaling.py new file: inpaint/model/anytext/ldm/modules/diffusionmodules/util.py new file: inpaint/model/anytext/ldm/modules/distributions/__init__.py new file: inpaint/model/anytext/ldm/modules/distributions/distributions.py new file: inpaint/model/anytext/ldm/modules/ema.py new file: inpaint/model/anytext/ldm/modules/encoders/__init__.py new file: inpaint/model/anytext/ldm/modules/encoders/modules.py new file: inpaint/model/anytext/ldm/util.py new file: inpaint/model/anytext/main.py new file: inpaint/model/anytext/ocr_recog/RNN.py new file: inpaint/model/anytext/ocr_recog/RecCTCHead.py new file: inpaint/model/anytext/ocr_recog/RecModel.py new file: inpaint/model/anytext/ocr_recog/RecMv1_enhance.py new file: inpaint/model/anytext/ocr_recog/RecSVTR.py new file: inpaint/model/anytext/ocr_recog/__init__.py new file: inpaint/model/anytext/ocr_recog/common.py new file: inpaint/model/anytext/ocr_recog/en_dict.txt new file: inpaint/model/anytext/ocr_recog/ppocr_keys_v1.txt new file: inpaint/model/anytext/utils.py new file: inpaint/model/base.py new file: inpaint/model/brushnet/__init__.py new file: inpaint/model/brushnet/brushnet.py new file: inpaint/model/brushnet/brushnet_unet_forward.py new file: inpaint/model/brushnet/brushnet_wrapper.py new file: inpaint/model/brushnet/pipeline_brushnet.py new file: inpaint/model/brushnet/unet_2d_blocks.py new file: inpaint/model/controlnet.py new file: inpaint/model/ddim_sampler.py new file: inpaint/model/fcf.py new file: inpaint/model/helper/__init__.py new file: inpaint/model/helper/controlnet_preprocess.py new file: inpaint/model/helper/cpu_text_encoder.py new file: inpaint/model/helper/g_diffuser_bot.py new file: inpaint/model/instruct_pix2pix.py new file: inpaint/model/kandinsky.py new file: inpaint/model/lama.py new file: inpaint/model/ldm.py new file: inpaint/model/manga.py new file: inpaint/model/mat.py new file: inpaint/model/mi_gan.py new file: inpaint/model/opencv2.py new file: inpaint/model/original_sd_configs/__init__.py new file: inpaint/model/original_sd_configs/sd_xl_base.yaml new file: inpaint/model/original_sd_configs/sd_xl_refiner.yaml new file: inpaint/model/original_sd_configs/v1-inference.yaml new file: inpaint/model/original_sd_configs/v2-inference-v.yaml new file: inpaint/model/paint_by_example.py new file: inpaint/model/plms_sampler.py new file: inpaint/model/power_paint/__init__.py new file: inpaint/model/power_paint/pipeline_powerpaint.py new file: inpaint/model/power_paint/power_paint.py new file: inpaint/model/power_paint/power_paint_v2.py new file: inpaint/model/power_paint/powerpaint_tokenizer.py
129 lines
4.1 KiB
Python
129 lines
4.1 KiB
Python
import json
|
|
from pathlib import Path
|
|
from typing import Dict, Optional
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from PIL import Image
|
|
from loguru import logger
|
|
from rich.console import Console
|
|
from rich.progress import (
|
|
Progress,
|
|
SpinnerColumn,
|
|
TimeElapsedColumn,
|
|
MofNCompleteColumn,
|
|
TextColumn,
|
|
BarColumn,
|
|
TaskProgressColumn,
|
|
)
|
|
|
|
from inpaint.helper import pil_to_bytes
|
|
from inpaint.model.utils import torch_gc
|
|
from inpaint.model_manager import ModelManager
|
|
from inpaint.schema import InpaintRequest
|
|
|
|
|
|
def glob_images(path: Path) -> Dict[str, Path]:
|
|
# png/jpg/jpeg
|
|
if path.is_file():
|
|
return {path.stem: path}
|
|
elif path.is_dir():
|
|
res = {}
|
|
for it in path.glob("*.*"):
|
|
if it.suffix.lower() in [".png", ".jpg", ".jpeg"]:
|
|
res[it.stem] = it
|
|
return res
|
|
|
|
|
|
def batch_inpaint(
|
|
model: str,
|
|
device,
|
|
image: Path,
|
|
mask: Path,
|
|
output: Path,
|
|
config: Optional[Path] = None,
|
|
concat: bool = False,
|
|
):
|
|
if image.is_dir() and output.is_file():
|
|
logger.error(
|
|
"invalid --output: when image is a directory, output should be a directory"
|
|
)
|
|
exit(-1)
|
|
output.mkdir(parents=True, exist_ok=True)
|
|
|
|
image_paths = glob_images(image)
|
|
mask_paths = glob_images(mask)
|
|
if len(image_paths) == 0:
|
|
logger.error("invalid --image: empty image folder")
|
|
exit(-1)
|
|
if len(mask_paths) == 0:
|
|
logger.error("invalid --mask: empty mask folder")
|
|
exit(-1)
|
|
|
|
if config is None:
|
|
inpaint_request = InpaintRequest()
|
|
logger.info(f"Using default config: {inpaint_request}")
|
|
else:
|
|
with open(config, "r", encoding="utf-8") as f:
|
|
inpaint_request = InpaintRequest(**json.load(f))
|
|
logger.info(f"Using config: {inpaint_request}")
|
|
|
|
model_manager = ModelManager(name=model, device=device)
|
|
first_mask = list(mask_paths.values())[0]
|
|
|
|
console = Console()
|
|
|
|
with Progress(
|
|
SpinnerColumn(),
|
|
TextColumn("[progress.description]{task.description}"),
|
|
BarColumn(),
|
|
TaskProgressColumn(),
|
|
MofNCompleteColumn(),
|
|
TimeElapsedColumn(),
|
|
console=console,
|
|
transient=False,
|
|
) as progress:
|
|
task = progress.add_task("Batch processing...", total=len(image_paths))
|
|
for stem, image_p in image_paths.items():
|
|
if stem not in mask_paths and mask.is_dir():
|
|
progress.log(f"mask for {image_p} not found")
|
|
progress.update(task, advance=1)
|
|
continue
|
|
mask_p = mask_paths.get(stem, first_mask)
|
|
|
|
infos = Image.open(image_p).info
|
|
|
|
img = np.array(Image.open(image_p).convert("RGB"))
|
|
mask_img = np.array(Image.open(mask_p).convert("L"))
|
|
|
|
if mask_img.shape[:2] != img.shape[:2]:
|
|
progress.log(
|
|
f"resize mask {mask_p.name} to image {image_p.name} size: {img.shape[:2]}"
|
|
)
|
|
mask_img = cv2.resize(
|
|
mask_img,
|
|
(img.shape[1], img.shape[0]),
|
|
interpolation=cv2.INTER_NEAREST,
|
|
)
|
|
mask_img[mask_img >= 127] = 255
|
|
mask_img[mask_img < 127] = 0
|
|
|
|
# bgr
|
|
inpaint_result = model_manager(img, mask_img, inpaint_request)
|
|
inpaint_result = cv2.cvtColor(inpaint_result, cv2.COLOR_BGR2RGB)
|
|
if concat:
|
|
mask_img = cv2.cvtColor(mask_img, cv2.COLOR_GRAY2RGB)
|
|
inpaint_result = cv2.hconcat([img, mask_img, inpaint_result])
|
|
|
|
img_bytes = pil_to_bytes(Image.fromarray(inpaint_result), "png", 100, infos)
|
|
save_p = output / f"{stem}.png"
|
|
with open(save_p, "wb") as fw:
|
|
fw.write(img_bytes)
|
|
|
|
progress.update(task, advance=1)
|
|
torch_gc()
|
|
# pid = psutil.Process().pid
|
|
# memory_info = psutil.Process(pid).memory_info()
|
|
# memory_in_mb = memory_info.rss / (1024 * 1024)
|
|
# print(f"原图大小:{img.shape},当前进程的内存占用:{memory_in_mb}MB")
|