kandinsky

lama_cleaner/tests/test_outpainting.py

@@ -17,6 +17,7 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 device = torch.device(device)
 
 
+@pytest.mark.parametrize("name", ["sd1.5"])
 @pytest.mark.parametrize("sd_device", ["mps"])
 @pytest.mark.parametrize(
     "rect",
@@ -30,16 +31,15 @@ device = torch.device(device)
         [-100, -100, 512 + 200, 512 + 200],
     ],
 )
-def test_sdxl_outpainting(sd_device, rect):
+def test_outpainting(name, sd_device, rect):
     def callback(i, t, latents):
         pass
 
     if sd_device == "cuda" and not torch.cuda.is_available():
         return
 
-    sd_steps = 50 if sd_device == "cuda" else 1
     model = ModelManager(
-        name="sd1.5",
+        name=name,
         device=torch.device(sd_device),
         hf_access_token="",
         sd_run_local=True,
@@ -50,21 +50,69 @@ def test_sdxl_outpainting(sd_device, rect):
     cfg = get_config(
         HDStrategy.ORIGINAL,
         prompt="a dog sitting on a bench in the park",
-        sd_steps=30,
+        sd_steps=50,
         use_croper=True,
         croper_is_outpainting=True,
         croper_x=rect[0],
         croper_y=rect[1],
         croper_width=rect[2],
         croper_height=rect[3],
-        sd_guidance_scale=14,
+        sd_guidance_scale=4,
         sd_sampler=SDSampler.dpm_plus_plus,
     )
 
     assert_equal(
         model,
         cfg,
-        f"sd15_outpainting_dpm++_{'_'.join(map(str, rect))}.png",
+        f"{name.replace('.', '_')}_outpainting_dpm++_{'_'.join(map(str, rect))}.png",
         img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
         mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
     )
+
+
+@pytest.mark.parametrize("name", ["kandinsky2.2"])
+@pytest.mark.parametrize("sd_device", ["mps"])
+@pytest.mark.parametrize(
+    "rect",
+    [
+        [-100, -100, 768, 768],
+    ],
+)
+def test_kandinsky_outpainting(name, sd_device, rect):
+    def callback(i, t, latents):
+        pass
+
+    if sd_device == "cuda" and not torch.cuda.is_available():
+        return
+
+    model = ModelManager(
+        name=name,
+        device=torch.device(sd_device),
+        hf_access_token="",
+        sd_run_local=True,
+        disable_nsfw=True,
+        sd_cpu_textencoder=False,
+        callback=callback,
+    )
+    cfg = get_config(
+        HDStrategy.ORIGINAL,
+        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",
+        sd_steps=50,
+        use_croper=True,
+        croper_is_outpainting=True,
+        croper_x=rect[0],
+        croper_y=rect[1],
+        croper_width=rect[2],
+        croper_height=rect[3],
+        sd_guidance_scale=4,
+        sd_sampler=SDSampler.dpm_plus_plus,
+    )
+
+    assert_equal(
+        model,
+        cfg,
+        f"{name.replace('.', '_')}_outpainting_dpm++_{'_'.join(map(str, rect))}.png",
+        img_p=current_dir / "cat.png",
+        mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
+    )