Update run.py

run.py

@@ -13,31 +13,27 @@ import os
 
 def load_hf_dataset(dataset_path, auth_token):
     dataset = load_dataset(dataset_path, token=auth_token)
-    
     video_paths = dataset
-    
     return video_paths
 
 def fast_caption(sys_prompt, usr_prompt, temp, top_p, max_tokens, model, key, endpoint, video_src, video_hf, video_hf_auth, parquet_index, video_od, video_od_auth, video_gd, video_gd_auth, frame_format, frame_limit):
+    progress_info = []
     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)
-    
         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
+        progress_info.append(f"Using model '{model}' with {len(frames)} frames extracted.")
+        return f"{caption}", "\n".join(progress_info), debug_image
     elif video_hf and video_hf_auth:
         current_file_path = os.path.abspath(__file__)
         current_directory = os.path.dirname(current_file_path)
-        print('begin video_hf')
-        # Process all videos in the dataset
+        progress_info.append('Begin processing Hugging Face dataset.')
         all_captions = []
         temp_parquet_file = hf_hub_download(
             repo_id=video_hf,
@@ -45,28 +41,22 @@ def fast_caption(sys_prompt, usr_prompt, temp, top_p, max_tokens, model, key, en
             repo_type="dataset",
             token=video_hf_auth,
         )
-        print(temp_parquet_file)
         parquet_file = pq.ParquetFile(temp_parquet_file)
-        
         for batch in parquet_file.iter_batches(batch_size=1):
             df = batch.to_pandas()
             video = df['video'][0]
-
             md5 = hashlib.md5(video).hexdigest()
-            print(md5)
             with tempfile.NamedTemporaryFile(dir=current_directory) as temp_file:
                 temp_file.write(video)
                 video_path = temp_file.name
-
                 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\n\n".join(all_captions), f"Processed {len(video_paths)} videos.", None
-
+                progress_info.append(f"Processed video with MD5: {md5}")
+        return "\n\n\n".join(all_captions), "\n".join(progress_info), None
     else:
         return "", "No video source selected.", None