Update model.py

model.py

@@ -103,10 +103,7 @@ class Model:
         control_strength: float,
         preprocessor_name: str,
     ) -> list[PIL.Image.Image]:
-        self.t5_model.model.to('cuda').to(torch.bfloat16)
-        self.load_gpt_weight('edge')
-        self.gpt_model.to('cuda').to(torch.bfloat16)
-        self.vq_model.to('cuda')
+        
         if isinstance(image, np.ndarray):
             image = Image.fromarray(image)
         origin_W, origin_H = image.size
@@ -125,9 +122,15 @@ class Model:
         elif preprocessor_name == 'No preprocess':
             condition_img = image
         print('get edge')
+        del self.preprocessor.model
+        torch.cuda.empty_cache()
         condition_img = condition_img.resize((512,512))
         W, H = condition_img.size
 
+        self.t5_model.model.to('cuda').to(torch.bfloat16)
+        self.load_gpt_weight('edge')
+        self.gpt_model.to('cuda').to(torch.bfloat16)
+        self.vq_model.to('cuda')
         condition_img = torch.from_numpy(np.array(condition_img)).unsqueeze(0).permute(0,3,1,2).repeat(1,1,1,1)
         condition_img = condition_img.to(self.device)
         condition_img = 2*(condition_img/255 - 0.5)
@@ -198,10 +201,7 @@ class Model:
         control_strength: float,
         preprocessor_name: str
     ) -> list[PIL.Image.Image]:
-        self.t5_model.model.to(self.device).to(torch.bfloat16)
-        self.load_gpt_weight('depth')
-        self.gpt_model.to('cuda').to(torch.bfloat16)
-        self.vq_model.to(self.device)
+        
         if isinstance(image, np.ndarray):
             image = Image.fromarray(image)
         origin_W, origin_H = image.size
@@ -216,9 +216,15 @@ class Model:
         elif preprocessor_name == 'No preprocess':
             condition_img = image
         print('get depth')
+        del self.preprocessor.model
+        torch.cuda.empty_cache()
         condition_img = condition_img.resize((512,512))
         W, H = condition_img.size
 
+        self.t5_model.model.to(self.device).to(torch.bfloat16)
+        self.load_gpt_weight('depth')
+        self.gpt_model.to('cuda').to(torch.bfloat16)
+        self.vq_model.to(self.device)
         condition_img = torch.from_numpy(np.array(condition_img)).unsqueeze(0).permute(0,3,1,2).repeat(1,1,1,1)
         condition_img = condition_img.to(self.device)
         condition_img = 2*(condition_img/255 - 0.5)