add NSFW gaurd

model/model_manager.py

@@ -8,6 +8,9 @@ from PIL import Image
 from .models import IMAGE_GENERATION_MODELS, IMAGE_EDITION_MODELS, VIDEO_GENERATION_MODELS, MUSEUM_UNSUPPORTED_MODELS, DESIRED_APPEAR_MODEL, load_pipeline
 from .fetch_museum_results import draw_from_imagen_museum, draw2_from_imagen_museum, draw_from_videogen_museum, draw2_from_videogen_museum
 
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+
 class ModelManager:
     def __init__(self):
         self.model_ig_list = IMAGE_GENERATION_MODELS
@@ -24,11 +27,28 @@ class ModelManager:
         else:
             pipe = self.loaded_models[model_name]
         return pipe
+
+    @spaces.GPU(duration=20)
+    def NSFW_filter(self, prompt):
+        model_id = "meta-llama/Meta-Llama-Guard-2-8B"
+        device = "cuda"
+        dtype = torch.bfloat16
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=dtype, device_map=device)
+        chat = [{"role": "user", "content": prompt}]
+        input_ids = tokenizer.apply_chat_template(chat, return_tensors="pt").to(device)
+        output = model.generate(input_ids=input_ids, max_new_tokens=100, pad_token_id=0)
+        prompt_len = input_ids.shape[-1]
+        result = tokenizer.decode(output[0][prompt_len:], skip_special_tokens=True)
+        return result
     
     @spaces.GPU(duration=120)
     def generate_image_ig(self, prompt, model_name):
-        pipe = self.load_model_pipe(model_name)
-        result = pipe(prompt=prompt)
+        if self.NSFW_filter(prompt) == 'safe':
+            pipe = self.load_model_pipe(model_name)
+            result = pipe(prompt=prompt)
+        else:
+            result = ''
         return result
 
     def generate_image_ig_api(self, prompt, model_name):
@@ -96,8 +116,11 @@ class ModelManager:
 
     @spaces.GPU(duration=200)
     def generate_image_ie(self, textbox_source, textbox_target, textbox_instruct, source_image, model_name):
-        pipe = self.load_model_pipe(model_name)
-        result = pipe(src_image = source_image, src_prompt = textbox_source, target_prompt = textbox_target, instruct_prompt = textbox_instruct)
+        if self.NSFW_filter(" ".join([textbox_source, textbox_target, textbox_instruct])) == 'safe':
+            pipe = self.load_model_pipe(model_name)
+            result = pipe(src_image = source_image, src_prompt = textbox_source, target_prompt = textbox_target, instruct_prompt = textbox_instruct)
+        else:
+            result = ''
         return result
 
     def generate_image_ie_museum(self, model_name):
@@ -161,8 +184,11 @@ class ModelManager:
 
     @spaces.GPU(duration=150)
     def generate_video_vg(self, prompt, model_name):
-        pipe = self.load_model_pipe(model_name)
-        result = pipe(prompt=prompt)
+        if self.NSFW_filter(prompt) == 'safe':
+            pipe = self.load_model_pipe(model_name)
+            result = pipe(prompt=prompt)
+        else:
+            result = ''
         return result
 
     def generate_video_vg_api(self, prompt, model_name):

serve/gradio_web.py

@@ -33,6 +33,12 @@ def build_side_by_side_ui_anony(models):
 - By using this service, you agree to the collection of your input and votes for research purposes. 
 - Your data will be anonymized and will not be used for commercial purposes.
 
+## 📑 User Guidelines for AI Content Generation
+- No NSFW Content: All user inputs must adhere to our content guidelines. Prompts containing NSFW (Not Safe For Work) content are strictly prohibited.
+- Automatic Filtering: Any input deemed inappropriate or NSFW will be automatically filtered out. Users will be prompted to re-enter a suitable request.
+- Content Guidelines: Please avoid using language or descriptions that could be considered offensive, explicit, or inappropriate. This includes, but is not limited to, sexual content, violence, and hate speech.
+- Respectful Use: We encourage all users to create content that is respectful and suitable for all audiences.
+
 ## 🏆 Arena Elo 
 Find out who is the 🥇conditional image generation models! More models are going to be supported.
 
@@ -206,6 +212,12 @@ def build_side_by_side_ui_named(models):
 - By using this service, you agree to the collection of your input and votes for research purposes. 
 - Your data will be anonymized and will not be used for commercial purposes.
 
