Stop using cpu_offload

model.py

@@ -2,6 +2,7 @@
 # The original license file is LICENSE.ControlNet in this repo.
 from __future__ import annotations
 
+import gc
 import pathlib
 import sys
 
@@ -24,7 +25,6 @@ from annotator.mlsd import apply_mlsd
 from annotator.openpose import apply_openpose
 from annotator.uniformer import apply_uniformer
 from annotator.util import HWC3, resize_image
-from share import *
 
 CONTROLNET_MODEL_IDS = {
     'canny': 'lllyasviel/sd-controlnet-canny',
@@ -47,6 +47,8 @@ class Model:
     def __init__(self,
                  base_model_id: str = 'runwayml/stable-diffusion-v1-5',
                  task_name: str = 'canny'):
+        self.device = torch.device(
+            'cuda:0' if torch.cuda.is_available() else 'cpu')
         self.base_model_id = ''
         self.task_name = ''
         self.pipe = self.load_pipe(base_model_id, task_name)
@@ -55,33 +57,41 @@ 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,
-                                                     torch_dtype=torch.float16)
+        controlnet = ControlNetModel.from_pretrained(model_id)
         pipe = StableDiffusionControlNetPipeline.from_pretrained(
-            base_model_id,
-            safety_checker=None,
-            controlnet=controlnet,
-            torch_dtype=torch.float16)
+            base_model_id, safety_checker=None, controlnet=controlnet)
         pipe.scheduler = UniPCMultistepScheduler.from_config(
             pipe.scheduler.config)
         pipe.enable_xformers_memory_efficient_attention()
-        pipe.enable_model_cpu_offload()
+        pipe.to(self.device)
+        torch.cuda.empty_cache()
+        gc.collect()
         self.base_model_id = base_model_id
         self.task_name = task_name
         return pipe
 
     def set_base_model(self, base_model_id: str) -> str:
-        self.pipe = self.load_pipe(base_model_id, self.task_name)
+        if not base_model_id or base_model_id == self.base_model_id:
+            return self.base_model_id
+        del self.pipe
+        torch.cuda.empty_cache()
+        gc.collect()
+        try:
+            self.pipe = self.load_pipe(base_model_id, self.task_name)
+        except Exception:
+            self.pipe = self.load_pipe(self.base_model_id, self.task_name)
         return self.base_model_id
 
     def load_controlnet_weight(self, task_name: str) -> None:
         if task_name == self.task_name:
             return
+        del self.pipe.controlnet
+        torch.cuda.empty_cache()
+        gc.collect()
         model_id = CONTROLNET_MODEL_IDS[task_name]
-        controlnet = ControlNetModel.from_pretrained(model_id,
-                                                     torch_dtype=torch.float16)
-        from accelerate import cpu_offload_with_hook
-        cpu_offload_with_hook(controlnet, torch.device('cuda:0'))
+        controlnet = ControlNetModel.from_pretrained(model_id).to(self.device)
+        torch.cuda.empty_cache()
+        gc.collect()
         self.pipe.controlnet = controlnet
         self.task_name = task_name