Merge pull request #63 from okotaku/feat/support_negative_prompt

[Feature] Support negative prompt for inference

diffengine/models/editors/ip_adapter/ip_adapter_xl.py

@@ -116,7 +116,7 @@ def _encode_image(self, image, num_images_per_prompt):
     def infer(self,
               prompt: List[str],
               example_image: List[Union[str, Image.Image]],
-              negative_prompt: Optional[List[str]] = None,
+              negative_prompt: Optional[str] = None,
               height: Optional[int] = None,
               width: Optional[int] = None) -> List[np.ndarray]:
         """Function invoked when calling the pipeline for generation.
@@ -126,8 +126,9 @@ def infer(self,
                 The prompt or prompts to guide the image generation.
             example_image (`List[Union[str, Image.Image]]`):
                 The image prompt or prompts to guide the image generation.
-            negative_prompt (`List[str]`):
+            negative_prompt (`Optional[str]`):
                 The prompt or prompts to guide the image generation.
+                Defaults to None.
             height (`int`, *optional*, defaults to
                 `self.unet.config.sample_size * self.vae_scale_factor`):
                 The height in pixels of the generated image.

diffengine/models/editors/stable_diffusion/stable_diffusion.py

@@ -112,13 +112,17 @@ def device(self):
     @torch.no_grad()
     def infer(self,
               prompt: List[str],
+              negative_prompt: Optional[str] = None,
               height: Optional[int] = None,
               width: Optional[int] = None) -> List[np.ndarray]:
         """Function invoked when calling the pipeline for generation.
 
         Args:
             prompt (`List[str]`):
                 The prompt or prompts to guide the image generation.
+            negative_prompt (`Optional[str]`):
+                The prompt or prompts to guide the image generation.
+                Defaults to None.
             height (`int`, *optional*, defaults to
                 `self.unet.config.sample_size * self.vae_scale_factor`):
                 The height in pixels of the generated image.
