Improve image transformation

image_transformation.py

@@ -30,7 +30,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_v11p_sd15_canny"
     description = IMAGE_TRANSFORMATION_DESCRIPTION
     inputs = ['image', 'text']
     outputs = ['image']
@@ -67,11 +67,14 @@ 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(device=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, 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()
 
@@ -87,12 +90,9 @@ class ImageTransformationTool(Tool):
         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,
+            prompt + added_prompt,
             canny_image,
-            negative_prompt="monochrome, lowres, bad anatomy, worst quality, low quality",
-            num_inference_steps=20,
-            generator=generator,
+            negative_prompt=negative_prompt,
+            num_inference_steps=25,
         ).images[0]