This commit is contained in:
Dragon Fire
2024-05-08 18:08:42 -04:00
parent 9e80e31a1d
commit 73d69a46d4
+3 -1
View File
@@ -78,8 +78,10 @@ module.exports = class Tensorflow {
const styleTensor = tf.node.decodeImage(styleImg, 3);
const loadedStyle = styleTensor.div(tf.scalar(255)).expandDims();
styleTensor.dispose();
const stylePrediction = await this.styleModel.predict(loadedStyle);
const reshapedStyle = loadedStyle.reshape([1, 1, 1, 100]);
loadedStyle.dispose();
const stylePrediction = await this.styleModel.predict(reshapedStyle);
reshapedStyle.dispose();
const stylizedImage = await this.transformerModel.predict([loadedImage, stylePrediction.squeeze()]);
loadedImage.dispose();
stylePrediction.dispose();