rename to iopaint

iopaint/tests/.gitignore

@@ -0,0 +1,2 @@
+*_result.png
+result/

iopaint/tests/test_adjust_mask.py

@@ -0,0 +1,15 @@
+import cv2
+from iopaint.helper import adjust_mask
+from iopaint.tests.utils import current_dir, save_dir
+
+mask_p = current_dir / "overture-creations-5sI6fQgYIuo_mask.png"
+
+
+def test_adjust_mask():
+    mask = cv2.imread(str(mask_p), cv2.IMREAD_GRAYSCALE)
+    res_mask = adjust_mask(mask, 0, "expand")
+    cv2.imwrite(str(save_dir / "adjust_mask_original.png"), res_mask)
+    res_mask = adjust_mask(mask, 40, "expand")
+    cv2.imwrite(str(save_dir / "adjust_mask_expand.png"), res_mask)
+    res_mask = adjust_mask(mask, 20, "shrink")
+    cv2.imwrite(str(save_dir / "adjust_mask_shrink.png"), res_mask)

iopaint/tests/test_controlnet.py

@@ -0,0 +1,117 @@
+import os
+
+from iopaint.const import SD_CONTROLNET_CHOICES
+from iopaint.tests.utils import current_dir, check_device, get_config, assert_equal
+
+os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
+from pathlib import Path
+
+import pytest
+import torch
+
+from iopaint.model_manager import ModelManager
+from iopaint.schema import HDStrategy, SDSampler
+
+
+model_name = "runwayml/stable-diffusion-inpainting"
+
+
+def convert_controlnet_method_name(name):
+    return name.replace("/", "--")
+
+
+@pytest.mark.parametrize("device", ["cuda", "mps", "cpu"])
+@pytest.mark.parametrize("controlnet_method", [SD_CONTROLNET_CHOICES[0]])
+def test_runway_sd_1_5(device, controlnet_method):
+    sd_steps = check_device(device)
+
+    model = ModelManager(
+        name=model_name,
+        device=torch.device(device),
+        disable_nsfw=True,
+        sd_cpu_textencoder=True,
+        enable_controlnet=True,
+        controlnet_method=controlnet_method,
+    )
+
+    cfg = get_config(
+        prompt="a fox sitting on a bench",
+        sd_steps=sd_steps,
+        enable_controlnet=True,
+        controlnet_conditioning_scale=0.5,
+        controlnet_method=controlnet_method,
+    )
+    name = f"device_{device}"
+
+    assert_equal(
+        model,
+        cfg,
+        f"sd_controlnet_{convert_controlnet_method_name(controlnet_method)}_{name}.png",
+        img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
+        mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
+    )
+
+
+@pytest.mark.parametrize("device", ["cuda", "mps", "cpu"])
+def test_controlnet_switch(device):
+    sd_steps = check_device(device)
+    model = ModelManager(
+        name=model_name,
+        device=torch.device(device),
+        disable_nsfw=True,
+        sd_cpu_textencoder=False,
+        cpu_offload=True,
+        enable_controlnet=True,
+        controlnet_method="lllyasviel/control_v11p_sd15_canny",
+    )
+    cfg = get_config(
+        prompt="a fox sitting on a bench",
+        sd_steps=sd_steps,
+        enable_controlnet=True,
+        controlnet_method="lllyasviel/control_v11f1p_sd15_depth",
+    )
+
+    assert_equal(
+        model,
+        cfg,
+        f"controlnet_switch_canny_to_depth_device_{device}.png",
+        img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
+        mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
+    )
+
+
+@pytest.mark.parametrize("device", ["cuda", "mps", "cpu"])
+@pytest.mark.parametrize(
+    "local_file", ["sd-v1-5-inpainting.ckpt", "v1-5-pruned-emaonly.safetensors"]
+)
+def test_local_file_path(device, local_file):
+    sd_steps = check_device(device)
+
+    controlnet_kwargs = dict(
+        enable_controlnet=True,
+        controlnet_method=SD_CONTROLNET_CHOICES[0],
+    )
+
+    model = ModelManager(
+        name=local_file,
+        device=torch.device(device),
+        disable_nsfw=True,
+        sd_cpu_textencoder=False,
+        cpu_offload=True,
+        **controlnet_kwargs,
+    )
+    cfg = get_config(
+        prompt="a fox sitting on a bench",
+        sd_steps=sd_steps,
+        **controlnet_kwargs,
+    )
+
+    name = f"device_{device}"
+
+    assert_equal(
+        model,
+        cfg,
+        f"{controlnet_kwargs['controlnet_method']}_local_model_{name}.png",
+        img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
+        mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
+    )

iopaint/tests/test_instruct_pix2pix.py

@@ -0,0 +1,41 @@
+from pathlib import Path
+
+import pytest
+import torch
+
+from iopaint.model_manager import ModelManager
+from iopaint.schema import HDStrategy
+from iopaint.tests.utils import get_config, check_device, assert_equal, current_dir
+
+model_name = "timbrooks/instruct-pix2pix"
+
+
+@pytest.mark.parametrize("device", ["cuda", "mps", "cpu"])
+@pytest.mark.parametrize("disable_nsfw", [True, False])
+@pytest.mark.parametrize("cpu_offload", [False, True])
+def test_instruct_pix2pix(device, disable_nsfw, cpu_offload):
+    sd_steps = check_device(device)
+    model = ModelManager(
+        name=model_name,
+        device=torch.device(device),
+        disable_nsfw=disable_nsfw,
+        sd_cpu_textencoder=False,
+        cpu_offload=cpu_offload,
+    )
+    cfg = get_config(
+        strategy=HDStrategy.ORIGINAL,
+        prompt="What if it were snowing?",
+        p2p_steps=sd_steps,
+        sd_scale=1.1,
+    )
+
+    name = f"device_{device}_disnsfw_{disable_nsfw}_cpu_offload_{cpu_offload}"
+
+    assert_equal(
+        model,
+        cfg,
+        f"instruct_pix2pix_{name}.png",
+        img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
+        mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
+        fx=1.3,
+    )

