renaming CogVideoX

model/fetch_museum_results/__init__.py

@@ -3,7 +3,7 @@ from .imagen_museum import fetch_indexes, fetch_indexes_no_csv
 import random
 
 ARENA_TO_IG_MUSEUM = {"LCM(v1.5/XL)":"LCM", "PlayGroundV2.5": "PlayGroundV2_5", "FLUX.1-schnell": "FLUX1schnell", "FLUX.1-dev": "FLUX1dev"}
-ARENA_TO_VG_MUSEUM = {"StableVideoDiffusion": "FastSVD"}
+ARENA_TO_VG_MUSEUM = {"StableVideoDiffusion": "FastSVD", "CogVideoX-2B": "CogVideoX"}
 
 def draw2_from_imagen_museum(task, model_name1, model_name2):
     task_name = TASK_DICT[task]

model/model_registry.py

@@ -410,9 +410,19 @@ register_model_info(
     "text2video_generation"
 )
 
+register_model_info(
+    ["videogenhub_CogVideoX_generation"],
+    "CogVideoX-2B",
+    "https://github.com/THUDM/CogVideo",
+    "Text-to-Video Diffusion Models with An Expert Transformer.",
+    "CogVideoX LICENSE",
+    "THUDM",
+    "text2video_generation"
+)
+
 register_model_info(
     ["fal_CogVideoX_text2video"],
-    "CogVideoX",
+    "CogVideoX-5B",
     "https://github.com/THUDM/CogVideo",
     "Text-to-Video Diffusion Models with An Expert Transformer.",
     "CogVideoX LICENSE",

model/models/__init__.py

@@ -21,13 +21,16 @@ VIDEO_GENERATION_MODELS = ['fal_AnimateDiff_text2video',
                            #'videogenhub_LaVie_generation', 
                            'videogenhub_VideoCrafter2_generation',
                            #'videogenhub_ModelScope_generation', 
-                           #'videogenhub_CogVideoX_generation', 
+                           'videogenhub_CogVideoX-2B_generation', 
                            'videogenhub_OpenSora12_generation',
                            #'videogenhub_OpenSora_generation', 
                            #'videogenhub_T2VTurbo_generation',
                            'fal_T2VTurbo_text2video',
                            'fal_StableVideoDiffusion_text2video',
-                           'fal_CogVideoX_text2video']
+                           'fal_CogVideoX-5B_text2video']
+
+MAP_NAMES_VIDEOGENHUB = {"CogVideoX-2B": "CogVideoX"}
+
 MUSEUM_UNSUPPORTED_MODELS = ['videogenhub_OpenSoraPlan_generation']
 DESIRED_APPEAR_MODEL = ['videogenhub_T2VTurbo_generation','fal_StableVideoDiffusion_text2video']
 
@@ -56,6 +59,8 @@ def load_pipeline(model_name):
     elif model_source == "fal":
         pipe = load_fal_model(model_name, model_type)
     elif model_source == "videogenhub":
+        if model_name in MAP_NAMES_VIDEOGENHUB.keys():
+            model_name = MAP_NAMES_VIDEOGENHUB[model_name]
         pipe = load_videogenhub_model(model_name)
     else:
         raise ValueError(f"Model source {model_source} not supported")

model/models/fal_api_models.py

@@ -62,7 +62,7 @@ class FalModel():
                 fal_model_name = 'fast-animatediff/turbo/text-to-video'
             elif self.model_name == 'StableVideoDiffusion':
                 fal_model_name = 'fast-svd/text-to-video'
-            elif self.model_name == 'CogVideoX':
+            elif self.model_name == 'CogVideoX-5B':
                 fal_model_name = 'cogvideox-5b'
             else:
                 raise NotImplementedError(f"text2video model of {self.model_name} in fal is not implemented yet")