Update image_transformation.py

image_transformation.py

@@ -17,9 +17,6 @@ if is_vision_available():
 if is_diffusers_available():
     from diffusers import ControlNetModel, StableDiffusionControlNetPipeline, UniPCMultistepScheduler
 
-if is_opencv_available():
-    import cv2
-
 
 IMAGE_TRANSFORMATION_DESCRIPTION = (
     "This is a tool that transforms an image according to a prompt. It takes two inputs: `image`, which should be "
@@ -30,7 +27,7 @@ IMAGE_TRANSFORMATION_DESCRIPTION = (
 
 class ImageTransformationTool(Tool):
     default_stable_diffusion_checkpoint = "runwayml/stable-diffusion-v1-5"
-    default_controlnet_checkpoint = "lllyasviel/sd-controlnet-canny"
+    default_controlnet_checkpoint = "lllyasviel/control_v11e_sd15_ip2p"
     description = IMAGE_TRANSFORMATION_DESCRIPTION
     inputs = ['image', 'text']
     outputs = ['image']
@@ -67,32 +64,20 @@ class ImageTransformationTool(Tool):
             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.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, added_prompt=", high quality, high resolution, beautiful, aesthetic, sharp"):
         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="monochrome, lowres, worst quality, low quality",
+            num_inference_steps=25,
         ).images[0]