Merge branch 'main' of https://huggingface.co/spaces/TIGER-Lab/GenAI-Arena

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/model_registry.py

@@ -321,6 +321,16 @@ register_model_info(
     "image_edition"
 )
 
+register_model_info(
+    ["imagenhub_AURORA_edition"],
+    "AURORA",
+    "https://aurora-editing.github.io/",
+    "AURORA (Action Reasoning Object Attribute) enables training an instruction-guided image editing model that can perform action and reasoning-centric edits.",
+    "MIT",
+    "McGill NLP",
+    "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_InfEdit_edition', 'imagenhub_CosXLEdit_edition', 'imagenhub_UltraEdit_edition', 'imagenhub_AURORA_edition']
 VIDEO_GENERATION_MODELS = ['fal_AnimateDiff_text2video',
                            'fal_AnimateDiffTurbo_text2video',
                            #'videogenhub_LaVie_generation', 

model/pre_download.py

@@ -1,16 +1,19 @@
 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()
+    else:
+        videogen_dl_error = ["None"]
     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 +36,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 +59,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
     """

requirements.txt

@@ -68,5 +68,5 @@ tensorboard
 timm
 bitsandbytes
 wandb
-pandarallel
 kaleido
+pandarallel