run.py

import gradio as gr
from huggingface.user import HFUser, GR_CONF
import os

Theme = gr.Theme.load(GR_CONF["theme"])
GR_CONF["theme"] = Theme
dataset_options = [
    "OpenVideo/pexels-raw",
    "OpenVideo/Sample-2k",
]
fetch_files = []
fetch_files_select_idx = 0
custom_css = """
    .class-for-small-font label {
        font-size: 12px; /* Adjust the font size as needed */
    }
"""

def login(token):
    u = HFUser.from_token(token)
    return u, u.name, gr.Column(visible=False)

def show_time(u, name):
    return u.ping(name), gr.Column(visible=True)

def list_dataset(u, repo):
    print(f"repo is: {repo}")
    files = u.list_dataset(repo)
    global fetch_files
    fetch_files = files
    return gr.Dropdown(value=files[0], choices=files), gr.Column(visible=True), gr.Column(visible=False)

def fetch_parquet(u, fname):
    _cache = u.fetch_file(fname)
    auto_tag_content_1 = ""
    auto_tag_content_2 = ""
    auto_tag_content_3 = ""
    if fname.endswith(".mp4"):
        directory_name = os.path.dirname(fname)
        file_name = os.path.basename(fname)
        file_name_without_extension = os.path.splitext(file_name)[0]
        gpt_4o_0621_tag = os.path.join(directory_name, "gpt-4o", "dev-20240621", file_name_without_extension+".txt")
        gpt_4o_0626_tag = os.path.join(directory_name, "gpt-4o", "dev-20240626", file_name_without_extension + ".txt")
        gemini_0626_tag = os.path.join(directory_name, "gemini-1.5-pro-latest", "dev-20240626", file_name_without_extension + ".txt")
        gpt_4o_0621_tag_cache = u.fetch_file(gpt_4o_0621_tag)
        with open(gpt_4o_0621_tag_cache, 'r') as file:
            auto_tag_content_1 = auto_tag_content_1 + file.read()
        gpt_4o_0626_tag_cache = u.fetch_file(gpt_4o_0626_tag)
        with open(gpt_4o_0626_tag_cache, 'r') as file:
            auto_tag_content_2 = auto_tag_content_2 + file.read()
        gemini_0626_tag_cache = u.fetch_file(gemini_0626_tag)
        with open(gemini_0626_tag_cache, 'r') as file:
            auto_tag_content_3 = auto_tag_content_3 + file.read()
    return _cache, gr.Label(value=auto_tag_content_1, label="gpt_4o_0621"), \
           gr.Label(value=auto_tag_content_2, label="gpt_4o_0626"), \
           gr.Label(value=auto_tag_content_3, label="gemini_0626")

def split_parquet(u, file, batch_size):
    batch_size = int(batch_size)
    if file.lower().endswith(".mp4"):
        print(f"file is: {file}")
        return file, [[file]], gr.Slider(value=0, maximum=batch_size-1), gr.Column(visible=True)
    file_slice = u.split_parquet(file, batch_size)
    return file_slice[0][0], file_slice, gr.Slider(value=0, maximum=batch_size-1), gr.Column(visible=True)

def select_video(chunks, epoch_idx , batch_idx):
    print(f"chunks: {chunks}, epoch_idx: {epoch_idx}, batch_idx: {batch_idx}")
    chunks_size = len(chunks)
    epoch_idx = int(epoch_idx)
    batch_idx = int(batch_idx)
    if epoch_idx >= chunks_size:
        epoch_idx = chunks_size-1
    if batch_idx >= len(chunks[epoch_idx]):
        batch_idx = len(chunks[epoch_idx])-1
    return chunks[epoch_idx][batch_idx]

def show_lables():
    return gr.Column(visible=True)

def next_chunks(video_chunks, epoch_idx):
    length = len(video_chunks)
    return (epoch_idx+1)%length, gr.Slider(value=0)

def next_mp4_chunks(video_chunks, epoch_idx):
    global fetch_files_select_idx
    fetch_files_select_idx = fetch_files_select_idx + 1
    return gr.Dropdown(value=fetch_files[fetch_files_select_idx + 1], choices=fetch_files)

with gr.Blocks(**GR_CONF) as Core:
    user = gr.State()
    epoch_idx = gr.State(0)
    video_chunks = gr.State()

    with gr.Row(variant="panel"):
        with gr.Column(scale=6):
            _video = gr.Video(height=720)
        with gr.Column(scale=2):
            with gr.Column() as Auth:
                _token = gr.Textbox(label="Huggingface Token")
                _auth = gr.Button("Auth", variant="primary", size="lg")

            with gr.Row() as UUID:
                name= gr.Textbox(label="Name", interactive=False, scale=1)
                time= gr.Textbox(label="Time", interactive=False, scale=1)

            with gr.Column(visible=False) as Repo:
                # raw_dataset= gr.Textbox("OpenVideo/pexels-raw", label="Raw Dataset")
                raw_dataset = gr.Dropdown(choices=dataset_options, label="Raw Dataset")
                _list = gr.Button("List", variant='secondary', size='sm')

            with gr.Column(visible=False) as Batch:
                file = gr.Dropdown(label="Parquet/MP4")
                with gr.Row():
                    _cache= gr.Textbox("Downloading", label="Cache")
                    batch_size= gr.Textbox("8", label="Batch")
                _fetch = gr.Button("Fetch", variant='primary', size='sm')

            with gr.Column(visible=False) as Pick:
                _pick = gr.Slider(0, 7, value=0, step=1, label="Batch", info="Choose between 1 and $BATCH")
                gr.Label()

    with gr.Row(variant="panel") as Auto_Tag:
        with gr.Column():
            auto_tag_1 = gr.Label()
        with gr.Column():
            auto_tag_2 = gr.Label()
        with gr.Column():
            auto_tag_3 = gr.Label()

    with gr.Row(variant="panel", visible=False) as Tag:
        _human_tag = gr.Textbox(label="Tag", scale=2)
        with gr.Column():
            submit = gr.Button("Submit", variant="primary", size="sm", scale=1)
            with gr.Row():
                rst = gr.Button("Reset", variant="stop", size="sm", scale=1)
                nxt = gr.Button("Next Batch", variant="secondary", size="sm", scale=1)
                nxt_mp4 = gr.Button("Next MP4", variant="secondary", size="sm", scale=1)

    _auth.click(fn=login, inputs=_token, outputs=[user, name,  Auth])
    name.change(fn=show_time, inputs=[user, name], outputs=[time, Repo])
    _list.click(fn=list_dataset, inputs=[user, raw_dataset], outputs=[file, Batch,  Repo])
    _fetch.click(fn=fetch_parquet, inputs=[user, file], outputs=[_cache, auto_tag_1, auto_tag_2, auto_tag_3] )
    file.change(fn=fetch_parquet, inputs=[user, file], outputs=[_cache, auto_tag_1, auto_tag_2, auto_tag_3] )
    _cache.change(fn=split_parquet, inputs=[user, _cache, batch_size], outputs=[_video, video_chunks, _pick, Pick])
    _pick.change(fn=select_video, inputs=[video_chunks, epoch_idx, _pick], outputs=_video)
    _video.change(fn=show_lables, outputs=Tag)
    nxt.click(fn=next_chunks, inputs=[video_chunks, epoch_idx], outputs=[epoch_idx, _pick])
    nxt_mp4.click(fn=next_mp4_chunks, inputs=[video_chunks, epoch_idx], outputs=[file])

if __name__ == "__main__":
    Core.launch()