iopaint/tests/test_load_img.py

@@ -0,0 +1,19 @@
+from iopaint.helper import load_img
+from iopaint.tests.utils import current_dir
+
+png_img_p = current_dir / "image.png"
+jpg_img_p = current_dir / "bunny.jpeg"
+
+
+def test_load_png_image():
+    with open(png_img_p, "rb") as f:
+        np_img, alpha_channel = load_img(f.read())
+    assert np_img.shape == (256, 256, 3)
+    assert alpha_channel.shape == (256, 256)
+
+
+def test_load_jpg_image():
+    with open(jpg_img_p, "rb") as f:
+        np_img, alpha_channel = load_img(f.read())
+    assert np_img.shape == (394, 448, 3)
+    assert alpha_channel is None

iopaint/tests/test_model.py

@@ -0,0 +1,160 @@
+import pytest
+import torch
+
+from iopaint.model_manager import ModelManager
+from iopaint.schema import HDStrategy, LDMSampler
+from iopaint.tests.utils import assert_equal, get_config, current_dir, check_device
+
+
+@pytest.mark.parametrize("device", ["cuda", "mps", "cpu"])
+@pytest.mark.parametrize(
+    "strategy", [HDStrategy.ORIGINAL, HDStrategy.RESIZE, HDStrategy.CROP]
+)
+def test_lama(device, strategy):
+    check_device(device)
+    model = ModelManager(name="lama", device=device)
+    assert_equal(
+        model,
+        get_config(strategy=strategy),
+        f"lama_{strategy[0].upper() + strategy[1:]}_result.png",
+    )
+
+    fx = 1.3
+    assert_equal(
+        model,
+        get_config(strategy=strategy),
+        f"lama_{strategy[0].upper() + strategy[1:]}_fx_{fx}_result.png",
+        fx=1.3,
+    )
+
+
+@pytest.mark.parametrize("device", ["cuda", "cpu"])
+@pytest.mark.parametrize(
+    "strategy", [HDStrategy.ORIGINAL, HDStrategy.RESIZE, HDStrategy.CROP]
+)
+@pytest.mark.parametrize("ldm_sampler", [LDMSampler.ddim, LDMSampler.plms])
+def test_ldm(device, strategy, ldm_sampler):
+    check_device(device)
+    model = ModelManager(name="ldm", device=device)
+    cfg = get_config(strategy=strategy, ldm_sampler=ldm_sampler)
+    assert_equal(
+        model, cfg, f"ldm_{strategy[0].upper() + strategy[1:]}_{ldm_sampler}_result.png"
+    )
+
+    fx = 1.3
+    assert_equal(
+        model,
+        cfg,
+        f"ldm_{strategy[0].upper() + strategy[1:]}_{ldm_sampler}_fx_{fx}_result.png",
+        fx=fx,
+    )
+
+
+@pytest.mark.parametrize("device", ["cuda", "cpu"])
+@pytest.mark.parametrize(
+    "strategy", [HDStrategy.ORIGINAL, HDStrategy.RESIZE, HDStrategy.CROP]
+)
+@pytest.mark.parametrize("zits_wireframe", [False, True])
+def test_zits(device, strategy, zits_wireframe):
+    check_device(device)
+    model = ModelManager(name="zits", device=device)
+    cfg = get_config(strategy=strategy, zits_wireframe=zits_wireframe)
+    assert_equal(
+        model,
+        cfg,
+        f"zits_{strategy[0].upper() + strategy[1:]}_wireframe_{zits_wireframe}_result.png",
+    )
+
+    fx = 1.3
+    assert_equal(
+        model,
+        cfg,
+        f"zits_{strategy.capitalize()}_wireframe_{zits_wireframe}_fx_{fx}_result.png",
+        fx=fx,
+    )
+
+
+@pytest.mark.parametrize("device", ["cuda", "cpu"])
+@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
+@pytest.mark.parametrize("no_half", [True, False])
+def test_mat(device, strategy, no_half):
+    check_device(device)
+    model = ModelManager(name="mat", device=device, no_half=no_half)
+    cfg = get_config(strategy=strategy)
+
+    assert_equal(
+        model,
+        cfg,
+        f"mat_{strategy.capitalize()}_result.png",
+    )
+
+
+@pytest.mark.parametrize("device", ["cuda", "cpu"])
+@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
+def test_fcf(device, strategy):
+    check_device(device)
+    model = ModelManager(name="fcf", device=device)
+    cfg = get_config(strategy=strategy)
+
+    assert_equal(model, cfg, f"fcf_{strategy.capitalize()}_result.png", fx=2, fy=2)
+    assert_equal(model, cfg, f"fcf_{strategy.capitalize()}_result.png", fx=3.8, fy=2)
+
+
+@pytest.mark.parametrize(
+    "strategy", [HDStrategy.ORIGINAL, HDStrategy.RESIZE, HDStrategy.CROP]
+)
+@pytest.mark.parametrize("cv2_flag", ["INPAINT_NS", "INPAINT_TELEA"])
+@pytest.mark.parametrize("cv2_radius", [3, 15])
+def test_cv2(strategy, cv2_flag, cv2_radius):
+    model = ModelManager(
+        name="cv2",
+        device=torch.device("cpu"),
+    )
+    cfg = get_config(strategy=strategy, cv2_flag=cv2_flag, cv2_radius=cv2_radius)
+    assert_equal(
+        model,
+        cfg,
+        f"cv2_{strategy.capitalize()}_{cv2_flag}_{cv2_radius}.png",
+        img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
+        mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
+    )
+
+
+@pytest.mark.parametrize("device", ["cuda", "cpu"])
+@pytest.mark.parametrize(
+    "strategy", [HDStrategy.ORIGINAL, HDStrategy.RESIZE, HDStrategy.CROP]
+)
+def test_manga(device, strategy):
+    check_device(device)
+    model = ModelManager(
+        name="manga",
+        device=torch.device(device),
+    )
+    cfg = get_config(strategy=strategy)
+    assert_equal(
+        model,
+        cfg,
+        f"manga_{strategy.capitalize()}.png",
+        img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
+        mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
+    )
+
+
+@pytest.mark.parametrize("device", ["cuda", "mps", "cpu"])
+@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
+def test_mi_gan(device, strategy):
+    check_device(device)
+    model = ModelManager(
+        name="migan",
+        device=torch.device(device),
+    )
+    cfg = get_config(strategy=strategy)
+    assert_equal(
+        model,
+        cfg,
+        f"migan_device_{device}.png",
+        img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
+        mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
+        fx=1.5,
+        fy=1.7
+    )

