temp override

model/model_manager.py

@@ -5,7 +5,7 @@ import requests
 import io, base64, json
 import spaces
 from PIL import Image
-from .models import IMAGE_GENERATION_MODELS, IMAGE_EDITION_MODELS, VIDEO_GENERATION_MODELS, MUSEUM_UNSUPPORTED_MODELS, load_pipeline
+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
 
 class ModelManager:
@@ -14,6 +14,7 @@ class ModelManager:
         self.model_ie_list = IMAGE_EDITION_MODELS
         self.model_vg_list = VIDEO_GENERATION_MODELS
         self.excluding_model_list = MUSEUM_UNSUPPORTED_MODELS
+        self.desired_model_list = DESIRED_APPEAR_MODEL
         self.loaded_models = {}
 
     def load_model_pipe(self, model_name):
@@ -196,6 +197,9 @@ class ModelManager:
         picking_list = [item for item in self.model_vg_list if item not in self.excluding_model_list]
         if model_A == "" and model_B == "":
             model_names = random.sample([model for model in picking_list], 2)
+
+            #override the random selection
+            model_names[random.choice([0, 1])] = random.choice(self.desired_model_list)
         else:
             model_names = [model_A, model_B]
 

model/models/__init__.py

@@ -16,17 +16,12 @@ IMAGE_EDITION_MODELS = ['imagenhub_CycleDiffusion_edition', 'imagenhub_Pix2PixZe
                         'imagenhub_InfEdit_edition', 'imagenhub_CosXLEdit_edition']
 VIDEO_GENERATION_MODELS = ['fal_AnimateDiff_text2video',
                            'fal_AnimateDiffTurbo_text2video',
-                           'fal_StableVideoDiffusion_text2video',
                            'videogenhub_LaVie_generation', 
                            'videogenhub_VideoCrafter2_generation',
                            'videogenhub_ModelScope_generation',
-                           'videogenhub_OpenSora_generation', 'videogenhub_T2VTurbo_generation']
-MUSEUM_UNSUPPORTED_MODELS = ['fal_AnimateDiff_text2video',
-                            'fal_AnimateDiffTurbo_text2video',
-                            'fal_StableVideoDiffusion_text2video',
-                            'videogenhub_LaVie_generation', 
-                            'videogenhub_VideoCrafter2_generation',
-                            'videogenhub_ModelScope_generation']
+                           'videogenhub_OpenSora_generation']
+MUSEUM_UNSUPPORTED_MODELS = ['videogenhub_OpenSoraPlan_generation']
+DESIRED_APPEAR_MODEL = ['videogenhub_T2VTurbo_generation','fal_StableVideoDiffusion_text2video']
 
 def load_pipeline(model_name):
     """