From c2c5cb5aafeb37f0ef2d3a44eb21778eafed42e1 Mon Sep 17 00:00:00 2001
From: xororz <xcandlrr@gmail.com>
Date: Sun, 14 Apr 2024 22:47:52 -0400
Subject: [PATCH] fix: resolve gpu memory leak

---
 src/upscale.ts | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/src/upscale.ts b/src/upscale.ts
index f4edcab..80e4a66 100644
--- a/src/upscale.ts
+++ b/src/upscale.ts
@@ -5,9 +5,13 @@ export default async function upscale(
   image: Image,
   model: any
 ): Promise<Image> {
-  let tensor = img2tensor(image);
-  let result = model.predict(tensor) as tf.Tensor;
-  let resultImage = await tensor2img(result);
+  const result = tf.tidy(() => {
+    const tensor = img2tensor(image);
+    const result = model.predict(tensor) as tf.Tensor;
+    return result;
+  });
+  const resultImage = await tensor2img(result);
+  tf.dispose(result);
   return resultImage;
 }
 
@@ -25,6 +29,7 @@ function img2tensor(image: Image): tf.Tensor {
 async function tensor2img(tensor: tf.Tensor): Promise<Image> {
   let [_, height, width, __] = tensor.shape;
   let arr = await tensor.data();
+  tensor.dispose();
   let clipped = new Uint8Array(
     arr.map((x) => {
       x = Math.min(1, Math.max(0, x));