let5see/IOPaint
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

wip

Browse Source
This commit is contained in:
Qing
2023-12-30 23:36:44 +08:00
parent 85c3397b97
commit c4abda3942
35 changed files with 969 additions and 854 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
lama_cleaner/model/base.py
View File

@@ -14,7 +14,7 @@ from lama_cleaner.helper import (
)
from lama_cleaner.model.helper.g_diffuser_bot import expand_image
from lama_cleaner.model.utils import get_scheduler
from lama_cleaner.schema import Config, HDStrategy, SDSampler
from lama_cleaner.schema import InpaintRequest, HDStrategy, SDSampler
class InpaintModel:
@@ -44,7 +44,7 @@ class InpaintModel:
return False
@abc.abstractmethod
def forward(self, image, mask, config: Config):
def forward(self, image, mask, config: InpaintRequest):
"""Input images and output images have same size
images: [H, W, C] RGB
masks: [H, W, 1] 255 为 masks 区域
@@ -56,7 +56,7 @@ class InpaintModel:
def download():
...
def _pad_forward(self, image, mask, config: Config):
def _pad_forward(self, image, mask, config: InpaintRequest):
origin_height, origin_width = image.shape[:2]
pad_image = pad_img_to_modulo(
image, mod=self.pad_mod, square=self.pad_to_square, min_size=self.min_size
@@ -74,7 +74,7 @@ class InpaintModel:
result, image, mask = self.forward_post_process(result, image, mask, config)
if config.sd_prevent_unmasked_area:
if config.sd_keep_unmasked_area:
mask = mask[:, :, np.newaxis]
result = result * (mask / 255) + image[:, :, ::-1] * (1 - (mask / 255))
return result
@@ -86,7 +86,7 @@ class InpaintModel:
return result, image, mask
@torch.no_grad()
def __call__(self, image, mask, config: Config):
def __call__(self, image, mask, config: InpaintRequest):
"""
images: [H, W, C] RGB, not normalized
masks: [H, W]
@@ -141,7 +141,7 @@ class InpaintModel:
return inpaint_result
def _crop_box(self, image, mask, box, config: Config):
def _crop_box(self, image, mask, box, config: InpaintRequest):
"""
Args:
@@ -233,7 +233,7 @@ class InpaintModel:
return result
def _apply_cropper(self, image, mask, config: Config):
def _apply_cropper(self, image, mask, config: InpaintRequest):
img_h, img_w = image.shape[:2]
l, t, w, h = (
config.croper_x,
@@ -253,7 +253,7 @@ class InpaintModel:
crop_mask = mask[t:b, l:r]
return crop_img, crop_mask, (l, t, r, b)
def _run_box(self, image, mask, box, config: Config):
def _run_box(self, image, mask, box, config: InpaintRequest):
"""
Args:
@@ -276,7 +276,7 @@ class DiffusionInpaintModel(InpaintModel):
super().__init__(device, **kwargs)
@torch.no_grad()
def __call__(self, image, mask, config: Config):
def __call__(self, image, mask, config: InpaintRequest):
"""
images: [H, W, C] RGB, not normalized
masks: [H, W]
@@ -295,7 +295,7 @@ class DiffusionInpaintModel(InpaintModel):
return inpaint_result
def _do_outpainting(self, image, config: Config):
def _do_outpainting(self, image, config: InpaintRequest):
# cropper 和 image 在同一个坐标系下，croper_x/y 可能为负数
# 从 image 中 crop 出 outpainting 区域
image_h, image_w = image.shape[:2]
@@ -368,7 +368,7 @@ class DiffusionInpaintModel(InpaintModel):
] = expanded_cropped_result_image
return outpainting_image
def _scaled_pad_forward(self, image, mask, config: Config):
def _scaled_pad_forward(self, image, mask, config: InpaintRequest):
longer_side_length = int(config.sd_scale * max(image.shape[:2]))
origin_size = image.shape[:2]
downsize_image = resize_max_size(image, size_limit=longer_side_length)
@@ -396,7 +396,7 @@ class DiffusionInpaintModel(InpaintModel):
# ]
return inpaint_result
def set_scheduler(self, config: Config):
def set_scheduler(self, config: InpaintRequest):
scheduler_config = self.model.scheduler.config
sd_sampler = config.sd_sampler
if config.sd_lcm_lora:

4
lama_cleaner/model/controlnet.py
View File

@@ -14,7 +14,7 @@ from lama_cleaner.model.helper.controlnet_preprocess import (
)
from lama_cleaner.model.helper.cpu_text_encoder import CPUTextEncoderWrapper
from lama_cleaner.model.utils import get_scheduler, handle_from_pretrained_exceptions
from lama_cleaner.schema import Config, ModelType
from lama_cleaner.schema import InpaintRequest, ModelType
class ControlNet(DiffusionInpaintModel):
@@ -130,7 +130,7 @@ class ControlNet(DiffusionInpaintModel):
raise NotImplementedError(f"{self.controlnet_method} not implemented")
return control_image
def forward(self, image, mask, config: Config):
def forward(self, image, mask, config: InpaintRequest):
"""Input image and output image have same size
image: [H, W, C] RGB
mask: [H, W, 1] 255 means area to repaint

