primitive anti nsfw using wdtagger

app.py

@@ -2,6 +2,7 @@ import spaces
 import os
 import gc
 import gradio as gr
+import gradio_client as grcl
 import numpy as np
 import torch
 import json
@@ -12,7 +13,7 @@ from PIL import Image, PngImagePlugin
 from datetime import datetime
 from diffusers.models import AutoencoderKL
 from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline
-
+GRAD_CLIENT = grcl.Client("https://yoinked-da-nsfw-checker.hf.space/")
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
@@ -33,7 +34,7 @@ MODEL = os.getenv(
     "OnomaAIResearch/Illustrious-xl-early-release-v0",
 )
 
-torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.deterministic = True # maybe disable this? seems 
 torch.backends.cudnn.benchmark = False
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
@@ -192,7 +193,19 @@ def generate(
         pipe.scheduler = backup_scheduler
         utils.free_memory()
 
-
+def genwrap(*args, **kwargs):
+    ipth, mtd = generate(*args, **kwargs)
+    r = GRAD_CLIENT(image=grcl.file(ipth), "chen-evangelion", 0.4, False, False, api_name="/classify"))
+    ratings = val[0]
+    rating = rating['confidences']
+    highestval, classtype = -1, "aa"
+    for o in rating:
+        if o['confidence'] > highestval:
+            highestval = o['confidence']
+            classtype = o['label']
+    if classtype not in ["general", "sensitive"]: #add "questionable" and "explicit" to enable nsfw, or just delete this func
+        return "https://upload.wikimedia.org/wikipedia/commons/b/bf/Bucephala-albeola-010.jpg", mtd
+    return ipth, mtd
 if torch.cuda.is_available():
     pipe = load_pipeline(MODEL)
     logger.info("Loaded on Device!")
@@ -369,7 +382,7 @@ with gr.Blocks(css="style.css", theme="NoCrypt/[email protected]") as demo:
         queue=False,
         api_name=False,
     ).then(
-        fn=generate,
+        fn=genwrap,
         inputs=[
             prompt,
             negative_prompt,