update

model/model_manager.py

@@ -12,13 +12,14 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
 class ModelManager:
-    def __init__(self):
+    def __init__(self, enable_nsfw=True):
         self.model_ig_list = IMAGE_GENERATION_MODELS
         self.model_ie_list = IMAGE_EDITION_MODELS
         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.enable_nsfw = enable_nsfw
+        self.load_guard(enable_nsfw)
         self.loaded_models = {}
 
     def load_model_pipe(self, model_name):
@@ -29,22 +30,33 @@ class ModelManager:
             pipe = self.loaded_models[model_name]
         return pipe
 
-    def load_guard(self):
+    def load_guard(self, enable_nsfw=True):
         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'])
-        self.guard = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=dtype, device_map=device, token=os.environ['HF_GUARD'])
+        if enable_nsfw:
+            self.guard_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'])
+        else:
+            self.guard_tokenizer = None
+            self.guard = None
         
-    @spaces.GPU(duration=30)
     def NSFW_filter(self, prompt):
         chat = [{"role": "user", "content": prompt}]
-        input_ids = self.tokenizer.apply_chat_template(chat, return_tensors="pt").to('cuda')
+        input_ids = self.guard_tokenizer.apply_chat_template(chat, return_tensors="pt").to('cuda')
         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 = self.tokenizer.decode(output[0][prompt_len:], skip_special_tokens=True)
-        return result
+        if self.guard:
+            @spaces.GPU(duration=30)
+            def _generate():
+                return self.guard.generate(input_ids=input_ids, max_new_tokens=100, pad_token_id=0)
+            output = _generate()
+            output = self.guard.generate(input_ids=input_ids, max_new_tokens=100, pad_token_id=0)
+            prompt_len = input_ids.shape[-1]
+            result = self.guard_tokenizer.decode(output[0][prompt_len:], skip_special_tokens=True)
+            return result
+        else:
+            # guard is disabled
+            return "safe"
     
     @spaces.GPU(duration=120)
     def generate_image_ig(self, prompt, model_name):