Update app.py

app.py

@@ -13,6 +13,24 @@ import torch
 from diffusers import DiffusionPipeline
 from typing import Tuple
 
+#Check for the Model Base..//
+
+
+
+bad_words = json.loads(os.getenv('BAD_WORDS', "[]"))
+bad_words_negative = json.loads(os.getenv('BAD_WORDS_NEGATIVE', "[]"))
+default_negative = os.getenv("default_negative","")
+
+def check_text(prompt, negative=""):
+    for i in bad_words:
+        if i in prompt:
+            return True
+    for i in bad_words_negative:
+        if i in negative:
+            return True
+    return False
+
+
 
 style_list = [
     {
@@ -131,7 +149,8 @@ def generate(
     use_resolution_binning: bool = True,
     progress=gr.Progress(track_tqdm=True),
 ):
-    
+    if check_text(prompt, negative_prompt):
+        raise ValueError("Prompt contains restricted words.")
     
     prompt, negative_prompt = apply_style(style, prompt, negative_prompt)
     seed = int(randomize_seed_fn(seed, randomize_seed))