+## 📑 User Guidelines for AI Content Generation
+- No NSFW Content: All user inputs must adhere to our content guidelines. Prompts containing NSFW (Not Safe For Work) content are strictly prohibited.
+- Automatic Filtering: Any input deemed inappropriate or NSFW will be automatically filtered out. Users will be prompted to re-enter a suitable request.
+- Content Guidelines: Please avoid using language or descriptions that could be considered offensive, explicit, or inappropriate. This includes, but is not limited to, sexual content, violence, and hate speech.
+- Respectful Use: We encourage all users to create content that is respectful and suitable for all audiences.
+
 ## 🤖 Choose two models to compare
 """
     model_list = models.model_ig_list

serve/gradio_web_image_editing.py

@@ -37,6 +37,12 @@ def build_side_by_side_ui_anony_ie(models):
 - By using this service, you agree to the collection of your input and votes for research purposes. 
 - Your data will be anonymized and will not be used for commercial purposes.
 
+## 📑 User Guidelines for AI Content Generation
+- No NSFW Content: All user inputs must adhere to our content guidelines. Prompts containing NSFW (Not Safe For Work) content are strictly prohibited.
+- Automatic Filtering: Any input deemed inappropriate or NSFW will be automatically filtered out. Users will be prompted to re-enter a suitable request.
+- Content Guidelines: Please avoid using language or descriptions that could be considered offensive, explicit, or inappropriate. This includes, but is not limited to, sexual content, violence, and hate speech.
+- Respectful Use: We encourage all users to create content that is respectful and suitable for all audiences.
+
 ## 🏆 Arena Elo 
 Find out who is the 🥇conditional image edition models!
 
@@ -244,6 +250,12 @@ def build_side_by_side_ui_named_ie(models):
 - By using this service, you agree to the collection of your input and votes for research purposes. 
 - Your data will be anonymized and will not be used for commercial purposes.
 
+## 📑 User Guidelines for AI Content Generation
+- No NSFW Content: All user inputs must adhere to our content guidelines. Prompts containing NSFW (Not Safe For Work) content are strictly prohibited.
+- Automatic Filtering: Any input deemed inappropriate or NSFW will be automatically filtered out. Users will be prompted to re-enter a suitable request.
+- Content Guidelines: Please avoid using language or descriptions that could be considered offensive, explicit, or inappropriate. This includes, but is not limited to, sexual content, violence, and hate speech.
+- Respectful Use: We encourage all users to create content that is respectful and suitable for all audiences.
+
 ## 🤖 Choose two models to compare
 """
 

serve/gradio_web_video_generation.py

@@ -34,6 +34,12 @@ def build_side_by_side_ui_anony_vg(models):
 - By using this service, you agree to the collection of your input and votes for research purposes. 
 - Your data will be anonymized and will not be used for commercial purposes.
 
+## 📑 User Guidelines for AI Content Generation
+- No NSFW Content: All user inputs must adhere to our content guidelines. Prompts containing NSFW (Not Safe For Work) content are strictly prohibited.
+- Automatic Filtering: Any input deemed inappropriate or NSFW will be automatically filtered out. Users will be prompted to re-enter a suitable request.
+- Content Guidelines: Please avoid using language or descriptions that could be considered offensive, explicit, or inappropriate. This includes, but is not limited to, sexual content, violence, and hate speech.
+- Respectful Use: We encourage all users to create content that is respectful and suitable for all audiences.
+
 ## 🏆 Arena Elo 
 Find out who is the 🥇conditional video generation models! More models are going to be supported.
 
@@ -210,6 +216,12 @@ def build_side_by_side_ui_named_vg(models):
 - By using this service, you agree to the collection of your input and votes for research purposes. 
 - Your data will be anonymized and will not be used for commercial purposes.
 
