Update model/fetch_museum_results/__init__.py

model/fetch_museum_results/__init__.py

@@ -2,8 +2,8 @@ from .imagen_museum import TASK_DICT, DOMAIN
 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"}
+ARENA_TO_IG_MUSEUM = {"LCM(v1.5/XL)":"LCM", "PlayGroundV2.5": "PlayGroundV2_5"}
+ARENA_TO_VG_MUSEUM = {"StableVideoDiffusion": "FastSVD"}
 
 def draw2_from_imagen_museum(task, model_name1, model_name2):
     task_name = TASK_DICT[task]
@@ -61,6 +61,9 @@ def draw2_from_videogen_museum(task, model_name1, model_name2):
     domain = "https://github.com/ChromAIca/VideoGenMuseum/raw/main/Museum/"
     baselink = domain + "VideoGenHub_Text-Guided_VG"
 
+    model_name1 = ARENA_TO_VG_MUSEUM[model_name1] if model_name1 in ARENA_TO_VG_MUSEUM else model_name1
+    model_name2 = ARENA_TO_VG_MUSEUM[model_name2] if model_name2 in ARENA_TO_VG_MUSEUM else model_name2
+
     matched_results = fetch_indexes_no_csv(baselink)
     r = random.Random()
     uid, value = r.choice(list(matched_results.items()))
@@ -77,6 +80,8 @@ def draw_from_videogen_museum(task, model_name):
     domain = "https://github.com/ChromAIca/VideoGenMuseum/raw/main/Museum/"
     baselink = domain + "VideoGenHub_Text-Guided_VG"
 
+    model_name = ARENA_TO_VG_MUSEUM[model_name] if model_name in ARENA_TO_VG_MUSEUM else model_name
+
     matched_results = fetch_indexes_no_csv(baselink)
     r = random.Random()
     uid, value = r.choice(list(matched_results.items()))