Update

app_normal.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_normal_image = gr.Checkbox(label='Is normal image',
+                                                  value=False)
                     num_samples = gr.Slider(label='Images',
                                             minimum=1,
                                             maximum=max_images,
@@ -74,6 +76,7 @@ def create_demo(process, max_images=12, default_num_images=3):
             guidance_scale,
             seed,
             bg_threshold,
+            is_normal_image,
         ]
         prompt.submit(fn=process, inputs=inputs, outputs=result)
         run_button.click(fn=process,

app_seg.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_segmentation_map = gr.Checkbox(
+                        label='Is segmentation map', value=False)
                     num_samples = gr.Slider(label='Images',
                                             minimum=1,
                                             maximum=max_images,
@@ -68,6 +70,7 @@ def create_demo(process, max_images=12, default_num_images=3):
             num_steps,
             guidance_scale,
             seed,
+            is_segmentation_map,
         ]
         prompt.submit(fn=process, inputs=inputs, outputs=result)
         run_button.click(fn=process,

model.py

@@ -494,14 +494,18 @@ class Model:
         input_image: np.ndarray,
         image_resolution: int,
         detect_resolution: int,
+        is_segmentation_map: bool,
     ) -> tuple[PIL.Image.Image, PIL.Image.Image]:
         input_image = HWC3(input_image)
-        control_image = apply_uniformer(
-            resize_image(input_image, detect_resolution))
-        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_segmentation_map:
+            control_image = apply_uniformer(
+                resize_image(input_image, detect_resolution))
+            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)
 
@@ -518,11 +522,13 @@ class Model:
         num_steps: int,
         guidance_scale: float,
         seed: int,
+        is_segmentation_map: bool,
     ) -> list[PIL.Image.Image]:
         control_image, vis_control_image = self.preprocess_seg(
             input_image=input_image,
             image_resolution=image_resolution,
             detect_resolution=detect_resolution,
+            is_segmentation_map=is_segmentation_map,
         )
         return self.process(
             task_name='seg',
@@ -597,17 +603,21 @@ class Model:
         input_image: np.ndarray,
         image_resolution: int,
         detect_resolution: int,
-        bg_threshold,
+        bg_threshold: float,
+        is_normal_image: bool,
     ) -> tuple[PIL.Image.Image, PIL.Image.Image]:
         input_image = HWC3(input_image)
-        _, control_image = apply_midas(resize_image(input_image,
-                                                    detect_resolution),
-                                       bg_th=bg_threshold)
-        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_normal_image:
+            _, control_image = apply_midas(resize_image(
+                input_image, detect_resolution),
+                                           bg_th=bg_threshold)
+            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)
 
@@ -624,13 +634,15 @@ class Model:
         num_steps: int,
         guidance_scale: float,
         seed: int,
-        bg_threshold,
+        bg_threshold: float,
+        is_normal_image: bool,
     ) -> list[PIL.Image.Image]:
         control_image, vis_control_image = self.preprocess_normal(
             input_image=input_image,
             image_resolution=image_resolution,
             detect_resolution=detect_resolution,
             bg_threshold=bg_threshold,
+            is_normal_image=is_normal_image,
         )
         return self.process(
             task_name='normal',