Update inference.py

inference.py

@@ -53,6 +53,7 @@ def get_args():
     parser.add_argument("--video_length", type=int, default=15, help="Length of synthesized video")
     parser.add_argument("--height", type=int, default=512, help="Height of synthesized video, and should be a multiple of 32")
     parser.add_argument("--width", type=int, default=512, help="Width of synthesized video, and should be a multiple of 32")
+    parser.add_argument("--fps", type=int, default=8, help="FPS for final output")
     parser.add_argument("--smoother_steps", nargs='+', default=[19, 20], type=int, help="Timesteps at which using interleaved-frame smoother")
     parser.add_argument("--is_long_video", action='store_true', help="Whether to use hierarchical sampler to produce long video")
     parser.add_argument("--seed", type=int, default=42, help="Random seed of generator")
@@ -110,7 +111,7 @@ if __name__ == "__main__":
 
     # Save condition video
     video_cond = [np.array(p).astype(np.uint8) for p in pil_annotation]
-    imageio.mimsave(os.path.join(args.output_path, f"{args.condition}_condition.mp4"), video_cond, fps=8)
+    imageio.mimsave(os.path.join(args.output_path, f"{args.condition}_condition.mp4"), video_cond, fps=args.fps)
 
     # Reduce memory (optional)
     del annotator; torch.cuda.empty_cache()