add input check

prompt_check.py

@@ -1,3 +1,7 @@
+from diffusers.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
 def is_unsafe_prompt(model, tokenizer, user_prompt):
     prompt = f"""You are a content safety filter. Your task is to analyze the following user input and determine if it contains unsafe content.
     
@@ -15,6 +19,8 @@ def is_unsafe_prompt(model, tokenizer, user_prompt):
     messages = [
         {"role": "user", "content": prompt}
     ]
+
+    logger.warning(str(messages))
     text = tokenizer.apply_chat_template(
         messages,
         tokenize=False,
@@ -32,6 +38,6 @@ def is_unsafe_prompt(model, tokenizer, user_prompt):
 
     content = tokenizer.decode(output_ids, skip_special_tokens=True).strip("\n")
 
-    print(content)
+    logger.warning(content)
 
     return "yes" in content.lower()