add video option in pipeline

Former-commit-id: c7f057ced06adf0ba3d6f5d4040c09693f267eed

pipeline.py

@@ -8,17 +8,19 @@ import ffmpeg
 
 parser = argparse.ArgumentParser()
 parser.add_argument("--link", help="youtube video link here", default=None, type=str, required=False)
-parser.add_argument("--local_path", help="local video path here", default=None, type=str, required=False)
+parser.add_argument("--video_file", help="local video path here", default=None, type=str, required=False)
+parser.add_argument("--audio_file", help="local audio path here", default=None, type=str, required=False)
 parser.add_argument("--srt_file", help="srt file input path here", default=None, type=str, required=False)  # New argument
 parser.add_argument("--download", help="download path", default='./downloads', type=str, required=False)
 parser.add_argument("--output_dir", help="translate result path", default='./results', type=str, required=False)
 parser.add_argument("--video_name", help="video name, if use video link as input, the name will auto-filled by youtube video name", default='placeholder', type=str, required=False)
 parser.add_argument("--model_name", help="model name only support text-davinci-003 and gpt-3.5-turbo", type=str, required=False, default="gpt-3.5-turbo")
-parser.add_argument("-ass", help="set script output to .ass file", action='store_true')
+parser.add_argument("-only_srt", help="set script output to only .srt file", action='store_true')
+parser.add_argument("-v", help="auto encode script with video", action='store_true')
 args = parser.parse_args()
 
 # input should be either video file or youtube video link.
-if args.link is None and args.local_path is None and args.srt_file is None:
+if args.link is None and args.video_file is None and args.srt_file is None:
     print("need video source or srt file")
     exit()
 
@@ -38,30 +40,47 @@ VIDEO_NAME = args.video_name
 model_name = args.model_name
 
 # get source audio
-if args.link is not None and args.local_path is None:
+if args.link is not None and args.video_file is None:
     # Download audio from YouTube
     video_link = args.link
     video = None
     audio = None
     try:
-        video = YouTube(video_link)
-        audio = video.streams.filter(only_audio=True, file_extension='mp4').first()
+        yt = YouTube(video_link)
+        video = yt.streams.filter(progressive=True, file_extension='mp4').order_by('resolution').desc().first()
+        if video:
+            video.download(f'{DOWNLOAD_PATH}/video')
+            print('Video download completed!')
+        else:
+            print("Error: Video stream not found")
+        audio = yt.streams.filter(only_audio=True, file_extension='mp4').first()
         if audio:
             audio.download(f'{DOWNLOAD_PATH}/audio')
-            print('Download Completed!')
+            print('Audio download completed!')
         else:
             print("Error: Audio stream not found")
     except Exception as e:
         print("Connection Error")
         print(e) 
+        exit()
+    
+    video_path = f'{DOWNLOAD_PATH}/video/{video.default_filename}'
+    # video_file = open(video_path, "rb")
     audio_path = '{}/audio/{}'.format(DOWNLOAD_PATH, audio.default_filename)
     audio_file = open(audio_path, "rb")
     if VIDEO_NAME == 'placeholder':
         VIDEO_NAME = audio.default_filename.split('.')[0]
-elif args.local_path is not None:
+elif args.video_file is not None:
     # Read from local
-    audio_file= open(args.local_path, "rb")
-    audio_path = args.local_path
+    # video_file = open(args.video_file, "rb")
+    video_path = args.video_file
+    if args.audio_file is not None:
+        audio_file= open(args.audio_file, "rb")
+        audio_path = args.audio_file
+    else:
+        os.system(f'ffmpeg -i {args.video_file} -f mp3 -ab 192000 -vn {DOWNLOAD_PATH}/audio/{VIDEO_NAME}.mp3')
+        audio_file= open(f'{DOWNLOAD_PATH}/audio/{VIDEO_NAME}.mp3', "rb")
+        audio_path = f'{DOWNLOAD_PATH}/audio/{VIDEO_NAME}.mp3'
 
 if not os.path.exists(f'{RESULT_PATH}/{VIDEO_NAME}'):
     os.mkdir(f'{RESULT_PATH}/{VIDEO_NAME}')
@@ -94,7 +113,7 @@ else:
         script_en = f.read()
         script_input = script_en
 
-if args.ass:
+if not args.only_srt:
     from srt2ass import srt2ass
     assSub_en = srt2ass(srt_file_en, "default", "No", "Modest")
     print('ASS subtitle saved as: ' + assSub_en)
@@ -145,6 +164,13 @@ for s in script_arr:
 
         with open(f"{RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_zh.srt", 'a+') as f:
             f.write(response['choices'][0]['text'].strip())
-if args.ass:
+
+if not args.only_srt:
     assSub_zh = srt2ass(f"{RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_zh.srt", "default", "No", "Modest")
     print('ASS subtitle saved as: ' + assSub_zh)
+
+if args.v:
+    if args.only_srt:
+        os.system(f'ffmpeg -i {video_path} -vf "subtitles={RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_zh.srt" {RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}.mp4')
+    else:
+        os.system(f'ffmpeg -i {video_path} -vf "subtitles={RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_zh.ass" {RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}.mp4')