Create main.py

main.py

@@ -0,0 +1,104 @@
+import argparse
+import os
+import cv2
+
+from dall_e import DalleImageGenerator
+from download_from_youtube import download_transcription, download_mp3, get_video_name
+from utils import clean_str, get_sqrt, read_and_preprocess_transcript, put_subtitles_on_frame
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Generate a Dall-e video clip from a youtube url')
+    parser.add_argument('url', help='URL of the a song to get the audio from')
+    parser.add_argument('--token', help='A Replicate API token',  default=None)
+    parser.add_argument('--song_name', help='The name of the song / output files', default=None)
+    parser.add_argument('--fps', help='FPS higher then 1 / sec_per_img to better control over timing', default=10)
+    parser.add_argument('--sec_per_img', help='How long to show each image', default=3)
+    parser.add_argument('--n_lines', help='Limit number of generated lines in the video to faster testing', default=None)
+    args = parser.parse_args()
+
+    dalle = DalleImageGenerator(token=args.token)
+    img_dim = 256 # Dall-e's output dim
+    resize_factor = 2 # upacale factor for frames
+
+    if args.song_name is None:
+        args.song_name = clean_str(get_video_name(args.url))
+
+    # Set paths
+    outputs_dir = f"data/{args.song_name}"
+    frames_dir = f"{outputs_dir}/frames"
+    os.makedirs(outputs_dir, exist_ok=True)
+    os.makedirs(frames_dir, exist_ok=True)
+    mp3_path = f"{outputs_dir}/audio.mp3"
+    transcript_path = f"{outputs_dir}/transcript.json"
+    vid_path = f"{outputs_dir}/frames.avi"
+    final_vid_path = f"{outputs_dir}/final.avi"
+
+    print("Getting audio file and transcript from youtube")
+    # Download data
+    download_transcription(args.url, transcript_path)
+    download_mp3(args.url, mp3_path)
+
+    transcript = read_and_preprocess_transcript(transcript_path, args.song_name, args.n_lines)
+
+    frames, video_duration = get_frames(dalle, transcript, resize_factor, args)
+
+    # Write video
+    video = cv2.VideoWriter(vid_path, 0, args.fps, (img_dim * resize_factor, img_dim * resize_factor))
+    for i, frame in enumerate(frames):
+        cv2.imwrite(f"{frames_dir}/frame-{i}.png", frame)
+        video.write(frame)
+    cv2.destroyAllWindows()
+    video.release()
+
+    # Mix video clip with audio
+    os.system(f"ffmpeg -ss 00:00:00 -t {video_duration} -i '{mp3_path}' -map 0:a -acodec libmp3lame '{f'data/{args.song_name}/tmp.mp3'}'")
+    os.system(f"ffmpeg -i '{vid_path}' -i '{f'data/{args.song_name}/tmp.mp3'}' -map 0 -map 1:a -c:v copy -shortest '{final_vid_path}'")
+    print(f"Final video available at: {final_vid_path}")
+
+
+def get_frames(dalle, transcript, resize_factor, args):
+    """For each line in the transcript prompt dall-e mini to get images and duplicate them in the correct FPS for writing as a video"""
+    print("Building video-clip")
+    frames = []
+    video_duration = 0
+    for line in transcript:
+        text = clean_str(line['text'])
+        # start = min(video_duration, start)
+        start = line['start']
+        duration = line['duration']
+
+        # Dall-e generatees grid_size**2 images
+        grid_size = max(get_sqrt(duration / args.sec_per_img), 1)
+
+        print(f"({start:.1f} - {start + duration:.1f}):")
+
+        print(f"* Generating {grid_size**2} images with prompt: '{text}'")
+        # Generate images
+        images = dalle.generate_images(text, grid_size, text_adherence=3)
+
+        # Write frames
+        segment_duration = 0
+        frames_per_image = int(duration * args.fps) // len(images)
+        for j in range(len(images)):
+            frame = cv2.cvtColor(images[j], cv2.COLOR_RGBA2BGR)
+            frame = put_subtitles_on_frame(frame, text, resize_factor)
+            print(f"* Writing image - {j} as {frames_per_image} frames")
+            for _ in range(frames_per_image):
+                frames.append(frame)
+                segment_duration += 1 / args.fps
+
+        # Write more frames from last image to fill the gap
+        if segment_duration < duration:
+            n_frames = int((duration - segment_duration) * args.fps)
+            print(f"* Writing image - {j} for {n_frames} frames")
+            for _ in range(n_frames):
+                frames.append(frame)
+                segment_duration += 1 / args.fps
+        video_duration += segment_duration
+
+    return frames, video_duration
+
+
+if __name__ == '__main__':
+    main()