Update app.py

app.py

@@ -34,7 +34,11 @@ hf_api = HfApi(endpoint="https://huggingface.co", token=hf_token)
 hf_api = HfApi(token=hf_token)
 repo_id = "Vchitect/VBench_sampled_video"
 
-model_names=['Gen-2','Gen-3']
+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']
@@ -44,20 +48,7 @@ with open("videos_by_dimension.json") as f:
 
 types = ['appearance_style', 'color', 'temporal_style', 'spatial_relationship', 'temporal_flickering', 'scene', 'multiple_objects', 'object_class', 'human_action', 'overall_consistency', 'subject_consistency']
 
-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_index = random.randint(0, len(model_names) - 1)
-    model_name =  model_names[random_index]
-
-    
-
+def get_video_path_local(model_name, type, prompt):
     video_path_subfolder = os.path.join(model_name, type)
     try:
         return hf_api.hf_hub_download(
@@ -83,18 +74,41 @@ def get_random_video():
             print(e)
     # video_path = dataset['train'][random_index]['video_path']
     print('error:', model_name, type, prompt)
-    return type
+    return None
 
-# Gradio 接口
-def display_video():
-    video_path = get_random_video()
-    return video_path
+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
+
+
+
+with gr.Blocks() as interface:
+    with gr.Row():
+        with gr.Column():
+            model_name_1_output = gr.Textbox(label="Model Name 1")
+            video_output_1 = gr.Video(label="Video 1")
+        with gr.Column():
+            model_name_2_output = gr.Textbox(label="Model Name 2")
+            video_output_2 = gr.Video(label="Video 2")
+    
+    type_output = gr.Textbox(label="Type")
+    prompt_output = gr.Textbox(label="Prompt")
 
-interface = gr.Interface(fn=display_video, 
-                         outputs=gr.Video(label="随机视频展示"),
-                         inputs=[], 
-                         title="随机视频展示",
-                         description="从 Vchitect/VBench_sampled_video 数据集中随机展示一个视频。")
+    display_button = gr.Button("Display Videos")
+    display_button.click(
+        fn=get_random_video,
+        outputs=[video_output_1, video_output_2, type_output, prompt_output, model_name_1_output, model_name_2_output]
+    )
 
-if __name__ == "__main__":
-    interface.launch()
+interface.launch()