iopaint/tests/test_model_md5.py

@@ -0,0 +1,16 @@
+def test_load_model():
+    from iopaint.plugins import InteractiveSeg
+    from iopaint.model_manager import ModelManager
+
+    interactive_seg_model = InteractiveSeg("vit_l", "cpu")
+
+    models = ["lama", "ldm", "zits", "mat", "fcf", "manga", "migan"]
+    for m in models:
+        ModelManager(
+            name=m,
+            device="cpu",
+            no_half=False,
+            disable_nsfw=False,
+            sd_cpu_textencoder=True,
+            cpu_offload=True,
+        )

iopaint/tests/test_model_switch.py

@@ -0,0 +1,70 @@
+import os
+
+from iopaint.schema import InpaintRequest
+
+os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
+
+import torch
+
+from iopaint.model_manager import ModelManager
+
+
+def test_model_switch():
+    model = ModelManager(
+        name="runwayml/stable-diffusion-inpainting",
+        enable_controlnet=True,
+        controlnet_method="lllyasviel/control_v11p_sd15_canny",
+        device=torch.device("mps"),
+        disable_nsfw=True,
+        sd_cpu_textencoder=True,
+        cpu_offload=False,
+    )
+
+    model.switch("lama")
+
+
+def test_controlnet_switch_onoff(caplog):
+    name = "runwayml/stable-diffusion-inpainting"
+    model = ModelManager(
+        name=name,
+        enable_controlnet=True,
+        controlnet_method="lllyasviel/control_v11p_sd15_canny",
+        device=torch.device("mps"),
+        disable_nsfw=True,
+        sd_cpu_textencoder=True,
+        cpu_offload=False,
+    )
+
+    model.switch_controlnet_method(
+        InpaintRequest(
+            name=name,
+            enable_controlnet=False,
+        )
+    )
+
+    assert "Disable controlnet" in caplog.text
+
+
+def test_switch_controlnet_method(caplog):
+    name = "runwayml/stable-diffusion-inpainting"
+    old_method = "lllyasviel/control_v11p_sd15_canny"
+    new_method = "lllyasviel/control_v11p_sd15_openpose"
+    model = ModelManager(
+        name=name,
+        enable_controlnet=True,
+        controlnet_method=old_method,
+        device=torch.device("mps"),
+        disable_nsfw=True,
+        sd_cpu_textencoder=True,
+        cpu_offload=False,
+    )
+
+    model.switch_controlnet_method(
+        InpaintRequest(
+            name=name,
+            enable_controlnet=True,
+            controlnet_method=new_method,
+        )
+    )
+
+    assert f"Switch Controlnet method from {old_method} to {new_method}" in caplog.text

iopaint/tests/test_outpainting.py

