Testing add on FluxEdit

model/model_manager.py

@@ -171,8 +171,22 @@ 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 picking_list], 2)
+            # Filter desired_model_list to only include models that exist in picking_list
+            valid_desired_models = [m for m in self.desired_model_list if m in picking_list]
+            
+            # 50% (or DESIRED_APPEAR_MODEL_CHANCE) chance to include exactly one model from valid desired_model_list
+            if valid_desired_models and random.random() < DESIRED_APPEAR_MODEL_CHANCE:
+                # Pick one model from valid desired list
+                desired_model = random.choice(valid_desired_models)
+                # Pick one model from regular list, excluding desired models
+                regular_model = random.choice([m for m in picking_list if m not in valid_desired_models])
+                # Randomly determine order
+                model_names = [desired_model, regular_model] if random.random() < 0.5 else [regular_model, desired_model]
+            else:
+                # Pick two models from the regular picking list
+                model_names = random.sample([model for model in picking_list], 2)
         else:
             model_names = [model_A, model_B]
 
@@ -185,8 +199,22 @@ class ModelManager:
     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 picking_list], 2)
+            # Filter desired_model_list to only include models that exist in picking_list
+            valid_desired_models = [m for m in self.desired_model_list if m in picking_list]
+            
+            # 50% (or DESIRED_APPEAR_MODEL_CHANCE) chance to include exactly one model from valid desired_model_list
+            if valid_desired_models and random.random() < DESIRED_APPEAR_MODEL_CHANCE:
+                # Pick one model from valid desired list
+                desired_model = random.choice(valid_desired_models)
+                # Pick one model from regular list, excluding desired models
+                regular_model = random.choice([m for m in picking_list if m not in valid_desired_models])
+                # Randomly determine order
+                model_names = [desired_model, regular_model] if random.random() < 0.5 else [regular_model, desired_model]
+            else:
+                # Pick two models from the regular picking list
+                model_names = random.sample([model for model in picking_list], 2)
         else:
             model_names = [model_A, model_B]
 
@@ -259,10 +287,25 @@ class ModelManager:
     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 picking_list], 2)
+            # Filter desired_model_list to only include models that exist in picking_list
+            valid_desired_models = [m for m in self.desired_model_list if m in picking_list]
+            
+            # 50% (or DESIRED_APPEAR_MODEL_CHANCE) chance to include exactly one model from valid desired_model_list
+            if valid_desired_models and random.random() < DESIRED_APPEAR_MODEL_CHANCE:
+                # Pick one model from valid desired list
+                desired_model = random.choice(valid_desired_models)
+                # Pick one model from regular list, excluding desired models
+                regular_model = random.choice([m for m in picking_list if m not in valid_desired_models])
+                # Randomly determine order
+                model_names = [desired_model, regular_model] if random.random() < 0.5 else [regular_model, desired_model]
+            else:
+                # Pick two models from the regular picking list
+                model_names = random.sample([model for model in picking_list], 2)
         else:
             model_names = [model_A, model_B]
+
         with concurrent.futures.ThreadPoolExecutor() as executor:
             futures = [executor.submit(self.generate_image_ie, textbox_source, textbox_target, textbox_instruct, source_image, model) for model in model_names]
             results = [future.result() for future in futures]
@@ -271,10 +314,25 @@ class ModelManager:
     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 picking_list], 2)
+            # Filter desired_model_list to only include models that exist in picking_list
+            valid_desired_models = [m for m in self.desired_model_list if m in picking_list]
+            
+            # 50% (or DESIRED_APPEAR_MODEL_CHANCE) chance to include exactly one model from valid desired_model_list
+            if valid_desired_models and random.random() < DESIRED_APPEAR_MODEL_CHANCE:
+                # Pick one model from valid desired list
+                desired_model = random.choice(valid_desired_models)
+                # Pick one model from regular list, excluding desired models
+                regular_model = random.choice([m for m in picking_list if m not in valid_desired_models])
+                # Randomly determine order
+                model_names = [desired_model, regular_model] if random.random() < 0.5 else [regular_model, desired_model]
+            else:
+                # Pick two models from the regular picking list
+                model_names = random.sample([model for model in picking_list], 2)
         else:
             model_names = [model_A, model_B]
