Update model.py

model.py

@@ -87,7 +87,7 @@ class Model:
     def load_t5(self):
         precision = torch.bfloat16
         t5_model = T5Embedder(
-            device="cuda",
+            device=self.device,
             local_cache=True,
             cache_dir='checkpoints/flan-t5-xl',
             dir_or_name='flan-t5-xl',
@@ -114,8 +114,8 @@ class Model:
         image = resize_image_to_16_multiple(image, 'canny')
         W, H = image.size
         print(W, H)
-        self.t5_model.model.to(self.device)
-        self.gpt_model_canny.to(self.device)
+        # self.t5_model.model.to(self.device)
+        # self.gpt_model_canny.to(self.device)
 
         condition_img = self.get_control_canny(np.array(image), low_threshold,
                                                high_threshold)
@@ -191,10 +191,10 @@ class Model:
         image = resize_image_to_16_multiple(image, 'depth')
         W, H = image.size
         print(W, H)
-        self.t5_model.model.to(self.device)
-        self.gpt_model_depth.to(self.device)
-        self.get_control_depth.model.to(self.device)
-        self.vq_model.to(self.device)
+        # self.t5_model.model.to(self.device)
+        # self.gpt_model_depth.to(self.device)
+        # self.get_control_depth.model.to(self.device)
+        # self.vq_model.to(self.device)
         image_tensor = torch.from_numpy(np.array(image)).to(self.device)
         condition_img = torch.from_numpy(
             self.get_control_depth(image_tensor)).unsqueeze(0)