TIGER-Lab/GenAI-Arena

1import gradio as gr2import os3from serve.gradio_web import *4from serve.gradio_web_image_editing import *5from serve.gradio_web_video_generation import *6from serve.leaderboard import build_leaderboard_tab7from model.model_manager import ModelManager8from pathlib import Path9from serve.constants import SERVER_PORT, ROOT_PATH, ELO_RESULTS_DIR10 11def build_combine_demo(models, elo_results_file, leaderboard_table_file):12 13    with gr.Blocks(14        title="Play with Open Vision Models",15        theme=gr.themes.Default(),16        css=block_css,17    ) as demo:18        with gr.Tabs() as tabs_combine:19            with gr.Tab("Image Generation", id=0):20                with gr.Tabs() as tabs_ig:21                    with gr.Tab("Generation Arena (battle)", id=0):22                        build_side_by_side_ui_anony(models)23 24                    with gr.Tab("Generation Arena (side-by-side)", id=1):25                        build_side_by_side_ui_named(models)26 27                    with gr.Tab("Generation Playground", id=2): #Direct Chat28                        build_single_model_ui(models, add_promotion_links=True)29                    if elo_results_file:30                        with gr.Tab("Generation Leaderboard", id=3):31                            build_leaderboard_tab(elo_results_file['t2i_generation'], leaderboard_table_file['t2i_generation'])32            33            with gr.Tab("Image Edition", id=5):34                with gr.Tabs() as tabs_ie:35                    with gr.Tab("Edition Arena (battle)", id=5):36                        build_side_by_side_ui_anony_ie(models)37 38                    with gr.Tab("Edition Arena (side-by-side)", id=6):39                        build_side_by_side_ui_named_ie(models)40 41                    with gr.Tab("Edition Playground", id=7): #Direct Chat42                        build_single_model_ui_ie(models, add_promotion_links=True)43                    if elo_results_file:44                        with gr.Tab("Edition Leaderboard", id=8):45                            build_leaderboard_tab(elo_results_file['image_editing'], leaderboard_table_file['image_editing'])46 47            with gr.Tab("Video Generation", id=10):48                with gr.Tabs() as tabs_vg:49                    with gr.Tab("Video Generation Arena (battle)", id=10):50                        build_side_by_side_ui_anony_vg(models)51 52                    with gr.Tab("Video Generation Arena (side-by-side)", id=11):53                        build_side_by_side_ui_named_vg(models)54 55                    with gr.Tab("Video Generation Playground", id=12): #Direct Chat56                        build_single_model_ui_vg(models, add_promotion_links=True)57                    if elo_results_file and 'video_generation' in elo_results_file:58                        with gr.Tab("Video Generation Leaderboard", id=13):59                            build_leaderboard_tab(elo_results_file['video_generation'], leaderboard_table_file['video_generation'])60            with gr.Tab("About Us", id=4): 61                build_about()62 63    return demo64 65 66def load_elo_results(elo_results_dir):67    from collections import defaultdict68    elo_results_file = defaultdict(lambda: None)69    leaderboard_table_file = defaultdict(lambda: None)70    if elo_results_dir is not None:71        elo_results_dir = Path(elo_results_dir)72        elo_results_file = {}73        leaderboard_table_file = {}74        for file in elo_results_dir.glob('elo_results_*.pkl'):75            if 't2i_generation' in file.name:76                elo_results_file['t2i_generation'] = file77            elif 'image_editing' in file.name:78                elo_results_file['image_editing'] = file79            elif 'video_generation' in file.name:80                elo_results_file['video_generation'] = file81            else:82                raise ValueError(f"Unknown file name: {file.name}")83        for file in elo_results_dir.glob('*_leaderboard.csv'):84            if 't2i_generation' in file.name:85                leaderboard_table_file['t2i_generation'] = file86            elif 'image_editing' in file.name:87                leaderboard_table_file['image_editing'] = file88            elif 'video_generation' in file.name:89                leaderboard_table_file['video_generation'] = file90            else:91                raise ValueError(f"Unknown file name: {file.name}")92            93    return elo_results_file, leaderboard_table_file94    95if __name__ == "__main__":96    server_port = int(SERVER_PORT)97    root_path = ROOT_PATH98    elo_results_dir = ELO_RESULTS_DIR99    # models = ModelManager(enable_nsfw=False, do_pre_download=True, do_debug_packages=True)100    models = ModelManager(enable_nsfw=False, do_pre_download=False, do_debug_packages=False)101    102    elo_results_file, leaderboard_table_file = load_elo_results(elo_results_dir)103    demo = build_combine_demo(models, elo_results_file, leaderboard_table_file)104    demo.queue(max_size=20).launch(server_port=server_port, root_path=ROOT_PATH)105