default use safe steps = 2

annotator/teed/__init__.py

@@ -36,7 +36,7 @@ class TEEDDector:
         if self.model is not None:
             self.model.cpu()
 
-    def __call__(self, image: np.ndarray, is_safe: bool = True) -> np.ndarray:
+    def __call__(self, image: np.ndarray, safe_steps: int = 2) -> np.ndarray:
 
         self.model.to(self.device)
 
@@ -49,7 +49,7 @@ class TEEDDector:
             edges = [cv2.resize(e, (W, H), interpolation=cv2.INTER_LINEAR) for e in edges]
             edges = np.stack(edges, axis=2)
             edge = 1 / (1 + np.exp(-np.mean(edges, axis=2).astype(np.float64)))
-            if is_safe:
-                edge = safe_step(edge)
+            if safe_steps != 0:
+                edge = safe_step(edge, safe_steps)
             edge = (edge * 255.0).clip(0, 255).astype(np.uint8)
             return edge

scripts/global_state.py

@@ -111,8 +111,7 @@ cn_preprocessor_modules = {
     "anime_face_segment": anime_face_segment,
     "densepose": functools.partial(densepose, cmap="viridis"),
     "densepose_parula": functools.partial(densepose, cmap="parula"),
-    "te_hed":functools.partial(te_hed, is_safe=False),
-    "te_hed_safe":functools.partial(te_hed, is_safe=True),
+    "te_hed":te_hed,
 }
 
 cn_preprocessor_unloadable = {
@@ -176,7 +175,6 @@ preprocessor_aliases = {
     "densepose": "densepose (pruple bg & purple torso)",
     "densepose_parula": "densepose_parula (black bg & blue torso)",
     "te_hed": "softedge_teed",
-    "te_hed_safe": "softedge_teed_safe",
 }
 
 ui_preprocessor_keys = ['none', preprocessor_aliases['invert']]

scripts/processor.py

@@ -692,13 +692,13 @@ def unload_densepose():
 
 model_te_hed = None
 
-def te_hed(img, res=512, is_safe=True,**kwargs):
+def te_hed(img, res=512, thr_a=2, **kwargs):
     img, remove_pad = resize_image_with_pad(img, res)
     global model_te_hed
     if model_te_hed is None:
         from annotator.teed import TEEDDector
         model_te_hed = TEEDDector()
-    result = model_te_hed(img, is_safe=is_safe)
+    result = model_te_hed(img, safe_steps=int(thr_a))
     return remove_pad(result), True
 
 def unload_te_hed():
@@ -1200,15 +1200,14 @@ preprocessor_sliders_config = {
             "value": 512,
             "min": 64,
             "max": 2048
-        } 
-    ],
-    "te_hed_safe": [
+        },
         {
-            "name": flag_preprocessor_resolution,
-            "value": 512,
-            "min": 64,
-            "max": 2048
-        } 
+            "name": "Safe Steps",
+            "min": 0,
+            "max": 10,
+            "value": 2,
+            "step": 1,
+        },
     ],
 }