mapping on imagenhub and videogenhub implementation

model/models/__init__.py

@@ -29,9 +29,10 @@ VIDEO_GENERATION_MODELS = ['fal_AnimateDiff_text2video',
                            'fal_StableVideoDiffusion_text2video',
                            'fal_CogVideoX-5B_text2video']
 
+MAP_NAMES_IMAGENHUB = {}
 MAP_NAMES_VIDEOGENHUB = {"CogVideoX-2B": "CogVideoX"}
 
-MUSEUM_UNSUPPORTED_MODELS = ['videogenhub_OpenSoraPlan_generation']
+MUSEUM_UNSUPPORTED_MODELS = ['fal_CogVideoX-5B_text2video']
 DESIRED_APPEAR_MODEL = ['videogenhub_T2VTurbo_generation','fal_StableVideoDiffusion_text2video']
 
 ALL_MODELS = IMAGE_GENERATION_MODELS + IMAGE_EDITION_MODELS + VIDEO_GENERATION_MODELS
@@ -53,14 +54,18 @@ def load_pipeline(model_name):
     # TODO ^Support model_name that contains _ in the name
     print("model_source, model_name, model_type =", model_source, model_name, model_type)
     if model_source == "imagenhub":
+        # This approach uses get() to try to fetch the value associated with model_name. If model_name isn't found in the dictionary, 
+        # it defaults to the current value of model_name,
+        model_name = MAP_NAMES_IMAGENHUB.get(model_name, model_name) 
         pipe = load_imagenhub_model(model_name, model_type)
     elif model_source == "playground":
         pipe = load_playground_model(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]
+        # This approach uses get() to try to fetch the value associated with model_name. If model_name isn't found in the dictionary, 
+        # it defaults to the current value of model_name,
+        model_name = MAP_NAMES_VIDEOGENHUB.get(model_name, model_name) 
         pipe = load_videogenhub_model(model_name)
     else:
         raise ValueError(f"Model source {model_source} not supported")