wip: controlnet

lama_cleaner/tests/test_controlnet.py

@@ -1,5 +1,7 @@
 import os
 
+from lama_cleaner.const import SD_CONTROLNET_CHOICES
+
 os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
 from pathlib import Path
 
@@ -22,7 +24,10 @@ device = torch.device(device)
 @pytest.mark.parametrize("sampler", [SDSampler.uni_pc])
 @pytest.mark.parametrize("cpu_textencoder", [True])
 @pytest.mark.parametrize("disable_nsfw", [True])
-def test_runway_sd_1_5(sd_device, strategy, sampler, cpu_textencoder, disable_nsfw):
+@pytest.mark.parametrize("sd_controlnet_method", SD_CONTROLNET_CHOICES)
+def test_runway_sd_1_5(
+    sd_device, strategy, sampler, cpu_textencoder, disable_nsfw, sd_controlnet_method
+):
     if sd_device == "cuda" and not torch.cuda.is_available():
         return
     if device == "mps" and not torch.backends.mps.is_available():
@@ -34,19 +39,20 @@ def test_runway_sd_1_5(sd_device, strategy, sampler, cpu_textencoder, disable_ns
         sd_controlnet=True,
         device=torch.device(sd_device),
         hf_access_token="",
-        sd_run_local=True,
+        sd_run_local=False,
         disable_nsfw=disable_nsfw,
         sd_cpu_textencoder=cpu_textencoder,
+        sd_controlnet_method=sd_controlnet_method,
     )
     cfg = get_config(strategy, prompt="a fox sitting on a bench", sd_steps=sd_steps)
     cfg.sd_sampler = sampler
 
-    name = f"device_{sd_device}_{sampler}_cpu_textencoder_{cpu_textencoder}_disnsfw_{disable_nsfw}"
+    name = f"device_{sd_device}_{sampler}_cpu_textencoder_disable_nsfw"
 
     assert_equal(
         model,
         cfg,
-        f"sd_controlnet_{name}.png",
+        f"sd_controlnet_{sd_controlnet_method}_{name}.png",
         img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
         mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
         fx=1.2,
@@ -68,11 +74,12 @@ def test_local_file_path(sd_device, sampler):
         sd_controlnet=True,
         device=torch.device(sd_device),
         hf_access_token="",
-        sd_run_local=True,
+        sd_run_local=False,
         disable_nsfw=True,
         sd_cpu_textencoder=False,
         cpu_offload=True,
         sd_local_model_path="/Users/cwq/data/models/sd-v1-5-inpainting.ckpt",
+        sd_controlnet_method="control_v11p_sd15_canny",
     )
     cfg = get_config(
         HDStrategy.ORIGINAL,
@@ -86,7 +93,48 @@ def test_local_file_path(sd_device, sampler):
     assert_equal(
         model,
         cfg,
-        f"sd_controlnet_local_model_{name}.png",
+        f"sd_controlnet_canny_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_local_file_path_controlnet_native_inpainting(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="sd1.5",
+        sd_controlnet=True,
+        device=torch.device(sd_device),
+        hf_access_token="",
+        sd_run_local=False,
+        disable_nsfw=True,
+        sd_cpu_textencoder=False,
+        cpu_offload=True,
+        sd_local_model_path="/Users/cwq/data/models/v1-5-pruned-emaonly.safetensors",
+        sd_controlnet_method="control_v11p_sd15_inpaint",
+    )
+    cfg = get_config(
+        HDStrategy.ORIGINAL,
+        prompt="a fox sitting on a bench",
+        sd_steps=sd_steps,
+        controlnet_conditioning_scale=1.0,
+        sd_strength=1.0
+    )
+    cfg.sd_sampler = sampler
+
+    name = f"device_{sd_device}_{sampler}"
+
+    assert_equal(
+        model,
+        cfg,
+        f"sd_controlnet_local_native_{name}.png",
         img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
         mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
     )

lama_cleaner/tests/test_instruct_pix2pix.py

@@ -20,7 +20,7 @@ def test_instruct_pix2pix(disable_nsfw, cpu_offload):
     model = ModelManager(name="instruct_pix2pix",
                          device=torch.device(device),
                          hf_access_token="",
-                         sd_run_local=True,
+                         sd_run_local=False,
                          disable_nsfw=disable_nsfw,
                          sd_cpu_textencoder=False,
                          cpu_offload=cpu_offload)
@@ -45,7 +45,7 @@ def test_instruct_pix2pix_snow(disable_nsfw, cpu_offload):
     model = ModelManager(name="instruct_pix2pix",
                          device=torch.device(device),
                          hf_access_token="",
-                         sd_run_local=True,
+                         sd_run_local=False,
                          disable_nsfw=disable_nsfw,
                          sd_cpu_textencoder=False,
                          cpu_offload=cpu_offload)