New models

arena_elo/elo_rating/clean_battle_data.py

@@ -58,6 +58,10 @@ IDENTITY_WORDS = [
 for i in range(len(IDENTITY_WORDS)):
     IDENTITY_WORDS[i] = IDENTITY_WORDS[i].lower()
 
+def parse_model_name(model_name):
+    model_source, *rest = model_name.split("_", 1)
+    model_type, model_name = rest[-1], "_".join(rest[:-1])
+    return model_source, model_name, model_type
 
 def remove_html(raw):
     if raw.startswith("<h3>"):
@@ -229,7 +233,8 @@ def clean_battle_data(
             valid = True
             for _model in models:
                 try:
-                    platform, model_name, task = _model.split("_")
+                    #platform, model_name, task = _model.split("_")
+                    platform, model_name, task = parse_model_name(_model)
                 except ValueError:
                     valid = False
                     break
@@ -240,7 +245,8 @@ def clean_battle_data(
                 ct_invalid += 1
                 continue
             for i, _model in enumerate(models):
-                platform, model_name, task = _model.split("_")
+                #platform, model_name, task = _model.split("_")
+                platform, model_name, task = parse_model_name(_model)
                 models[i] = model_name
                 
             # if not all(x.startswith("imagenhub_") and x.endswith("_edition") for x in models):
@@ -253,7 +259,8 @@ def clean_battle_data(
             valid = True
             for _model in models:
                 try:
-                    platform, model_name, task = _model.split("_")
+                    #platform, model_name, task = _model.split("_")
+                    platform, model_name, task = parse_model_name(_model)
                 except ValueError:
                     valid = False
                     break
@@ -264,7 +271,8 @@ def clean_battle_data(
                 ct_invalid += 1
                 continue
             for i, _model in enumerate(models):
-                platform, model_name, task = _model.split("_")
+                #platform, model_name, task = _model.split("_")
+                platform, model_name, task = parse_model_name(_model)
                 models[i] = model_name
             # if not all("playground" in x.lower() or (x.startswith("imagenhub_") and x.endswith("_generation")) for x in models):
             #     print(f"Invalid model names: {models}")
@@ -280,7 +288,8 @@ def clean_battle_data(
             valid = True
             for _model in models:
                 try:
-                    platform, model_name, task = _model.split("_")
+                    #platform, model_name, task = _model.split("_")
+                    platform, model_name, task = parse_model_name(_model)
                 except ValueError:
                     valid = False
                     break
@@ -291,7 +300,8 @@ def clean_battle_data(
                 ct_invalid += 1
                 continue
             for i, _model in enumerate(models):
-                platform, model_name, task = _model.split("_")
+                #platform, model_name, task = _model.split("_")
+                platform, model_name, task = parse_model_name(_model)
                 models[i] = model_name
 
         else:

model/model_registry.py

@@ -143,6 +143,28 @@ register_model_info(
     "Kolors is a large-scale text-to-image generation model based on latent diffusion",
 )
 
+register_model_info(
+    ["fal_AuraFlow_text2image"],
+    "AuraFlow",
+    "https://huggingface.co/fal/AuraFlow",
+    "Opensourced flow-based text-to-image generation model.",
+)
+
+register_model_info(
+    ["fal_FluxTimestep_text2image"],
+    "FLUX.1-schnell",
+    "https://huggingface.co/docs/diffusers/main/en/api/pipelines/flux",
+    "Flux is a series of text-to-image generation models based on diffusion transformers. Timestep-distilled version.",
+)
+
+register_model_info(
+    ["fal_FluxGuidance_text2image"],
+    "FLUX.1-dev",
+    "https://huggingface.co/docs/diffusers/main/en/api/pipelines/flux",
+    "Flux is a series of text-to-image generation models based on diffusion transformers. Guidance-distilled version.",
+)
+
+
 # regist image edition models
 register_model_info(
     ["imagenhub_CycleDiffusion_edition"],

model/models/__init__.py

@@ -9,7 +9,8 @@ from .videogenhub_models import load_videogenhub_model
 #                             'playground_PlayGroundV2_generation', 'playground_PlayGroundV2.5_generation']
 IMAGE_GENERATION_MODELS = ['imagenhub_SDXLTurbo_generation','imagenhub_SDXL_generation', 'imagenhub_PixArtAlpha_generation', 'imagenhub_PixArtSigma_generation',
                            'imagenhub_OpenJourney_generation','imagenhub_SDXLLightning_generation', 'imagenhub_StableCascade_generation', 'imagenhub_HunyuanDiT_generation',
-                           'playground_PlayGroundV2.5_generation', 'imagenhub_Kolors_generation', 'imagenhub_SD3_generation']  # 'playground_PlayGroundV2_generation'
+                           'playground_PlayGroundV2.5_generation', 'imagenhub_Kolors_generation', 'imagenhub_SD3_generation',
+                           'fal_AuraFlow_text2image', 'fal_FluxTimestep_text2image', 'fal_FluxGuidance_text2image']  # 'playground_PlayGroundV2_generation'
 IMAGE_EDITION_MODELS = ['imagenhub_CycleDiffusion_edition', 'imagenhub_Pix2PixZero_edition', 'imagenhub_Prompt2prompt_edition',
                         'imagenhub_SDEdit_edition', 'imagenhub_InstructPix2Pix_edition', 
                         'imagenhub_MagicBrush_edition', 'imagenhub_PNP_edition', 
@@ -32,7 +33,9 @@ def load_pipeline(model_name):
         the name is the name of the model used to load the model
         the type is the type of the model, either generation or edition
     """
-    model_source, model_name, model_type = model_name.split("_")
+    #model_source, model_name, model_type = model_name.split("_")
+    model_source, *rest = model_name.split("_", 1)
+    model_type, model_name = rest[-1], "_".join(rest[:-1])
     if model_source == "imagenhub":
         pipe = load_imagenhub_model(model_name, model_type)
     elif model_source == "playground":

model/models/fal_api_models.py

@@ -6,7 +6,8 @@ import os
 import base64
 
 FAL_MODEl_NAME_MAP = {"SDXL": "fast-sdxl", "SDXLTurbo": "fast-turbo-diffusion", "SDXLLightning": "fast-lightning-sdxl",
-                      "LCM(v1.5/XL)": "fast-lcm-diffusion", "PixArtSigma": "pixart-sigma", "StableCascade": "stable-cascade"}
+                      "LCM(v1.5/XL)": "fast-lcm-diffusion", "PixArtSigma": "pixart-sigma", "StableCascade": "stable-cascade",
+                      "AuraFlow": "aura-flow", "FluxTimestep": "flux/schnell", "FluxGuidance": "flux/dev"}
 
 class FalModel():
     def __init__(self, model_name, model_type):