main.py

import gradio as gr
import os
from huggingface_hub import hf_hub_download
from pathlib import Path
from transformers import GPT2Config, GPT2LMHeadModel, GPT2Tokenizer

config_class, model_class, tokenizer_class = GPT2Config, GPT2LMHeadModel, GPT2Tokenizer
model = model_class.from_pretrained('gpt2')
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")

def search_index(query):
    # 示例返回，实际中应根据查询来搜索索引
    return "example_uuid"


# 下载视频并返回路径的函数
def download_video(uuid):
    dataset_name = "quchenyuan/360x_dataset"
    dataset_path = "360_dataset/binocular/"
    video_filename = f"{uuid}.mp4"

    # 确保存储目录存在
    storage_dir = Path("videos")
    storage_dir.mkdir(exist_ok=True)

    storage_limit = 40*1024 * 1024 * 1024
    current_storage = sum(f.stat().st_size for f in storage_dir.glob('*') if f.is_file())
    if current_storage + os.path.getsize(video_filename) > storage_limit:
        oldest_file = min(storage_dir.glob('*'), key=os.path.getmtime)
        oldest_file.unlink()

    downloaded_file_path = hf_hub_download(dataset_name, dataset_path + video_filename)

    return str(storage_dir / video_filename)


# Gradio 接口函数
def search_and_show_video(query):
    uuid = search_index(query)
    video_path = download_video(uuid)
    return video_path



if __name__ == "__main__":
    with gr.Blocks() as demo:
        with gr.Column():
            with gr.Row():
                search_input = gr.Textbox(label="输入查询")
            with gr.Row():
                with gr.Column():
                    video_output_1 = gr.Video(label="匹配的视频")
                with gr.Column():
                    video_output_2 = gr.Video(label="匹配的视频")
                with gr.Column():
                    video_output_3 = gr.Video(label="匹配的视频")
            with gr.Row():
                submit_button = gr.Button(label="搜索")

        submit_button.click(search_and_show_video, search_input, outputs=[video_output_1, video_output_2, video_output_3])

    # 运行 Gradio 应用
    demo.launch()