Update model.py

model.py

@@ -74,8 +74,8 @@ class Model:
 
     def load_gpt(self, condition_type='canny'):
         gpt_ckpt = models[condition_type]
-        precision = torch.bfloat16
-        # precision = torch.float32
+        # precision = torch.bfloat16
+        precision = torch.float32
         latent_size = 768 // 16
         gpt_model = GPT_models["GPT-XL"](
             block_size=latent_size**2,
@@ -93,8 +93,8 @@ class Model:
         return gpt_model
 
     def load_t5(self):
-        precision = torch.bfloat16
-        # precision = torch.float32
+        # precision = torch.bfloat16
+        precision = torch.float32
         t5_model = T5Embedder(
             device=self.device,
             local_cache=True,
@@ -124,8 +124,8 @@ class Model:
         W, H = image.size
         print(W, H)
         # self.gpt_model_depth.to('cpu')
-        self.t5_model.model.to('cuda')
-        self.gpt_model_canny.to('cuda')
+        self.t5_model.model.to('cuda').to(torch.bfloat16)
+        self.gpt_model_canny.to('cuda').to(torch.bfloat16)
         self.vq_model.to('cuda')
         # print("after cuda", self.gpt_model_canny.adapter.model.embeddings.patch_embeddings.projection.weight)