app.py

import gradio as gr
import os
# os.system("pip install -r requirements.txt -U")
# os.system("pip uninstall -y apex")
# os.system("pip uninstall -y flash-attn")
# os.system("FLASH_ATTENTION_FORCE_BUILD=TRUE pip install flash-attn")
from serve.gradio_web import *
from serve.leaderboard import build_leaderboard_tab
from model.model_manager import ModelManager
from pathlib import Path
from serve.constants import SERVER_PORT, ROOT_PATH, ELO_RESULTS_DIR   

def build_combine_demo(models, elo_results_file, leaderboard_table_file):

    with gr.Blocks(
        title="Play with Open Vision Models",
        theme=gr.themes.Default(),
        css=block_css,
    ) as demo:
        with gr.Tabs() as tabs_combine:
            with gr.Tab("Image Generation", id=0):
                with gr.Tabs() as tabs_ig:
                    with gr.Tab("Generation Arena (battle)", id=0):
                        build_side_by_side_ui_anony(models)

                    with gr.Tab("Generation Arena (side-by-side)", id=1):
                        build_side_by_side_ui_named(models)

                    if elo_results_file:
                        with gr.Tab("Generation Leaderboard", id=2):
                            build_leaderboard_tab(elo_results_file['t2i_generation'], leaderboard_table_file['t2i_generation'])
            
            with gr.Tab("About Us", id=3): 
                build_about()

    return demo


def load_elo_results(elo_results_dir):
    from collections import defaultdict
    elo_results_file = defaultdict(lambda: None)
    leaderboard_table_file = defaultdict(lambda: None)
    if elo_results_dir is not None:
        elo_results_dir = Path(elo_results_dir)
        elo_results_file = {}
        leaderboard_table_file = {}
        for file in elo_results_dir.glob('elo_results*.pkl'):
            elo_results_file['t2i_generation'] = file
        for file in elo_results_dir.glob('*leaderboard.csv'):
            leaderboard_table_file['t2i_generation'] = file
            
    return elo_results_file, leaderboard_table_file
    
if __name__ == "__main__":

    server_port = int(SERVER_PORT)
    root_path = ROOT_PATH
    elo_results_dir = ELO_RESULTS_DIR
    
    models = ModelManager()
    
    elo_results_file, leaderboard_table_file = load_elo_results(elo_results_dir)
    demo = build_combine_demo(models, elo_results_file, leaderboard_table_file)
    demo.queue(max_size=20).launch(server_port=server_port, root_path=ROOT_PATH, show_error=True)
    demo.launch(server_port=server_port, root_path=ROOT_PATH, show_error=True)