+            
         with concurrent.futures.ThreadPoolExecutor() as executor:
             model_1 = model_names[0].split('_')[1]
             model_2 = model_names[1].split('_')[1]
@@ -313,8 +371,22 @@ class ModelManager:
     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 picking_list], 2)
+            # Filter desired_model_list to only include models that exist in picking_list
+            valid_desired_models = [m for m in self.desired_model_list if m in picking_list]
+            
+            # 50% (or DESIRED_APPEAR_MODEL_CHANCE) chance to include exactly one model from valid desired_model_list
+            if valid_desired_models and random.random() < DESIRED_APPEAR_MODEL_CHANCE:
+                # Pick one model from valid desired list
+                desired_model = random.choice(valid_desired_models)
+                # Pick one model from regular list, excluding desired models
+                regular_model = random.choice([m for m in picking_list if m not in valid_desired_models])
+                # Randomly determine order
+                model_names = [desired_model, regular_model] if random.random() < 0.5 else [regular_model, desired_model]
+            else:
+                # Pick two models from the regular picking list
+                model_names = random.sample([model for model in picking_list], 2)
         else:
             model_names = [model_A, model_B]
 

model/model_registry.py

@@ -331,6 +331,16 @@ register_model_info(
     "image_edition"
 )
 
+register_model_info(
+    ["imagenhub_FluxEdit_edition"],
+    "FluxEdit",
+    "https://github.com/sayakpaul/flux-image-editing",
+    "Flux Control trained on OmniEdit dataset",
+    "Apache 2.0",
+    "HuggingFace",
+    "image_edition"
+)
+
 register_model_info(
     ["fal_stable-cascade_text2image"],
     "StableCascade",

model/models/__init__.py

@@ -15,7 +15,7 @@ IMAGE_GENERATION_MODELS = ['imagenhub_SDXLTurbo_generation','imagenhub_SDXL_gene
 IMAGE_EDITION_MODELS = ['imagenhub_CycleDiffusion_edition', 'imagenhub_Pix2PixZero_edition', 'imagenhub_Prompt2prompt_edition',
                         'imagenhub_SDEdit_edition', 'imagenhub_InstructPix2Pix_edition', 
                         'imagenhub_MagicBrush_edition', 'imagenhub_PNP_edition', 
-                        'imagenhub_InfEdit_edition', 'imagenhub_CosXLEdit_edition', 'imagenhub_UltraEdit_edition', 'imagenhub_AURORA_edition']
+                        'imagenhub_InfEdit_edition', 'imagenhub_CosXLEdit_edition', 'imagenhub_UltraEdit_edition', 'imagenhub_AURORA_edition', 'imagenhub_FluxEdit_edition']
 VIDEO_GENERATION_MODELS = ['fal_AnimateDiff_text2video',
                            'fal_AnimateDiffTurbo_text2video',
                            #'videogenhub_LaVie_generation', 
@@ -37,7 +37,7 @@ MAP_NAMES_IMAGENHUB = {}
 MAP_NAMES_VIDEOGENHUB = {"CogVideoX-2B": "CogVideoX", "CogVideoX-5B": "CogVideoX5B"}
 
 MUSEUM_UNSUPPORTED_MODELS = []
-DESIRED_APPEAR_MODEL = []
+DESIRED_APPEAR_MODEL = ['imagenhub_FluxEdit_edition', 'imagenhub_AURORA_edition']
 DESIRED_APPEAR_MODEL_CHANCE = 0.5
 
 ALL_MODELS = IMAGE_GENERATION_MODELS + IMAGE_EDITION_MODELS + VIDEO_GENERATION_MODELS