@@ -0,0 +1,137 @@
+import os
+
+from iopaint.tests.utils import current_dir, check_device
+
+os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
+from pathlib import Path
+
+import pytest
+import torch
+
+from iopaint.model_manager import ModelManager
+from iopaint.schema import HDStrategy, SDSampler
+from iopaint.tests.test_model import get_config, assert_equal
+
+
+@pytest.mark.parametrize("name", ["runwayml/stable-diffusion-inpainting"])
+@pytest.mark.parametrize("device", ["cuda", "mps", "cpu"])
+@pytest.mark.parametrize(
+    "rect",
+    [
+        [0, -100, 512, 512 - 128 + 100],
+        [0, 128, 512, 512 - 128 + 100],
+        [128, 0, 512 - 128 + 100, 512],
+        [-100, 0, 512 - 128 + 100, 512],
+        [0, 0, 512, 512 + 200],
+        [0, 0, 512 + 200, 512],
+        [-100, -100, 512 + 200, 512 + 200],
+    ],
+)
+def test_outpainting(name, device, rect):
+    sd_steps = check_device(device)
+
+    model = ModelManager(
+        name=name,
+        device=torch.device(device),
+        disable_nsfw=True,
+        sd_cpu_textencoder=False,
+    )
+    cfg = get_config(
+        prompt="a dog sitting on a bench in the park",
+        sd_steps=sd_steps,
+        use_extender=True,
+        extender_x=rect[0],
+        extender_y=rect[1],
+        extender_width=rect[2],
+        extender_height=rect[3],
+        sd_guidance_scale=8.0,
+        sd_sampler=SDSampler.dpm_plus_plus_2m,
+    )
+
+    assert_equal(
+        model,
+        cfg,
+        f"{name.replace('/', '--')}_outpainting_{'_'.join(map(str, rect))}_device_{device}.png",
+        img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
+        mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
+    )
+
+
+@pytest.mark.parametrize("name", ["kandinsky-community/kandinsky-2-2-decoder-inpaint"])
+@pytest.mark.parametrize("device", ["cuda", "mps", "cpu"])
+@pytest.mark.parametrize(
+    "rect",
+    [
+        [-128, -128, 768, 768],
+    ],
+)
+def test_kandinsky_outpainting(name, device, rect):
+    sd_steps = check_device(device)
+
+    model = ModelManager(
+        name=name,
+        device=torch.device(device),
+        disable_nsfw=True,
+        sd_cpu_textencoder=False,
+    )
+    cfg = get_config(
+        prompt="a cat",
+        negative_prompt="lowres, text, error, cropped, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, out of frame, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck, username, watermark, signature",
+        sd_steps=sd_steps,
+        use_extender=True,
+        extender_x=rect[0],
+        extender_y=rect[1],
+        extender_width=rect[2],
+        extender_height=rect[3],
+        sd_guidance_scale=7,
+        sd_sampler=SDSampler.dpm_plus_plus_2m,
+    )
+
+    assert_equal(
+        model,
+        cfg,
+        f"{name.replace('/', '--')}_outpainting_{'_'.join(map(str, rect))}_device_{device}.png",
+        img_p=current_dir / "cat.png",
+        mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
+        fx=1,
+        fy=1,
+    )
+
+
+@pytest.mark.parametrize("name", ["Sanster/PowerPaint-V1-stable-diffusion-inpainting"])
+@pytest.mark.parametrize("device", ["cuda", "mps", "cpu"])
+@pytest.mark.parametrize(
+    "rect",
+    [
+        [-100, -100, 512 + 200, 512 + 200],
+    ],
+)
+def test_powerpaint_outpainting(name, device, rect):
+    sd_steps = check_device(device)
+
+    model = ModelManager(
+        name=name,
+        device=torch.device(device),
+        disable_nsfw=True,
+        sd_cpu_textencoder=False,
+    )
+    cfg = get_config(
+        prompt="a dog sitting on a bench in the park",
+        sd_steps=sd_steps,
+        use_extender=True,
+        extender_x=rect[0],
+        extender_y=rect[1],
+        extender_width=rect[2],
+        extender_height=rect[3],
+        sd_guidance_scale=8.0,
+        sd_sampler=SDSampler.dpm_plus_plus_2m,
+        powerpaint_task="outpainting",
+    )
+
+    assert_equal(
+        model,
+        cfg,
+        f"{name.replace('/', '--')}_outpainting_{'_'.join(map(str, rect))}_device_{device}.png",
+        img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
+        mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
+    )

iopaint/tests/test_paint_by_example.py

@@ -0,0 +1,55 @@
+import cv2
+import pytest
+from PIL import Image
+
+from iopaint.model_manager import ModelManager
+from iopaint.schema import HDStrategy
+from iopaint.tests.utils import (
+    current_dir,
+    get_config,
+    get_data,
+    save_dir,
+    check_device,
+)
+
+model_name = "Fantasy-Studio/Paint-by-Example"
+
+
+def assert_equal(
+    model,
+    config,
+    save_name: str,
+    fx: float = 1,
+    fy: float = 1,
+    img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
+    mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
+    example_p=current_dir / "bunny.jpeg",
+):
+    img, mask = get_data(fx=fx, fy=fy, img_p=img_p, mask_p=mask_p)
+
+    example_image = cv2.imread(str(example_p))
+    example_image = cv2.cvtColor(example_image, cv2.COLOR_BGRA2RGB)
+    example_image = cv2.resize(
+        example_image, None, fx=fx, fy=fy, interpolation=cv2.INTER_AREA
+    )
+
+    print(f"Input image shape: {img.shape}, example_image: {example_image.shape}")
+    config.paint_by_example_example_image = Image.fromarray(example_image)
+    res = model(img, mask, config)
+    cv2.imwrite(str(save_dir / save_name), res)
+
+
+@pytest.mark.parametrize("device", ["cuda", "mps", "cpu"])
+def test_paint_by_example(device):
+    sd_steps = check_device(device)
+    model = ModelManager(name=model_name, device=device, disable_nsfw=True)
+    cfg = get_config(strategy=HDStrategy.ORIGINAL, sd_steps=sd_steps)
+    assert_equal(
+        model,
+        cfg,
+        f"paint_by_example_device_{device}.png",
+        img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
+        mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
+        fy=0.9,
+        fx=1.3,
+    )

