Update run.py

support hfdata

run.py

@@ -2,20 +2,47 @@
 import gradio as gr
 from utils import VideoProcessor, AzureAPI, GoogleAPI, AnthropicAPI, OpenAIAPI
 from constraint import SYS_PROMPT, USER_PROMPT
+from datasets import load_dataset
 
-def fast_caption(sys_prompt, usr_prompt, temp, top_p, max_tokens, model, key, endpoint, video, frame_format, frame_limit):
-    processor = VideoProcessor(frame_format=frame_format, frame_limit=frame_limit)
-    frames = processor._decode(video)
+def load_hf_dataset(dataset_path, auth_token):
+    dataset = load_dataset(dataset_path, use_auth_token=auth_token)
     
-    base64_list = processor.to_base64_list(frames)
-    debug_image = processor.concatenate(frames)
+    video_paths = dataset
+    
+    return video_paths
 
-    if not key or not endpoint:
-        return "", f"API key or endpoint is missing. Processed {len(frames)} frames.", debug_image
+def fast_caption(sys_prompt, usr_prompt, temp, top_p, max_tokens, model, key, endpoint, video_src, video_hf, video_hf_auth, video_od, video_od_auth, video_gd, video_gd_auth, frame_format, frame_limit):
+    if video_src:
+        video = video_src
+        processor = VideoProcessor(frame_format=frame_format, frame_limit=frame_limit)
+        frames = processor._decode(video)
+        
+        base64_list = processor.to_base64_list(frames)
+        debug_image = processor.concatenate(frames)
     
-    api = AzureAPI(key=key, endpoint=endpoint, model=model, temp=temp, top_p=top_p, max_tokens=max_tokens)
-    caption = api.get_caption(sys_prompt, usr_prompt, base64_list)
-    return f"{caption}", f"Using model '{model}' with {len(frames)} frames extracted.", debug_image
+        if not key or not endpoint:
+            return "", f"API key or endpoint is missing. Processed {len(frames)} frames.", debug_image
+        
+        api = AzureAPI(key=key, endpoint=endpoint, model=model, temp=temp, top_p=top_p, max_tokens=max_tokens)
+        caption = api.get_caption(sys_prompt, usr_prompt, base64_list)
+        return f"{caption}", f"Using model '{model}' with {len(frames)} frames extracted.", debug_image
+    elif video_hf and video_hf_auth:
+        # Handle Hugging Face dataset
+        video_paths = load_hf_dataset(video_hf, video_hf_auth)
+        # Process all videos in the dataset
+        all_captions = []
+        for video_path in video_paths:
+            if video_path.endswith('.mp4'):  # 假设我们只处理.mp4文件
+                processor = VideoProcessor(frame_format=frame_format, frame_limit=frame_limit)
+                frames = processor._decode(video_path)
+                base64_list = processor.to_base64_list(frames)
+                api = AzureAPI(key=key, endpoint=endpoint, model=model, temp=temp, top_p=top_p, max_tokens=max_tokens)
+                caption = api.get_caption(sys_prompt, usr_prompt, base64_list)
+                all_captions.append(caption)
+        return "\n".join(all_captions), f"Processed {len(video_paths)} videos.", None
+    # ... (Handle other sources)
+    else:
+        return "", "No video source selected.", None
 
 with gr.Blocks() as Core:
     with gr.Row(variant="panel"):
@@ -82,7 +109,7 @@ with gr.Blocks() as Core:
                 caption_button = gr.Button("Caption", variant="primary", size="lg")
         caption_button.click(
             fast_caption, 
-            inputs=[sys_prompt, usr_prompt, temp, top_p, max_tokens, model, key, endpoint, video_src, frame_format, frame_limit], 
+            inputs=[sys_prompt, usr_prompt, temp, top_p, max_tokens, model, key, endpoint, video_src, video_hf, video_hf_auth, video_od, video_od_auth, video_gd, video_gd_auth, frame_format, frame_limit], 
             outputs=[result, info, frame]
         )