simple fix to disable museum-unsupported models

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, load_pipeline
+from .models import IMAGE_GENERATION_MODELS, IMAGE_EDITION_MODELS, VIDEO_GENERATION_MODELS, MUSEUM_UNSUPPORTED_MODELS, 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:
@@ -13,6 +13,7 @@ class ModelManager:
         self.model_ig_list = IMAGE_GENERATION_MODELS
         self.model_ie_list = IMAGE_EDITION_MODELS
         self.model_vg_list = VIDEO_GENERATION_MODELS
+        self.excluding_model_list = MUSEUM_UNSUPPORTED_MODELS
         self.loaded_models = {}
 
     def load_model_pipe(self, model_name):
@@ -44,8 +45,10 @@ class ModelManager:
 
 
     def generate_image_ig_parallel_anony(self, prompt, model_A, model_B):
+        # Using list comprehension to get the difference between two lists
+        picking_list = [item for item in self.model_ig_list if item not in self.excluding_model_list]
         if model_A == "" and model_B == "":
-            model_names = random.sample([model for model in self.model_ig_list], 2)
+            model_names = random.sample([model for model in picking_list], 2)
         else:
             model_names = [model_A, model_B]
 
@@ -56,8 +59,10 @@ class ModelManager:
         return results[0], results[1], model_names[0], model_names[1]
 
     def generate_image_ig_museum_parallel_anony(self, model_A, model_B):
+        # Using list comprehension to get the difference between two lists
+        picking_list = [item for item in self.model_ig_list if item not in self.excluding_model_list]
         if model_A == "" and model_B == "":
-            model_names = random.sample([model for model in self.model_ig_list], 2)
+            model_names = random.sample([model for model in picking_list], 2)
         else:
             model_names = [model_A, model_B]
 
@@ -125,8 +130,10 @@ class ModelManager:
         return image_links[0], image_links[1], image_links[2], prompt_list[0], prompt_list[1], prompt_list[2]
 
     def generate_image_ie_parallel_anony(self, textbox_source, textbox_target, textbox_instruct, source_image, model_A, model_B):
+        # Using list comprehension to get the difference between two lists
+        picking_list = [item for item in self.model_ie_list if item not in self.excluding_model_list]
         if model_A == "" and model_B == "":
-            model_names = random.sample([model for model in self.model_ie_list], 2)
+            model_names = random.sample([model for model in picking_list], 2)
         else:
             model_names = [model_A, model_B]
         with concurrent.futures.ThreadPoolExecutor() as executor:
@@ -135,8 +142,10 @@ class ModelManager:
         return results[0], results[1], model_names[0], model_names[1]
 
     def generate_image_ie_museum_parallel_anony(self, model_A, model_B):
+        # Using list comprehension to get the difference between two lists
+        picking_list = [item for item in self.model_ie_list if item not in self.excluding_model_list]
         if model_A == "" and model_B == "":
-            model_names = random.sample([model for model in self.model_ie_list], 2)
+            model_names = random.sample([model for model in picking_list], 2)
         else:
             model_names = [model_A, model_B]
         with concurrent.futures.ThreadPoolExecutor() as executor:
@@ -169,8 +178,10 @@ class ModelManager:
         return video_link, prompt
 
     def generate_video_vg_parallel_anony(self, prompt, model_A, model_B):
+        # Using list comprehension to get the difference between two lists
+        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 self.model_vg_list], 2)
+            model_names = random.sample([model for model in picking_list], 2)
         else:
             model_names = [model_A, model_B]
 
@@ -181,8 +192,10 @@ class ModelManager:
         return results[0], results[1], model_names[0], model_names[1]
 
     def generate_video_vg_museum_parallel_anony(self, model_A, model_B):
+        # Using list comprehension to get the difference between two lists
+        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 self.model_vg_list], 2)
+            model_names = random.sample([model for model in picking_list], 2)
         else:
             model_names = [model_A, model_B]
 

model/models/__init__.py

@@ -21,7 +21,7 @@ VIDEO_GENERATION_MODELS = ['fal_AnimateDiff_text2video',
                            'videogenhub_VideoCrafter2_generation',
                            'videogenhub_ModelScope_generation',
                            'videogenhub_OpenSora_generation']
-
+MUSEUM_UNSUPPORTED_MODELS = ['videogenhub_T2VTurbo_generation', 'videogenhub_OpenSoraPlan_generation']
 
 def load_pipeline(model_name):
     """