iopaint/tests/test_plugins.py

@@ -0,0 +1,120 @@
+import hashlib
+import os
+import time
+from PIL import Image
+
+from iopaint.helper import encode_pil_to_base64, gen_frontend_mask
+from iopaint.plugins.anime_seg import AnimeSeg
+from iopaint.schema import RunPluginRequest
+from iopaint.tests.utils import check_device, current_dir, save_dir
+
+os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
+
+import cv2
+import pytest
+
+from iopaint.plugins import (
+    RemoveBG,
+    RealESRGANUpscaler,
+    GFPGANPlugin,
+    RestoreFormerPlugin,
+    InteractiveSeg,
+)
+
+img_p = current_dir / "bunny.jpeg"
+img_bytes = open(img_p, "rb").read()
+bgr_img = cv2.imread(str(img_p))
+rgb_img = cv2.cvtColor(bgr_img, cv2.COLOR_BGR2RGB)
+rgb_img_base64 = encode_pil_to_base64(Image.fromarray(rgb_img), 100, {})
+bgr_img_base64 = encode_pil_to_base64(Image.fromarray(bgr_img), 100, {})
+
+
+def _save(img, name):
+    cv2.imwrite(str(save_dir / name), img)
+
+
+def test_remove_bg():
+    model = RemoveBG()
+    rgba_np_img = model.gen_image(
+        rgb_img, RunPluginRequest(name=RemoveBG.name, image=rgb_img_base64)
+    )
+    res = cv2.cvtColor(rgba_np_img, cv2.COLOR_RGBA2BGRA)
+    _save(res, "test_remove_bg.png")
+
+    bgr_np_img = model.gen_mask(
+        rgb_img, RunPluginRequest(name=RemoveBG.name, image=rgb_img_base64)
+    )
+
+    res_mask = gen_frontend_mask(bgr_np_img)
+    _save(res_mask, "test_remove_bg_frontend_mask.png")
+
+    assert len(bgr_np_img.shape) == 2
+    _save(bgr_np_img, "test_remove_bg_mask.jpeg")
+
+
+def test_anime_seg():
+    model = AnimeSeg()
+    img = cv2.imread(str(current_dir / "anime_test.png"))
+    img_base64 = encode_pil_to_base64(Image.fromarray(img), 100, {})
+    res = model.gen_image(img, RunPluginRequest(name=AnimeSeg.name, image=img_base64))
+    assert len(res.shape) == 3
+    assert res.shape[-1] == 4
+    _save(res, "test_anime_seg.png")
+
+    res = model.gen_mask(img, RunPluginRequest(name=AnimeSeg.name, image=img_base64))
+    assert len(res.shape) == 2
+    _save(res, "test_anime_seg_mask.png")
+
+
+@pytest.mark.parametrize("device", ["cuda", "cpu", "mps"])
+def test_upscale(device):
+    check_device(device)
+    model = RealESRGANUpscaler("realesr-general-x4v3", device)
+    res = model.gen_image(
+        rgb_img,
+        RunPluginRequest(name=RealESRGANUpscaler.name, image=rgb_img_base64, scale=2),
+    )
+    _save(res, f"test_upscale_x2_{device}.png")
+
+    res = model.gen_image(
+        rgb_img,
+        RunPluginRequest(name=RealESRGANUpscaler.name, image=rgb_img_base64, scale=4),
+    )
+    _save(res, f"test_upscale_x4_{device}.png")
+
+
+@pytest.mark.parametrize("device", ["cuda", "cpu", "mps"])
+def test_gfpgan(device):
+    check_device(device)
+    model = GFPGANPlugin(device)
+    res = model.gen_image(
+        rgb_img, RunPluginRequest(name=GFPGANPlugin.name, image=rgb_img_base64)
+    )
+    _save(res, f"test_gfpgan_{device}.png")
+
+
+@pytest.mark.parametrize("device", ["cuda", "cpu", "mps"])
+def test_restoreformer(device):
+    check_device(device)
+    model = RestoreFormerPlugin(device)
+    res = model.gen_image(
+        rgb_img, RunPluginRequest(name=RestoreFormerPlugin.name, image=rgb_img_base64)
+    )
+    _save(res, f"test_restoreformer_{device}.png")
+
+
+@pytest.mark.parametrize("device", ["cuda", "cpu", "mps"])
+def test_segment_anything(device):
+    check_device(device)
+    model = InteractiveSeg("vit_l", device)
+    new_mask = model.gen_mask(
+        rgb_img,
+        RunPluginRequest(
+            name=InteractiveSeg.name,
+            image=rgb_img_base64,
+            clicks=([[448 // 2, 394 // 2, 1]]),
+        ),
+    )
+
+    save_name = f"test_segment_anything_{device}.png"
+    _save(new_mask, save_name)

iopaint/tests/test_save_exif.py

@@ -0,0 +1,59 @@
+import io
+import tempfile
+from pathlib import Path
+from typing import List
+
+from PIL import Image
+
+from iopaint.helper import pil_to_bytes, load_img
+
+current_dir = Path(__file__).parent.absolute().resolve()
+
+
+def print_exif(exif):
+    for k, v in exif.items():
+        print(f"{k}: {v}")
+
+
+def extra_info(img_p: Path):
+    ext = img_p.suffix.strip(".")
+    img_bytes = img_p.read_bytes()
+    np_img, _, infos = load_img(img_bytes, False, True)
+    res_pil_bytes = pil_to_bytes(Image.fromarray(np_img), ext=ext, infos=infos)
+    res_img = Image.open(io.BytesIO(res_pil_bytes))
+    return infos, res_img.info, res_pil_bytes
+
+
+def assert_keys(keys: List[str], infos, res_infos):
+    for k in keys:
+        assert k in infos
+        assert k in res_infos
+        assert infos[k] == res_infos[k]
+
+
+def run_test(file_path, keys):
+    infos, res_infos, res_pil_bytes = extra_info(file_path)
+    assert_keys(keys, infos, res_infos)
+    with tempfile.NamedTemporaryFile("wb", suffix=file_path.suffix) as temp_file:
+        temp_file.write(res_pil_bytes)
+        temp_file.flush()
+        infos, res_infos, res_pil_bytes = extra_info(Path(temp_file.name))
+        assert_keys(keys, infos, res_infos)
+
+
+def test_png_icc_profile_png():
+    run_test(current_dir / "icc_profile_test.png", ["icc_profile", "exif"])
+
+
+def test_png_icc_profile_jpeg():
+    run_test(current_dir / "icc_profile_test.jpg", ["icc_profile", "exif"])
+
+
+def test_jpeg():
+    jpg_img_p = current_dir / "bunny.jpeg"
+    run_test(jpg_img_p, ["dpi", "exif"])
+
+
+def test_png_parameter():
+    jpg_img_p = current_dir / "png_parameter_test.png"
+    run_test(jpg_img_p, ["parameters"])

iopaint/tests/test_sd_model.py

@@ -0,0 +1,237 @@
+import os
+
+from loguru import logger
+
+from iopaint.tests.utils import check_device, get_config, assert_equal
+
+os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
+from pathlib import Path
+
+import pytest
+import torch
+
+from iopaint.model_manager import ModelManager
+from iopaint.schema import HDStrategy, SDSampler, FREEUConfig
+
+current_dir = Path(__file__).parent.absolute().resolve()
+save_dir = current_dir / "result"
+save_dir.mkdir(exist_ok=True, parents=True)
+
+
+@pytest.mark.parametrize("device", ["cuda", "mps"])
+def test_runway_sd_1_5_all_samplers(device):
+    sd_steps = check_device(device)
+    model = ModelManager(
+        name="runwayml/stable-diffusion-inpainting",
+        device=torch.device(device),
+        disable_nsfw=True,
+        sd_cpu_textencoder=False,
+    )
+
+    all_samplers = [member.value for member in SDSampler.__members__.values()]
+    print(all_samplers)
+    for sampler in all_samplers:
+        print(f"Testing sampler {sampler}")
+        if (
+            sampler
+            in [SDSampler.dpm2_karras, SDSampler.dpm2_a_karras, SDSampler.lms_karras]
+            and device == "mps"
+        ):
+            # diffusers 0.25.0 still has bug on these sampler on mps, wait main branch released to fix it
+            logger.warning(
+                "skip dpm2_karras on mps, diffusers does not support it on mps. TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64. Please use float32 instead."
+            )
+            continue
+        cfg = get_config(
+            strategy=HDStrategy.ORIGINAL,
+            prompt="a fox sitting on a bench",
+            sd_steps=sd_steps,
+            sd_sampler=sampler,
+        )
+
+        name = f"device_{device}_{sampler}"
+
+        assert_equal(
+            model,
+            cfg,
+            f"runway_sd_{name}.png",
+            img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
+            mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
+        )
+
+
+@pytest.mark.parametrize("device", ["cuda", "mps", "cpu"])
+@pytest.mark.parametrize("sampler", [SDSampler.lcm])
+def test_runway_sd_lcm_lora(device, sampler):
+    check_device(device)
+
+    sd_steps = 5
+    model = ModelManager(
+        name="runwayml/stable-diffusion-inpainting",
+        device=torch.device(device),
+        disable_nsfw=True,
+        sd_cpu_textencoder=False,
+    )
+    cfg = get_config(
+        strategy=HDStrategy.ORIGINAL,
+        prompt="face of a fox, sitting on a bench",
+        sd_steps=sd_steps,
+        sd_guidance_scale=2,
+        sd_lcm_lora=True,
+    )
+    cfg.sd_sampler = sampler
+
+    assert_equal(
+        model,
+        cfg,
+        f"runway_sd_1_5_lcm_lora_device_{device}.png",
+        img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
+        mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
+    )
+
+
+@pytest.mark.parametrize("device", ["cuda", "mps", "cpu"])
+@pytest.mark.parametrize("sampler", [SDSampler.ddim])
+def test_runway_sd_freeu(device, sampler):
+    sd_steps = check_device(device)
+    model = ModelManager(
+        name="runwayml/stable-diffusion-inpainting",
+        device=torch.device(device),
+        disable_nsfw=True,
+        sd_cpu_textencoder=False,
+    )
+    cfg = get_config(
+        strategy=HDStrategy.ORIGINAL,
+        prompt="face of a fox, sitting on a bench",
+        sd_steps=sd_steps,
+        sd_guidance_scale=7.5,
+        sd_freeu=True,
+        sd_freeu_config=FREEUConfig(),
+    )
+    cfg.sd_sampler = sampler
+
+    assert_equal(
+        model,
+        cfg,
+        f"runway_sd_1_5_freeu_device_{device}.png",
+        img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
+        mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
+    )
+
+
+@pytest.mark.parametrize("device", ["cuda", "mps"])
+@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
+@pytest.mark.parametrize("sampler", [SDSampler.ddim])
+def test_runway_sd_sd_strength(device, strategy, sampler):
+    sd_steps = check_device(device)
+    model = ModelManager(
+        name="runwayml/stable-diffusion-inpainting",
+        device=torch.device(device),
+        disable_nsfw=True,
+        sd_cpu_textencoder=False,
+    )
+    cfg = get_config(
+        strategy=strategy,
+        prompt="a fox sitting on a bench",
+        sd_steps=sd_steps,
+        sd_strength=0.8,
+    )
+    cfg.sd_sampler = sampler
+
+    assert_equal(
+        model,
+        cfg,
+        f"runway_sd_strength_0.8_device_{device}.png",
+        img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
+        mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
+    )
+
+
+@pytest.mark.parametrize("device", ["cuda", "mps", "cpu"])
+@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
+@pytest.mark.parametrize("sampler", [SDSampler.ddim])
+def test_runway_norm_sd_model(device, strategy, sampler):
+    sd_steps = check_device(device)
+    model = ModelManager(
+        name="runwayml/stable-diffusion-v1-5",
+        device=torch.device(device),
+        disable_nsfw=True,
+        sd_cpu_textencoder=False,
+    )
+    cfg = get_config(
+        strategy=strategy, prompt="face of a fox, sitting on a bench", sd_steps=sd_steps
+    )
+    cfg.sd_sampler = sampler
+
+    assert_equal(
+        model,
+        cfg,
+        f"runway_{device}_norm_sd_model_device_{device}.png",
+        img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
+        mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
+    )
+
+
+@pytest.mark.parametrize("device", ["cuda"])
+@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
+@pytest.mark.parametrize("sampler", [SDSampler.dpm_plus_plus_2m])
+def test_runway_sd_1_5_cpu_offload(device, strategy, sampler):
+    sd_steps = check_device(device)
+    model = ModelManager(
+        name="runwayml/stable-diffusion-inpainting",
+        device=torch.device(device),
+        disable_nsfw=True,
+        sd_cpu_textencoder=False,
+        cpu_offload=True,
+    )
+    cfg = get_config(
+        strategy=strategy, prompt="a fox sitting on a bench", sd_steps=sd_steps
+    )
+    cfg.sd_sampler = sampler
+
+    name = f"device_{device}_{sampler}"
+
+    assert_equal(
+        model,
+        cfg,
+        f"runway_sd_{strategy.capitalize()}_{name}_cpu_offload.png",
+        img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
+        mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
+    )
+
+
+@pytest.mark.parametrize("device", ["cuda", "mps", "cpu"])
+@pytest.mark.parametrize("sampler", [SDSampler.ddim])
+@pytest.mark.parametrize(
+    "name",
+    [
+        "sd-v1-5-inpainting.ckpt",
+        "sd-v1-5-inpainting.safetensors",
+        "v1-5-pruned-emaonly.safetensors",
+    ],
+)
+def test_local_file_path(device, sampler, name):
+    sd_steps = check_device(device)
+    model = ModelManager(
+        name=name,
+        device=torch.device(device),
+        disable_nsfw=True,
+        sd_cpu_textencoder=False,
+        cpu_offload=False,
+    )
+    cfg = get_config(
+        strategy=HDStrategy.ORIGINAL,
+        prompt="a fox sitting on a bench",
+        sd_steps=sd_steps,
+    )
+    cfg.sd_sampler = sampler
+
+    name = f"device_{device}_{sampler}_{name}"
+
+    assert_equal(
+        model,
+        cfg,
+        f"sd_local_model_{name}.png",
+        img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
+        mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
+    )

iopaint/tests/test_sdxl.py

@@ -0,0 +1,140 @@
+import os
+
+from iopaint.tests.utils import check_device, current_dir
+
+os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
+
+import pytest
+import torch
+
+from iopaint.model_manager import ModelManager
+from iopaint.schema import HDStrategy, SDSampler, FREEUConfig
+from iopaint.tests.test_model import get_config, assert_equal
+
+
+@pytest.mark.parametrize("device", ["cuda", "mps"])
+@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
+@pytest.mark.parametrize("sampler", [SDSampler.ddim])
+def test_sdxl(device, strategy, sampler):
+    sd_steps = check_device(device)
+
+    model = ModelManager(
+        name="diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
+        device=torch.device(device),
+        disable_nsfw=True,
+        sd_cpu_textencoder=False,
+    )
+    cfg = get_config(
+        strategy=strategy,
+        prompt="face of a fox, sitting on a bench",
+        sd_steps=sd_steps,
+        sd_strength=1.0,
+        sd_guidance_scale=7.0,
+    )
+    cfg.sd_sampler = sampler
+
+    assert_equal(
+        model,
+        cfg,
+        f"sdxl_device_{device}.png",
+        img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
+        mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
+        fx=2,
+        fy=2,
+    )
+
+
+@pytest.mark.parametrize("device", ["cuda", "mps"])
+@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
+@pytest.mark.parametrize("sampler", [SDSampler.ddim])
+def test_sdxl_lcm_lora_and_freeu(device, strategy, sampler):
+    sd_steps = check_device(device)
+
+    model = ModelManager(
+        name="diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
+        device=torch.device(device),
+        disable_nsfw=True,
+        sd_cpu_textencoder=False,
+    )
+    cfg = get_config(
+        strategy=strategy,
+        prompt="face of a fox, sitting on a bench",
+        sd_steps=sd_steps,
+        sd_strength=1.0,
+        sd_guidance_scale=2.0,
+        sd_lcm_lora=True,
+    )
+    cfg.sd_sampler = sampler
+
+    name = f"device_{device}_{sampler}"
+
+    assert_equal(
+        model,
+        cfg,
+        f"sdxl_{name}_lcm_lora.png",
+        img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
+        mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
+        fx=2,
+        fy=2,
+    )
+
+    cfg = get_config(
+        strategy=strategy,
+        prompt="face of a fox, sitting on a bench",
+        sd_steps=sd_steps,
+        sd_guidance_scale=7.5,
+        sd_freeu=True,
+        sd_freeu_config=FREEUConfig(),
+    )
+
+    assert_equal(
+        model,
+        cfg,
+        f"sdxl_{name}_freeu_device_{device}.png",
+        img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
+        mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
+        fx=2,
+        fy=2,
+    )
+
+
+@pytest.mark.parametrize("device", ["cuda", "mps"])
+@pytest.mark.parametrize(
+    "rect",
+    [
+        [-128, -128, 1024, 1024],
+    ],
+)
+def test_sdxl_outpainting(device, rect):
+    sd_steps = check_device(device)
+
+    model = ModelManager(
+        name="diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
+        device=torch.device(device),
+        disable_nsfw=True,
+        sd_cpu_textencoder=False,
+    )
+
+    cfg = get_config(
+        strategy=HDStrategy.ORIGINAL,
+        prompt="a dog sitting on a bench in the park",
+        sd_steps=sd_steps,
+        use_extender=True,
+        extender_x=rect[0],
+        extender_y=rect[1],
+        extender_width=rect[2],
+        extender_height=rect[3],
+        sd_strength=1.0,
+        sd_guidance_scale=8.0,
+        sd_sampler=SDSampler.ddim,
+    )
+
+    assert_equal(
+        model,
+        cfg,
+        f"sdxl_outpainting_dog_ddim_{'_'.join(map(str, rect))}_device_{device}.png",
+        img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
+        mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
+        fx=1.5,
+        fy=1.5,
+    )

iopaint/tests/utils.py

@@ -0,0 +1,77 @@
+from pathlib import Path
+import cv2
+import pytest
+import torch
+
+from iopaint.helper import encode_pil_to_base64
+from iopaint.schema import LDMSampler, HDStrategy, InpaintRequest, SDSampler
+from PIL import Image
+
+current_dir = Path(__file__).parent.absolute().resolve()
+save_dir = current_dir / "result"
+save_dir.mkdir(exist_ok=True, parents=True)
+
+
+def check_device(device: str) -> int:
+    if device == "cuda" and not torch.cuda.is_available():
+        pytest.skip("CUDA is not available, skip test on cuda")
+    if device == "mps" and not torch.backends.mps.is_available():
+        pytest.skip("mps is not available, skip test on mps")
+    steps = 1 if device == "cpu" else 20
+    return steps
+
+
+def assert_equal(
+    model,
+    config: InpaintRequest,
+    gt_name,
+    fx: float = 1,
+    fy: float = 1,
+    img_p=current_dir / "image.png",
+    mask_p=current_dir / "mask.png",
+):
+    img, mask = get_data(fx=fx, fy=fy, img_p=img_p, mask_p=mask_p)
+    print(f"Input image shape: {img.shape}")
+    res = model(img, mask, config)
+    ok = cv2.imwrite(
+        str(save_dir / gt_name),
+        res,
+        [int(cv2.IMWRITE_JPEG_QUALITY), 100, int(cv2.IMWRITE_PNG_COMPRESSION), 0],
+    )
+    assert ok, save_dir / gt_name
+
+    """
+    Note that JPEG is lossy compression, so even if it is the highest quality 100, 
+    when the saved images is reloaded, a difference occurs with the original pixel value. 
+    If you want to save the original images as it is, save it as PNG or BMP.
+    """
+    # gt = cv2.imread(str(current_dir / gt_name), cv2.IMREAD_UNCHANGED)
+    # assert np.array_equal(res, gt)
+
+
+def get_data(
+    fx: float = 1,
+    fy: float = 1.0,
+    img_p=current_dir / "image.png",
+    mask_p=current_dir / "mask.png",
+):
+    img = cv2.imread(str(img_p))
+    img = cv2.cvtColor(img, cv2.COLOR_BGRA2RGB)
+    mask = cv2.imread(str(mask_p), cv2.IMREAD_GRAYSCALE)
+    img = cv2.resize(img, None, fx=fx, fy=fy, interpolation=cv2.INTER_AREA)
+    mask = cv2.resize(mask, None, fx=fx, fy=fy, interpolation=cv2.INTER_NEAREST)
+    return img, mask
+
+
+def get_config(**kwargs):
+    data = dict(
+        sd_sampler=kwargs.get("sd_sampler", SDSampler.uni_pc),
+        ldm_steps=1,
+        ldm_sampler=LDMSampler.plms,
+        hd_strategy=kwargs.get("strategy", HDStrategy.ORIGINAL),
+        hd_strategy_crop_margin=32,
+        hd_strategy_crop_trigger_size=200,
+        hd_strategy_resize_limit=200,
+    )
+    data.update(**kwargs)
+    return InpaintRequest(image="", mask="", **data)