Initial commit

Minor modifications to TTS-arena code

app.py

@@ -0,0 +1,429 @@
+# An arena for 3D generations with code inspired from TTS arena
+
+import gradio as gr
+import pandas as pd
+from langdetect import detect
+from datasets import load_dataset
+import threading, time, uuid, sqlite3, shutil, os, random, asyncio, threading
+from pathlib import Path
+from huggingface_hub import CommitScheduler, delete_file, hf_hub_download
+from gradio_client import Client
+from detoxify import Detoxify
+import os
+import tempfile
+
+toxicity = Detoxify('original')
+
+####################################
+# Constants
+####################################
+AVAILABLE_MODELS = {
+    'TripoSR': 'TripSR',
+    'Shape-E': 'shap-e',
+}
+
+SPACE_ID = os.getenv('SPACE_ID')
+MAX_SAMPLE_TXT_LENGTH = 300
+MIN_SAMPLE_TXT_LENGTH = 10
+DB_DATASET_ID = os.getenv('DATASET_ID')
+DB_NAME = "database.db"
+
+# If /data available => means local storage is enabled => let's use it!
+DB_PATH = f"/data/{DB_NAME}" if os.path.isdir("/data") else DB_NAME
+print(f"Using {DB_PATH}")
+
+####################################
+# Functions
+####################################
+
+def create_db_if_missing():
+    conn = get_db()
+    cursor = conn.cursor()
+    cursor.execute('''
+        CREATE TABLE IF NOT EXISTS model (
+            name TEXT UNIQUE,
+            upvote INTEGER,
+            downvote INTEGER
+        );
+    ''')
+    cursor.execute('''
+        CREATE TABLE IF NOT EXISTS vote (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            username TEXT,
+            model TEXT,
+            vote INTEGER,
+            timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+        );
+    ''')
+    cursor.execute('''
+        CREATE TABLE IF NOT EXISTS votelog (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            username TEXT,
+            chosen TEXT,
+            rejected TEXT,
+            timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+        );
+    ''')
+    
+def get_db():
+    return sqlite3.connect(DB_PATH)
+
+####################################
+# Space initialization
+####################################
+
+# Download existing DB
+if not os.path.isfile(DB_PATH):
+    print("Downloading DB...")
+    try:
+        cache_path = hf_hub_download(repo_id=DB_DATASET_ID, repo_type='dataset', filename=DB_NAME)
+        shutil.copyfile(cache_path, DB_PATH)
+        print("Downloaded DB")
+    except Exception as e:
+        print("Error while downloading DB:", e)
+
+# Create DB table (if doesn't exist)
+create_db_if_missing()
+    
+# Sync local DB with remote repo every 5 minute (only if a change is detected)
+scheduler = CommitScheduler(
+    repo_id=DB_DATASET_ID,
+    repo_type="dataset",
+    folder_path=Path(DB_PATH).parent,
+    every=5,
+    allow_patterns=DB_NAME,
+)
+
+####################################
+# Router API
+####################################
+router = Client("RamAnanth1/3D-Arena-Router", hf_token=os.getenv('HF_TOKEN'))
+####################################
+# Gradio app
+####################################
+MUST_BE_LOGGEDIN = "Please login with Hugging Face to participate in the 3D Arena."
+DESCR = """
+# ⚔️3D Arena: Benchmarking Image-to-3D models
+
+Vote to help the community find the best Image-to-3D model!
+""".strip()
+
+INSTR = """
+## 🗳️ Vote
+
+* Input image to generate a 3D reconstruction.
+* View the responses of the models, one after the other.
+* Vote on which model made a better reconstruction.
+* _Note: Model names are revealed after the vote is cast._
+
+Note: It may take up to 60 seconds to get a response.
+""".strip()
+request = ''
+if SPACE_ID:
+    request = f"""
+### Request a model
+
+Please [create a Discussion](https://huggingface.co/spaces/{SPACE_ID}/discussions/new) to request a model.
+"""
+ABOUT = f"""
+## 📄 About
+
+The 3D Arena evaluates leading 3D generation model. It is inspired by LMsys's [Chatbot Arena](https://chat.lmsys.org/) and [TTS-Arena](https://huggingface.co/spaces/TTS-AGI/TTS-Arena).
+
+### The Arena
+
+The leaderboard allows a user to input an image, for which a 3D reconstruction be synthesized by two models. After viewing each sample, the user can vote on which model works better. Due to the risks of human bias and abuse, model names are revealed only after a vote is submitted.
+
+{request}
+
+
+""".strip()
+LDESC = """
+## 🏆 Leaderboard
+
+Vote to help the community find the best Image-to-3D model!
+
+The leaderboard displays models in descending order of how suitable the models are (based on votes cast by the community).
+
+Important: In order to help keep results fair, the leaderboard hides results by default until the number of votes passes a threshold. Tick the `Reveal preliminary results` to show models without sufficient votes. Please note that preliminary results may be inaccurate.
+""".strip()
+
+def del_db(txt):
+    if not txt.lower() == 'delete db':
+        raise gr.Error('You did not enter "delete db"')
+
+    # Delete local + remote
+    os.remove(DB_PATH)
+    delete_file(path_in_repo=DB_NAME, repo_id=DB_DATASET_ID, repo_type='dataset')
+
+    # Recreate
+    create_db_if_missing()
+    return 'Delete DB'
+
+theme = gr.themes.Monochrome(
+    primary_hue="indigo",
+    secondary_hue="blue",
+    neutral_hue="slate",
+    radius_size=gr.themes.sizes.radius_sm,
+    font=[gr.themes.GoogleFont("Open Sans"), "ui-sans-serif", "system-ui", "sans-serif"],
+)
+css = ".generating {visibility: hidden}"
+
+model_names = {
+    'TripoSR': 'TripoSR',
+    'Shap-E': 'Shap-E',
+}
+model_licenses = {
+    'TripoSR': 'MIT License',
+    'Shap-E': 'MIT License'
+}
+model_links = {
+    'TripoSR': 'https://github.com/VAST-AI-Research/TripoSR',
+    'Shap-E': 'https://github.com/openai/shap-e',
+}
+
+def model_license(name):
+    print(name)
+    for k, v in AVAILABLE_MODELS.items():
+        if k == name:
+            if v in model_licenses:
+                return model_licenses[v]
+    print('---')
+    return 'Unknown'
+def get_leaderboard(reveal_prelim = False):
+    conn = get_db()
+    cursor = conn.cursor()
+    sql = 'SELECT name, upvote, downvote FROM model'
+    # if not reveal_prelim: sql += ' WHERE EXISTS (SELECT 1 FROM model WHERE (upvote + downvote) > 750)'
+    if not reveal_prelim: sql += ' WHERE (upvote + downvote) > 500'
+    cursor.execute(sql)
+    data = cursor.fetchall()
+    df = pd.DataFrame(data, columns=['name', 'upvote', 'downvote'])
+    # df['license'] = df['name'].map(model_license)
+    df['name'] = df['name'].replace(model_names)
+    df['votes'] = df['upvote'] + df['downvote']
+    # df['score'] = round((df['upvote'] / df['votes']) * 100, 2) # Percentage score
+
+    ## ELO SCORE
+    df['score'] = 1200
+    for i in range(len(df)):
+        for j in range(len(df)):
+            if i != j:
+                expected_a = 1 / (1 + 10 ** ((df['score'][j] - df['score'][i]) / 400))
+                expected_b = 1 / (1 + 10 ** ((df['score'][i] - df['score'][j]) / 400))
+                actual_a = df['upvote'][i] / df['votes'][i]
+                actual_b = df['upvote'][j] / df['votes'][j]
+                df.at[i, 'score'] += 32 * (actual_a - expected_a)
+                df.at[j, 'score'] += 32 * (actual_b - expected_b)
+    df['score'] = round(df['score'])
+    ## ELO SCORE
+    df = df.sort_values(by='score', ascending=False)
+    df['order'] = ['#' + str(i + 1) for i in range(len(df))]
+    # df = df[['name', 'score', 'upvote', 'votes']]
+    # df = df[['order', 'name', 'score', 'license', 'votes']]
+    df = df[['order', 'name', 'score', 'votes']]
+    return df
+
+def mkuuid(uid):
+    if not uid:
+        uid = uuid.uuid4()
+    return uid
+
+def upvote_model(model, uname):
+    conn = get_db()
+    cursor = conn.cursor()
+    cursor.execute('UPDATE model SET upvote = upvote + 1 WHERE name = ?', (model,))
+    if cursor.rowcount == 0:
+        cursor.execute('INSERT OR REPLACE INTO model (name, upvote, downvote) VALUES (?, 1, 0)', (model,))
+    cursor.execute('INSERT INTO vote (username, model, vote) VALUES (?, ?, ?)', (uname, model, 1,))
+    with scheduler.lock:
+        conn.commit()
+    cursor.close()
+
+def downvote_model(model, uname):
+    conn = get_db()
+    cursor = conn.cursor()
+    cursor.execute('UPDATE model SET downvote = downvote + 1 WHERE name = ?', (model,))
+    if cursor.rowcount == 0:
+        cursor.execute('INSERT OR REPLACE INTO model (name, upvote, downvote) VALUES (?, 0, 1)', (model,))
+    cursor.execute('INSERT INTO vote (username, model, vote) VALUES (?, ?, ?)', (uname, model, -1,))
+    with scheduler.lock:
+        conn.commit()
+    cursor.close()
+
+def a_is_better(model1, model2, userid):
+    userid = mkuuid(userid)
+    if model1 and model2:
+        conn = get_db()
+        cursor = conn.cursor()
+        cursor.execute('INSERT INTO votelog (username, chosen, rejected) VALUES (?, ?, ?)', (str(userid), model1, model2,))
+        with scheduler.lock:
+            conn.commit()
+            cursor.close()
+        upvote_model(model1, str(userid))
+        downvote_model(model2, str(userid))
+    return reload(model1, model2, userid, chose_a=True)
+def b_is_better(model1, model2, userid):
+    userid = mkuuid(userid)
+    if model1 and model2:
+        conn = get_db()
+        cursor = conn.cursor()
+        cursor.execute('INSERT INTO votelog (username, chosen, rejected) VALUES (?, ?, ?)', (str(userid), model2, model1,))
+        with scheduler.lock:
+            conn.commit()
+            cursor.close()
+        upvote_model(model2, str(userid))
+        downvote_model(model1, str(userid))
+    return reload(model1, model2, userid, chose_b=True)
+def both_bad(model1, model2, userid):
+    userid = mkuuid(userid)
+    if model1 and model2:
+        downvote_model(model1, str(userid))
+        downvote_model(model2, str(userid))
+    return reload(model1, model2, userid)
+def both_good(model1, model2, userid):
+    userid = mkuuid(userid)
+    if model1 and model2:
+        upvote_model(model1, str(userid))
+        upvote_model(model2, str(userid))
+    return reload(model1, model2, userid)
+def reload(chosenmodel1=None, chosenmodel2=None, userid=None, chose_a=False, chose_b=False):
+    out = [
+        gr.update(interactive=False, visible=False),
+        gr.update(interactive=False, visible=False)
+    ]
+    if chose_a == True:
+        out.append(gr.update(value=f'Your vote: {chosenmodel1}', interactive=False, visible=True))
+        out.append(gr.update(value=f'{chosenmodel2}', interactive=False, visible=True))
+    else:
+        out.append(gr.update(value=f'{chosenmodel1}', interactive=False, visible=True))
+        out.append(gr.update(value=f'Your vote: {chosenmodel2}', interactive=False, visible=True))
+    out.append(gr.update(visible=True))
+    return out
+
+with gr.Blocks() as leaderboard:
+    gr.Markdown(LDESC)
+    # df = gr.Dataframe(interactive=False, value=get_leaderboard())
+    df = gr.Dataframe(interactive=False, min_width=0, wrap=True, column_widths=[30, 200, 50, 50])
+    with gr.Row():
+        reveal_prelim = gr.Checkbox(label="Reveal preliminary results", info="Show all models, including models with very few human ratings.", scale=1)
+        reloadbtn = gr.Button("Refresh", scale=3)
+    reveal_prelim.input(get_leaderboard, inputs=[reveal_prelim], outputs=[df])
+    leaderboard.load(get_leaderboard, inputs=[reveal_prelim], outputs=[df])
+    reloadbtn.click(get_leaderboard, inputs=[reveal_prelim], outputs=[df])
+    
+def synthandreturn(text):
+    text = text.strip()
+    if len(text) > MAX_SAMPLE_TXT_LENGTH:
+        raise gr.Error(f'You exceeded the limit of {MAX_SAMPLE_TXT_LENGTH} characters')
+    if len(text) < MIN_SAMPLE_TXT_LENGTH:
+        raise gr.Error(f'Please input a text longer than {MIN_SAMPLE_TXT_LENGTH} characters')
+    if (
+        # test toxicity 
+        toxicity.predict(text)['toxicity'] > 0.8
+    ):
+        print(f'Detected toxic content! "{text}"')
+        raise gr.Error('Your text failed the toxicity test')
+    if not text:
+        raise gr.Error(f'You did not enter any text')
+    # Check language
+    try:
+        if not detect(text) == "en":
+            gr.Warning('Warning: The input text may not be in English')
+    except:
+        pass
+    # Get two random models
+    mdl1, mdl2 = random.sample(list(AVAILABLE_MODELS.keys()), 2)
+    log_text(text)
+    print("[debug] Using", mdl1, mdl2)
+    def predict_and_update_result(text, model, result_storage):
+        try:
+            if model in AVAILABLE_MODELS:
+                result = router.predict(text, AVAILABLE_MODELS[model].lower(), api_name="/synthesize")
+            else:
+                result = router.predict(text, model.lower(), api_name="/synthesize")
+        except:
+            raise gr.Error('Unable to call API, please try again :)')
+        print('Done with', model)
+        try:
+            with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as f:
+                audio = AudioSegment.from_file(result)
+                current_sr = audio.frame_rate
+                if current_sr > 24000:
+                    audio = audio.set_frame_rate(24000)
+                try:
+                    print('Trying to normalize audio')
+                    audio = match_target_amplitude(audio, -20)
+                except:
+                    print('[WARN] Unable to normalize audio')
+                audio.export(f.name, format="wav")
+                os.unlink(result)
+                result = f.name
+        except:
+            pass
+
+        result_storage[model] = result
+    
+    results = {}
+    thread1 = threading.Thread(target=predict_and_update_result, args=(text, mdl1, results))
+    thread2 = threading.Thread(target=predict_and_update_result, args=(text, mdl2, results))
+    thread1.start()
+    thread2.start()
+    thread1.join()
+    thread2.join()
+    
+    return (
+        text,
+        "Synthesize",
+        gr.update(visible=True), # r2
+        mdl1, # model1
+        mdl2, # model2
+        gr.update(visible=True, value=results[mdl1]), # aud1
+        gr.update(visible=True, value=results[mdl2]), # aud2
+        gr.update(visible=True, interactive=True),
+        gr.update(visible=True, interactive=True),
+        gr.update(visible=False),
+        gr.update(visible=False),
+        gr.update(visible=False), #nxt round btn
+    )
+    
+def clear_stuff():
+    return "", "Synthesize", gr.update(visible=False), '', '', gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+
+with gr.Blocks() as vote:
+    useridstate = gr.State()
+    gr.Markdown(INSTR)
+    with gr.Group():
+        with gr.Row():
+            text = gr.Textbox(container=False, show_label=False, placeholder="Enter text to synthesize", lines=1, max_lines=1, scale=9999999, min_width=0)
+        btn = gr.Button("Synthesize", variant='primary')
+    model1 = gr.Textbox(interactive=False, lines=1, max_lines=1, visible=False)
+    model2 = gr.Textbox(interactive=False, lines=1, max_lines=1, visible=False)
+    with gr.Row(visible=False) as r2:
+        with gr.Column():
+            with gr.Group():
+                aud1 = gr.Audio(interactive=False, show_label=False, show_download_button=False, show_share_button=False, waveform_options={'waveform_progress_color': '#3C82F6'})
+                abetter = gr.Button("A is better", variant='primary')
+                prevmodel1 = gr.Textbox(interactive=False, show_label=False, container=False, value="Vote to reveal model A", text_align="center", lines=1, max_lines=1, visible=False)
+        with gr.Column():
+            with gr.Group():
+                aud2 = gr.Audio(interactive=False, show_label=False, show_download_button=False, show_share_button=False, waveform_options={'waveform_progress_color': '#3C82F6'})
+                bbetter = gr.Button("B is better", variant='primary')
+                prevmodel2 = gr.Textbox(interactive=False, show_label=False, container=False, value="Vote to reveal model B", text_align="center", lines=1, max_lines=1, visible=False)
+    nxtroundbtn = gr.Button('Next round', visible=False)
+    outputs = [text, btn, r2, model1, model2, aud1, aud2, abetter, bbetter, prevmodel1, prevmodel2, nxtroundbtn]
+    btn.click(synthandreturn, inputs=[text], outputs=outputs)
+    nxtroundbtn.click(clear_stuff, outputs=outputs)
+
+    
+    nxt_outputs = [abetter, bbetter, prevmodel1, prevmodel2, nxtroundbtn]
+    abetter.click(a_is_better, outputs=nxt_outputs, inputs=[model1, model2, useridstate])
+    bbetter.click(b_is_better, outputs=nxt_outputs, inputs=[model1, model2, useridstate])
+
+with gr.Blocks() as about:
+    gr.Markdown(ABOUT)
+with gr.Blocks(theme=theme, css=css, title="3D Arena") as demo:
+    gr.Markdown(DESCR)
+    gr.TabbedInterface([vote, leaderboard, about], ['🗳️ Vote', '🏆 Leaderboard', '📄 About'])
+    
+demo.queue(api_open=False, default_concurrency_limit=40).launch(show_api=False)