From 640d9cd52b1aa4450bba80db6ba2036aea4a1b5f Mon Sep 17 00:00:00 2001
From: k4yt3x <i@k4yt3x.com>
Date: Mon, 21 Oct 2024 00:00:00 +0000
Subject: [PATCH] fix(encoder): fixed memory leak in the encoder

Signed-off-by: k4yt3x <i@k4yt3x.com>
---
 Makefile        | 44 ++++++++++++++++++++++++++++++++------------
 src/encoder.cpp | 11 ++++++++---
 2 files changed, 40 insertions(+), 15 deletions(-)

diff --git a/Makefile b/Makefile
index d9b02e5..b017947 100644
--- a/Makefile
+++ b/Makefile
@@ -1,9 +1,15 @@
-.PHONY: build static debug windows test-realesrgan test-libplacebo leakcheck clean
+.PHONY: build static debug clean \
+	test-realesrgan test-libplacebo \
+	memcheck-realesrgan memcheck-libplacebo \
+	heaptrack-realesrgan heaptrack-libplacebo
 
 BINDIR=build
 CC=clang
 CXX=clang++
 
+TEST_VIDEO=data/standard-test.mp4
+TEST_OUTPUT=data/output.mp4
+
 build:
 	cmake -S . -B $(BINDIR) \
 		-DCMAKE_EXPORT_COMPILE_COMMANDS=ON \
@@ -48,21 +54,24 @@ debian:
 		libspdlog-dev \
 		libopencv-dev
 	cmake -B /tmp/build -S . -DUSE_SYSTEM_NCNN=OFF \
-		-DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++ \
+		-DCMAKE_C_COMPILER=$(CC) -DCMAKE_CXX_COMPILER=$(CXX) \
 		-DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=/tmp/install \
 		-DINSTALL_BIN_DESTINATION=. -DINSTALL_INCLUDE_DESTINATION=include \
 		-DINSTALL_LIB_DESTINATION=. -DINSTALL_MODEL_DESTINATION=.
 	cmake --build /tmp/build --config Release --target install --parallel
 
+clean:
+	rm -rf $(BINDIR)
+
 test-realesrgan:
-	LD_LIBRARY_PATH=$(BINDIR) $(BINDIR)/video2x -i data/standard-test.mp4 -o data/output.mp4 \
-		-f realesrgan -r 4 --model realesr-animevideov3
+	LD_LIBRARY_PATH=$(BINDIR) $(BINDIR)/video2x -i $(TEST_VIDEO) -o $(TEST_OUTPUT) \
+		-f realesrgan -r 4 -m realesr-animevideov3
 
 test-libplacebo:
-	LD_LIBRARY_PATH=$(BINDIR) $(BINDIR)/video2x -i data/standard-test.mp4 -o data/output.mp4 \
+	LD_LIBRARY_PATH=$(BINDIR) $(BINDIR)/video2x -i $(TEST_VIDEO) -o $(TEST_OUTPUT) \
 		-f libplacebo -w 1920 -h 1080 -s anime4k-mode-a
 
-leakcheck-realesrgan:
+memcheck-realesrgan:
 	LD_LIBRARY_PATH=$(BINDIR) valgrind \
 		--tool=memcheck \
 		--leak-check=full \
@@ -71,11 +80,11 @@ leakcheck-realesrgan:
 		--show-reachable=yes \
 		--verbose --log-file="valgrind.log" \
 		$(BINDIR)/video2x \
-		-i data/standard-test.mp4 -o data/output.mp4 \
-		-f realesrgan -r 2 --model realesr-animevideov3 \
+		-i $(TEST_VIDEO) -o $(TEST_OUTPUT) \
+		-f realesrgan -r 2 -m realesr-animevideov3 \
 		-p veryfast -b 1000000 -q 30
 
-leakcheck-libplacebo:
+memcheck-libplacebo:
 	LD_LIBRARY_PATH=$(BINDIR) valgrind \
 		--tool=memcheck \
 		--leak-check=full \
@@ -84,9 +93,20 @@ leakcheck-libplacebo:
 		--show-reachable=yes \
 		--verbose --log-file="valgrind.log" \
 		$(BINDIR)/video2x \
-		-i data/standard-test.mp4 -o data/output.mp4 \
+		-i $(TEST_VIDEO) -o $(TEST_OUTPUT) \
 		-f libplacebo -w 1920 -h 1080 -s anime4k-mode-a \
 		-p veryfast -b 1000000 -q 30
 
-clean:
-	rm -rf $(BINDIR)
+heaptrack-realesrgan:
+	LD_LIBRARY_PATH=$(BINDIR) HEAPTRACK_ENABLE_DEBUGINFOD=1 heaptrack \
+		$(BINDIR)/video2x \
+		-i $(TEST_VIDEO) -o $(TEST_OUTPUT) \
+		-f realesrgan -r 4 -m realesr-animevideov3 \
+		-p veryfast -b 1000000 -q 30
+
+heaptrack-libplacebo:
+	LD_LIBRARY_PATH=$(BINDIR) HEAPTRACK_ENABLE_DEBUGINFOD=1 heaptrack \
+		$(BINDIR)/video2x \
+		-i $(TEST_VIDEO) -o $(TEST_OUTPUT) \
+		-f libplacebo -w 1920 -h 1080 -s anime4k-mode-a \
+		-p veryfast -b 1000000 -q 30
diff --git a/src/encoder.cpp b/src/encoder.cpp
index 5c35a5f..0486e75 100644
--- a/src/encoder.cpp
+++ b/src/encoder.cpp
@@ -183,18 +183,18 @@ int encode_and_write_frame(
     AVFormatContext *ofmt_ctx,
     int vstream_idx
 ) {
+    AVFrame *converted_frame = nullptr;
     int ret;
 
     // Convert the frame to the encoder's pixel format if needed
     if (frame->format != enc_ctx->pix_fmt) {
-        AVFrame *converted_frame = convert_avframe_pix_fmt(frame, enc_ctx->pix_fmt);
+        converted_frame = convert_avframe_pix_fmt(frame, enc_ctx->pix_fmt);
         if (!converted_frame) {
             spdlog::error("Error converting frame to encoder's pixel format");
             return AVERROR_EXTERNAL;
         }
 
         converted_frame->pts = frame->pts;
-        frame = converted_frame;
     }
 
     AVPacket *enc_pkt = av_packet_alloc();
@@ -203,7 +203,12 @@ int encode_and_write_frame(
         return AVERROR(ENOMEM);
     }
 
-    ret = avcodec_send_frame(enc_ctx, frame);
+    if (converted_frame != nullptr) {
+        ret = avcodec_send_frame(enc_ctx, converted_frame);
+        av_frame_free(&converted_frame);
+    } else {
+        ret = avcodec_send_frame(enc_ctx, frame);
+    }
     if (ret < 0) {
         spdlog::error("Error sending frame to encoder");
         av_packet_free(&enc_pkt);