update NSFW

model/model_manager.py

@@ -30,7 +30,7 @@ class ModelManager:
         return pipe
 
     def load_guard(self):
-        model_id = "meta-llama/Meta-Llama-Guard-2-8B"
+        model_id = "meta-llama/Llama-Guard-3-8B"
         device = "cuda" if torch.cuda.is_available() else "cpu"
         dtype = torch.bfloat16
         self.tokenizer = AutoTokenizer.from_pretrained(model_id, token=os.environ['HF_GUARD'])
@@ -48,7 +48,7 @@ class ModelManager:
     
     @spaces.GPU(duration=120)
     def generate_image_ig(self, prompt, model_name):
-        if self.NSFW_filter(prompt) == 'safe':
+        if 'unsafe' not in self.NSFW_filter(prompt):
             print('The prompt is safe')
             pipe = self.load_model_pipe(model_name)
             result = pipe(prompt=prompt)
@@ -57,7 +57,7 @@ class ModelManager:
         return result
 
     def generate_image_ig_api(self, prompt, model_name):
-        if self.NSFW_filter(prompt) == 'safe':
+        if 'unsafe' not in self.NSFW_filter(prompt):
             print('The prompt is safe')
             pipe = self.load_model_pipe(model_name)
             result = pipe(prompt=prompt)
@@ -125,7 +125,7 @@ class ModelManager:
 
     @spaces.GPU(duration=200)
     def generate_image_ie(self, textbox_source, textbox_target, textbox_instruct, source_image, model_name):
-        # if self.NSFW_filter(" ".join([textbox_source, textbox_target, textbox_instruct])) == 'safe':
+        # if 'unsafe' not in self.NSFW_filter(" ".join([textbox_source, textbox_target, textbox_instruct])):
         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:
@@ -193,7 +193,7 @@ class ModelManager:
 
     @spaces.GPU(duration=150)
     def generate_video_vg(self, prompt, model_name):
-        # if self.NSFW_filter(prompt) == 'safe':
+        # if 'unsafe' not in self.NSFW_filter(prompt):
         pipe = self.load_model_pipe(model_name)
         result = pipe(prompt=prompt)
         # else:
@@ -201,7 +201,7 @@ class ModelManager:
         return result
 
     def generate_video_vg_api(self, prompt, model_name):
-        # if self.NSFW_filter(prompt) == 'safe':
+        # if 'unsafe' not in self.NSFW_filter(prompt):
         pipe = self.load_model_pipe(model_name)
         result = pipe(prompt=prompt)
         # else: