From 22531b75c4f898b8910d501d1392c2f059aeb867 Mon Sep 17 00:00:00 2001
From: Dragon Fire <danielbodendahl@gmail.com>
Date: Wed, 8 May 2024 18:32:00 -0400
Subject: [PATCH] Fix

---
 structures/Tensorflow.js | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/structures/Tensorflow.js b/structures/Tensorflow.js
index fe4fe878..ea0f44b4 100644
--- a/structures/Tensorflow.js
+++ b/structures/Tensorflow.js
@@ -80,12 +80,13 @@ module.exports = class Tensorflow {
 		styleTensor.dispose();
 		const stylePrediction = await this.styleModel.predict(loadedStyle);
 		loadedStyle.dispose();
-		const stylizedImage = await this.transformerModel.predict([loadedImage, stylePrediction]);
+		const stylizedImage = (await this.transformerModel.predict([loadedImage, stylePrediction])).squeeze();
 		loadedImage.dispose();
 		stylePrediction.dispose();
-		stylizedImage.mul(255).toInt();
-		const buffer = await tf.node.encodeJpeg(stylizedImage.squeeze());
+		const encodedImg = stylizedImage.mul(255).toInt();
 		stylizedImage.dispose();
+		const buffer = await tf.node.encodeJpeg(encodedImg);
+		encodedImg.dispose();
 		return Buffer.from(buffer);
 	}
 };