From 80d394d64ac43e1a11201de24e6fa777b4d094e9 Mon Sep 17 00:00:00 2001 From: root Date: Wed, 6 Mar 2024 21:04:29 +0800 Subject: [PATCH] [fix]: disable gradient computation to save GPU memory when reconstructing a video Former-commit-id: b7bf763e7f8800959ed82bb89f41684db216dd8e --- opensora/models/ae/videobase/vqvae/videogpt/rec_video.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/opensora/models/ae/videobase/vqvae/videogpt/rec_video.py b/opensora/models/ae/videobase/vqvae/videogpt/rec_video.py index 0353334d2..5bc2d2f19 100644 --- a/opensora/models/ae/videobase/vqvae/videogpt/rec_video.py +++ b/opensora/models/ae/videobase/vqvae/videogpt/rec_video.py @@ -89,8 +89,9 @@ def main(args): x_vae = preprocess(read_video(video_path, num_frames, sample_rate), resolution, crop_size) x_vae = x_vae.to(device) - encodings, embeddings = vqvae.encode(x_vae, include_embeddings=True) - video_recon = vqvae.decode(encodings) + with torch.no_grad(): + encodings, embeddings = vqvae.encode(x_vae, include_embeddings=True) + video_recon = vqvae.decode(encodings) # custom_to_video(x_vae[0], fps=sample_fps/sample_rate, output_file='origin_input.mp4') custom_to_video(video_recon[0], fps=sample_fps/sample_rate, output_file=args.rec_path)