update

lama_cleaner/tests/test_instruct_pix2pix.py

@@ -4,20 +4,17 @@ import pytest
 import torch
 
 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.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"
 
 
+@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(disable_nsfw, cpu_offload):
-    sd_steps = 50 if device == "cuda" else 20
+def test_instruct_pix2pix(device, disable_nsfw, cpu_offload):
+    sd_steps = check_device(device)
     model = ModelManager(
         name=model_name,
         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",
         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",
-    )