fix: commute inference to gpu and touches

handler.py

@@ -11,12 +11,12 @@ import cv2
 import controlnet_hinter
 
 # set device
-# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-# if device.type != 'cuda':
-#     raise ValueError("need to run on GPU")
-# # set mixed precision dtype
-# dtype = torch.bfloat16 if torch.cuda.get_device_capability()[
-#     0] == 8 else torch.float16
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+if device.type != 'cuda':
+    raise ValueError("need to run on GPU")
+# set mixed precision dtype
+dtype = torch.bfloat16 if torch.cuda.get_device_capability()[
+    0] == 8 else torch.float16
 
 # for the moment, support only canny edge
 SDXLCONTROLNET_MAPPING = {
@@ -24,7 +24,7 @@ SDXLCONTROLNET_MAPPING = {
         "model_id": "diffusers/controlnet-canny-sdxl-1.0",
         "hinter": controlnet_hinter.hint_canny
     },
-     "pose": {
+    "pose": {
         "model_id": "lllyasviel/sd-controlnet-openpose",
         "hinter": controlnet_hinter.hint_openpose
     },
@@ -96,7 +96,7 @@ class EndpointHandler():
         # define default controlnet id and load controlnet
         self.control_type = "canny_edge"
         self.controlnet = ControlNetModel.from_pretrained(
-            SDXLCONTROLNET_MAPPING[self.control_type]["model_id"], torch_dtype=torch.float16, use_safetensors=True)
+            SDXLCONTROLNET_MAPPING[self.control_type]["model_id"], torch_dtype=dtype).to(device)
 
         # Load StableDiffusionControlNetPipeline
         self.sdxl_id = "stabilityai/stable-diffusion-xl-base-1.0"
@@ -104,10 +104,8 @@ class EndpointHandler():
         self.pipe = StableDiffusionXLControlNetPipeline.from_pretrained(self.sdxl_id,
                                                                         controlnet=self.controlnet,
                                                                         # vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16, use_safetensors=True),
-                                                                        torch_dtype=torch.float16,
-                                                                        use_safetensors=True)
-        self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
-        self.pipe.enable_model_cpu_offload()
+                                                                        torch_dtype=dtype,
+                                                                        safety_checker=None).to(device)
         self.generator = torch.Generator(device="cpu").manual_seed(3)
 
     def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
@@ -127,12 +125,16 @@ class EndpointHandler():
         if num_of_images is None:
             num_of_images = 1
 
+        if controlnet_type is not "canny_edge" and controlnet_type is not None:
+            return {"error": "Please provide a valid controlnet type. Only canny_edge is supported at the moment."}
+
         # Check if a new controlnet is provided
         if controlnet_type is not None and controlnet_type != self.control_type:
             print(
                 f"changing controlnet from {self.control_type} to {controlnet_type} using {SDXLCONTROLNET_MAPPING[controlnet_type]['model_id']} model")
             self.control_type = controlnet_type
-            self.controlnet = ControlNetModel.from_pretrained(SDXLCONTROLNET_MAPPING[self.control_type]["model_id"], torch_dtype=torch.float16, use_safetensors=True)
+            self.controlnet = ControlNetModel.from_pretrained(
+                SDXLCONTROLNET_MAPPING[self.control_type]["model_id"], torch_dtype=dtype).to(device)
             self.pipe.controlnet = self.controlnet
 
         # hyperparamters
@@ -148,8 +150,6 @@ class EndpointHandler():
         image = self.decode_base64_image(image)
         control_image = SDXLCONTROLNET_MAPPING[self.control_type]["hinter"](
             image, width=1024, height=1024)
-    
-        self.generator = torch.manual_seed(1)
 
         # run inference pipeline
         images = self.pipe(