app.py

import os
import json
import shutil
import gradio as gr
import random
from huggingface_hub import Repository,HfApi
from huggingface_hub import snapshot_download
# from datasets import load_dataset
from datasets import config

hf_token = os.environ['hf_token']  # 确保环境变量中有你的令牌

local_dir = "VBench_sampled_video"  # 本地文件夹路径
# dataset = load_dataset("Vchitect/VBench_sampled_video")
# print(os.listdir("~/.cache/huggingface/datasets/Vchitect___VBench_sampled_video/"))
# root = "~/.cache/huggingface/datasets/Vchitect___VBench_sampled_video/"
# print(config.HF_DATASETS_CACHE)
# root = config.HF_DATASETS_CACHE
# print(root)
def print_directory_contents(path, indent=0):
    # 打印当前目录的内容
    try:
        for item in os.listdir(path):
            item_path = os.path.join(path, item)
            print('    ' * indent + item)  # 使用缩进打印文件或文件夹
            if os.path.isdir(item_path):  # 如果是目录，则递归调用
                print_directory_contents(item_path, indent + 1)
    except PermissionError:
        print('    ' * indent + "[权限错误，无法访问该目录]")

# 拉取数据集
os.makedirs(local_dir, exist_ok=True)
hf_api = HfApi(endpoint="https://huggingface.co", token=hf_token)
hf_api = HfApi(token=hf_token)
repo_id = "Vchitect/VBench_sampled_video"

model_names=[]
for i in hf_api.list_repo_tree('Vchitect/VBench_sampled_video',repo_type='dataset'):
    model_name = i.path
    if '.git' not in model_name and '.md' not in model_name:
        model_names.append(model_name)

with open("videos_by_dimension.json") as f:
    dimension = json.load(f)['videos_by_dimension']
    for key in dimension:
        new_item = []
        for item in dimension[key]:
            new_item.append(os.path.basename(item))
        dimension[key] = new_item

# with open("all_videos.json") as f:
    # all_videos = json.load(f)

types = ['appearance_style', 'color', 'temporal_style', 'spatial_relationship', 'temporal_flickering', 'scene', 'multiple_objects', 'object_class', 'human_action', 'overall_consistency', 'subject_consistency']

def get_video_path_local(model_name, type, prompt):
    if 'Show-1' in model_name:
        video_path_subfolder = os.path.join(model_name, type, 'super2')
    elif 'videocrafter-1' in model_name:
        video_path_subfolder = os.path.join(model_name, type, '1024x576')
    else:
        video_path_subfolder = os.path.join(model_name, type)

    if model_name == 'cogvideo':
        prompt = prompt.replace(".mp4",".gif")
    
    try:
        return hf_api.hf_hub_download(
            repo_id = repo_id,
            filename = prompt,
            subfolder = video_path_subfolder,
            repo_type = "dataset",
            local_dir = local_dir
        )
    except Exception as e:
        print(f"[PATH]{video_path_subfolder}/{prompt} NOT in hf repo, try {model_name}",e)
        video_path_subfolder = model_name
        try:
            return hf_api.hf_hub_download(
                repo_id = repo_id,
                filename = prompt,
                subfolder = video_path_subfolder,
                repo_type = 'dataset',
                local_dir = local_dir
            )
        except Exception as e:
            print(f"[PATH]{video_path_subfolder}/{prompt} NOT in hf repo, try {model_name}",e)
            print(e)
    # video_path = dataset['train'][random_index]['video_path']
    print('error:', model_name, type, prompt)
    return None

def get_random_video():
    # 随机选择一个索引
    random_index = random.randint(0, len(types) - 1)
    type = types[random_index]
    # 随机选择一个Prompt
    random_index = random.randint(0, len(dimension[type]) - 1)
    prompt = dimension[type][random_index]
    prompt = os.path.basename(prompt)
    # 随机选择两个不同的模型名称
    random_model_names = random.sample(model_names, 2)
    model_name_1, model_name_2 = random_model_names
    video_path1 = get_video_path_local(model_name_1, type, prompt)
    video_path2 = get_video_path_local(model_name_2, type, prompt)
    return video_path1, video_path2, model_name_1, model_name_2, type, prompt

def update_prompt_options(type, value=None):
    if value:
        return gr.update(choices=dimension[type], value=value if dimension[type] else None)
    else:    
        return gr.update(choices=dimension[type], value=dimension[type][0] if dimension[type] else None)


