trying preload

app.py

@@ -7,6 +7,7 @@ from serve.leaderboard import build_leaderboard_tab
 from model.model_manager import ModelManager
 from pathlib import Path
 from serve.constants import SERVER_PORT, ROOT_PATH, ELO_RESULTS_DIR
+from model.pre_download import pre_download_all_models
 
 def build_combine_demo(models, elo_results_file, leaderboard_table_file):
 
@@ -97,6 +98,8 @@ if __name__ == "__main__":
     root_path = ROOT_PATH
     elo_results_dir = ELO_RESULTS_DIR
     models = ModelManager()
+
+    pre_download_all_models()
     
     elo_results_file, leaderboard_table_file = load_elo_results(elo_results_dir)
     demo = build_combine_demo(models, elo_results_file, leaderboard_table_file)

model/pre_download.py

@@ -0,0 +1,77 @@
+from .models import IMAGE_GENERATION_MODELS, IMAGE_EDITION_MODELS, VIDEO_GENERATION_MODELS
+
+def pre_download_all_models():
+    """
+    Pre-download all models to avoid download delay during the first user request
+    """
+    imagen_dl_error = pre_download_image_models()
+    imagedit_dl_error = pre_download_image_models()
+    videogen_dl_error = pre_download_video_models()
+    print("All models downloaded.")
+    print("Models that encountered download error:", "Image Generation:", imagen_dl_error, "Image Edition:", imagedit_dl_error, "Video Generation:", videogen_dl_error)
+
+def pre_download_image_models():
+    """
+    Pre-download image models to avoid download delay during the first user request
+    """
+    import imagen_hub
+    errored_models = []
+    for model_string in IMAGE_GENERATION_MODELS:
+        print("Loading image generation model:", model_name)
+        model_lib, model_name, model_type = model_string.split("_")
+
+        if model_lib == "imagenhub":
+            try:
+                temp_model = imagen_hub.get_model(model_name) # Forcing model to download weight files
+                del temp_model
+            except Exception as e:
+                print(f"Failed to load model {model_name} \n {e}")
+                errored_models.append(model_string)
+                continue
+        else:
+            pass
+    return errored_models
+
+def pre_download_image_models():
+    """
+    Pre-download image models to avoid download delay during the first user request
+    """
+    import imagen_hub
+    errored_models = []
+    for model_string in IMAGE_EDITION_MODELS:
+        print("Loading image edition model:", model_name)
+        model_lib, model_name, model_type = model_string.split("_")
+
+        if model_lib == "imagenhub":
+            try:
+                temp_model = imagen_hub.get_model(model_name) # Forcing model to download weight files
+                del temp_model
+            except Exception as e:
+                print(f"Failed to load model {model_name} \n {e}")
+                errored_models.append(model_string)
+                continue
+        else:
+            pass
+    return errored_models
+
+def pre_download_video_models():
+    """
+    Pre-download video models to avoid download delay during the first user request
+    """
+    import videogen_hub
+    errored_models = []
+    for model_string in VIDEO_GENERATION_MODELS:
+        print("Loading video generation model:", model_name)
+        model_lib, model_name, model_type = model_string.split("_")
+
+        if model_lib == "videogenhub":
+            try:
+                temp_model = videogen_hub.get_model(model_name) # Forcing model to download weight files
+                del temp_model
+            except Exception as e:
+                print(f"Failed to load model {model_name} \n {e}")
+                errored_models.append(model_string)
+                continue
+        else:
+            pass
+    return errored_models