Merge branch 'main' of https://huggingface.co/spaces/TIGER-Lab/GenAI-Arena

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,23 +29,27 @@ 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('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 = 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':
+            print('The prompt is safe')
             pipe = self.load_model_pipe(model_name)
             result = pipe(prompt=prompt)
         else:
@@ -53,6 +58,7 @@ class ModelManager:
 
     def generate_image_ig_api(self, prompt, model_name):
         if self.NSFW_filter(prompt) == 'safe':
+            print('The prompt is safe')
             pipe = self.load_model_pipe(model_name)
             result = pipe(prompt=prompt)
         else:
@@ -119,11 +125,11 @@ 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':
-            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 = ''
+        # 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):
@@ -187,19 +193,19 @@ class ModelManager:
 
     @spaces.GPU(duration=150)
     def generate_video_vg(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_video_vg_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_video_vg_museum(self, model_name):

model/model_registry.py

@@ -285,6 +285,20 @@ register_model_info(
     "https://github.com/hpcaitech/Open-Sora",
     "A community-driven opensource implementation of Sora.",
 )
+
+register_model_info(
+    ["videogenhub_OpenSora12_generation"],
+    "OpenSora v1.2",
+    "https://github.com/hpcaitech/Open-Sora",
+    "A community-driven opensource implementation of Sora. v1.2",
+)
+
+register_model_info(
+    ["videogenhub_CogVideoX_generation"],
+    "CogVideoX",
+    "https://github.com/THUDM/CogVideo",
+    "Text-to-Video Diffusion Models with An Expert Transformer.",
+)
     
 register_model_info(
     ["videogenhub_T2VTurbo_generation"],

model/models/__init__.py

@@ -19,7 +19,7 @@ VIDEO_GENERATION_MODELS = ['fal_AnimateDiff_text2video',
                            'fal_AnimateDiffTurbo_text2video',
                            'videogenhub_LaVie_generation', 
                            'videogenhub_VideoCrafter2_generation',
-                           'videogenhub_ModelScope_generation',
+                           'videogenhub_ModelScope_generation', 'videogenhub_CogVideoX_generation', 'videogenhub_OpenSora12_generation',
                            'videogenhub_OpenSora_generation', 'videogenhub_T2VTurbo_generation','fal_StableVideoDiffusion_text2video']
 MUSEUM_UNSUPPORTED_MODELS = ['videogenhub_OpenSoraPlan_generation']
 DESIRED_APPEAR_MODEL = ['videogenhub_T2VTurbo_generation','fal_StableVideoDiffusion_text2video']

requirements.txt

@@ -7,7 +7,7 @@ h5py
 xformers~=0.0.20
 numpy>=1.23.5
 pandas<2.0.0
-peft
+peft>=0.12
 torch==2.2
 torchvision 
 torchaudio
@@ -28,7 +28,6 @@ setuptools>=59.5.0
 transformers
 torchmetrics>=0.6.0
 lpips
-dreamsim
 image-reward
 kornia>=0.6
 diffusers>=0.18.0
@@ -49,7 +48,7 @@ statsmodels
 plotly
 git+https://github.com/TIGER-AI-Lab/ImagenHub.git#egg=imagen-hub
 fal_client
--e git+https://github.com/TIGER-AI-Lab/VideoGenHub.git@arena#egg=videogen-hub
+git+https://github.com/TIGER-AI-Lab/VideoGenHub.git@arena#egg=videogen-hub
 open_clip_torch
 decord
 huggingface_hub