Update model/model_manager.py

model/model_manager.py

@@ -18,6 +18,7 @@ class ModelManager:
         self.model_vg_list = VIDEO_GENERATION_MODELS
         self.excluding_model_list = MUSEUM_UNSUPPORTED_MODELS
         self.desired_model_list = DESIRED_APPEAR_MODEL
+        self.load_guard()
         self.loaded_models = {}
 
     def load_model_pipe(self, model_name):
@@ -28,35 +29,38 @@ class ModelManager:
             pipe = self.loaded_models[model_name]
         return pipe
 
-    @spaces.GPU(duration=20)
-    def NSFW_filter(self, prompt):
+    def load_guard(self)
         model_id = "meta-llama/Meta-Llama-Guard-2-8B"
-        device = "cuda"
+        device = "cuda" if torch.cuda.is_available() else "cpu"
         dtype = torch.bfloat16
-        tokenizer = AutoTokenizer.from_pretrained(model_id, token=os.environ['HF_GUARD'])
-        model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=dtype, device_map=device, token=os.environ['HF_GUARD'])
+        self.tokenizer = AutoTokenizer.from_pretrained(model_id, token=os.environ['HF_GUARD'])
+        self.guard = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=dtype, device_map=device, token=os.environ['HF_GUARD'])
+        
+    @spaces.GPU(duration=30)
+    def NSFW_filter(self, prompt):
         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)
+        input_ids = self.tokenizer.apply_chat_template(chat, return_tensors="pt").to(device)
+        self.guard.cuda()
+        output = self.guard.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)
+        result = self.tokenizer.decode(output[0][prompt_len:], skip_special_tokens=True)
         return result
     
     @spaces.GPU(duration=120)
     def generate_image_ig(self, prompt, model_name):
-        #if self.NSFW_filter(prompt) == 'safe':
-        pipe = self.load_model_pipe(model_name)
-        result = pipe(prompt=prompt)
-        # else:
-        #    result = ''
+        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):
-        # if self.NSFW_filter(prompt) == 'safe':
-        pipe = self.load_model_pipe(model_name)
-        result = pipe(prompt=prompt)
-        # else:
-        #    result = ''
+        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_museum(self, model_name):