update tests for maskedim2impipe

singnet · Sep 26, 2024 · 38c02dc · 38c02dc
1 parent 1e72335
commit 38c02dc
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/multigen/pipes.py b/multigen/pipes.py
@@ -747,6 +747,8 @@ def __init__(self, model_id, pipe: Optional[StableDiffusionControlNetPipeline] =
                     self.pipe = self._class.from_pipe(self.pipe, controlnet=cnets)
             for cnet in cnets:
                 cnet.to(self.pipe.dtype)
+                if 'offload_device' not in args:
+                    cnet.to(self.pipe.device)
         else:
             # don't load anything, just reuse pipe
             super().__init__(model_id=model_id, pipe=pipe, **args)

diff --git a/tests/pipe_test.py b/tests/pipe_test.py
@@ -162,9 +162,9 @@ def test_maskedimg2img_basic(self):
         scheduler = "EulerAncestralDiscreteScheduler"
         seed = 49045438434843
         blur = 48
-        param_3_3 = dict(original_image=img, image_painted=img_paint, strength=0.96,
+        param_3_3 = dict(image=img, image_painted=img_paint, strength=0.96,
                scheduler=scheduler, clip_skip=0, blur=blur, blur_compose=3, steps=5, guidance_scale=3.3)
-        param_7_6 = dict(original_image=img, image_painted=img_paint, strength=0.96,
+        param_7_6 = dict(image=img, image_painted=img_paint, strength=0.96,
                scheduler=scheduler, clip_skip=0, blur=blur, blur_compose=3, steps=5, guidance_scale=7.6)
         pipe.setup(**param_3_3)
         self.assertEqual(3.3, pipe.pipe_params['guidance_scale'])