This commit is contained in:
Qing 2023-12-28 10:48:52 +08:00
parent f0b852725f
commit 9cc9bd7a88
19 changed files with 345 additions and 482 deletions

1
.gitignore vendored
View File

@ -8,3 +8,4 @@ build
dist/ dist/
lama_cleaner.egg-info/ lama_cleaner.egg-info/
venv/ venv/
tmp/

View File

@ -72,7 +72,7 @@ def start(
enable_anime_seg: bool = Option(False, help=ANIMESEG_HELP), enable_anime_seg: bool = Option(False, help=ANIMESEG_HELP),
enable_realesrgan: bool = Option(False), enable_realesrgan: bool = Option(False),
realesrgan_device: Device = Option(Device.cpu), realesrgan_device: Device = Option(Device.cpu),
realesrgan_model: str = Option(RealESRGANModel.realesr_general_x4v3), realesrgan_model: RealESRGANModel = Option(RealESRGANModel.realesr_general_x4v3),
enable_gfpgan: bool = Option(False), enable_gfpgan: bool = Option(False),
gfpgan_device: Device = Option(Device.cpu), gfpgan_device: Device = Option(Device.cpu),
enable_restoreformer: bool = Option(False), enable_restoreformer: bool = Option(False),

View File

@ -37,7 +37,7 @@ class ModelManager:
def init_model(self, name: str, device, **kwargs): def init_model(self, name: str, device, **kwargs):
logger.info(f"Loading model: {name}") logger.info(f"Loading model: {name}")
if name not in self.available_models: if name not in self.available_models:
raise NotImplementedError(f"Unsupported model: {name}") raise NotImplementedError(f"Unsupported model: {name}. Available models: {self.available_models.keys()}")
model_info = self.available_models[name] model_info = self.available_models[name]
kwargs = { kwargs = {

View File

@ -6,7 +6,7 @@ from .realesrgan import RealESRGANUpscaler
from .gfpgan_plugin import GFPGANPlugin from .gfpgan_plugin import GFPGANPlugin
from .restoreformer import RestoreFormerPlugin from .restoreformer import RestoreFormerPlugin
from .anime_seg import AnimeSeg from .anime_seg import AnimeSeg
from ..const import InteractiveSegModel, Device from ..const import InteractiveSegModel, Device, RealESRGANModel
def build_plugins( def build_plugins(
@ -18,7 +18,7 @@ def build_plugins(
enable_anime_seg: bool, enable_anime_seg: bool,
enable_realesrgan: bool, enable_realesrgan: bool,
realesrgan_device: Device, realesrgan_device: Device,
realesrgan_model: str, realesrgan_model: RealESRGANModel,
enable_gfpgan: bool, enable_gfpgan: bool,
gfpgan_device: Device, gfpgan_device: Device,
enable_restoreformer: bool, enable_restoreformer: bool,

View File

@ -507,7 +507,7 @@ def start(
enable_anime_seg: bool, enable_anime_seg: bool,
enable_realesrgan: bool, enable_realesrgan: bool,
realesrgan_device: Device, realesrgan_device: Device,
realesrgan_model: str, realesrgan_model: RealESRGANModel,
enable_gfpgan: bool, enable_gfpgan: bool,
gfpgan_device: Device, gfpgan_device: Device,
enable_restoreformer: bool, enable_restoreformer: bool,
@ -525,6 +525,7 @@ def start(
output_dir, "lama_cleaner_thumbnails" output_dir, "lama_cleaner_thumbnails"
) )
file_manager.output_dir = output_dir file_manager.output_dir = output_dir
global_config.file_manager = file_manager
else: else:
global_config.input_image_path = input global_config.input_image_path = input

View File

@ -1,6 +1,7 @@
import os import os
from lama_cleaner.const import SD_CONTROLNET_CHOICES from lama_cleaner.const import SD_CONTROLNET_CHOICES
from lama_cleaner.tests.utils import current_dir, check_device, get_config, assert_equal
os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1" os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
from pathlib import Path from pathlib import Path
@ -10,178 +11,107 @@ import torch
from lama_cleaner.model_manager import ModelManager from lama_cleaner.model_manager import ModelManager
from lama_cleaner.schema import HDStrategy, SDSampler from lama_cleaner.schema import HDStrategy, SDSampler
from lama_cleaner.tests.test_model import get_config, assert_equal
current_dir = Path(__file__).parent.absolute().resolve()
save_dir = current_dir / "result"
save_dir.mkdir(exist_ok=True, parents=True)
device = "cuda" if torch.cuda.is_available() else "cpu"
device = torch.device(device)
model_name = "runwayml/stable-diffusion-inpainting" model_name = "runwayml/stable-diffusion-inpainting"
@pytest.mark.parametrize("sd_device", ["cuda", "mps"]) def convert_controlnet_method_name(name):
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL]) return name.replace("/", "--")
@pytest.mark.parametrize("sampler", [SDSampler.uni_pc])
@pytest.mark.parametrize("cpu_textencoder", [True])
@pytest.mark.parametrize("disable_nsfw", [True]) @pytest.mark.parametrize("device", ["cuda", "mps", "cpu"])
@pytest.mark.parametrize("sd_controlnet_method", SD_CONTROLNET_CHOICES) @pytest.mark.parametrize("controlnet_method", [SD_CONTROLNET_CHOICES[0]])
def test_runway_sd_1_5( def test_runway_sd_1_5(device, controlnet_method):
sd_device, strategy, sampler, cpu_textencoder, disable_nsfw, sd_controlnet_method sd_steps = check_device(device)
):
if sd_device == "cuda" and not torch.cuda.is_available():
return
if device == "mps" and not torch.backends.mps.is_available():
return
sd_steps = 1 if sd_device == "cpu" else 30
model = ModelManager( model = ModelManager(
name=model_name, name=model_name,
sd_controlnet=True, device=torch.device(device),
device=torch.device(sd_device), disable_nsfw=True,
disable_nsfw=disable_nsfw, sd_cpu_textencoder=True,
sd_cpu_textencoder=cpu_textencoder, enable_controlnet=True,
sd_controlnet_method=sd_controlnet_method, controlnet_method=controlnet_method,
) )
controlnet_conditioning_scale = {
"control_v11p_sd15_canny": 0.4,
"control_v11p_sd15_openpose": 0.4,
"control_v11p_sd15_inpaint": 1.0,
"control_v11f1p_sd15_depth": 1.0,
}[sd_controlnet_method]
cfg = get_config( cfg = get_config(
strategy,
prompt="a fox sitting on a bench", prompt="a fox sitting on a bench",
sd_steps=sd_steps, sd_steps=sd_steps,
controlnet_conditioning_scale=controlnet_conditioning_scale, enable_controlnet=True,
controlnet_method=sd_controlnet_method, controlnet_conditioning_scale=0.5,
controlnet_method=controlnet_method,
) )
cfg.sd_sampler = sampler name = f"device_{device}"
name = f"device_{sd_device}_{sampler}_cpu_textencoder_disable_nsfw"
assert_equal( assert_equal(
model, model,
cfg, cfg,
f"sd_controlnet_{sd_controlnet_method}_{name}.png", f"sd_controlnet_{convert_controlnet_method_name(controlnet_method)}_{name}.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png", img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png", mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
fx=1.2,
fy=1.2,
) )
@pytest.mark.parametrize("sd_device", ["cuda", "mps"]) @pytest.mark.parametrize("device", ["cuda", "mps", "cpu"])
@pytest.mark.parametrize("sampler", [SDSampler.uni_pc]) def test_controlnet_switch(device):
def test_local_file_path(sd_device, sampler): sd_steps = check_device(device)
if sd_device == "cuda" and not torch.cuda.is_available():
return
if device == "mps" and not torch.backends.mps.is_available():
return
sd_steps = 1 if sd_device == "cpu" else 30
model = ModelManager( model = ModelManager(
name=model_name, name=model_name,
sd_controlnet=True, device=torch.device(device),
device=torch.device(sd_device),
disable_nsfw=True, disable_nsfw=True,
sd_cpu_textencoder=False, sd_cpu_textencoder=False,
cpu_offload=True, cpu_offload=True,
sd_controlnet_method="control_v11p_sd15_canny", enable_controlnet=True,
controlnet_method="lllyasviel/control_v11p_sd15_canny",
) )
cfg = get_config( cfg = get_config(
HDStrategy.ORIGINAL,
prompt="a fox sitting on a bench", prompt="a fox sitting on a bench",
sd_steps=sd_steps, sd_steps=sd_steps,
controlnet_method="control_v11p_sd15_canny", enable_controlnet=True,
controlnet_method="lllyasviel/control_v11f1p_sd15_depth",
) )
cfg.sd_sampler = sampler
name = f"device_{sd_device}_{sampler}"
assert_equal( assert_equal(
model, model,
cfg, cfg,
f"sd_controlnet_canny_local_model_{name}.png", f"controlnet_switch_canny_to_depth_device_{device}.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png", img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png", mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
) )
@pytest.mark.parametrize("sd_device", ["cuda", "mps"]) @pytest.mark.parametrize("device", ["cuda", "mps", "cpu"])
@pytest.mark.parametrize("sampler", [SDSampler.uni_pc]) @pytest.mark.parametrize(
def test_local_file_path_controlnet_native_inpainting(sd_device, sampler): "local_file", ["sd-v1-5-inpainting.ckpt", "v1-5-pruned-emaonly.safetensors"]
if sd_device == "cuda" and not torch.cuda.is_available(): )
return def test_local_file_path(device, local_file):
if device == "mps" and not torch.backends.mps.is_available(): sd_steps = check_device(device)
return
controlnet_kwargs = dict(
enable_controlnet=True,
controlnet_method=SD_CONTROLNET_CHOICES[0],
)
sd_steps = 1 if sd_device == "cpu" else 30
model = ModelManager( model = ModelManager(
name=model_name, name=local_file,
sd_controlnet=True, device=torch.device(device),
device=torch.device(sd_device),
disable_nsfw=True, disable_nsfw=True,
sd_cpu_textencoder=False, sd_cpu_textencoder=False,
cpu_offload=True, cpu_offload=True,
sd_local_model_path="/Users/cwq/data/models/v1-5-pruned-emaonly.safetensors", **controlnet_kwargs,
sd_controlnet_method="control_v11p_sd15_inpaint",
) )
cfg = get_config( cfg = get_config(
HDStrategy.ORIGINAL,
prompt="a fox sitting on a bench", prompt="a fox sitting on a bench",
sd_steps=sd_steps, sd_steps=sd_steps,
controlnet_conditioning_scale=1.0, **controlnet_kwargs,
sd_strength=1.0,
controlnet_method="control_v11p_sd15_inpaint",
) )
cfg.sd_sampler = sampler
name = f"device_{sd_device}_{sampler}" name = f"device_{device}"
assert_equal( assert_equal(
model, model,
cfg, cfg,
f"sd_controlnet_local_native_{name}.png", 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",
)
@pytest.mark.parametrize("sd_device", ["cuda", "mps"])
@pytest.mark.parametrize("sampler", [SDSampler.uni_pc])
def test_controlnet_switch(sd_device, sampler):
if sd_device == "cuda" and not torch.cuda.is_available():
return
if device == "mps" and not torch.backends.mps.is_available():
return
sd_steps = 1 if sd_device == "cpu" else 30
model = ModelManager(
name=model_name,
sd_controlnet=True,
device=torch.device(sd_device),
disable_nsfw=True,
sd_cpu_textencoder=False,
cpu_offload=True,
sd_controlnet_method="control_v11p_sd15_canny",
)
cfg = get_config(
HDStrategy.ORIGINAL,
prompt="a fox sitting on a bench",
sd_steps=sd_steps,
controlnet_method="control_v11p_sd15_inpaint",
)
cfg.sd_sampler = sampler
name = f"device_{sd_device}_{sampler}"
assert_equal(
model,
cfg,
f"sd_controlnet_switch_to_inpaint_local_model_{name}.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png", img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png", mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
) )

View File

@ -4,20 +4,17 @@ import pytest
import torch import torch
from lama_cleaner.model_manager import ModelManager from lama_cleaner.model_manager import ModelManager
from lama_cleaner.tests.test_model import get_config, assert_equal
from lama_cleaner.schema import HDStrategy from lama_cleaner.schema import HDStrategy
from lama_cleaner.tests.utils import get_config, check_device, assert_equal, current_dir
current_dir = Path(__file__).parent.absolute().resolve()
save_dir = current_dir / "result"
save_dir.mkdir(exist_ok=True, parents=True)
device = "cuda" if torch.cuda.is_available() else "mps"
model_name = "timbrooks/instruct-pix2pix" model_name = "timbrooks/instruct-pix2pix"
@pytest.mark.parametrize("device", ["cuda", "mps", "cpu"])
@pytest.mark.parametrize("disable_nsfw", [True, False]) @pytest.mark.parametrize("disable_nsfw", [True, False])
@pytest.mark.parametrize("cpu_offload", [False, True]) @pytest.mark.parametrize("cpu_offload", [False, True])
def test_instruct_pix2pix(disable_nsfw, cpu_offload): def test_instruct_pix2pix(device, disable_nsfw, cpu_offload):
sd_steps = 50 if device == "cuda" else 20 sd_steps = check_device(device)
model = ModelManager( model = ModelManager(
name=model_name, name=model_name,
device=torch.device(device), device=torch.device(device),
@ -42,31 +39,3 @@ def test_instruct_pix2pix(disable_nsfw, cpu_offload):
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png", mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
fx=1.3, fx=1.3,
) )
@pytest.mark.parametrize("disable_nsfw", [False])
@pytest.mark.parametrize("cpu_offload", [False])
def test_instruct_pix2pix_snow(disable_nsfw, cpu_offload):
sd_steps = 50 if device == "cuda" else 20
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,
)
name = f"snow"
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",
)

View File

@ -1,8 +1,6 @@
from pathlib import Path
from lama_cleaner.helper import load_img from lama_cleaner.helper import load_img
from lama_cleaner.tests.utils import current_dir
current_dir = Path(__file__).parent.absolute().resolve()
png_img_p = current_dir / "image.png" png_img_p = current_dir / "image.png"
jpg_img_p = current_dir / "bunny.jpeg" jpg_img_p = current_dir / "bunny.jpeg"

View File

@ -1,100 +1,42 @@
from pathlib import Path
import cv2
import pytest import pytest
import torch import torch
from lama_cleaner.model_manager import ModelManager from lama_cleaner.model_manager import ModelManager
from lama_cleaner.schema import Config, HDStrategy, LDMSampler, SDSampler from lama_cleaner.schema import HDStrategy, LDMSampler
from lama_cleaner.tests.utils import assert_equal, get_config, current_dir, check_device
current_dir = Path(__file__).parent.absolute().resolve()
save_dir = current_dir / "result"
save_dir.mkdir(exist_ok=True, parents=True)
device = "cuda" if torch.cuda.is_available() else "cpu"
device = torch.device(device)
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(strategy, **kwargs):
data = dict(
ldm_steps=1,
ldm_sampler=LDMSampler.plms,
hd_strategy=strategy,
hd_strategy_crop_margin=32,
hd_strategy_crop_trigger_size=200,
hd_strategy_resize_limit=200,
)
data.update(**kwargs)
return Config(**data)
def assert_equal(
model,
config,
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)
cv2.imwrite(
str(save_dir / gt_name),
res,
[int(cv2.IMWRITE_JPEG_QUALITY), 100, int(cv2.IMWRITE_PNG_COMPRESSION), 0],
)
"""
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)
@pytest.mark.parametrize("device", ["cuda", "mps", "cpu"])
@pytest.mark.parametrize( @pytest.mark.parametrize(
"strategy", [HDStrategy.ORIGINAL, HDStrategy.RESIZE, HDStrategy.CROP] "strategy", [HDStrategy.ORIGINAL, HDStrategy.RESIZE, HDStrategy.CROP]
) )
def test_lama(strategy): def test_lama(device, strategy):
check_device(device)
model = ModelManager(name="lama", device=device) model = ModelManager(name="lama", device=device)
assert_equal( assert_equal(
model, model,
get_config(strategy), get_config(strategy=strategy),
f"lama_{strategy[0].upper() + strategy[1:]}_result.png", f"lama_{strategy[0].upper() + strategy[1:]}_result.png",
) )
fx = 1.3 fx = 1.3
assert_equal( assert_equal(
model, model,
get_config(strategy), get_config(strategy=strategy),
f"lama_{strategy[0].upper() + strategy[1:]}_fx_{fx}_result.png", f"lama_{strategy[0].upper() + strategy[1:]}_fx_{fx}_result.png",
fx=1.3, fx=1.3,
) )
@pytest.mark.parametrize("device", ["cuda", "cpu"])
@pytest.mark.parametrize( @pytest.mark.parametrize(
"strategy", [HDStrategy.ORIGINAL, HDStrategy.RESIZE, HDStrategy.CROP] "strategy", [HDStrategy.ORIGINAL, HDStrategy.RESIZE, HDStrategy.CROP]
) )
@pytest.mark.parametrize("ldm_sampler", [LDMSampler.ddim, LDMSampler.plms]) @pytest.mark.parametrize("ldm_sampler", [LDMSampler.ddim, LDMSampler.plms])
def test_ldm(strategy, ldm_sampler): def test_ldm(device, strategy, ldm_sampler):
check_device(device)
model = ModelManager(name="ldm", device=device) model = ModelManager(name="ldm", device=device)
cfg = get_config(strategy, ldm_sampler=ldm_sampler) cfg = get_config(strategy=strategy, ldm_sampler=ldm_sampler)
assert_equal( assert_equal(
model, cfg, f"ldm_{strategy[0].upper() + strategy[1:]}_{ldm_sampler}_result.png" model, cfg, f"ldm_{strategy[0].upper() + strategy[1:]}_{ldm_sampler}_result.png"
) )
@ -108,15 +50,15 @@ def test_ldm(strategy, ldm_sampler):
) )
@pytest.mark.parametrize("device", ["cuda", "cpu"])
@pytest.mark.parametrize( @pytest.mark.parametrize(
"strategy", [HDStrategy.ORIGINAL, HDStrategy.RESIZE, HDStrategy.CROP] "strategy", [HDStrategy.ORIGINAL, HDStrategy.RESIZE, HDStrategy.CROP]
) )
@pytest.mark.parametrize("zits_wireframe", [False, True]) @pytest.mark.parametrize("zits_wireframe", [False, True])
def test_zits(strategy, zits_wireframe): def test_zits(device, strategy, zits_wireframe):
check_device(device)
model = ModelManager(name="zits", device=device) model = ModelManager(name="zits", device=device)
cfg = get_config(strategy, zits_wireframe=zits_wireframe) cfg = get_config(strategy=strategy, zits_wireframe=zits_wireframe)
# os.environ['ZITS_DEBUG_LINE_PATH'] = str(current_dir / 'zits_debug_line.jpg')
# os.environ['ZITS_DEBUG_EDGE_PATH'] = str(current_dir / 'zits_debug_edge.jpg')
assert_equal( assert_equal(
model, model,
cfg, cfg,
@ -132,13 +74,14 @@ def test_zits(strategy, zits_wireframe):
) )
@pytest.mark.parametrize("device", ["cuda", "cpu"])
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL]) @pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
@pytest.mark.parametrize("no_half", [True, False]) @pytest.mark.parametrize("no_half", [True, False])
def test_mat(strategy, no_half): def test_mat(device, strategy, no_half):
check_device(device)
model = ModelManager(name="mat", device=device, no_half=no_half) model = ModelManager(name="mat", device=device, no_half=no_half)
cfg = get_config(strategy) cfg = get_config(strategy=strategy)
for _ in range(10):
assert_equal( assert_equal(
model, model,
cfg, cfg,
@ -146,13 +89,14 @@ def test_mat(strategy, no_half):
) )
@pytest.mark.parametrize("device", ["cuda", "cpu"])
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL]) @pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
def test_fcf(strategy): def test_fcf(device, strategy):
check_device(device)
model = ModelManager(name="fcf", device=device) model = ModelManager(name="fcf", device=device)
cfg = get_config(strategy) 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=2, fy=2)
assert_equal(model, cfg, f"fcf_{strategy.capitalize()}_result.png", fx=3.8, fy=2) assert_equal(model, cfg, f"fcf_{strategy.capitalize()}_result.png", fx=3.8, fy=2)
@ -164,9 +108,9 @@ def test_fcf(strategy):
def test_cv2(strategy, cv2_flag, cv2_radius): def test_cv2(strategy, cv2_flag, cv2_radius):
model = ModelManager( model = ModelManager(
name="cv2", name="cv2",
device=torch.device(device), device=torch.device("cpu"),
) )
cfg = get_config(strategy, cv2_flag=cv2_flag, cv2_radius=cv2_radius) cfg = get_config(strategy=strategy, cv2_flag=cv2_flag, cv2_radius=cv2_radius)
assert_equal( assert_equal(
model, model,
cfg, cfg,
@ -176,15 +120,17 @@ def test_cv2(strategy, cv2_flag, cv2_radius):
) )
@pytest.mark.parametrize("device", ["cuda", "cpu"])
@pytest.mark.parametrize( @pytest.mark.parametrize(
"strategy", [HDStrategy.ORIGINAL, HDStrategy.RESIZE, HDStrategy.CROP] "strategy", [HDStrategy.ORIGINAL, HDStrategy.RESIZE, HDStrategy.CROP]
) )
def test_manga(strategy): def test_manga(device, strategy):
check_device(device)
model = ModelManager( model = ModelManager(
name="manga", name="manga",
device=torch.device(device), device=torch.device(device),
) )
cfg = get_config(strategy) cfg = get_config(strategy=strategy)
assert_equal( assert_equal(
model, model,
cfg, cfg,
@ -194,17 +140,21 @@ def test_manga(strategy):
) )
@pytest.mark.parametrize("device", ["cuda", "mps", "cpu"])
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL]) @pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
def test_mi_gan(strategy): def test_mi_gan(device, strategy):
check_device(device)
model = ModelManager( model = ModelManager(
name="migan", name="migan",
device=torch.device(device), device=torch.device(device),
) )
cfg = get_config(strategy) cfg = get_config(strategy=strategy)
assert_equal( assert_equal(
model, model,
cfg, cfg,
f"migan_{strategy.capitalize()}.png", f"migan_device_{device}.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png", img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png", mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
fx=1.5,
fy=1.7
) )

View File

@ -2,16 +2,9 @@ def test_load_model():
from lama_cleaner.plugins import InteractiveSeg from lama_cleaner.plugins import InteractiveSeg
from lama_cleaner.model_manager import ModelManager from lama_cleaner.model_manager import ModelManager
interactive_seg_model = InteractiveSeg('vit_l', 'cpu') interactive_seg_model = InteractiveSeg("vit_l", "cpu")
models = [ models = ["lama", "ldm", "zits", "mat", "fcf", "manga", "migan"]
"lama",
"ldm",
"zits",
"mat",
"fcf",
"manga",
]
for m in models: for m in models:
ModelManager( ModelManager(
name=m, name=m,
@ -21,25 +14,3 @@ def test_load_model():
sd_cpu_textencoder=True, sd_cpu_textencoder=True,
cpu_offload=True, cpu_offload=True,
) )
# def create_empty_file(tmp_dir, name):
# tmp_model_dir = os.path.join(tmp_dir, "torch", "hub", "checkpoints")
# Path(tmp_model_dir).mkdir(exist_ok=True, parents=True)
# path = os.path.join(tmp_model_dir, name)
# with open(path, "w") as f:
# f.write("1")
#
#
# def test_load_model_error():
# MODELS = [
# ("big-lama.pt", "e3aa4aaa15225a33ec84f9f4bc47e500"),
# ("cond_stage_model_encode.pt", "23239fc9081956a3e70de56472b3f296"),
# ("cond_stage_model_decode.pt", "fe419cd15a750d37a4733589d0d3585c"),
# ("diffusion.pt", "b0afda12bf790c03aba2a7431f11d22d"),
# ]
# with tempfile.TemporaryDirectory() as tmp_dir:
# os.environ["XDG_CACHE_HOME"] = tmp_dir
# for name, md5 in MODELS:
# create_empty_file(tmp_dir, name)
# test_load_model()

View File

@ -1,4 +1,3 @@
import logging
import os import os
from lama_cleaner.schema import Config from lama_cleaner.schema import Config
@ -13,8 +12,8 @@ from lama_cleaner.model_manager import ModelManager
def test_model_switch(): def test_model_switch():
model = ModelManager( model = ModelManager(
name="runwayml/stable-diffusion-inpainting", name="runwayml/stable-diffusion-inpainting",
sd_controlnet=True, enable_controlnet=True,
sd_controlnet_method="lllyasviel/control_v11p_sd15_canny", controlnet_method="lllyasviel/control_v11p_sd15_canny",
device=torch.device("mps"), device=torch.device("mps"),
disable_nsfw=True, disable_nsfw=True,
sd_cpu_textencoder=True, sd_cpu_textencoder=True,
@ -29,8 +28,8 @@ def test_controlnet_switch_onoff(caplog):
name = "runwayml/stable-diffusion-inpainting" name = "runwayml/stable-diffusion-inpainting"
model = ModelManager( model = ModelManager(
name=name, name=name,
sd_controlnet=True, enable_controlnet=True,
sd_controlnet_method="lllyasviel/control_v11p_sd15_canny", controlnet_method="lllyasviel/control_v11p_sd15_canny",
device=torch.device("mps"), device=torch.device("mps"),
disable_nsfw=True, disable_nsfw=True,
sd_cpu_textencoder=True, sd_cpu_textencoder=True,
@ -41,21 +40,21 @@ def test_controlnet_switch_onoff(caplog):
model.switch_controlnet_method( model.switch_controlnet_method(
Config( Config(
name=name, name=name,
controlnet_enabled=False, enable_controlnet=False,
) )
) )
assert "Disable controlnet" in caplog.text assert "Disable controlnet" in caplog.text
def test_controlnet_switch_method(caplog): def test_switch_controlnet_method(caplog):
name = "runwayml/stable-diffusion-inpainting" name = "runwayml/stable-diffusion-inpainting"
old_method = "lllyasviel/control_v11p_sd15_canny" old_method = "lllyasviel/control_v11p_sd15_canny"
new_method = "lllyasviel/control_v11p_sd15_openpose" new_method = "lllyasviel/control_v11p_sd15_openpose"
model = ModelManager( model = ModelManager(
name=name, name=name,
sd_controlnet=True, enable_controlnet=True,
sd_controlnet_method=old_method, controlnet_method=old_method,
device=torch.device("mps"), device=torch.device("mps"),
disable_nsfw=True, disable_nsfw=True,
sd_cpu_textencoder=True, sd_cpu_textencoder=True,
@ -66,7 +65,7 @@ def test_controlnet_switch_method(caplog):
model.switch_controlnet_method( model.switch_controlnet_method(
Config( Config(
name=name, name=name,
controlnet_enabled=True, enable_controlnet=True,
controlnet_method=new_method, controlnet_method=new_method,
) )
) )

View File

@ -1,5 +1,7 @@
import os import os
from lama_cleaner.tests.utils import current_dir, check_device
os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1" os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
from pathlib import Path from pathlib import Path
@ -10,15 +12,9 @@ from lama_cleaner.model_manager import ModelManager
from lama_cleaner.schema import HDStrategy, SDSampler from lama_cleaner.schema import HDStrategy, SDSampler
from lama_cleaner.tests.test_model import get_config, assert_equal from lama_cleaner.tests.test_model import get_config, assert_equal
current_dir = Path(__file__).parent.absolute().resolve()
save_dir = current_dir / "result"
save_dir.mkdir(exist_ok=True, parents=True)
device = "cuda" if torch.cuda.is_available() else "cpu"
device = torch.device(device)
@pytest.mark.parametrize("name", ["runwayml/stable-diffusion-inpainting"]) @pytest.mark.parametrize("name", ["runwayml/stable-diffusion-inpainting"])
@pytest.mark.parametrize("sd_device", ["mps"]) @pytest.mark.parametrize("device", ["cuda", "mps", "cpu"])
@pytest.mark.parametrize( @pytest.mark.parametrize(
"rect", "rect",
[ [
@ -31,24 +27,22 @@ device = torch.device(device)
[-100, -100, 512 + 200, 512 + 200], [-100, -100, 512 + 200, 512 + 200],
], ],
) )
def test_outpainting(name, sd_device, rect): def test_outpainting(name, device, rect):
sd_steps = check_device(device)
def callback(i, t, latents): def callback(i, t, latents):
pass pass
if sd_device == "cuda" and not torch.cuda.is_available():
return
model = ModelManager( model = ModelManager(
name=name, name=name,
device=torch.device(sd_device), device=torch.device(device),
disable_nsfw=True, disable_nsfw=True,
sd_cpu_textencoder=False, sd_cpu_textencoder=False,
callback=callback, callback=callback,
) )
cfg = get_config( cfg = get_config(
HDStrategy.ORIGINAL,
prompt="a dog sitting on a bench in the park", prompt="a dog sitting on a bench in the park",
sd_steps=50, sd_steps=sd_steps,
use_extender=True, use_extender=True,
extender_x=rect[0], extender_x=rect[0],
extender_y=rect[1], extender_y=rect[1],
@ -61,39 +55,37 @@ def test_outpainting(name, sd_device, rect):
assert_equal( assert_equal(
model, model,
cfg, cfg,
f"{name.replace('/', '--')}_outpainting_dpm++_{'_'.join(map(str, rect))}.png", f"{name.replace('/', '--')}_outpainting_{'_'.join(map(str, rect))}_device_{device}.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png", img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png", mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
) )
@pytest.mark.parametrize("name", ["kandinsky-community/kandinsky-2-2-decoder-inpaint"]) @pytest.mark.parametrize("name", ["kandinsky-community/kandinsky-2-2-decoder-inpaint"])
@pytest.mark.parametrize("sd_device", ["mps"]) @pytest.mark.parametrize("device", ["cuda", "mps", "cpu"])
@pytest.mark.parametrize( @pytest.mark.parametrize(
"rect", "rect",
[ [
[-128, -128, 768, 768], [-128, -128, 768, 768],
], ],
) )
def test_kandinsky_outpainting(name, sd_device, rect): def test_kandinsky_outpainting(name, device, rect):
sd_steps = check_device(device)
def callback(i, t, latents): def callback(i, t, latents):
pass pass
if sd_device == "cuda" and not torch.cuda.is_available():
return
model = ModelManager( model = ModelManager(
name=name, name=name,
device=torch.device(sd_device), device=torch.device(device),
disable_nsfw=True, disable_nsfw=True,
sd_cpu_textencoder=False, sd_cpu_textencoder=False,
callback=callback, callback=callback,
) )
cfg = get_config( cfg = get_config(
HDStrategy.ORIGINAL,
prompt="a cat", 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", 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=50, sd_steps=sd_steps,
use_extender=True, use_extender=True,
extender_x=rect[0], extender_x=rect[0],
extender_y=rect[1], extender_y=rect[1],
@ -106,9 +98,52 @@ def test_kandinsky_outpainting(name, sd_device, rect):
assert_equal( assert_equal(
model, model,
cfg, cfg,
f"{name.replace('/', '--')}_outpainting_dpm++_{'_'.join(map(str, rect))}.png", f"{name.replace('/', '--')}_outpainting_{'_'.join(map(str, rect))}_device_{device}.png",
img_p=current_dir / "cat.png", img_p=current_dir / "cat.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png", mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
fx=1, fx=1,
fy=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)
def callback(i, t, latents):
pass
model = ModelManager(
name=name,
device=torch.device(device),
disable_nsfw=True,
sd_cpu_textencoder=False,
callback=callback,
)
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,
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",
)

View File

@ -1,26 +1,24 @@
from pathlib import Path
import cv2 import cv2
import pytest import pytest
import torch
from PIL import Image from PIL import Image
from lama_cleaner.model_manager import ModelManager from lama_cleaner.model_manager import ModelManager
from lama_cleaner.schema import HDStrategy from lama_cleaner.schema import HDStrategy
from lama_cleaner.tests.test_model import get_config, get_data from lama_cleaner.tests.utils import (
current_dir,
get_config,
get_data,
save_dir,
check_device,
)
current_dir = Path(__file__).parent.absolute().resolve()
save_dir = current_dir / "result"
save_dir.mkdir(exist_ok=True, parents=True)
device = "cuda" if torch.cuda.is_available() else "mps"
device = torch.device(device)
model_name = "Fantasy-Studio/Paint-by-Example" model_name = "Fantasy-Studio/Paint-by-Example"
def assert_equal( def assert_equal(
model, model,
config, config,
gt_name, save_name: str,
fx: float = 1, fx: float = 1,
fy: float = 1, fy: float = 1,
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png", img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
@ -38,46 +36,18 @@ def assert_equal(
print(f"Input image shape: {img.shape}, example_image: {example_image.shape}") print(f"Input image shape: {img.shape}, example_image: {example_image.shape}")
config.paint_by_example_example_image = Image.fromarray(example_image) config.paint_by_example_example_image = Image.fromarray(example_image)
res = model(img, mask, config) res = model(img, mask, config)
cv2.imwrite(str(save_dir / gt_name), res) cv2.imwrite(str(save_dir / save_name), res)
def test_paint_by_example(): @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) model = ModelManager(name=model_name, device=device, disable_nsfw=True)
cfg = get_config(HDStrategy.ORIGINAL, sd_steps=30) cfg = get_config(strategy=HDStrategy.ORIGINAL, sd_steps=sd_steps)
assert_equal( assert_equal(
model, model,
cfg, cfg,
f"paint_by_example.png", 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,
)
def test_paint_by_example_cpu_offload():
model = ModelManager(
name=model_name, device=device, cpu_offload=True, disable_nsfw=False
)
cfg = get_config(HDStrategy.ORIGINAL, sd_steps=30)
assert_equal(
model,
cfg,
f"paint_by_example_cpu_offload.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
)
def test_paint_by_example_cpu_offload_cpu_device():
model = ModelManager(
name=model_name, device=torch.device("cpu"), cpu_offload=True, disable_nsfw=True
)
cfg = get_config(HDStrategy.ORIGINAL, sd_steps=1)
assert_equal(
model,
cfg,
f"paint_by_example_cpu_offload_cpu_device.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png", img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png", mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
fy=0.9, fy=0.9,

View File

@ -3,13 +3,12 @@ import os
import time import time
from lama_cleaner.plugins.anime_seg import AnimeSeg from lama_cleaner.plugins.anime_seg import AnimeSeg
from lama_cleaner.tests.utils import check_device, current_dir, save_dir
os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1" os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
from pathlib import Path
import cv2 import cv2
import pytest import pytest
import torch.cuda
from lama_cleaner.plugins import ( from lama_cleaner.plugins import (
RemoveBG, RemoveBG,
@ -19,9 +18,6 @@ from lama_cleaner.plugins import (
InteractiveSeg, InteractiveSeg,
) )
current_dir = Path(__file__).parent.absolute().resolve()
save_dir = current_dir / "result"
save_dir.mkdir(exist_ok=True, parents=True)
img_p = current_dir / "bunny.jpeg" img_p = current_dir / "bunny.jpeg"
img_bytes = open(img_p, "rb").read() img_bytes = open(img_p, "rb").read()
bgr_img = cv2.imread(str(img_p)) bgr_img = cv2.imread(str(img_p))
@ -50,11 +46,7 @@ def test_anime_seg():
@pytest.mark.parametrize("device", ["cuda", "cpu", "mps"]) @pytest.mark.parametrize("device", ["cuda", "cpu", "mps"])
def test_upscale(device): def test_upscale(device):
if device == "cuda" and not torch.cuda.is_available(): check_device(device)
return
if device == "mps" and not torch.backends.mps.is_available():
return
model = RealESRGANUpscaler("realesr-general-x4v3", device) model = RealESRGANUpscaler("realesr-general-x4v3", device)
res = model.forward(bgr_img, 2) res = model.forward(bgr_img, 2)
_save(res, f"test_upscale_x2_{device}.png") _save(res, f"test_upscale_x2_{device}.png")
@ -65,10 +57,7 @@ def test_upscale(device):
@pytest.mark.parametrize("device", ["cuda", "cpu", "mps"]) @pytest.mark.parametrize("device", ["cuda", "cpu", "mps"])
def test_gfpgan(device): def test_gfpgan(device):
if device == "cuda" and not torch.cuda.is_available(): check_device(device)
return
if device == "mps" and not torch.backends.mps.is_available():
return
model = GFPGANPlugin(device) model = GFPGANPlugin(device)
res = model(rgb_img, None, None) res = model(rgb_img, None, None)
_save(res, f"test_gfpgan_{device}.png") _save(res, f"test_gfpgan_{device}.png")
@ -76,10 +65,7 @@ def test_gfpgan(device):
@pytest.mark.parametrize("device", ["cuda", "cpu", "mps"]) @pytest.mark.parametrize("device", ["cuda", "cpu", "mps"])
def test_restoreformer(device): def test_restoreformer(device):
if device == "cuda" and not torch.cuda.is_available(): check_device(device)
return
if device == "mps" and not torch.backends.mps.is_available():
return
model = RestoreFormerPlugin(device) model = RestoreFormerPlugin(device)
res = model(rgb_img, None, None) res = model(rgb_img, None, None)
_save(res, f"test_restoreformer_{device}.png") _save(res, f"test_restoreformer_{device}.png")
@ -87,10 +73,7 @@ def test_restoreformer(device):
@pytest.mark.parametrize("device", ["cuda", "cpu", "mps"]) @pytest.mark.parametrize("device", ["cuda", "cpu", "mps"])
def test_segment_anything(device): def test_segment_anything(device):
if device == "cuda" and not torch.cuda.is_available(): check_device(device)
return
if device == "mps" and not torch.backends.mps.is_available():
return
img_md5 = hashlib.md5(img_bytes).hexdigest() img_md5 = hashlib.md5(img_bytes).hexdigest()
model = InteractiveSeg("vit_l", device) model = InteractiveSeg("vit_l", device)
new_mask = model.forward(rgb_img, [[448 // 2, 394 // 2, 1]], img_md5) new_mask = model.forward(rgb_img, [[448 // 2, 394 // 2, 1]], img_md5)

View File

@ -1,5 +1,7 @@
import os import os
from lama_cleaner.tests.utils import check_device, get_config, assert_equal
os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1" os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
from pathlib import Path from pathlib import Path
@ -8,14 +10,13 @@ import torch
from lama_cleaner.model_manager import ModelManager from lama_cleaner.model_manager import ModelManager
from lama_cleaner.schema import HDStrategy, SDSampler, FREEUConfig from lama_cleaner.schema import HDStrategy, SDSampler, FREEUConfig
from lama_cleaner.tests.test_model import get_config, assert_equal
current_dir = Path(__file__).parent.absolute().resolve() current_dir = Path(__file__).parent.absolute().resolve()
save_dir = current_dir / "result" save_dir = current_dir / "result"
save_dir.mkdir(exist_ok=True, parents=True) save_dir.mkdir(exist_ok=True, parents=True)
@pytest.mark.parametrize("sd_device", ["cuda", "mps"]) @pytest.mark.parametrize("device", ["cuda", "mps"])
@pytest.mark.parametrize( @pytest.mark.parametrize(
"sampler", "sampler",
[ [
@ -28,25 +29,24 @@ save_dir.mkdir(exist_ok=True, parents=True)
], ],
) )
def test_runway_sd_1_5_all_samplers( def test_runway_sd_1_5_all_samplers(
sd_device, device,
sampler, sampler,
): ):
if sd_device == "cuda" and not torch.cuda.is_available(): sd_steps = check_device(device)
return
sd_steps = 30
model = ModelManager( model = ModelManager(
name="runwayml/stable-diffusion-inpainting", name="runwayml/stable-diffusion-inpainting",
device=torch.device(sd_device), device=torch.device(device),
disable_nsfw=True, disable_nsfw=True,
sd_cpu_textencoder=False, sd_cpu_textencoder=False,
) )
cfg = get_config( cfg = get_config(
HDStrategy.ORIGINAL, prompt="a fox sitting on a bench", sd_steps=sd_steps strategy=HDStrategy.ORIGINAL,
prompt="a fox sitting on a bench",
sd_steps=sd_steps,
) )
cfg.sd_sampler = sampler cfg.sd_sampler = sampler
name = f"device_{sd_device}_{sampler}" name = f"device_{device}_{sampler}"
assert_equal( assert_equal(
model, model,
@ -57,22 +57,20 @@ def test_runway_sd_1_5_all_samplers(
) )
@pytest.mark.parametrize("sd_device", ["cuda", "mps"]) @pytest.mark.parametrize("device", ["cuda", "mps", "cpu"])
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
@pytest.mark.parametrize("sampler", [SDSampler.lcm]) @pytest.mark.parametrize("sampler", [SDSampler.lcm])
def test_runway_sd_lcm_lora(sd_device, strategy, sampler): def test_runway_sd_lcm_lora(device, sampler):
if sd_device == "cuda" and not torch.cuda.is_available(): check_device(device)
return
sd_steps = 5 sd_steps = 5
model = ModelManager( model = ModelManager(
name="runwayml/stable-diffusion-inpainting", name="runwayml/stable-diffusion-inpainting",
device=torch.device(sd_device), device=torch.device(device),
disable_nsfw=True, disable_nsfw=True,
sd_cpu_textencoder=False, sd_cpu_textencoder=False,
) )
cfg = get_config( cfg = get_config(
strategy, strategy=HDStrategy.ORIGINAL,
prompt="face of a fox, sitting on a bench", prompt="face of a fox, sitting on a bench",
sd_steps=sd_steps, sd_steps=sd_steps,
sd_guidance_scale=2, sd_guidance_scale=2,
@ -83,28 +81,24 @@ def test_runway_sd_lcm_lora(sd_device, strategy, sampler):
assert_equal( assert_equal(
model, model,
cfg, cfg,
f"runway_sd_1_5_lcm_lora.png", f"runway_sd_1_5_lcm_lora_device_{device}.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png", img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png", mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
) )
@pytest.mark.parametrize("sd_device", ["cuda", "mps"]) @pytest.mark.parametrize("device", ["cuda", "mps", "cpu"])
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
@pytest.mark.parametrize("sampler", [SDSampler.ddim]) @pytest.mark.parametrize("sampler", [SDSampler.ddim])
def test_runway_sd_freeu(sd_device, strategy, sampler): def test_runway_sd_freeu(device, sampler):
if sd_device == "cuda" and not torch.cuda.is_available(): sd_steps = check_device(device)
return
sd_steps = 30
model = ModelManager( model = ModelManager(
name="runwayml/stable-diffusion-inpainting", name="runwayml/stable-diffusion-inpainting",
device=torch.device(sd_device), device=torch.device(device),
disable_nsfw=True, disable_nsfw=True,
sd_cpu_textencoder=False, sd_cpu_textencoder=False,
) )
cfg = get_config( cfg = get_config(
strategy, strategy=HDStrategy.ORIGINAL,
prompt="face of a fox, sitting on a bench", prompt="face of a fox, sitting on a bench",
sd_steps=sd_steps, sd_steps=sd_steps,
sd_guidance_scale=7.5, sd_guidance_scale=7.5,
@ -116,85 +110,83 @@ def test_runway_sd_freeu(sd_device, strategy, sampler):
assert_equal( assert_equal(
model, model,
cfg, cfg,
f"runway_sd_1_5_freeu.png", f"runway_sd_1_5_freeu_device_{device}.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png", img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png", mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
) )
@pytest.mark.parametrize("sd_device", ["cuda", "mps"]) @pytest.mark.parametrize("device", ["cuda", "mps"])
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL]) @pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
@pytest.mark.parametrize("sampler", [SDSampler.ddim]) @pytest.mark.parametrize("sampler", [SDSampler.ddim])
def test_runway_sd_sd_strength(sd_device, strategy, sampler): def test_runway_sd_sd_strength(device, strategy, sampler):
if sd_device == "cuda" and not torch.cuda.is_available(): sd_steps = check_device(device)
return
sd_steps = 30
model = ModelManager( model = ModelManager(
name="runwayml/stable-diffusion-inpainting", name="runwayml/stable-diffusion-inpainting",
device=torch.device(sd_device), device=torch.device(device),
disable_nsfw=True, disable_nsfw=True,
sd_cpu_textencoder=False, sd_cpu_textencoder=False,
) )
cfg = get_config( cfg = get_config(
strategy, prompt="a fox sitting on a bench", sd_steps=sd_steps, sd_strength=0.8 strategy=strategy,
prompt="a fox sitting on a bench",
sd_steps=sd_steps,
sd_strength=0.8,
) )
cfg.sd_sampler = sampler cfg.sd_sampler = sampler
assert_equal( assert_equal(
model, model,
cfg, cfg,
f"runway_sd_strength_0.8.png", f"runway_sd_strength_0.8_device_{device}.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png", img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png", mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
) )
@pytest.mark.parametrize("sd_device", ["cuda", "mps"]) @pytest.mark.parametrize("device", ["cuda", "mps", "cpu"])
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL]) @pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
@pytest.mark.parametrize("sampler", [SDSampler.ddim]) @pytest.mark.parametrize("sampler", [SDSampler.ddim])
def test_runway_norm_sd_model(sd_device, strategy, sampler): def test_runway_norm_sd_model(device, strategy, sampler):
if sd_device == "cuda" and not torch.cuda.is_available(): sd_steps = check_device(device)
return
sd_steps = 30
model = ModelManager( model = ModelManager(
name="runwayml/stable-diffusion-v1-5", name="runwayml/stable-diffusion-v1-5",
device=torch.device(sd_device), device=torch.device(device),
disable_nsfw=True, disable_nsfw=True,
sd_cpu_textencoder=False, sd_cpu_textencoder=False,
) )
cfg = get_config(strategy, prompt="face of a fox, sitting on a bench", sd_steps=sd_steps) cfg = get_config(
strategy=strategy, prompt="face of a fox, sitting on a bench", sd_steps=sd_steps
)
cfg.sd_sampler = sampler cfg.sd_sampler = sampler
assert_equal( assert_equal(
model, model,
cfg, cfg,
f"runway_{sd_device}_norm_sd_model.png", f"runway_{device}_norm_sd_model_device_{device}.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png", img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png", mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
) )
@pytest.mark.parametrize("sd_device", ["cuda"]) @pytest.mark.parametrize("device", ["cuda"])
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL]) @pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
@pytest.mark.parametrize("sampler", [SDSampler.k_euler_a]) @pytest.mark.parametrize("sampler", [SDSampler.k_euler_a])
def test_runway_sd_1_5_cpu_offload(sd_device, strategy, sampler): def test_runway_sd_1_5_cpu_offload(device, strategy, sampler):
if sd_device == "cuda" and not torch.cuda.is_available(): sd_steps = check_device(device)
return
sd_steps = 30
model = ModelManager( model = ModelManager(
name="runwayml/stable-diffusion-inpainting", name="runwayml/stable-diffusion-inpainting",
device=torch.device(sd_device), device=torch.device(device),
disable_nsfw=True, disable_nsfw=True,
sd_cpu_textencoder=False, sd_cpu_textencoder=False,
cpu_offload=True, cpu_offload=True,
) )
cfg = get_config(strategy, prompt="a fox sitting on a bench", sd_steps=sd_steps) cfg = get_config(
strategy=strategy, prompt="a fox sitting on a bench", sd_steps=sd_steps
)
cfg.sd_sampler = sampler cfg.sd_sampler = sampler
name = f"device_{sd_device}_{sampler}" name = f"device_{device}_{sampler}"
assert_equal( assert_equal(
model, model,
@ -205,7 +197,7 @@ def test_runway_sd_1_5_cpu_offload(sd_device, strategy, sampler):
) )
@pytest.mark.parametrize("sd_device", ["cuda", "mps"]) @pytest.mark.parametrize("device", ["cuda", "mps", "cpu"])
@pytest.mark.parametrize("sampler", [SDSampler.ddim]) @pytest.mark.parametrize("sampler", [SDSampler.ddim])
@pytest.mark.parametrize( @pytest.mark.parametrize(
"name", "name",
@ -215,26 +207,23 @@ def test_runway_sd_1_5_cpu_offload(sd_device, strategy, sampler):
"v1-5-pruned-emaonly.safetensors", "v1-5-pruned-emaonly.safetensors",
], ],
) )
def test_local_file_path(sd_device, sampler, name): def test_local_file_path(device, sampler, name):
if sd_device == "cuda" and not torch.cuda.is_available(): sd_steps = check_device(device)
return
sd_steps = 30
model = ModelManager( model = ModelManager(
name=name, name=name,
device=torch.device(sd_device), device=torch.device(device),
disable_nsfw=True, disable_nsfw=True,
sd_cpu_textencoder=False, sd_cpu_textencoder=False,
cpu_offload=False, cpu_offload=False,
) )
cfg = get_config( cfg = get_config(
HDStrategy.ORIGINAL, strategy=HDStrategy.ORIGINAL,
prompt="a fox sitting on a bench", prompt="a fox sitting on a bench",
sd_steps=sd_steps, sd_steps=sd_steps,
) )
cfg.sd_sampler = sampler cfg.sd_sampler = sampler
name = f"device_{sd_device}_{sampler}_{name}" name = f"device_{device}_{sampler}_{name}"
assert_equal( assert_equal(
model, model,

View File

@ -1,7 +1,8 @@
import os import os
from lama_cleaner.tests.utils import check_device, current_dir
os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1" os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
from pathlib import Path
import pytest import pytest
import torch import torch
@ -10,31 +11,25 @@ from lama_cleaner.model_manager import ModelManager
from lama_cleaner.schema import HDStrategy, SDSampler, FREEUConfig from lama_cleaner.schema import HDStrategy, SDSampler, FREEUConfig
from lama_cleaner.tests.test_model import get_config, assert_equal from lama_cleaner.tests.test_model import get_config, assert_equal
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"])
@pytest.mark.parametrize("sd_device", ["cuda", "mps"])
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL]) @pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
@pytest.mark.parametrize("sampler", [SDSampler.ddim]) @pytest.mark.parametrize("sampler", [SDSampler.ddim])
def test_sdxl(sd_device, strategy, sampler): def test_sdxl(device, strategy, sampler):
sd_steps = check_device(device)
def callback(i, t, latents): def callback(i, t, latents):
pass pass
if sd_device == "cuda" and not torch.cuda.is_available():
return
sd_steps = 20
model = ModelManager( model = ModelManager(
name="diffusers/stable-diffusion-xl-1.0-inpainting-0.1", name="diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
device=torch.device(sd_device), device=torch.device(device),
disable_nsfw=True, disable_nsfw=True,
sd_cpu_textencoder=False, sd_cpu_textencoder=False,
callback=callback, callback=callback,
) )
cfg = get_config( cfg = get_config(
strategy, strategy=strategy,
prompt="face of a fox, sitting on a bench", prompt="face of a fox, sitting on a bench",
sd_steps=sd_steps, sd_steps=sd_steps,
sd_strength=1.0, sd_strength=1.0,
@ -42,12 +37,10 @@ def test_sdxl(sd_device, strategy, sampler):
) )
cfg.sd_sampler = sampler cfg.sd_sampler = sampler
name = f"device_{sd_device}_{sampler}"
assert_equal( assert_equal(
model, model,
cfg, cfg,
f"sdxl_{name}.png", f"sdxl_device_{device}.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png", img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png", mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
fx=2, fx=2,
@ -55,26 +48,24 @@ def test_sdxl(sd_device, strategy, sampler):
) )
@pytest.mark.parametrize("sd_device", ["cuda", "mps"]) @pytest.mark.parametrize("device", ["cuda", "mps"])
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL]) @pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
@pytest.mark.parametrize("sampler", [SDSampler.ddim]) @pytest.mark.parametrize("sampler", [SDSampler.ddim])
def test_sdxl_lcm_lora_and_freeu(sd_device, strategy, sampler): def test_sdxl_lcm_lora_and_freeu(device, strategy, sampler):
sd_steps = check_device(device)
def callback(i, t, latents): def callback(i, t, latents):
pass pass
if sd_device == "cuda" and not torch.cuda.is_available():
return
sd_steps = 5
model = ModelManager( model = ModelManager(
name="diffusers/stable-diffusion-xl-1.0-inpainting-0.1", name="diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
device=torch.device(sd_device), device=torch.device(device),
disable_nsfw=True, disable_nsfw=True,
sd_cpu_textencoder=False, sd_cpu_textencoder=False,
callback=callback, callback=callback,
) )
cfg = get_config( cfg = get_config(
strategy, strategy=strategy,
prompt="face of a fox, sitting on a bench", prompt="face of a fox, sitting on a bench",
sd_steps=sd_steps, sd_steps=sd_steps,
sd_strength=1.0, sd_strength=1.0,
@ -83,7 +74,7 @@ def test_sdxl_lcm_lora_and_freeu(sd_device, strategy, sampler):
) )
cfg.sd_sampler = sampler cfg.sd_sampler = sampler
name = f"device_{sd_device}_{sampler}" name = f"device_{device}_{sampler}"
assert_equal( assert_equal(
model, model,
@ -96,7 +87,7 @@ def test_sdxl_lcm_lora_and_freeu(sd_device, strategy, sampler):
) )
cfg = get_config( cfg = get_config(
strategy, strategy=strategy,
prompt="face of a fox, sitting on a bench", prompt="face of a fox, sitting on a bench",
sd_steps=sd_steps, sd_steps=sd_steps,
sd_guidance_scale=7.5, sd_guidance_scale=7.5,
@ -107,7 +98,7 @@ def test_sdxl_lcm_lora_and_freeu(sd_device, strategy, sampler):
assert_equal( assert_equal(
model, model,
cfg, cfg,
f"sdxl_{name}_freeu.png", f"sdxl_{name}_freeu_device_{device}.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png", img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png", mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
fx=2, fx=2,
@ -115,28 +106,27 @@ def test_sdxl_lcm_lora_and_freeu(sd_device, strategy, sampler):
) )
@pytest.mark.parametrize("sd_device", ["mps"]) @pytest.mark.parametrize("device", ["cuda", "mps"])
@pytest.mark.parametrize( @pytest.mark.parametrize(
"rect", "rect",
[ [
[-128, -128, 1024, 1024], [-128, -128, 1024, 1024],
], ],
) )
def test_sdxl_outpainting(sd_device, rect): def test_sdxl_outpainting(device, rect):
if sd_device == "cuda" and not torch.cuda.is_available(): sd_steps = check_device(device)
return
model = ModelManager( model = ModelManager(
name="diffusers/stable-diffusion-xl-1.0-inpainting-0.1", name="diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
device=torch.device(sd_device), device=torch.device(device),
disable_nsfw=True, disable_nsfw=True,
sd_cpu_textencoder=False, sd_cpu_textencoder=False,
) )
cfg = get_config( cfg = get_config(
HDStrategy.ORIGINAL, strategy=HDStrategy.ORIGINAL,
prompt="a dog sitting on a bench in the park", prompt="a dog sitting on a bench in the park",
sd_steps=20, sd_steps=sd_steps,
use_extender=True, use_extender=True,
extender_x=rect[0], extender_x=rect[0],
extender_y=rect[1], extender_y=rect[1],
@ -150,7 +140,7 @@ def test_sdxl_outpainting(sd_device, rect):
assert_equal( assert_equal(
model, model,
cfg, cfg,
f"sdxl_outpainting_dog_ddim_{'_'.join(map(str, rect))}.png", f"sdxl_outpainting_dog_ddim_{'_'.join(map(str, rect))}_device_{device}.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png", img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png", mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
fx=1.5, fx=1.5,

View File

@ -0,0 +1,75 @@
from pathlib import Path
import cv2
import pytest
import torch
from lama_cleaner.schema import LDMSampler, HDStrategy, Config, SDSampler
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,
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 Config(**data)

View File

@ -18,7 +18,7 @@ const SheetOverlay = React.forwardRef<
>(({ className, ...props }, ref) => ( >(({ className, ...props }, ref) => (
<SheetPrimitive.Overlay <SheetPrimitive.Overlay
className={cn( className={cn(
"fixed inset-0 z-50 bg-background/80 backdrop-blur-sm data-[state=open]:animate-in data-[state=closed]:animate-out data-[state=closed]:fade-out-0 data-[state=open]:fade-in-0", "fixed inset-0 z-50 bg-background/80 backdrop-blur-sm data-[state=closed]:fade-out-0 data-[state=open]:fade-in-0",
className className
)} )}
{...props} {...props}
@ -28,7 +28,7 @@ const SheetOverlay = React.forwardRef<
SheetOverlay.displayName = SheetPrimitive.Overlay.displayName SheetOverlay.displayName = SheetPrimitive.Overlay.displayName
const sheetVariants = cva( const sheetVariants = cva(
"fixed z-50 gap-4 bg-background p-6 shadow-lg transition ease-in-out data-[state=open]:animate-in data-[state=closed]:animate-out data-[state=closed]:duration-200 data-[state=open]:duration-300", "fixed z-50 gap-4 bg-background p-6 shadow-lg transition ease-in-out data-[state=closed]:duration-200 data-[state=open]:duration-300",
{ {
variants: { variants: {
side: { side: {

View File

@ -109,6 +109,7 @@ type ServerConfig = {
enableAutoSaving: boolean enableAutoSaving: boolean
enableControlnet: boolean enableControlnet: boolean
controlnetMethod: string controlnetMethod: string
disableModelSwitch: boolean
isDesktop: boolean isDesktop: boolean
} }
@ -279,6 +280,7 @@ const defaultValues: AppState = {
enableAutoSaving: false, enableAutoSaving: false,
enableControlnet: false, enableControlnet: false,
controlnetMethod: "lllyasviel/control_v11p_sd15_canny", controlnetMethod: "lllyasviel/control_v11p_sd15_canny",
disableModelSwitch: false,
isDesktop: false, isDesktop: false,
}, },
settings: { settings: {