Update model.py

model.py

@@ -52,8 +52,8 @@ class Model:
         self.gpt_model_canny = self.load_gpt(condition_type='canny')
         self.gpt_model_depth = self.load_gpt(condition_type='depth')
         self.get_control_canny = CannyDetector()
-        # self.get_control_depth = MidasDetector('cuda')
-        self.get_control_depth =  MidasDetector.from_pretrained("lllyasviel/Annotators")
+        self.get_control_depth = MidasDetector('cuda')
+        # self.get_control_depth =  MidasDetector.from_pretrained("lllyasviel/Annotators")
 
     def to(self, device):
         self.gpt_model_canny.to('cuda')
@@ -206,12 +206,12 @@ class Model:
         # condition_img = condition_img.unsqueeze(0).repeat(2, 3, 1, 1)
         # condition_img = condition_img.to(self.device)
         # condition_img = 2 * (condition_img / 255 - 0.5)
-
-        control_image = self.get_control_depth(
-                image=image,
-                image_resolution=512,
-                detect_resolution=512,
-            )
+        condition_img = 2 * (image_tensor / 255 - 0.5)
+        # control_image = self.get_control_depth(
+        #         image=image,
+        #         image_resolution=512,
+        #         detect_resolution=512,
+        #     )
         
         prompts = [prompt] * 2
         caption_embs, emb_masks = self.t5_model.get_text_embeddings(prompts)