Improve image transformation

image_transformation.py

@@ -15,7 +15,7 @@ if is_vision_available():
     from PIL import Image
 
 if is_diffusers_available():
-    from diffusers import ControlNetModel, StableDiffusionControlNetPipeline, UniPCMultistepScheduler
+    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/sd-controlnet-canny"
+    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,37 +55,20 @@ 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.stable_diffusion_checkpoint, controlnet=self.controlnet
-        )
-        self.pipeline.scheduler = UniPCMultistepScheduler.from_config(self.pipeline.scheduler.config)
-        self.pipeline.enable_model_cpu_offload()
+        self.pipeline = DiffusionPipeline.from_pretrained(self.stable_diffusion)
+
+        self.pipeline.to(self.device)
+        if self.device.type == "cuda":
+            self.pipeline.to(torch_dtype=torch.float16)
 
         self.is_initialized = True
 
-    def __call__(self, image, prompt):
+    def __call__(self, image, prompt, negative_prompt="low quality, bad quality, deformed, low resolution", added_prompt=" , highest quality, highly realistic, very high resolution"):
         if not self.is_initialized:
             self.setup()
 
-        initial_prompt = "super-hero character, best quality, extremely detailed"
-        prompt = initial_prompt + prompt
-
-        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)
-
-        generator = torch.Generator(device="cpu").manual_seed(2)
-
         return self.pipeline(
-            prompt,
-            canny_image,
-            negative_prompt="monochrome, lowres, bad anatomy, worst quality, low quality",
-            num_inference_steps=20,
-            generator=generator,
+            prompt + added_prompt,
+            image,
+            negative_prompt=negative_prompt,
         ).images[0]