Use mixed precision

model.py

@@ -57,9 +57,13 @@ class Model:
         if base_model_id == self.base_model_id and task_name == self.task_name:
             return self.pipe
         model_id = CONTROLNET_MODEL_IDS[task_name]
-        controlnet = ControlNetModel.from_pretrained(model_id)
+        controlnet = ControlNetModel.from_pretrained(model_id,
+                                                     torch_dtype=torch.float16)
         pipe = StableDiffusionControlNetPipeline.from_pretrained(
-            base_model_id, safety_checker=None, controlnet=controlnet)
+            base_model_id,
+            safety_checker=None,
+            controlnet=controlnet,
+            torch_dtype=torch.float16)
         pipe.scheduler = UniPCMultistepScheduler.from_config(
             pipe.scheduler.config)
         pipe.enable_xformers_memory_efficient_attention()
@@ -89,7 +93,9 @@ class Model:
         torch.cuda.empty_cache()
         gc.collect()
         model_id = CONTROLNET_MODEL_IDS[task_name]
-        controlnet = ControlNetModel.from_pretrained(model_id).to(self.device)
+        controlnet = ControlNetModel.from_pretrained(model_id,
+                                                     torch_dtype=torch.float16)
+        controlnet.to(self.device)
         torch.cuda.empty_cache()
         gc.collect()
         self.pipe.controlnet = controlnet
@@ -102,6 +108,7 @@ class Model:
             prompt = f'{prompt}, {additional_prompt}'
         return prompt
 
+    @torch.autocast('cuda')
     def run_pipe(
         self,
         prompt: str,