Update inference.py

inference.py

@@ -101,8 +101,8 @@ if __name__ == "__main__":
     video = read_video(video_path=args.video_path, video_length=args.video_length, width=args.width, height=args.height)
 
     # Save source video
-    original_pixels = rearrange(video, "(b f) c h w -> b c f h w", b=1)
-    save_videos_grid(original_pixels, os.path.join(args.output_path, "source_video.mp4"), rescale=True)
+    # original_pixels = rearrange(video, "(b f) c h w -> b c f h w", b=1)
+    # save_videos_grid(original_pixels, os.path.join(args.output_path, "source_video.mp4"), rescale=True)
 
 
     # Step 2. Parse a video to conditional frames
@@ -111,11 +111,11 @@ if __name__ == "__main__":
         pil_annotation = [pil_annot[0] for pil_annot in pil_annotation]
 
     # 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=args.fps)
+    #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=args.fps)
 
     # Reduce memory (optional)
-    del annotator; torch.cuda.empty_cache()
+    #del annotator; torch.cuda.empty_cache()
 
     # Step 3. inference
 
@@ -132,4 +132,6 @@ if __name__ == "__main__":
                     generator=generator, guidance_scale=12.5, negative_prompt=NEG_PROMPT,
                     width=args.width, height=args.height
                 ).videos
-    save_videos_grid(sample, f"{args.output_path}/{args.temp_chunk_path}.mp4")
+    save_videos_grid(sample, f"{args.output_path}/{args.temp_chunk_path}.mp4")
+    del pipe
+    torch.cuda.empty_cache()