Added NSFW filter for Stable Diffusion, which can be controlled with …

…a checkbox
Hukasx0 · Nov 17, 2023 · 7ecb438 · 7ecb438
1 parent 05a51bc
commit 7ecb438
Show file tree

Hide file tree

Showing 4 changed files with 73 additions and 20 deletions.
diff --git a/app/main-mistral-webui.py b/app/main-mistral-webui.py
@@ -9,6 +9,7 @@
 
 llm = None
 sd = None
+safety_checker_sd = None
 
 folder_path = 'models'
 model_url = 'https://huggingface.co/TheBloke/Mistral-7B-Instruct-v0.1-GGUF/resolve/main/mistral-7b-instruct-v0.1.Q4_K_M.gguf'
@@ -157,7 +158,7 @@ def generate_example_messages(character_name, character_summary, character_perso
     print(output)
     return output
 
-def generate_character_avatar(character_name, character_summary, topic, negative_prompt, avatar_prompt):
+def generate_character_avatar(character_name, character_summary, topic, negative_prompt, avatar_prompt, nsfw_filter):
     example_dialogue = """
 <s>[INST] create a prompt that lists the appearance characteristics of a character whose summary is Jamie Hale is a savvy and accomplished businessman who has carved a name for himself in the world of corporate success. With his sharp mind, impeccable sense of style, and unwavering determination, he has risen to the top of the business world. Jamie stands at 6 feet tall with a confident and commanding presence. He exudes charisma and carries himself with an air of authority that draws people to him.
 Jamie's appearance is always polished and professional. He is often seen in tailored suits that accentuate his well-maintained physique. His dark, well-groomed hair and neatly trimmed beard add to his refined image. His piercing blue eyes exude a sense of intense focus and ambition. Topic: business [/INST]
@@ -168,6 +169,7 @@ def generate_character_avatar(character_name, character_summary, topic, negative
     """
     sd_prompt = input_none(avatar_prompt) or llm(example_dialogue+f"\n[INST] create a prompt that lists the appearance characteristics of a character whose summary is {character_summary}. Topic: {topic if input_none(topic) else 'any theme'} [/INST]\n")
     print(sd_prompt)
+    sd_filter(nsfw_filter)
     return image_generate(character_name, sd_prompt, input_none(negative_prompt))
 
 def image_generate(character_name, prompt, negative_prompt):
@@ -188,6 +190,14 @@ def image_generate(character_name, prompt, negative_prompt):
 
     return generated_image
 
+def sd_filter(enable):
+  if enable:
+    sd.safety_checker = safety_checker_sd
+    sd.requires_safety_checker = True
+  else:
+    sd.safety_checker = None
+    sd.requires_safety_checker = False
+
 def input_none(text):
   user_input = text
   if user_input == "":
@@ -268,8 +278,10 @@ def export_character_card(name, summary, personality, scenario, greeting_message
           negative_prompt = gr.Textbox(placeholder="negative prompt for stable diffusion (optional)", label="negative prompt")
           avatar_prompt = gr.Textbox(placeholder="prompt for generating character avatar (If not provided, LLM will generate prompt from character description)", label="stable diffusion prompt")
           avatar_button = gr.Button("Generate avatar with stable diffusion (set character name first)")
+          potential_nsfw_checkbox = gr.Checkbox(label="Block potential NSFW image (Upon detection of this content, a black image will be returned)", value=True,
+                                               interactive=True)
           avatar_button.click(generate_character_avatar,
-                    inputs=[name, summary, topic, negative_prompt, avatar_prompt],
+                    inputs=[name, summary, topic, negative_prompt, avatar_prompt, potential_nsfw_checkbox],
                     outputs=image_input)
   with gr.Tab("Export character"):
     with gr.Column():
@@ -284,5 +296,6 @@ def export_character_card(name, summary, personality, scenario, greeting_message
         export_card_button.click(export_character_card, inputs=[name, summary, personality, scenario, greeting_message, example_messages], outputs=export_image)
         export_json_button.click(export_as_json, inputs=[name, summary, personality, scenario, greeting_message, example_messages], outputs=export_json_textbox)
 
+safety_checker_sd = sd.safety_checker
 
 webui.launch(debug=True)
diff --git a/app/main-zephyr-webui.py b/app/main-zephyr-webui.py
@@ -9,6 +9,7 @@
 
 llm = None
 sd = None
+safety_checker_sd = None
 
 folder_path = 'models'
 model_url = 'https://huggingface.co/TheBloke/zephyr-7B-beta-GGUF/resolve/main/zephyr-7b-beta.Q4_K_M.gguf'
@@ -201,7 +202,7 @@ def generate_example_messages(character_name, character_summary, character_perso
     print(output)
     return output
 
-def generate_character_avatar(character_name, character_summary, topic, negative_prompt, avatar_prompt):
+def generate_character_avatar(character_name, character_summary, topic, negative_prompt, avatar_prompt, nsfw_filter):
     example_dialogue = """
 <|system|>
 You are a text generation tool, in the response you are supposed to give only descriptions of the appearance, what the character looks like, describe the character simply and unambiguously
@@ -215,6 +216,7 @@ def generate_character_avatar(character_name, character_summary, topic, negative
 """
     sd_prompt = input_none(avatar_prompt) or llm(example_dialogue+f"\n<|user|> create a prompt that lists the appearance characteristics of a character whose summary is {character_summary}. Topic: {topic} </s>\n<|assistant|> ").strip()
     print(sd_prompt)
+    sd_filter(nsfw_filter)
     return image_generate(character_name, sd_prompt, input_none(negative_prompt))
 
 def image_generate(character_name, prompt, negative_prompt):
@@ -234,6 +236,14 @@ def image_generate(character_name, prompt, negative_prompt):
     print("Generated character avatar")
     return generated_image
 
+def sd_filter(enable):
+  if enable:
+    sd.safety_checker = safety_checker_sd
+    sd.requires_safety_checker = True
+  else:
+    sd.safety_checker = None
+    sd.requires_safety_checker = False
+
 def input_none(text):
   user_input = text
   if user_input == "":
@@ -314,8 +324,10 @@ def export_character_card(name, summary, personality, scenario, greeting_message
           negative_prompt = gr.Textbox(placeholder="negative prompt for stable diffusion (optional)", label="negative prompt")
           avatar_prompt = gr.Textbox(placeholder="prompt for generating character avatar (If not provided, LLM will generate prompt from character description)", label="stable diffusion prompt")
           avatar_button = gr.Button("Generate avatar with stable diffusion (set character name first)")
+          potential_nsfw_checkbox = gr.Checkbox(label="Block potential NSFW image (Upon detection of this content, a black image will be returned)", value=True,
+                                               interactive=True)
           avatar_button.click(generate_character_avatar,
-                    inputs=[name, summary, topic, negative_prompt, avatar_prompt],
+                    inputs=[name, summary, topic, negative_prompt, avatar_prompt, potential_nsfw_checkbox],
                     outputs=image_input)
   with gr.Tab("Export character"):
     with gr.Column():
@@ -330,5 +342,6 @@ def export_character_card(name, summary, personality, scenario, greeting_message
         export_card_button.click(export_character_card, inputs=[name, summary, personality, scenario, greeting_message, example_messages], outputs=export_image)
         export_json_button.click(export_as_json, inputs=[name, summary, personality, scenario, greeting_message, example_messages], outputs=export_json_textbox)
 
+safety_checker_sd = sd.safety_checker
 
 webui.launch(debug=True)
diff --git a/colab/character_factory_colab_mistral.ipynb b/colab/character_factory_colab_mistral.ipynb
@@ -11,12 +11,12 @@
     },
     {
       "cell_type": "markdown",
-      "metadata": {
-        "id": "D-ALE6QCD4mC"
-      },
       "source": [
         "## This version of the script uses the Mistral 7B instruct 0.1 model"
-      ]
+      ],
+      "metadata": {
+        "id": "D-ALE6QCD4mC"
+      }
     },
     {
       "cell_type": "code",
@@ -77,6 +77,7 @@
       "source": [
         "llm = None\n",
         "sd = None\n",
+        "safety_checker_sd = None\n",
         "\n",
         "google_drive_path = \"/content/drive/My Drive/Character_factory/\"\n",
         "folder_path = 'models'\n",
@@ -238,7 +239,7 @@
         "    print(output)\n",
         "    return output\n",
         "\n",
-        "def generate_character_avatar(character_name, character_summary, topic, negative_prompt, avatar_prompt):\n",
+        "def generate_character_avatar(character_name, character_summary, topic, negative_prompt, avatar_prompt, nsfw_filter):\n",
         "    example_dialogue = \"\"\"\n",
         "<s>[INST] create a prompt that lists the appearance characteristics of a character whose summary is Jamie Hale is a savvy and accomplished businessman who has carved a name for himself in the world of corporate success. With his sharp mind, impeccable sense of style, and unwavering determination, he has risen to the top of the business world. Jamie stands at 6 feet tall with a confident and commanding presence. He exudes charisma and carries himself with an air of authority that draws people to him.\n",
         "Jamie's appearance is always polished and professional. He is often seen in tailored suits that accentuate his well-maintained physique. His dark, well-groomed hair and neatly trimmed beard add to his refined image. His piercing blue eyes exude a sense of intense focus and ambition. Topic: business [/INST]\n",
@@ -249,6 +250,7 @@
         "    \"\"\"\n",
         "    sd_prompt = input_none(avatar_prompt) or llm(example_dialogue+f\"\\n[INST] create a prompt that lists the appearance characteristics of a character whose summary is {character_summary}. Topic: {topic if input_none(topic) else 'any theme'} [/INST]\\n\")\n",
         "    print(sd_prompt)\n",
+        "    sd_filter(nsfw_filter)\n",
         "    return image_generate(character_name, sd_prompt, input_none(negative_prompt))\n",
         "\n",
         "def image_generate(character_name, prompt, negative_prompt):\n",
@@ -269,6 +271,14 @@
         "\n",
         "    return generated_image\n",
         "\n",
+        "def sd_filter(enable):\n",
+        "  if enable:\n",
+        "    sd.safety_checker = safety_checker_sd\n",
+        "    sd.requires_safety_checker = True\n",
+        "  else:\n",
+        "    sd.safety_checker = None\n",
+        "    sd.requires_safety_checker = False\n",
+        "\n",
         "def input_none(text):\n",
         "  user_input = text\n",
         "  if user_input == \"\":\n",
@@ -279,18 +289,18 @@
     },
     {
       "cell_type": "markdown",
-      "metadata": {
-        "id": "MsVEXd4l_Glw"
-      },
       "source": [
         "## Start WebUI (Alpha)"
-      ]
+      ],
+      "metadata": {
+        "id": "mVWvxB6NO62K"
+      }
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "nvlXjC5lxp5U"
+        "id": "oyPGMmQcsoeN"
       },
       "outputs": [],
       "source": [
@@ -365,8 +375,10 @@
         "          negative_prompt = gr.Textbox(placeholder=\"negative prompt for stable diffusion (optional)\", label=\"negative prompt\")\n",
         "          avatar_prompt = gr.Textbox(placeholder=\"prompt for generating character avatar (If not provided, LLM will generate prompt from character description)\", label=\"stable diffusion prompt\")\n",
         "          avatar_button = gr.Button(\"Generate avatar with stable diffusion (set character name first)\")\n",
+        "          potential_nsfw_checkbox = gr.Checkbox(label=\"Block potential NSFW image (Upon detection of this content, a black image will be returned)\", value=True,\n",
+        "                                               interactive=True)\n",
         "          avatar_button.click(generate_character_avatar,\n",
-        "                    inputs=[name, summary, topic, negative_prompt, avatar_prompt],\n",
+        "                    inputs=[name, summary, topic, negative_prompt, avatar_prompt, potential_nsfw_checkbox],\n",
         "                    outputs=image_input)\n",
         "  with gr.Tab(\"Export character\"):\n",
         "    with gr.Column():\n",
@@ -381,6 +393,7 @@
         "        export_card_button.click(export_character_card, inputs=[name, summary, personality, scenario, greeting_message, example_messages], outputs=export_image)\n",
         "        export_json_button.click(export_as_json, inputs=[name, summary, personality, scenario, greeting_message, example_messages], outputs=export_json_textbox)\n",
         "\n",
+        "safety_checker_sd = sd.safety_checker\n",
         "\n",
         "webui.launch(debug=True)"
       ]
@@ -389,6 +402,7 @@
   "metadata": {
     "accelerator": "GPU",
     "colab": {
+      "gpuType": "T4",
       "provenance": []
     },
     "kernelspec": {
@@ -401,4 +415,4 @@
   },
   "nbformat": 4,
   "nbformat_minor": 0
-}
+}
diff --git a/colab/character_factory_colab_zephyr.ipynb b/colab/character_factory_colab_zephyr.ipynb
@@ -88,6 +88,7 @@
       "source": [
         "llm = None\n",
         "sd = None\n",
+        "safety_checker_sd = None\n",
         "\n",
         "google_drive_path = \"/content/drive/My Drive/Character_factory/\"\n",
         "folder_path = 'models'\n",
@@ -293,7 +294,7 @@
         "    print(output)\n",
         "    return output\n",
         "\n",
-        "def generate_character_avatar(character_name, character_summary, topic, negative_prompt, avatar_prompt):\n",
+        "def generate_character_avatar(character_name, character_summary, topic, negative_prompt, avatar_prompt, nsfw_filter):\n",
         "    example_dialogue = \"\"\"\n",
         "<|system|>\n",
         "You are a text generation tool, in the response you are supposed to give only descriptions of the appearance, what the character looks like, describe the character simply and unambiguously\n",
@@ -307,6 +308,7 @@
         "\"\"\"\n",
         "    sd_prompt = input_none(avatar_prompt) or llm(example_dialogue+f\"\\n<|user|> create a prompt that lists the appearance characteristics of a character whose summary is {character_summary}. Topic: {topic} </s>\\n<|assistant|> \").strip()\n",
         "    print(sd_prompt)\n",
+        "    sd_filter(nsfw_filter)\n",
         "    return image_generate(character_name, sd_prompt, input_none(negative_prompt))\n",
         "\n",
         "def image_generate(character_name, prompt, negative_prompt):\n",
@@ -326,6 +328,14 @@
         "    print(\"Generated character avatar\")\n",
         "    return generated_image\n",
         "\n",
+        "def sd_filter(enable):\n",
+        "  if enable:\n",
+        "    sd.safety_checker = safety_checker_sd\n",
+        "    sd.requires_safety_checker = True\n",
+        "  else:\n",
+        "    sd.safety_checker = None\n",
+        "    sd.requires_safety_checker = False\n",
+        "\n",
         "def input_none(text):\n",
         "  user_input = text\n",
         "  if user_input == \"\":\n",
@@ -345,7 +355,7 @@
         "## Start WebUI (Alpha)"
       ],
       "metadata": {
-        "id": "R46uXoLBC1qs"
+        "id": "_KdNbqNVPJg7"
       }
     },
     {
@@ -422,8 +432,10 @@
         "          negative_prompt = gr.Textbox(placeholder=\"negative prompt for stable diffusion (optional)\", label=\"negative prompt\")\n",
         "          avatar_prompt = gr.Textbox(placeholder=\"prompt for generating character avatar (If not provided, LLM will generate prompt from character description)\", label=\"stable diffusion prompt\")\n",
         "          avatar_button = gr.Button(\"Generate avatar with stable diffusion (set character name first)\")\n",
+        "          potential_nsfw_checkbox = gr.Checkbox(label=\"Block potential NSFW image (Upon detection of this content, a black image will be returned)\", value=True,\n",
+        "                                               interactive=True)\n",
         "          avatar_button.click(generate_character_avatar,\n",
-        "                    inputs=[name, summary, topic, negative_prompt, avatar_prompt],\n",
+        "                    inputs=[name, summary, topic, negative_prompt, avatar_prompt, potential_nsfw_checkbox],\n",
         "                    outputs=image_input)\n",
         "  with gr.Tab(\"Export character\"):\n",
         "    with gr.Column():\n",
@@ -438,11 +450,12 @@
         "        export_card_button.click(export_character_card, inputs=[name, summary, personality, scenario, greeting_message, example_messages], outputs=export_image)\n",
         "        export_json_button.click(export_as_json, inputs=[name, summary, personality, scenario, greeting_message, example_messages], outputs=export_json_textbox)\n",
         "\n",
+        "safety_checker_sd = sd.safety_checker\n",
         "\n",
         "webui.launch(debug=True)"
       ],
       "metadata": {
-        "id": "so5isBKlCyv6"
+        "id": "oyPGMmQcsoeN"
       },
       "execution_count": null,
       "outputs": []