diff --git a/lama_cleaner/model/mat.py b/lama_cleaner/model/mat.py index a709660..36a8a8e 100644 --- a/lama_cleaner/model/mat.py +++ b/lama_cleaner/model/mat.py @@ -1880,7 +1880,10 @@ class MAT(InpaintModel): seed = 240 # pick up a random number set_seed(seed) - self.torch_dtype = torch.float16 + fp16 = not kwargs.get("no_half", False) + use_gpu = "cuda" in str(device) and torch.cuda.is_available() + self.torch_dtype = torch.float16 if use_gpu and fp16 else torch.float32 + G = Generator( z_dim=512, c_dim=0, diff --git a/lama_cleaner/tests/test_model.py b/lama_cleaner/tests/test_model.py index de8a000..d1af506 100644 --- a/lama_cleaner/tests/test_model.py +++ b/lama_cleaner/tests/test_model.py @@ -69,9 +69,73 @@ def assert_equal( # assert np.array_equal(res, gt) -@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL]) -def test_mat(strategy): - model = ModelManager(name="mat", device=device) +@pytest.mark.parametrize( + "strategy", [HDStrategy.ORIGINAL, HDStrategy.RESIZE, HDStrategy.CROP] +) +def test_lama(strategy): + model = ModelManager(name="lama", device=device) + assert_equal( + model, + get_config(strategy), + f"lama_{strategy[0].upper() + strategy[1:]}_result.png", + ) + + fx = 1.3 + assert_equal( + model, + get_config(strategy), + f"lama_{strategy[0].upper() + strategy[1:]}_fx_{fx}_result.png", + fx=1.3, + ) + + +@pytest.mark.parametrize( + "strategy", [HDStrategy.ORIGINAL, HDStrategy.RESIZE, HDStrategy.CROP] +) +@pytest.mark.parametrize("ldm_sampler", [LDMSampler.ddim, LDMSampler.plms]) +def test_ldm(strategy, ldm_sampler): + model = ModelManager(name="ldm", device=device) + cfg = get_config(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( + "strategy", [HDStrategy.ORIGINAL, HDStrategy.RESIZE, HDStrategy.CROP] +) +@pytest.mark.parametrize("zits_wireframe", [False, True]) +def test_zits(strategy, zits_wireframe): + model = ModelManager(name="zits", device=device) + cfg = get_config(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( + 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("strategy", [HDStrategy.ORIGINAL], "no_half", [True, False]) +def test_mat(strategy, no_half): + model = ModelManager(name="mat", device=device, no_half=no_half) cfg = get_config(strategy) for _ in range(10): @@ -79,3 +143,50 @@ def test_mat(strategy): model, cfg, f"mat_{strategy.capitalize()}_result.png", ) +@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL]) +def test_fcf(strategy): + model = ModelManager(name="fcf", device=device) + cfg = get_config(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(device), + ) + cfg = get_config(strategy, cv2_flag=cv2_flag, cv2_radius=cv2_radius) + assert_equal( + model, + cfg, + f"sd_{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( + "strategy", [HDStrategy.ORIGINAL, HDStrategy.RESIZE, HDStrategy.CROP] +) +def test_manga(strategy): + model = ModelManager( + name="manga", + device=torch.device(device), + ) + cfg = get_config(strategy) + assert_equal( + model, + cfg, + f"sd_{strategy.capitalize()}.png", + img_p=current_dir / "overture-creations-5sI6fQgYIuo.png", + mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png", + ) +>>>>>>> 944494a (mat support float16)