update pytorch tests

VoltaML · Nov 26, 2023 · f43bc33 · f43bc33
1 parent 02e42f1
commit f43bc33
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/tests/inference/test_pytorch.py b/tests/inference/test_pytorch.py
@@ -187,15 +187,18 @@ def test_inpaint(pipe: PyTorchStableDiffusion):
     pipe.generate(job)
 
 
-def test_controlnet(pipe: PyTorchStableDiffusion):
+@pytest.mark.parametrize(
+    "scheduler", [KarrasDiffusionSchedulers.UniPCMultistepScheduler, "dpmpp_2m"]
+)
+def test_controlnet(pipe: PyTorchStableDiffusion, scheduler):
     "Generate an image with ControlNet Image to Image"
 
     job = ControlNetQueueEntry(
         data=ControlNetData(
             image=generate_random_image_base64(),
             prompt="This is a test",
             id="test",
-            scheduler=KarrasDiffusionSchedulers.UniPCMultistepScheduler,
+            scheduler=scheduler,
             controlnet="lllyasviel/sd-controlnet-canny",
             width=128,
             height=128,