Update app.py

app.py

@@ -18,6 +18,7 @@ import requests
 import torch
 from diffusers import DiffusionPipeline
 import importlib
+import re
 from urllib.parse import urlparse
 
 random.seed(time.time())
@@ -28,7 +29,24 @@ USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1"
 ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"
 
 NUM_IMAGES_PER_PROMPT = 1
+# Define the regular expression
+child_related_regex = re.compile(
+    r'(child|children|kid|kids|baby|babies|toddler|infant|juvenile|minor|underage|preteen|adolescent|youngster|youth|kindergarten|preschool|'
+    r'([1-9]|1[0-7])\s*year(s)?\s*old|'  # Matches 1 to 17 years old
+    r'little|small|tiny|short|new\s*born\s*(boy|girl|bro|brother|sis|sister))',
+    re.IGNORECASE
+)
 
+# Function to remove child-related content from a prompt
+def remove_child_related_content(prompt):
+    cleaned_prompt = re.sub(child_related_regex, '', prompt)
+    return cleaned_prompt.strip()
+
+# Function to check if a prompt contains child-related content
+def contains_child_related_content(prompt):
+    if child_related_regex.search(prompt):
+        return True
+    return False
 
 cfg = json.load(open("app.conf"))
 
@@ -254,10 +272,10 @@ def generate(p, progress=gr.Progress(track_tqdm=True)):
     inference_steps = cfg.get("inference_steps", 30)
     randomize_seed = True
     guidance_scale = cfg.get("guidance_scale", 7.5)
+    p = remove_child_related_content(p)
     prompt_str = cfg.get("prompt", "{prompt}").replace("{prompt}", p)
     seed = int(randomize_seed_fn(seed, randomize_seed))
     generator = torch.Generator(pipe.device).manual_seed(seed)
-        
     images = pipe(
         prompt=prompt_str,
         negative_prompt=negative_prompt,