🐛 Fix unload

annotator/densepose/__init__.py

@@ -10,9 +10,9 @@
 
 N_PART_LABELS = 24
 result_visualizer = DensePoseMaskedColormapResultsVisualizer(
-    alpha=1, 
-    data_extractor=_extract_i_from_iuvarr, 
-    segm_extractor=_extract_i_from_iuvarr, 
+    alpha=1,
+    data_extractor=_extract_i_from_iuvarr,
+    segm_extractor=_extract_i_from_iuvarr,
     val_scale = 255.0 / N_PART_LABELS
 )
 remote_torchscript_path = "https://huggingface.co/LayerNorm/DensePose-TorchScript-with-hint-image/resolve/main/densepose_r50_fpn_dl.torchscript"

scripts/global_state.py

@@ -108,8 +108,8 @@ def unified_preprocessor(preprocessor_name: str, *args, **kwargs):
     "recolor_intensity": recolor_intensity,
     "blur_gaussian": blur_gaussian,
     "anime_face_segment": anime_face_segment,
-    "densepose": densepose,
-    "densepose_parula": densepose_parula
+    "densepose": functools.partial(densepose, cmap="viridis"),
+    "densepose_parula": functools.partial(densepose, cmap="parula"),
 }
 
 cn_preprocessor_unloadable = {
@@ -144,8 +144,8 @@ def unified_preprocessor(preprocessor_name: str, *args, **kwargs):
     "lineart_anime_denoise": unload_lineart_anime_denoise,
     "inpaint_only+lama": unload_lama_inpaint,
     "anime_face_segment": unload_anime_face_segment,
-    "densepose": unload_densepose_model,
-    "densepose_parula": unload_densepose_model,
+    "densepose": unload_densepose,
+    "densepose_parula": unload_densepose,
     "depth_hand_refiner": g_hand_refiner_model.unload,
 }
 

scripts/processor.py

@@ -647,31 +647,18 @@ def unload_anime_face_segment():
         model_anime_face_segment.unload_model()
 
 
-model_densepose = None
 
-
-def densepose(img, res=512, **kwargs):
+def densepose(img, res=512, cmap="viridis", **kwargs):
     img, remove_pad = resize_image_with_pad(img, res)
-    global model_densepose
-    if model_hed is None:
-        from annotator.densepose import apply_densepose
-        model_densepose = apply_densepose
-    result = model_densepose(img)
+    from annotator.densepose import apply_densepose
+    result = apply_densepose(img, cmap=cmap)
     return remove_pad(result), True
 
-def densepose_parula(img, res=512, **kwargs):
-    img, remove_pad = resize_image_with_pad(img, res)
-    global model_densepose
-    if model_hed is None:
-        from annotator.densepose import apply_densepose
-        model_densepose = apply_densepose
-    result = model_densepose(img, cmap="parula")
-    return remove_pad(result), True
 
-def unload_densepose_model():
-    global model_densepose
-    if model_densepose is not None:
-        model_densepose.unload_model()
+def unload_densepose():
+    from annotator.densepose import unload_model
+    unload_model()
+
 
 class InsightFaceModel:
     def __init__(self):