From 38c02dcb671e7ca40a5bafc6ec88703899c1c120 Mon Sep 17 00:00:00 2001 From: Anatoly Belikov Date: Thu, 26 Sep 2024 21:33:52 +0300 Subject: [PATCH] update tests for maskedim2impipe --- multigen/pipes.py | 2 ++ tests/pipe_test.py | 4 ++-- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/multigen/pipes.py b/multigen/pipes.py index df7d5e2..7e9a846 100755 --- 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 index 6690a69..8b6890e 100644 --- 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'])