get samplers from backend

lama_cleaner/tests/test_outpainting.py

@@ -49,7 +49,7 @@ def test_outpainting(name, device, rect):
         extender_width=rect[2],
         extender_height=rect[3],
         sd_guidance_scale=8.0,
-        sd_sampler=SDSampler.dpm_plus_plus,
+        sd_sampler=SDSampler.dpm_plus_plus_2m,
     )
 
     assert_equal(
@@ -92,7 +92,7 @@ def test_kandinsky_outpainting(name, device, rect):
         extender_width=rect[2],
         extender_height=rect[3],
         sd_guidance_scale=7,
-        sd_sampler=SDSampler.dpm_plus_plus,
+        sd_sampler=SDSampler.dpm_plus_plus_2m,
     )
 
     assert_equal(
@@ -136,7 +136,7 @@ def test_powerpaint_outpainting(name, device, rect):
         extender_width=rect[2],
         extender_height=rect[3],
         sd_guidance_scale=8.0,
-        sd_sampler=SDSampler.dpm_plus_plus,
+        sd_sampler=SDSampler.dpm_plus_plus_2m,
         powerpaint_task="outpainting",
     )
 

lama_cleaner/tests/test_sd_model.py

@@ -1,5 +1,7 @@
 import os
 
+from loguru import logger
+
 from lama_cleaner.tests.utils import check_device, get_config, assert_equal
 
 os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
@@ -17,21 +19,7 @@ save_dir.mkdir(exist_ok=True, parents=True)
 
 
 @pytest.mark.parametrize("device", ["cuda", "mps"])
-@pytest.mark.parametrize(
-    "sampler",
-    [
-        SDSampler.ddim,
-        SDSampler.pndm,
-        SDSampler.k_lms,
-        SDSampler.k_euler,
-        SDSampler.k_euler_a,
-        SDSampler.lcm,
-    ],
-)
-def test_runway_sd_1_5_all_samplers(
-    device,
-    sampler,
-):
+def test_runway_sd_1_5_all_samplers(device):
     sd_steps = check_device(device)
     model = ModelManager(
         name="runwayml/stable-diffusion-inpainting",
@@ -39,22 +27,37 @@ def test_runway_sd_1_5_all_samplers(
         disable_nsfw=True,
         sd_cpu_textencoder=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}"
+    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,
+        )
 
-    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",
-    )
+        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"])
@@ -171,7 +174,7 @@ def test_runway_norm_sd_model(device, strategy, sampler):
 
 @pytest.mark.parametrize("device", ["cuda"])
 @pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
-@pytest.mark.parametrize("sampler", [SDSampler.k_euler_a])
+@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(

lama_cleaner/tests/utils.py

@@ -3,7 +3,9 @@ import cv2
 import pytest
 import torch
 
+from lama_cleaner.helper import encode_pil_to_base64
 from lama_cleaner.schema import LDMSampler, HDStrategy, InpaintRequest, SDSampler
+from PIL import Image
 
 current_dir = Path(__file__).parent.absolute().resolve()
 save_dir = current_dir / "result"
@@ -21,7 +23,7 @@ def check_device(device: str) -> int:
 
 def assert_equal(
     model,
-    config,
+    config: InpaintRequest,
     gt_name,
     fx: float = 1,
     fy: float = 1,
@@ -29,6 +31,8 @@ def assert_equal(
     mask_p=current_dir / "mask.png",
 ):
     img, mask = get_data(fx=fx, fy=fy, img_p=img_p, mask_p=mask_p)
+    config.image = encode_pil_to_base64(Image.fromarray(img), 95, {})[0]
+    config.mask = encode_pil_to_base64(Image.fromarray(mask), 95, {})[0]
     print(f"Input image shape: {img.shape}")
     res = model(img, mask, config)
     ok = cv2.imwrite(
@@ -72,4 +76,4 @@ def get_config(**kwargs):
         hd_strategy_resize_limit=200,
     )
     data.update(**kwargs)
-    return InpaintRequest(**data)
+    return InpaintRequest(image="", mask="", **data)