Fix bug on model preloadling

model/model_manager.py

@@ -7,7 +7,7 @@ import spaces
 from PIL import Image
 from .models import IMAGE_GENERATION_MODELS, IMAGE_EDITION_MODELS, VIDEO_GENERATION_MODELS, MUSEUM_UNSUPPORTED_MODELS, DESIRED_APPEAR_MODEL, load_pipeline
 from .fetch_museum_results import draw_from_imagen_museum, draw2_from_imagen_museum, draw_from_videogen_museum, draw2_from_videogen_museum
-from .pre_download import pre_download_all_models, pre_download_video_models
+from .pre_download import pre_download_all_models, pre_download_image_models_gen, pre_download_image_models_edit, pre_download_video_models_gen
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 import re
@@ -82,7 +82,7 @@ class ModelManager:
         self.load_guard(enable_nsfw)
         self.loaded_models = {}
         if do_pre_download:
-            pre_download_all_models()
+            pre_download_all_models(include_video=False)
         if do_debug_packages:
             debug_packages()
 

model/pre_download.py

@@ -1,16 +1,17 @@
 from .models import IMAGE_GENERATION_MODELS, IMAGE_EDITION_MODELS, VIDEO_GENERATION_MODELS
 
-def pre_download_all_models():
+def pre_download_all_models(include_video=True):
     """
     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()
+    imagen_dl_error = pre_download_image_models_gen()
+    imagedit_dl_error = pre_download_image_models_edit()
+    if include_video:
+        videogen_dl_error = pre_download_video_models_gen()
     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():
+def pre_download_image_models_gen():
     """
     Pre-download image models to avoid download delay during the first user request
     """
@@ -33,7 +34,7 @@ def pre_download_image_models():
             pass
     return errored_models
 
-def pre_download_image_models():
+def pre_download_image_models_edit():
     """
     Pre-download image models to avoid download delay during the first user request
     """
@@ -56,7 +57,7 @@ def pre_download_image_models():
             pass
     return errored_models
 
-def pre_download_video_models():
+def pre_download_video_models_gen():
     """
     Pre-download video models to avoid download delay during the first user request
     """