try pix2pix

image_transformation.py

@@ -15,7 +15,7 @@ if is_vision_available():
     from PIL import Image
 
 if is_diffusers_available():
-    from diffusers import ControlNetModel, StableDiffusionControlNetPipeline, DPMSolverMultistepScheduler
+    from diffusers import DiffusionPipeline
 
 if is_opencv_available():
     import cv2
@@ -29,8 +29,7 @@ IMAGE_TRANSFORMATION_DESCRIPTION = (
 
 
 class ImageTransformationTool(Tool):
-    default_stable_diffusion_checkpoint = "runwayml/stable-diffusion-v1-5"
-    default_controlnet_checkpoint = "lllyasviel/control_v11e_sd15_ip2p"
+    default_stable_diffusion_checkpoint = "timbrooks/instruct-pix2pix"
     description = IMAGE_TRANSFORMATION_DESCRIPTION
     inputs = ['image', 'text']
     outputs = ['image']
@@ -47,13 +46,7 @@ class ImageTransformationTool(Tool):
 
         super().__init__()
 
-        if controlnet is None:
-            controlnet = self.default_controlnet_checkpoint
-        self.controlnet_checkpoint = controlnet
-
-        if stable_diffusion is None:
-            stable_diffusion = self.default_stable_diffusion_checkpoint
-        self.stable_diffusion_checkpoint = stable_diffusion
+        self.stable_diffusion = self.default_stable_diffusion_checkpoint
 
         self.device = device
         self.hub_kwargs = hub_kwargs
@@ -62,11 +55,9 @@ class ImageTransformationTool(Tool):
         if self.device is None:
             self.device = get_default_device()
 
-        self.controlnet = ControlNetModel.from_pretrained(self.controlnet_checkpoint)
-        self.pipeline = StableDiffusionControlNetPipeline.from_pretrained(
+        self.pipeline = DiffusionPipeline.from_pretrained(
             self.stable_diffusion_checkpoint, controlnet=self.controlnet
         )
-        self.pipeline.scheduler = DPMSolverMultistepScheduler.from_config(self.pipeline.scheduler.config)
 
         self.pipeline.to(self.device)
         if self.device.type == "cuda":
@@ -78,18 +69,8 @@ class ImageTransformationTool(Tool):
         if not self.is_initialized:
             self.setup()
 
-        low_threshold = 100
-        high_threshold = 200
-
-        image = np.array(image)
-        image = cv2.Canny(image, low_threshold, high_threshold)
-        image = image[:, :, None]
-        image = np.concatenate([image, image, image], axis=2)
-        canny_image = Image.fromarray(image)
-
         return self.pipeline(
             prompt + added_prompt,
-            canny_image,
+            image,
             negative_prompt=negative_prompt,
-            num_inference_steps=25,
         ).images[0]