Merge pull request Mikubill#497 from kft334/main

Fix detect endpoint
ch8os · Mar 5, 2023 · 84a2b22 · 84a2b22
2 parents dcd0832 + ac8007c
commit 84a2b22
Showing 1 changed file with 14 additions and 11 deletions.
diff --git a/scripts/api.py b/scripts/api.py
@@ -273,7 +273,8 @@ async def detect(
             "normal_map", 
             "openpose", 
             "segmentation",
-            "binary"
+            "binary",
+            "color"
         ]
 
         if controlnet_module not in available_modules:
@@ -289,25 +290,27 @@ async def detect(
             img = to_base64_nparray(input_image)
 
             if controlnet_module == "canny":
-                results.append(canny(img, controlnet_processor_res, controlnet_threshold_a, controlnet_threshold_b))
+                results.append(canny(img, controlnet_processor_res, controlnet_threshold_a, controlnet_threshold_b)[0])
             elif controlnet_module == "hed":
-                results.append(hed(img, controlnet_processor_res))
+                results.append(hed(img, controlnet_processor_res)[0])
             elif controlnet_module == "mlsd":
-                results.append(mlsd(img, controlnet_processor_res, controlnet_threshold_a, controlnet_threshold_b))
+                results.append(mlsd(img, controlnet_processor_res, controlnet_threshold_a, controlnet_threshold_b)[0])
             elif controlnet_module == "depth":
-                results.append(midas(img, controlnet_processor_res, np.pi * 2.0))
+                results.append(midas(img, controlnet_processor_res, np.pi * 2.0)[0])
             elif controlnet_module == "normal_map":
-                results.append(midas_normal(img, controlnet_processor_res, np.pi * 2.0, controlnet_threshold_a))
+                results.append(midas_normal(img, controlnet_processor_res, np.pi * 2.0, controlnet_threshold_a)[0])
             elif controlnet_module == "depth_leres":
-                results.append(leres(img, controlnet_processor_res, np.pi * 2.0, controlnet_threshold_a, controlnet_threshold_b))
+                results.append(leres(img, controlnet_processor_res, np.pi * 2.0, controlnet_threshold_a, controlnet_threshold_b)[0])
             elif controlnet_module == "openpose":
-                results.append(openpose(img, controlnet_processor_res, False))
+                results.append(openpose(img, controlnet_processor_res, False)[0])
             elif controlnet_module == "fake_scribble":
-                results.append(fake_scribble(img, controlnet_processor_res))
+                results.append(fake_scribble(img, controlnet_processor_res)[0])
             elif controlnet_module == "segmentation":
-                results.append(uniformer(img, controlnet_processor_res))
+                results.append(uniformer(img, controlnet_processor_res)[0])
             elif controlnet_module == "binary":
-                results.append(binary(img, controlnet_processor_res, controlnet_threshold_a))
+                results.append(binary(img, controlnet_processor_res, controlnet_threshold_a)[0])
+            elif controlnet_module == "color":
+                results.append(color(img, controlnet_processor_res)[0])
 
         if controlnet_module == "hed":
             unload_hed()