Merge pull request #73 from noskill/moretests

gscale -> guidance_scale

multigen/pipes.py

@@ -370,7 +370,7 @@ def __init__(self, model_id, pipe: Optional[StableDiffusionImg2ImgPipeline] = No
         super().__init__(model_id=model_id, pipe=pipe, **args)
         self._input_image = None
 
-    def setup(self, fimage, image=None, strength=0.75, gscale=7.5, scale=None, **args):
+    def setup(self, fimage, image=None, strength=0.75, guidance_scale=7.5, scale=None, **args):
         """
         Setup pipeline for generation.
 
@@ -380,7 +380,7 @@ def setup(self, fimage, image=None, strength=0.75, gscale=7.5, scale=None, **arg
             strength (float, *optional*):
                 Strength image modification. Defaults to 0.75. A lower strength values keep result close to the input image.
                  value of 1 means input image more or less ignored.
-            gscale (float, *optional*): Guidance scale for the model. Defaults to 7.5.
+            guidance_scale (float, *optional*): Guidance scale for the model. Defaults to 7.5.
             scale (list or float, *optional*): Scale factor for the input image. Defaults to None.
             **args: Additional arguments passed to BasePipe setup method.
         """
@@ -390,7 +390,7 @@ def setup(self, fimage, image=None, strength=0.75, gscale=7.5, scale=None, **arg
         self._input_image = self.scale_image(self._input_image, scale)
         self.pipe_params.update({
             "strength": strength,
-            "guidance_scale": gscale
+            "guidance_scale": guidance_scale
         })
 
     def scale_image(self, image, scale):
@@ -458,7 +458,7 @@ class MaskedIm2ImPipe(Im2ImPipe):
     _classxl = MaskedStableDiffusionXLImg2ImgPipeline
     _autopipeline = DiffusionPipeline
 
-    def __init__(self, *args, pipe: Optional[StableDiffusionImg2ImgPipeline] = None, lpw=False, **kwargs):
+    def __init__(self, *args, pipe: Optional[StableDiffusionImg2ImgPipeline] = None, lpw=False, model_type=None, **kwargs):
         """
         Initialize a MaskedIm2ImPipe instance.
 
@@ -467,7 +467,7 @@ def __init__(self, *args, pipe: Optional[StableDiffusionImg2ImgPipeline] = None,
             pipe (StableDiffusionImg2ImgPipeline, *optional*): The underlying pipeline. Defaults to None.
             **kwargs: Additional keyword arguments passed to Im2ImPipe constructor.
         """
-        super().__init__(*args, pipe=pipe, lpw=lpw, **kwargs)
+        super().__init__(*args, pipe=pipe, lpw=lpw, model_type=model_type, **kwargs)
         # convert loaded pipeline if necessary
         if not isinstance(self.pipe, (self._class, self._classxl)):
             self.pipe = self._from_pipe(self.pipe, **kwargs)
@@ -718,7 +718,6 @@ def setup(self, fimage, width=None, height=None,
             "height": image.size[1] if height is None else height,
             "controlnet_conditioning_scale": cscales,
             "guess_mode": guess_mode,
-            "num_inference_steps": 20
         })
 
     def get_default_cond_scales(self):

multigen/worker.py

@@ -113,7 +113,8 @@ def _update(sess, job, gs):
             device = pipe.pipe.device
             self.logger.debug(f'running job on {device}')
             if device.type == 'cuda':
-                self._gpu_jobs[device.index] = model_id
+                with self._lock:
+                    self._gpu_jobs[device.index] = model_id
             class_name = str(pipe.__class__)
             self.logger.debug(f'got pipeline {class_name}')
 

tests/mech_beard_sigm.png

@@ -0,0 +1 @@
+../examples/mech_beard_sigm.png

tests/mech_beard_sigm_mask.png

@@ -0,0 +1 @@
+../examples/mech_beard_sigm_mask.png

tests/pipe_test.py

@@ -121,9 +121,49 @@ def test_loader(self):
     def test_img2img_basic(self):
         pipe = Im2ImPipe(self.get_model(), model_type=self.model_type())
         im = self.get_ref_image()
-        pipe.setup(im, strength=0.7, steps=5)
-        result = pipe.gen(dict(prompt="cube planet cartoon style"))
-        result.save('test_img2img_basic.png')
+        seed = 49045438434843
+        pipe.setup(im, strength=0.7, steps=5, guidance_scale=3.3)
+        self.assertEqual(3.3, pipe.pipe_params['guidance_scale'])
+        image = pipe.gen(dict(prompt="cube planet cartoon style", generator=torch.cuda.manual_seed(seed)))
+        image.save('test_img2img_basic.png')
+        pipe.setup(im, strength=0.7, steps=5, guidance_scale=7.6)
+        image1 = pipe.gen(dict(prompt="cube planet cartoon style", generator=torch.cuda.manual_seed(seed)))
+        diff = self.compute_diff(image1, image)
+        # check that difference is large
+        self.assertGreater(diff, 1000)
+        pipe.setup(im, strength=0.7, steps=5, guidance_scale=3.3)
+        image2 = pipe.gen(dict(prompt="cube planet cartoon style", generator=torch.cuda.manual_seed(seed)))
+        diff = self.compute_diff(image2, image)
+        # check that difference is small
+        self.assertLess(diff, 1)
+
+    def test_maskedimg2img_basic(self):
+        pipe = MaskedIm2ImPipe(self.get_model(), model_type=self.model_type())
+        img = PIL.Image.open("./mech_beard_sigm.png")
+        # read image with mask painted over
+        img_paint = numpy.array(PIL.Image.open("./mech_beard_sigm_mask.png"))
+
+        scheduler = "EulerAncestralDiscreteScheduler"
+        seed = 49045438434843
+        blur = 48
+        param_3_3 = dict(original_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,
+               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'])
+        image = pipe.gen(dict(prompt="cube planet cartoon style", generator=torch.cuda.manual_seed(seed)))
+        image.save('test_img2img_basic.png')
+        pipe.setup(**param_7_6)
+        image1 = pipe.gen(dict(prompt="cube planet cartoon style", generator=torch.cuda.manual_seed(seed)))
+        diff = self.compute_diff(image1, image)
+        # check that difference is large
+        self.assertGreater(diff, 1000)
+        pipe.setup(**param_3_3)
+        image2 = pipe.gen(dict(prompt="cube planet cartoon style", generator=torch.cuda.manual_seed(seed)))
+        diff = self.compute_diff(image2, image)
+        # check that difference is small
+        self.assertLess(diff, 1)
 
     @unittest.skipIf(not can_run_lpw(), "can't run on tiny version of SD")
     def test_lpw(self):