Update model_new.py

model_new.py

@@ -180,21 +180,33 @@ class Model:
         top_k: int,
         top_p: int,
         seed: int,
+        control_strength: float,
+        preprocessor_name: str
     ) -> list[PIL.Image.Image]:
-        image = resize_image_to_16_multiple(image, 'depth')
-        W, H = image.size
-        print(W, H)
         self.gpt_model_canny.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.vq_model.to(self.device)
-        image_tensor = torch.from_numpy(np.array(image)).to(self.device)
+        if isinstance(image, np.ndarray):
+            image = Image.fromarray(image)
+        origin_W, origin_H = image.size
+        # print(image)
+        if preprocessor_name == 'depth':
+            self.preprocessor.load("Depth")
+            condition_img = self.preprocessor(
+                    image=image,
+                    image_resolution=512,
+                    detect_resolution=512,
+                )
+        elif preprocessor_name == 'No preprocess':
+            condition_img = image
+        condition_img = condition_img.resize((512,512))
+        W, H = condition_img.size
 
-        condition_img = 2 * (image_tensor / 255 - 0.5)
-        print(condition_img.shape)
-        condition_img = condition_img.permute(2,0,1).unsqueeze(0).repeat(2, 1, 1, 1)
-        
+        condition_img = torch.from_numpy(np.array(condition_img)).unsqueeze(0).permute(0,3,1,2).repeat(2,1,1,1)
+        condition_img = condition_img.to(self.device)
+        condition_img = 2*(condition_img/255 - 0.5)
         prompts = [prompt] * 2
         caption_embs, emb_masks = self.t5_model.get_text_embeddings(prompts)