Support raw pose image and depth image

app_depth.py

@@ -13,6 +13,8 @@ def create_demo(process, max_images=12, default_num_images=3):
                 prompt = gr.Textbox(label='Prompt')
                 run_button = gr.Button(label='Run')
                 with gr.Accordion('Advanced options', open=False):
+                    is_depth_image = gr.Checkbox(label='Is depth image',
+                                                 value=False)
                     num_samples = gr.Slider(label='Images',
                                             minimum=1,
                                             maximum=max_images,
@@ -67,6 +69,7 @@ def create_demo(process, max_images=12, default_num_images=3):
             num_steps,
             guidance_scale,
             seed,
+            is_depth_image,
         ]
         prompt.submit(fn=process, inputs=inputs, outputs=result)
         run_button.click(fn=process,

app_pose.py

@@ -13,6 +13,8 @@ def create_demo(process, max_images=12, default_num_images=3):
                 prompt = gr.Textbox(label='Prompt')
                 run_button = gr.Button(label='Run')
                 with gr.Accordion('Advanced options', open=False):
+                    is_pose_image = gr.Checkbox(label='Is pose image',
+                                                value=False)
                     num_samples = gr.Slider(label='Images',
                                             minimum=1,
                                             maximum=max_images,
@@ -67,6 +69,7 @@ def create_demo(process, max_images=12, default_num_images=3):
             num_steps,
             guidance_scale,
             seed,
+            is_pose_image,
         ]
         prompt.submit(fn=process, inputs=inputs, outputs=result)
         run_button.click(fn=process,

model.py

@@ -438,16 +438,19 @@ class Model:
         input_image: np.ndarray,
         image_resolution: int,
         detect_resolution: int,
+        is_pose_image: bool,
     ) -> tuple[PIL.Image.Image, PIL.Image.Image]:
         input_image = HWC3(input_image)
-        control_image, _ = apply_openpose(
-            resize_image(input_image, detect_resolution))
-        control_image = HWC3(control_image)
-        image = resize_image(input_image, image_resolution)
-        H, W = image.shape[:2]
-
-        control_image = cv2.resize(control_image, (W, H),
-                                   interpolation=cv2.INTER_NEAREST)
+        if not is_pose_image:
+            control_image, _ = apply_openpose(
+                resize_image(input_image, detect_resolution))
+            control_image = HWC3(control_image)
+            image = resize_image(input_image, image_resolution)
+            H, W = image.shape[:2]
+            control_image = cv2.resize(control_image, (W, H),
+                                       interpolation=cv2.INTER_NEAREST)
+        else:
+            control_image = input_image
 
         return PIL.Image.fromarray(control_image), PIL.Image.fromarray(
             control_image)
@@ -465,11 +468,13 @@ class Model:
         num_steps: int,
         guidance_scale: float,
         seed: int,
+        is_pose_image: bool,
     ) -> list[PIL.Image.Image]:
         control_image, vis_control_image = self.preprocess_pose(
             input_image=input_image,
             image_resolution=image_resolution,
             detect_resolution=detect_resolution,
+            is_pose_image=is_pose_image,
         )
         return self.process(
             task_name='pose',
@@ -537,15 +542,19 @@ class Model:
         input_image: np.ndarray,
         image_resolution: int,
         detect_resolution: int,
+        is_depth_image: bool,
     ) -> tuple[PIL.Image.Image, PIL.Image.Image]:
         input_image = HWC3(input_image)
-        control_image, _ = apply_midas(
-            resize_image(input_image, detect_resolution))
-        control_image = HWC3(control_image)
-        image = resize_image(input_image, image_resolution)
-        H, W = image.shape[:2]
-        control_image = cv2.resize(control_image, (W, H),
-                                   interpolation=cv2.INTER_LINEAR)
+        if not is_depth_image:
+            control_image, _ = apply_midas(
+                resize_image(input_image, detect_resolution))
+            control_image = HWC3(control_image)
+            image = resize_image(input_image, image_resolution)
+            H, W = image.shape[:2]
+            control_image = cv2.resize(control_image, (W, H),
+                                       interpolation=cv2.INTER_LINEAR)
+        else:
+            control_image = input_image
         return PIL.Image.fromarray(control_image), PIL.Image.fromarray(
             control_image)
 
@@ -562,11 +571,13 @@ class Model:
         num_steps: int,
         guidance_scale: float,
         seed: int,
+        is_depth_image: bool,
     ) -> list[PIL.Image.Image]:
         control_image, vis_control_image = self.preprocess_depth(
             input_image=input_image,
             image_resolution=image_resolution,
             detect_resolution=detect_resolution,
+            is_depth_image=is_depth_image,
         )
         return self.process(
             task_name='depth',