Update run.py

run.py

@@ -18,53 +18,52 @@ def load_hf_dataset(dataset_path, auth_token):
 
 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 = []
-    with tempfile.NamedTemporaryFile(mode='w', delete=False, newline='') as csv_file:
-        csv_filename = csv_file.name
-        fieldnames = ['md5', 'caption']
-        writer = csv.DictWriter(csv_file, fieldnames=fieldnames)
-        writer.writeheader()
-
-        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)
-            progress_info.append(f"Using model '{model}' with {len(frames)} frames extracted.")
-            writer.writerow({'md5': 'single_video', 'caption': caption})
-            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)
-            progress_info.append('Begin processing Hugging Face dataset.')
-            temp_parquet_file = hf_hub_download(
-                repo_id=video_hf,
-                filename='data/' + str(parquet_index).zfill(6) + '.parquet',
-                repo_type="dataset",
-                token=video_hf_auth,
-            )
-            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()
-                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)
-                    writer.writerow({'md5': md5, 'caption': caption})
-                    progress_info.append(f"Processed video with MD5: {md5}")
-            return csv_filename, "\n".join(progress_info), None
-        else:
-            return "", "No video source selected.", None
+    with tempfile.TemporaryDirectory() as temp_dir:
+        csv_filename = os.path.join(temp_dir, str(parquet_index) + '_caption.csv')
+        with open(csv_filename, mode='w', newline='') as csv_file:
+            fieldnames = ['md5', 'caption']
+            writer = csv.DictWriter(csv_file, fieldnames=fieldnames)
+            writer.writeheader()
+    
+            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)
+                progress_info.append(f"Using model '{model}' with {len(frames)} frames extracted.")
+                writer.writerow({'md5': 'single_video', 'caption': caption})
+                return f"{caption}", "\n".join(progress_info), debug_image
+            elif video_hf and video_hf_auth:
+                progress_info.append('Begin processing Hugging Face dataset.')
+                temp_parquet_file = hf_hub_download(
+                    repo_id=video_hf,
+                    filename='data/' + str(parquet_index).zfill(6) + '.parquet',
+                    repo_type="dataset",
+                    token=video_hf_auth,
+                )
+                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()
+                    with tempfile.NamedTemporaryFile(dir=temp_dir) 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)
+                        writer.writerow({'md5': md5, 'caption': caption})
+                        progress_info.append(f"Processed video with MD5: {md5}")
+                return csv_filename, "\n".join(progress_info), None
+            else:
+                return "", "No video source selected.", None
 
 with gr.Blocks() as Core:
     with gr.Row(variant="panel"):