Update app.py

app.py

@@ -6,7 +6,7 @@ import random
 # import spaces #[uncomment to use ZeroGPU]
 from diffusers import DiffusionPipeline
 import torch
-
+targets = {"pussy", "boobs", "breasts", "vagina", "penis", "sex", "oral", "anal", "butt", "ass"}
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model_repo_id = "John6666/spicy-realism-nsfw-mix-v30-sdxl"  # Replace to the model you would like to use
 
@@ -20,8 +20,15 @@ print(pipe)
 pipe = pipe.to(device)
 
 MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
+MAX_IMAGE_SIZE = 2048
+def safe_infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_steps):
+    # check if prompt contains any banned words
+    if any(word in prompt.lower() for word in targets):
+        print("Found at least one banned word!")
+        return "Refused due to safety.", seed
 
+    # otherwise run normal inference
+    return infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_steps)
 
 # @spaces.GPU #[uncomment to use ZeroGPU]
 def infer(
@@ -139,7 +146,9 @@ with gr.Blocks(css=css) as demo:
         gr.Examples(examples=examples, inputs=[prompt])
     gr.on(
         triggers=[run_button.click, prompt.submit],
-        fn=infer,
+        if any(word in prompt for word in targets):
+            print("Found at least one!")
+        fn=safe_infer,
         inputs=[
             prompt,
             negative_prompt,