+## 📑 User Guidelines for AI Content Generation
+- No NSFW Content: All user inputs must adhere to our content guidelines. Prompts containing NSFW (Not Safe For Work) content are strictly prohibited.
+- Automatic Filtering: Any input deemed inappropriate or NSFW will be automatically filtered out. Users will be prompted to re-enter a suitable request.
+- Content Guidelines: Please avoid using language or descriptions that could be considered offensive, explicit, or inappropriate. This includes, but is not limited to, sexual content, violence, and hate speech.
+- Respectful Use: We encourage all users to create content that is respectful and suitable for all audiences.
+
 ## 🤖 Choose two models to compare
 """
     model_list = models.model_vg_list

serve/vote_utils.py

@@ -514,6 +514,8 @@ def generate_ig(gen_func, state, text, model_name, request: gr.Request):
     ig_logger.info(f"generate. ip: {ip}")
     start_tstamp = time.time()
     generated_image = gen_func(text, model_name)
+    if generated_image == '':
+        raise gr.Warning("Input prompt is blocked by the NSFW filter, please input safe content and try again!")
     state.prompt = text
     state.output = generated_image
     state.model_name = model_name
@@ -596,6 +598,8 @@ def generate_igm(gen_func, state0, state1, text, model_name0, model_name1, reque
     model_name0 = re.sub(r"### Model A: ", "", model_name0)
     model_name1 = re.sub(r"### Model B: ", "", model_name1)
     generated_image0, generated_image1 = gen_func(text, model_name0, model_name1)
+    if generated_image0 == '' and generated_image1 == '':
+        raise gr.Warning("Input prompt is blocked by the NSFW filter, please input safe content and try again!")
     state0.prompt = text
     state1.prompt = text
     state0.output = generated_image0
@@ -712,6 +716,8 @@ def generate_igm_annoy(gen_func, state0, state1, text, model_name0, model_name1,
     model_name0 = ""
     model_name1 = ""
     generated_image0, generated_image1, model_name0, model_name1 = gen_func(text, model_name0, model_name1)
+    if generated_image0 == '' and generated_image1 == '':
+        raise gr.Warning("Input prompt is blocked by the NSFW filter, please input safe content and try again!")
     state0.prompt = text
     state1.prompt = text
     state0.output = generated_image0
@@ -831,6 +837,8 @@ def generate_ie(gen_func, state, source_text, target_text, instruct_text, source
     ig_logger.info(f"generate. ip: {ip}")
     start_tstamp = time.time()
     generated_image = gen_func(source_text, target_text, instruct_text, source_image, model_name)
+    if generated_image == '':
+        raise gr.Warning("Input prompt is blocked by the NSFW filter, please input safe content and try again!")
     state.source_prompt = source_text
     state.target_prompt = target_text
     state.instruct_prompt = instruct_text
@@ -933,6 +941,8 @@ def generate_iem(gen_func, state0, state1, source_text, target_text, instruct_te
     model_name0 = re.sub(r"### Model A: ", "", model_name0)
     model_name1 = re.sub(r"### Model B: ", "", model_name1)
     generated_image0, generated_image1 = gen_func(source_text, target_text, instruct_text, source_image, model_name0, model_name1)
+    if generated_image0 == '' and generated_image1 == '':
+        raise gr.Warning("Input prompt is blocked by the NSFW filter, please input safe content and try again!")
     state0.source_prompt = source_text
     state0.target_prompt = target_text
     state0.instruct_prompt = instruct_text
@@ -1074,6 +1084,8 @@ def generate_iem_annoy(gen_func, state0, state1, source_text, target_text, instr
     model_name0 = ""
     model_name1 = ""
     generated_image0, generated_image1, model_name0, model_name1 = gen_func(source_text, target_text, instruct_text, source_image, model_name0, model_name1)
+    if generated_image0 == '' and generated_image1 == '':
+        raise gr.Warning("Input prompt is blocked by the NSFW filter, please input safe content and try again!")
     state0.source_prompt = source_text
     state0.target_prompt = target_text
     state0.instruct_prompt = instruct_text
@@ -1205,6 +1217,8 @@ def generate_vg(gen_func, state, text, model_name, request: gr.Request):
     vg_logger.info(f"generate. ip: {ip}")
     start_tstamp = time.time()
     generated_video = gen_func(text, model_name)
+    if generated_video == '':
+        raise gr.Warning("Input prompt is blocked by the NSFW filter, please input safe content and try again!")
     state.prompt = text
     state.output = generated_video
     state.model_name = model_name
@@ -1299,6 +1313,8 @@ def generate_vgm(gen_func, state0, state1, text, model_name0, model_name1, reque
     model_name0 = re.sub(r"### Model A: ", "", model_name0)
     model_name1 = re.sub(r"### Model B: ", "", model_name1)
     generated_video0, generated_video1 = gen_func(text, model_name0, model_name1)
+    if generated_video0 == '' and generated_video1 == '':
+        raise gr.Warning("Input prompt is blocked by the NSFW filter, please input safe content and try again!")
     state0.prompt = text
     state1.prompt = text
     state0.output = generated_video0
@@ -1437,6 +1453,8 @@ def generate_vgm_annoy(gen_func, state0, state1, text, model_name0, model_name1,
     model_name0 = ""
     model_name1 = ""
     generated_video0, generated_video1, model_name0, model_name1 = gen_func(text, model_name0, model_name1)
+    if generated_video0 == '' and generated_video1 == '':
+        raise gr.Warning("Input prompt is blocked by the NSFW filter, please input safe content and try again!")
     state0.prompt = text
     state1.prompt = text
     state0.output = generated_video0