enable automatic download of nsfw detector for face swap

Signed-off-by: Markus Schuettler <[email protected]>
TNG · Feb 3, 2025 · e4e50e7 · e4e50e7
1 parent 65779a9
commit e4e50e7
Show file tree

Hide file tree

Showing 8 changed files with 469 additions and 453 deletions.
diff --git a/WebUI/external/workflows/CopyFace.json b/WebUI/external/workflows/CopyFace.json
@@ -6,41 +6,33 @@
       "https://github.com/Gourieff/Assets/raw/main/Insightface/insightface-0.7.3-cp311-cp311-win_amd64.whl",
       "onnxruntime"
     ],
-    "customNodes": ["Gourieff/comfyui-reactor@ee13ff754f35d2af284f0cf9d81ccce4defe1e5e"],
+    "customNodes": ["Gourieff/comfyui-reactor@d2318ad140582c6d0b68c51df342319b502006ed"],
     "requiredModels": [
       {
-        "type": "defaultCheckpoint",
-        "model": "RunDiffusion/Juggernaut-XL-v9/unet/diffusion_pytorch_model.fp16.safetensors"
-      },
-      {
-        "type": "defaultCheckpoint",
-        "model": "RunDiffusion/Juggernaut-XL-v9/text_encoder/model.fp16.safetensors"
-      },
-      {
-        "type": "defaultCheckpoint",
-        "model": "RunDiffusion/Juggernaut-XL-v9/text_encoder_2/model.fp16.safetensors"
+        "type": "faceswap",
+        "model": "Aitrepreneur/insightface/inswapper_128.onnx",
+        "additionalLicenceLink": "https://huggingface.co/datasets/Gourieff/ReActor"
       },
       {
-        "type": "defaultCheckpoint",
-        "model": "RunDiffusion/Juggernaut-XL-v9/vae/diffusion_pytorch_model.fp16.safetensors"
+        "type": "facerestore",
+        "model": "gmk123/GFPGAN/GFPGANv1.4.pth"
       },
       {
-        "type": "defaultLora",
-        "model": "latent-consistency/lcm-lora-sdxl/pytorch_lora_weights.safetensors"
+        "type": "nsfwdetector",
+        "model": "AdamCodd/vit-base-nsfw-detector/config.json"
       },
       {
-        "type": "faceswap",
-        "model": "Aitrepreneur/insightface/inswapper_128.onnx",
-        "additionalLicenceLink": "https://huggingface.co/datasets/Gourieff/ReActor"
+        "type": "nsfwdetector",
+        "model": "AdamCodd/vit-base-nsfw-detector/model.safetensors"
       },
       {
-        "type": "facerestore",
-        "model": "gmk123/GFPGAN/GFPGANv1.4.pth"
+        "type": "nsfwdetector",
+        "model": "AdamCodd/vit-base-nsfw-detector/preprocessor_config.json"
       }
     ]
   },
-  "tags": ["SDXL", "Fast-LCM", "non-com"],
-  "requirements": ["high-vram"],
+  "tags": ["non-com"],
+  "requirements": [],
   "inputs": [
     {
       "nodeTitle": "ReActor 🌌 Fast Face Swap",
@@ -79,11 +71,10 @@
     }
   ],
   "defaultSettings": {
-    "resolution": "896x896",
-    "inferenceSteps": 4
+    "batchSize": 1
   },
   "displayedSettings": [],
-  "modifiableSettings": ["inferenceSteps", "seed", "batchSize", "imagePreview", "resolution"],
+  "modifiableSettings": [],
   "comfyUiApiWorkflow": {
     "9": {
       "inputs": {

diff --git a/WebUI/external/workflows/FaceSwapHD.json b/WebUI/external/workflows/FaceSwapHD.json
@@ -6,7 +6,7 @@
       "https://github.com/Gourieff/Assets/raw/main/Insightface/insightface-0.7.3-cp311-cp311-win_amd64.whl",
       "onnxruntime"
     ],
-    "customNodes": ["Gourieff/comfyui-reactor@ee13ff754f35d2af284f0cf9d81ccce4defe1e5e"],
+    "customNodes": ["Gourieff/comfyui-reactor@d2318ad140582c6d0b68c51df342319b502006ed"],
     "requiredModels": [
       {
         "type": "defaultCheckpoint",
@@ -36,6 +36,18 @@
       {
         "type": "facerestore",
         "model": "gmk123/GFPGAN/GFPGANv1.4.pth"
+      },
+      {
+        "type": "nsfwdetector",
+        "model": "AdamCodd/vit-base-nsfw-detector/config.json"
+      },
+      {
+        "type": "nsfwdetector",
+        "model": "AdamCodd/vit-base-nsfw-detector/model.safetensors"
+      },
+      {
+        "type": "nsfwdetector",
+        "model": "AdamCodd/vit-base-nsfw-detector/preprocessor_config.json"
       }
     ]
   },

diff --git a/WebUI/src/assets/js/const.ts b/WebUI/src/assets/js/const.ts
@@ -15,3 +15,4 @@ export const MODEL_TYPE_COMFY_DEFAULT_LORA = 104
 export const MODEL_TYPE_COMFY_CONTROL_NET = 105
 export const MODEL_TYPE_FACESWAP = 106
 export const MODEL_TYPE_FACERESTORE = 107
+export const MODEL_TYPE_NSFW_DETECTOR = 108
diff --git a/WebUI/src/assets/js/store/imageGeneration.ts b/WebUI/src/assets/js/store/imageGeneration.ts
@@ -674,6 +674,8 @@ export const useImageGeneration = defineStore(
               return Const.MODEL_TYPE_FACESWAP
             case 'facerestore':
               return Const.MODEL_TYPE_FACERESTORE
+            case 'nsfwdetector':
+              return Const.MODEL_TYPE_NSFW_DETECTOR
             case 'defaultCheckpoint':
               return Const.MODEL_TYPE_COMFY_DEFAULT_CHECKPOINT
             case 'defaultLora':

diff --git a/WebUI/src/components/SettingsImageGeneration.vue b/WebUI/src/components/SettingsImageGeneration.vue
@@ -8,7 +8,7 @@
           :disabled="!modifiable('resolution')"
         />
       </div>
-      <div class="flex items-center gap-5">
+      <div v-if="modifiableOrDisplayed('imagePreview')" class="flex items-center gap-5">
         <p>{{ languages.SETTINGS_MODEL_IMAGE_PREVIEW }}</p>
         <button
           class="v-checkbox-control flex-none w-5 h-5"

diff --git a/service/aipg_utils.py b/service/aipg_utils.py
@@ -63,6 +63,8 @@ def check_comfyui_model_exists(type, repo_id) -> bool:
     model_dir = service_config.comfy_ui_model_paths.get(model_type)
     if model_type == 'faceswap' or model_type == 'facerestore':
         dir_to_look_for = os.path.join(model_dir, flat_repo_local_dir_name(repo_id))
+    elif model_type == 'nsfwdetector':
+        dir_to_look_for = os.path.join(model_dir, 'vit-base-nsfw-detector', extract_model_id_pathsegments(repo_id))
     else:
         dir_to_look_for = os.path.join(model_dir, repo_local_root_dir_name(repo_id), extract_model_id_pathsegments(repo_id))
     return os.path.exists(dir_to_look_for)
@@ -168,6 +170,8 @@ def convert_model_type(type: int):
         return "faceswap"
     elif type == 107:
         return "facerestore"
+    elif type == 108:
+        return "nsfwdetector"
     else:
         raise Exception(f"unknown model type value {type}")