Merge pull request #31 from jkremser/use_stop_list

stable-diff: Add also a stoplist for the prompt
kedify · Jun 19, 2024 · decc812 · decc812
2 parents c91b721 + e84ab59
commit decc812
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/samples/stable-diffusion/webui/app/prompt.tsx b/samples/stable-diffusion/webui/app/prompt.tsx
@@ -30,8 +30,18 @@ export const Prompt: React.FC<PromptProps> = ({ sendToQ }) => {
       return sendToQ(img, inParallel, count);
     },
   });
+  const safetyCheck = (prompt: string) => {
+    const stopList = ["boobs", "tits", "vagina", "penis", "dick", "naked", "nude", "fuck", "coitus", "sex", "shit", "blood", "murder", "dead"];
+    return stopList.some(w => prompt.toLowerCase().includes(w));
+  }
+
   const generateOnClick = () => {
     if (prompt) {
+      if (!safetyCheck(prompt)) {
+        toast.warning("Please don't use this tool for creating NSFW content.");
+        focusPrompt();
+        return
+      }
       toast.success("Task has been added to job queue 🤖");
       mutation.mutate([prompt, inParallel, count]);
       setPrompt('');