Update inference_manager.py

inference_manager.py

@@ -4,6 +4,7 @@ import json
 import time
 import copy
 import torch
+import random
 from diffusers import AutoPipelineForText2Image, StableDiffusionPipeline,DiffusionPipeline, StableDiffusionXLPipeline, AutoencoderKL, AutoencoderTiny, UNet2DConditionModel
 from huggingface_hub import hf_hub_download, snapshot_download
 from pathlib import Path
@@ -21,7 +22,9 @@ import glob
 import traceback
 from insightface.app import FaceAnalysis
 import cv2
+import re
 import gradio as gr
+from PIL import Image
 
 #from onediffx import compile_pipe, save_pipe, load_pipe
 
@@ -308,7 +311,47 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     return seed
+
+
+child_related_regex = re.compile(
+    r'(child|children|kid|kids|baby|babies|toddler|infant|juvenile|minor|underage|preteen|adolescent|youngster|youth|son|daughter|young|kindergarten|preschool|'
+    r'([1-9]|1[0-7])[\s_\-|\.\,]*year(s)?[\s_\-|\.\,]*old|'  # Matches 1 to 17 years old with various separators
+    r'little|small|tiny|short|young|new[\s_\-|\.\,]*born[\s_\-|\.\,]*(boy|girl|male|man|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
+
+def save_image(img):
+    path = "./tmp/"
+    # Ensure the Hugging Face path exists locally
+    if not os.path.exists(path):
+        os.makedirs(path)
+    # Generate a unique filename
+    unique_name = str(uuid.uuid4()) + ".webp"
+    unique_name = os.path.join(path, unique_name)
+    
+    # Convert the image to WebP format
+    webp_img = img.convert("RGB")  # Ensure the image is in RGB mode
+    
+    # Save the image in WebP format with high quality
+    webp_img.save(unique_name, "WEBP", quality=90)
     
+    # Open the saved WebP file and return it as a PIL Image object
+    with Image.open(unique_name) as webp_file:
+        webp_image = webp_file.copy()
+    
+    return webp_image, unique_name
+
 class ModelManager:
     def __init__(self, model_directory):
         """
@@ -506,7 +549,7 @@ class ModelManager:
         cfg = model.cfg
         p = inference_params.get("prompt")
         negative_prompt = inference_params.get("negative_prompt", cfg.get("negative_prompt", ""))
-        inference_steps = inference_params.get("steps", cfg.get("inference_steps", 30))
+        steps = inference_params.get("steps", cfg.get("inference_steps", 30))
         guidance_scale = inference_params.get("guidance_scale", cfg.get("guidance_scale", 7))
         width = inference_params.get("width", cfg.get("width", 512))
         height = inference_params.get("height", cfg.get("height", 512))
@@ -534,7 +577,7 @@ class ModelManager:
             width=width,
             height=height,
             guidance_scale=guidance_scale,
-            num_inference_steps=inference_steps,
+            num_inference_steps=steps,
             generator=generator,
             num_images_per_prompt=1,
             output_type="pil",