Update generate.py

generate.py

@@ -26,7 +26,8 @@ def main():
     parser.add_argument("--frames", type=int, default=25, help="frame count")
     parser.add_argument("--fps", type=int, default=24, help="output fps")
     parser.add_argument("--seed", type=int, default=None, help="random seed")
-    parser.add_argument("--crf", type=int, default=18, help="video quality (0-51, lower=better, 18=visually lossless)")
+    parser.add_argument("--crf", type=int, default=10, help="video quality (0-51, lower=better)")
+    parser.add_argument("--prores", action="store_true", help="use prores codec (large files, best quality)")
 
     args = parser.parse_args()
 
@@ -88,18 +89,31 @@ def main():
         frame = np.array(frame, dtype=np.uint8)
         frames.append(frame)
 
-    # export with high quality h264
-    writer = imageio.get_writer(
-        args.output,
-        fps=args.fps,
-        codec='libx264',
-        quality=None,
-        pixelformat='yuv420p',
-        output_params=['-crf', str(args.crf), '-preset', 'slow']
-    )
+    # export video
+    if args.prores:
+        output_path = args.output.replace('.mp4', '.mov') if args.output.endswith('.mp4') else args.output
+        writer = imageio.get_writer(
+            output_path,
+            fps=args.fps,
+            codec='prores_ks',
+            pixelformat='yuv422p10le',
+            output_params=['-profile:v', '3']  # prores hq
+        )
+    else:
+        output_path = args.output
+        writer = imageio.get_writer(
+            output_path,
+            fps=args.fps,
+            codec='libx264',
+            quality=None,
+            pixelformat='yuv420p',
+            output_params=['-crf', str(args.crf), '-preset', 'slow']
+        )
+
     for frame in frames:
         writer.append_data(frame)
     writer.close()
+    args.output = output_path
 
     print(f"\nsaved to: {args.output}")
     print(f"seed: {args.seed}")