def display_videos(type, prompt, model_name_1, model_name_2):
    video_path1 = get_video_path_local(model_name_1, type, prompt)
    video_path2 = get_video_path_local(model_name_2, type, prompt)
    return video_path1, video_path2

def record_user_feedback_a(model_name1, model_name2, type, prompt):
    # 0 means model A  better, 1 means model B better 
    hf_api.hf_hub_download(
            repo_id = "Vchitect/VBench_human_annotation",
            filename = "arena_feedback.csv",
            repo_type = "dataset",
            local_dir = './'
        )
    with open("arena_feedback.csv",'a') as f:
        f.write(f"{model_name1}\t{model_name2}\t{type}\t{prompt}\t{0}\n")
    hf_api.upload_file(
        path_or_fileobj="arena_feedback.csv",
        path_in_repo="arena_feedback.csv",
        repo_id="Vchitect/VBench_human_annotation",
        token=hf_token,
        repo_type="dataset",
        commit_message="[From VBench Arena] user feedback",
    )
    return gr.update(visible=False),gr.update(visible=False)
    
def record_user_feedback_b(model_name1, model_name2, type, prompt):
    # 0 means model A  better, 1 means model B better 
    hf_api.hf_hub_download(
            repo_id = "Vchitect/VBench_human_annotation",
            filename = "arena_feedback.csv",
            repo_type = "dataset",
            local_dir = './'
        )
    with open("arena_feedback.csv",'a') as f:
        f.write(f"{model_name1}\t{model_name2}\t{type}\t{prompt}\t{1}\n")  
    hf_api.upload_file(
        path_or_fileobj="arena_feedback.csv",
        path_in_repo="arena_feedback.csv",
        repo_id="Vchitect/VBench_human_annotation",
        token=hf_token,
        repo_type="dataset",
        commit_message="[From VBench Arena] user feedback",
    )
    return gr.update(visible=False),gr.update(visible=False)

def show_feedback_button():
    return gr.update(visible=True),gr.update(visible=True)

with gr.Blocks() as interface:
    gr.Markdown("# VBench Video Arena")
    gr.Markdown("""
**Random 2 videos** for randomly picking two models to compare random with the same dimension and prompt
**Play Selection** is used for the user to select the model, dimension, prompt in the drop-down box, and display them 
If you are interested, you can also leave your comments.""")
    
    type_output = gr.Dropdown(label="Type", choices=types, value=types[0])
    prompt_output = gr.Dropdown(label="Prompt", choices=dimension[types[0]], value=dimension[types[0]][0])
    prompt_placeholder = gr.State()
    with gr.Row():
        random_button = gr.Button("🎲 Random 2 videos")
        display_button = gr.Button("🎇 Play Selection")
   
    with gr.Row():
        with gr.Column():
            model_name_1_output = gr.Dropdown(label="Model Name 1", choices=model_names, value=model_names[0])
            video_output_1 = gr.Video(label="Video 1")
        with gr.Column():
            model_name_2_output = gr.Dropdown(label="Model Name 2", choices=model_names, value=model_names[1])
            video_output_2 = gr.Video(label="Video 2")
    with gr.Row():
        feed0 = gr.Button("👈 Model A is better",visible=False)
        feed1 = gr.Button("👉 Model B is better",visible=False)
    type_output.change(fn=update_prompt_options, inputs=[type_output], outputs=[prompt_output])
    
    
    random_button.click(
        fn=get_random_video,
        outputs=[video_output_1, video_output_2,model_name_1_output, model_name_2_output, type_output, prompt_placeholder]
    ).then(fn=update_prompt_options,
          inputs=[type_output],
          outputs=[prompt_output]
    ).then(fn=update_prompt_options,
          inputs=[type_output,prompt_placeholder],
          outputs=[prompt_output]
    ).then(
        fn= show_feedback_button,
        outputs=[feed0, feed1]
    )


    display_button.click(
        fn=display_videos,
        inputs=[type_output, prompt_output, model_name_1_output, model_name_2_output],
        outputs=[video_output_1, video_output_2]
    )

    feed0.click(
        fn = record_user_feedback_a,
        inputs=[model_name_1_output, model_name_2_output, type_output, prompt_placeholder],
        outputs=[feed0, feed1]
    )
    feed1.click(
        fn = record_user_feedback_b,
        inputs=[model_name_1_output, model_name_2_output, type_output, prompt_placeholder],
        outputs=[feed0, feed1]
    )

interface.launch()