@@ -138,7 +142,10 @@ def infer(self,
         images = []
         for p in prompt:
             image = pipeline(
-                p, num_inference_steps=50, height=height,
+                p,
+                negative_prompt=negative_prompt,
+                num_inference_steps=50,
+                height=height,
                 width=width).images[0]
             images.append(np.array(image))
 

diffengine/models/editors/stable_diffusion_controlnet/stable_diffusion_controlnet.py

@@ -105,6 +105,7 @@ def prepare_model(self):
     def infer(self,
               prompt: List[str],
               condition_image: List[Union[str, Image.Image]],
+              negative_prompt: Optional[str] = None,
               height: Optional[int] = None,
               width: Optional[int] = None) -> List[np.ndarray]:
         """Function invoked when calling the pipeline for generation.
@@ -114,6 +115,9 @@ def infer(self,
                 The prompt or prompts to guide the image generation.
             condition_image (`List[Union[str, Image.Image]]`):
                 The condition image for ControlNet.
+            negative_prompt (`Optional[str]`):
+                The prompt or prompts to guide the image generation.
+                Defaults to None.
             height (`int`, *optional*, defaults to
                 `self.unet.config.sample_size * self.vae_scale_factor`):
                 The height in pixels of the generated image.
@@ -138,7 +142,11 @@ def infer(self,
                 img = load_image(img)
             img = img.convert('RGB')
             image = pipeline(
-                p, img, num_inference_steps=50, height=height,
+                p,
+                img,
+                negative_prompt=negative_prompt,
+                num_inference_steps=50,
+                height=height,
                 width=width).images[0]
             images.append(np.array(image))
 

diffengine/models/editors/stable_diffusion_xl/stable_diffusion_xl.py

@@ -159,13 +159,17 @@ def device(self):
     @torch.no_grad()
     def infer(self,
               prompt: List[str],
+              negative_prompt: Optional[str] = None,
               height: Optional[int] = None,
               width: Optional[int] = None) -> List[np.ndarray]:
         """Function invoked when calling the pipeline for generation.
 
         Args:
             prompt (`List[str]`):
                 The prompt or prompts to guide the image generation.
+            negative_prompt (`Optional[str]`):
+                The prompt or prompts to guide the image generation.
+                Defaults to None.
             height (`int`, *optional*, defaults to
                 `self.unet.config.sample_size * self.vae_scale_factor`):
                 The height in pixels of the generated image.
@@ -198,7 +202,10 @@ def infer(self,
         images = []
         for p in prompt:
             image = pipeline(
-                p, num_inference_steps=50, height=height,
+                p,
+                negative_prompt=negative_prompt,
+                num_inference_steps=50,
+                height=height,
                 width=width).images[0]
             images.append(np.array(image))
 

diffengine/models/editors/stable_diffusion_xl_controlnet/stable_diffusion_xl_controlnet.py

@@ -106,6 +106,7 @@ def prepare_model(self):
     def infer(self,
               prompt: List[str],
               condition_image: List[Union[str, Image.Image]],
+              negative_prompt: Optional[str] = None,
               height: Optional[int] = None,
               width: Optional[int] = None) -> List[np.ndarray]:
         """Function invoked when calling the pipeline for generation.
@@ -115,6 +116,9 @@ def infer(self,
                 The prompt or prompts to guide the image generation.
             condition_image (`List[Union[str, Image.Image]]`):
                 The condition image for ControlNet.
+            negative_prompt (`Optional[str]`):
+                The prompt or prompts to guide the image generation.
+                Defaults to None.
             height (`int`, *optional*, defaults to
                 `self.unet.config.sample_size * self.vae_scale_factor`):
                 The height in pixels of the generated image.
@@ -143,7 +147,12 @@ def infer(self,
                 img = load_image(img)
             img = img.convert('RGB')
             image = pipeline(
-                p, p, img, num_inference_steps=50, height=height,
+                p,
+                p,
+                img,
+                negative_prompt=negative_prompt,
+                num_inference_steps=50,
+                height=height,
                 width=width).images[0]
             images.append(np.array(image))
 

tests/test_models/test_editors/test_ip_adapter/test_ip_adapter.py

@@ -45,6 +45,16 @@ def test_infer(self):
         # test device
         assert StableDiffuser.device.type == 'cpu'
 
+        # test infer with negative_prompt
+        result = StableDiffuser.infer(
+            ['an insect robot preparing a delicious meal'],
+            ['tests/testdata/color.jpg'],
+            negative_prompt='noise',
+            height=64,
+            width=64)
+        assert len(result) == 1
+        assert result[0].shape == (64, 64, 3)
+
     def test_train_step(self):
         # test load with loss module
         StableDiffuser = IPAdapterXL(

tests/test_models/test_editors/test_ip_adapter/test_ip_adapter_plus.py

@@ -46,6 +46,16 @@ def test_infer(self):
         # test device
         assert StableDiffuser.device.type == 'cpu'
 
+        # test infer with negative_prompt
+        result = StableDiffuser.infer(
+            ['an insect robot preparing a delicious meal'],
+            ['tests/testdata/color.jpg'],
+            negative_prompt='noise',
+            height=64,
+            width=64)
+        assert len(result) == 1
+        assert result[0].shape == (64, 64, 3)
+
     def test_train_step(self):
         # test load with loss module
         StableDiffuser = IPAdapterXLPlus(

tests/test_models/test_editors/test_stable_diffusion/test_stable_diffusion.py

@@ -26,6 +26,15 @@ def test_infer(self):
         # test device
         assert StableDiffuser.device.type == 'cpu'
 
+        # test infer with negative_prompt
+        result = StableDiffuser.infer(
+            ['an insect robot preparing a delicious meal'],
+            negative_prompt='noise',
+            height=64,
+            width=64)
+        assert len(result) == 1
+        assert result[0].shape == (64, 64, 3)
+
     def test_train_step(self):
         # test load with loss module
         StableDiffuser = StableDiffusion(

tests/test_models/test_editors/test_stable_diffusion_controlnet/test_stable_diffusion_controlnet.py

@@ -49,6 +49,16 @@ def test_infer(self):
         # test device
         assert StableDiffuser.device.type == 'cpu'
 
+        # test infer with negative_prompt
+        result = StableDiffuser.infer(
+            ['an insect robot preparing a delicious meal'],
+            ['tests/testdata/color.jpg'],
+            negative_prompt='noise',
+            height=64,
+            width=64)
+        assert len(result) == 1
+        assert result[0].shape == (64, 64, 3)
+
         # test controlnet small
         StableDiffuser = StableDiffusionControlNet(
             'hf-internal-testing/tiny-stable-diffusion-pipe',

tests/test_models/test_editors/test_stable_diffusion_xl/test_stable_diffusion_xl.py

@@ -26,6 +26,15 @@ def test_infer(self):
         # test device
         assert StableDiffuser.device.type == 'cpu'
 
+        # test infer with negative_prompt
+        result = StableDiffuser.infer(
+            ['an insect robot preparing a delicious meal'],
+            negative_prompt='noise',
+            height=64,
+            width=64)
+        assert len(result) == 1
+        assert result[0].shape == (64, 64, 3)
+
     def test_infer_with_pre_compute_embs(self):
         StableDiffuser = StableDiffusionXL(
             'hf-internal-testing/tiny-stable-diffusion-xl-pipe',

tests/test_models/test_editors/test_stable_diffusion_xl_controlnet/test_stable_diffusion_xl_controlnet.py

@@ -49,6 +49,16 @@ def test_infer(self):
         # test device
         assert StableDiffuser.device.type == 'cpu'
 
+        # test infer with negative_prompt
+        result = StableDiffuser.infer(
+            ['an insect robot preparing a delicious meal'],
+            ['tests/testdata/color.jpg'],
+            negative_prompt='noise',
+            height=64,
+            width=64)
+        assert len(result) == 1
+        assert result[0].shape == (64, 64, 3)
+
         # test controlnet small
         StableDiffuser = StableDiffusionXLControlNet(
             'hf-internal-testing/tiny-stable-diffusion-xl-pipe',