6
lama_cleaner/model/fcf.py
View File

@@ -6,7 +6,7 @@ import torch
import numpy as np
import torch.fft as fft
from lama_cleaner.schema import Config
from lama_cleaner.schema import InpaintRequest
from lama_cleaner.helper import (
load_model,
@@ -1665,7 +1665,7 @@ class FcF(InpaintModel):
return os.path.exists(get_cache_path_by_url(FCF_MODEL_URL))
@torch.no_grad()
def __call__(self, image, mask, config: Config):
def __call__(self, image, mask, config: InpaintRequest):
"""
images: [H, W, C] RGB, not normalized
masks: [H, W]
@@ -1705,7 +1705,7 @@ class FcF(InpaintModel):
return inpaint_result
def forward(self, image, mask, config: Config):
def forward(self, image, mask, config: InpaintRequest):
"""Input images and output images have same size
images: [H, W, C] RGB
masks: [H, W] mask area == 255

4
lama_cleaner/model/instruct_pix2pix.py
View File

@@ -4,7 +4,7 @@ import torch
from loguru import logger
from lama_cleaner.model.base import DiffusionInpaintModel
from lama_cleaner.schema import Config
from lama_cleaner.schema import InpaintRequest
class InstructPix2Pix(DiffusionInpaintModel):
@@ -40,7 +40,7 @@ class InstructPix2Pix(DiffusionInpaintModel):
else:
self.model = self.model.to(device)
def forward(self, image, mask, config: Config):
def forward(self, image, mask, config: InpaintRequest):
"""Input image and output image have same size
image: [H, W, C] RGB
mask: [H, W, 1] 255 means area to repaint

4
lama_cleaner/model/kandinsky.py
View File

@@ -5,7 +5,7 @@ import torch
from lama_cleaner.model.base import DiffusionInpaintModel
from lama_cleaner.model.utils import get_scheduler
from lama_cleaner.schema import Config
from lama_cleaner.schema import InpaintRequest
class Kandinsky(DiffusionInpaintModel):
@@ -29,7 +29,7 @@ class Kandinsky(DiffusionInpaintModel):
self.callback = kwargs.pop("callback", None)
def forward(self, image, mask, config: Config):
def forward(self, image, mask, config: InpaintRequest):
"""Input image and output image have same size
image: [H, W, C] RGB
mask: [H, W, 1] 255 means area to repaint

4
lama_cleaner/model/lama.py
View File

@@ -11,7 +11,7 @@ from lama_cleaner.helper import (
download_model,
)
from lama_cleaner.model.base import InpaintModel
from lama_cleaner.schema import Config
from lama_cleaner.schema import InpaintRequest
LAMA_MODEL_URL = os.environ.get(
"LAMA_MODEL_URL",
@@ -36,7 +36,7 @@ class LaMa(InpaintModel):
def is_downloaded() -> bool:
return os.path.exists(get_cache_path_by_url(LAMA_MODEL_URL))
def forward(self, image, mask, config: Config):
def forward(self, image, mask, config: InpaintRequest):
"""Input image and output image have same size
image: [H, W, C] RGB
mask: [H, W]

4
lama_cleaner/model/ldm.py
View File

@@ -7,7 +7,7 @@ from loguru import logger
from lama_cleaner.model.base import InpaintModel
from lama_cleaner.model.ddim_sampler import DDIMSampler
from lama_cleaner.model.plms_sampler import PLMSSampler
from lama_cleaner.schema import Config, LDMSampler
from lama_cleaner.schema import InpaintRequest, LDMSampler
torch.manual_seed(42)
import torch.nn as nn
@@ -277,7 +277,7 @@ class LDM(InpaintModel):
return all([os.path.exists(it) for it in model_paths])
@torch.cuda.amp.autocast()
def forward(self, image, mask, config: Config):
def forward(self, image, mask, config: InpaintRequest):
"""
image: [H, W, C] RGB
mask: [H, W, 1]

4
lama_cleaner/model/manga.py
View File

@@ -9,7 +9,7 @@ from loguru import logger
from lama_cleaner.helper import get_cache_path_by_url, load_jit_model, download_model
from lama_cleaner.model.base import InpaintModel
from lama_cleaner.schema import Config
from lama_cleaner.schema import InpaintRequest
MANGA_INPAINTOR_MODEL_URL = os.environ.get(
@@ -56,7 +56,7 @@ class Manga(InpaintModel):
]
return all([os.path.exists(it) for it in model_paths])
def forward(self, image, mask, config: Config):
def forward(self, image, mask, config: InpaintRequest):
"""
image: [H, W, C] RGB
mask: [H, W, 1]

4
lama_cleaner/model/mat.py
View File

@@ -28,7 +28,7 @@ from lama_cleaner.model.utils import (
normalize_2nd_moment,
set_seed,
)
from lama_cleaner.schema import Config
from lama_cleaner.schema import InpaintRequest
class ModulatedConv2d(nn.Module):
@@ -1912,7 +1912,7 @@ class MAT(InpaintModel):
def is_downloaded() -> bool:
return os.path.exists(get_cache_path_by_url(MAT_MODEL_URL))
def forward(self, image, mask, config: Config):
def forward(self, image, mask, config: InpaintRequest):
"""Input images and output images have same size
images: [H, W, C] RGB
masks: [H, W] mask area == 255

6
lama_cleaner/model/mi_gan.py
View File

@@ -3,7 +3,6 @@ import os
import cv2
import torch
from lama_cleaner.const import Config
from lama_cleaner.helper import (
load_jit_model,
download_model,
@@ -13,6 +12,7 @@ from lama_cleaner.helper import (
norm_img,
)
from lama_cleaner.model.base import InpaintModel
from lama_cleaner.schema import InpaintRequest
MIGAN_MODEL_URL = os.environ.get(
"MIGAN_MODEL_URL",
@@ -40,7 +40,7 @@ class MIGAN(InpaintModel):
return os.path.exists(get_cache_path_by_url(MIGAN_MODEL_URL))
@torch.no_grad()
def __call__(self, image, mask, config: Config):
def __call__(self, image, mask, config: InpaintRequest):
"""
images: [H, W, C] RGB, not normalized
masks: [H, W]
@@ -80,7 +80,7 @@ class MIGAN(InpaintModel):
return inpaint_result
def forward(self, image, mask, config: Config):
def forward(self, image, mask, config: InpaintRequest):
"""Input images and output images have same size
images: [H, W, C] RGB
masks: [H, W] mask area == 255

4
lama_cleaner/model/opencv2.py
View File

@@ -1,6 +1,6 @@
import cv2
from lama_cleaner.model.base import InpaintModel
from lama_cleaner.schema import Config
from lama_cleaner.schema import InpaintRequest
flag_map = {"INPAINT_NS": cv2.INPAINT_NS, "INPAINT_TELEA": cv2.INPAINT_TELEA}
@@ -14,7 +14,7 @@ class OpenCV2(InpaintModel):
def is_downloaded() -> bool:
return True
def forward(self, image, mask, config: Config):
def forward(self, image, mask, config: InpaintRequest):
"""Input image and output image have same size
image: [H, W, C] RGB
mask: [H, W, 1]

12
lama_cleaner/model/paint_by_example.py
View File

@@ -4,8 +4,9 @@ import cv2
import torch
from loguru import logger
from lama_cleaner.helper import decode_base64_to_image
from lama_cleaner.model.base import DiffusionInpaintModel
from lama_cleaner.schema import Config
from lama_cleaner.schema import InpaintRequest
class PaintByExample(DiffusionInpaintModel):
@@ -38,16 +39,21 @@ class PaintByExample(DiffusionInpaintModel):
else:
self.model = self.model.to(device)
def forward(self, image, mask, config: Config):
def forward(self, image, mask, config: InpaintRequest):
"""Input image and output image have same size
image: [H, W, C] RGB
mask: [H, W, 1] 255 means area to repaint
return: BGR IMAGE
"""
if config.paint_by_example_example_image is None:
raise ValueError("paint_by_example_example_image is required")
example_image, _, _ = decode_base64_to_image(
config.paint_by_example_example_image
)
output = self.model(
image=PIL.Image.fromarray(image),
mask_image=PIL.Image.fromarray(mask[:, :, -1], mode="L"),
example_image=config.paint_by_example_example_image,
example_image=PIL.Image.fromarray(example_image),
num_inference_steps=config.sd_steps,
guidance_scale=config.sd_guidance_scale,
negative_prompt="out of frame, lowres, error, cropped, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, out of frame, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, disfigured, gross proportions, malformed limbs, watermark, signature",

4
lama_cleaner/model/power_paint/power_paint.py
View File

@@ -7,7 +7,7 @@ from loguru import logger
from lama_cleaner.model.base import DiffusionInpaintModel
from lama_cleaner.model.helper.cpu_text_encoder import CPUTextEncoderWrapper
from lama_cleaner.model.utils import handle_from_pretrained_exceptions
from lama_cleaner.schema import Config
from lama_cleaner.schema import InpaintRequest
from .powerpaint_tokenizer import add_task_to_prompt
@@ -58,7 +58,7 @@ class PowerPaint(DiffusionInpaintModel):
self.callback = kwargs.pop("callback", None)
def forward(self, image, mask, config: Config):
def forward(self, image, mask, config: InpaintRequest):
"""Input image and output image have same size
image: [H, W, C] RGB
mask: [H, W, 1] 255 means area to repaint

4
lama_cleaner/model/sd.py
View File

@@ -6,7 +6,7 @@ from loguru import logger
from lama_cleaner.model.base import DiffusionInpaintModel
from lama_cleaner.model.helper.cpu_text_encoder import CPUTextEncoderWrapper
from lama_cleaner.model.utils import handle_from_pretrained_exceptions
from lama_cleaner.schema import Config, ModelType
from lama_cleaner.schema import InpaintRequest, ModelType
class SD(DiffusionInpaintModel):
@@ -64,7 +64,7 @@ class SD(DiffusionInpaintModel):
self.callback = kwargs.pop("callback", None)
def forward(self, image, mask, config: Config):
def forward(self, image, mask, config: InpaintRequest):
"""Input image and output image have same size
image: [H, W, C] RGB
mask: [H, W, 1] 255 means area to repaint

4
lama_cleaner/model/sdxl.py
View File

@@ -9,7 +9,7 @@ from loguru import logger
from lama_cleaner.model.base import DiffusionInpaintModel
from lama_cleaner.model.utils import handle_from_pretrained_exceptions
from lama_cleaner.schema import Config, ModelType
from lama_cleaner.schema import InpaintRequest, ModelType
class SDXL(DiffusionInpaintModel):
@@ -60,7 +60,7 @@ class SDXL(DiffusionInpaintModel):
self.callback = kwargs.pop("callback", None)
def forward(self, image, mask, config: Config):
def forward(self, image, mask, config: InpaintRequest):
"""Input image and output image have same size
image: [H, W, C] RGB
mask: [H, W, 1] 255 means area to repaint

4
lama_cleaner/model/zits.py
View File

@@ -6,7 +6,7 @@ import torch
import torch.nn.functional as F
from lama_cleaner.helper import get_cache_path_by_url, load_jit_model, download_model
from lama_cleaner.schema import Config
from lama_cleaner.schema import InpaintRequest
import numpy as np
from lama_cleaner.model.base import InpaintModel
@@ -343,7 +343,7 @@ class ZITS(InpaintModel):
items["line"] = line_pred.detach()
@torch.no_grad()
def forward(self, image, mask, config: Config):
def forward(self, image, mask, config: InpaintRequest):
"""Input images and output images have same size
images: [H, W, C] RGB
masks: [H, W]