Update model.py

model.py

@@ -64,10 +64,14 @@ class Model:
         return gpt_model
 
     def load_gpt_weight(self, condition_type='edge'):
+        torch.cuda.empty_cache()
+        gc.collect()
         gpt_ckpt = models[condition_type]
         model_weight = load_file(gpt_ckpt)
         self.gpt_model.load_state_dict(model_weight, strict=False)
         self.gpt_model.eval()
+        torch.cuda.empty_cache()
+        gc.collect()
         # print("gpt model is loaded")
         
     def load_t5(self):
@@ -193,10 +197,12 @@ class Model:
         control_strength: float,
         preprocessor_name: str
     ) -> list[PIL.Image.Image]:
-        self.gpt_model_edge.to('cpu')
+        # self.gpt_model_edge.to('cpu')
         self.t5_model.model.to(self.device)
-        self.gpt_model_depth.to(self.device)
-        self.get_control_depth.model.to(self.device)
+        # self.gpt_model_depth.to(self.device)
+        self.load_gpt_weight('depth')
+        self.gpt_model.to('cuda').to(torch.bfloat16)
+        # self.get_control_depth.model.to(self.device)
         self.vq_model.to(self.device)
         if isinstance(image, np.ndarray):
             image = Image.fromarray(image)
@@ -237,7 +243,7 @@ class Model:
         qzshape = [len(c_indices), 8, H // 16, W // 16]
         t1 = time.time()
         index_sample = generate(
-            self.gpt_model_depth,
+            self.gpt_model,
             c_indices,
             (H // 16) * (W // 16),
             c_emb_masks,