Update app.py

app.py

@@ -3,6 +3,7 @@
 import os
 import random
 import uuid
+import json
 
 import gradio as gr
 import numpy as np
@@ -11,6 +12,13 @@ import spaces
 import torch
 from diffusers import DiffusionPipeline
 
+bad_words = json.loads(os.getenv('BAD_WORDS', "[]"))
+def check_text(prompt):
+    for i in bad_words:
+        if i in prompt:
+            return True
+    return False
+
 DESCRIPTION = """# RealVisXL_V3.0"""
 if not torch.cuda.is_available():
     DESCRIPTION += "\n<p>Running on CPU 🥶 This demo may not work on CPU.</p>"
@@ -56,8 +64,33 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     return seed
 
 
+def generate(prompt: str,
+    negative_prompt: str = "",
+    use_negative_prompt: bool = False,
+    seed: int = 0,
+    width: int = 1024,
+    height: int = 1024,
+    guidance_scale: float = 3,
+    randomize_seed: bool = False,
+    use_resolution_binning: bool = True,
+    progress=gr.Progress(track_tqdm=True)):
+    if check_text(prompt):
+        return 'NSFW Detection'
+    return _generate(
+        prompt,
+        negative_prompt,
+        use_negative_prompt,
+        seed,
+        width,
+        height,
+        guidance_scale,
+        randomize_seed,
+        use_resolution_binning,
+        progress
+    )
+
 @spaces.GPU(enable_queue=True)
-def generate(
+def _generate(
     prompt: str,
     negative_prompt: str = "",
     use_negative_prompt: bool = False,
@@ -90,7 +123,6 @@ def generate(
     ).images
 
     image_paths = [save_image(img) for img in images]
-    print(image_paths)
     return image_paths, seed