From 081465d576775278a837e7519cc93faa14fd5ddd Mon Sep 17 00:00:00 2001
From: Haorui Jiang <143785706+HiramJiang@users.noreply.github.com>
Date: Wed, 25 Mar 2026 13:21:09 -0400
Subject: [PATCH 1/2] fix: register MultiplyBeta custom layer for ESRGAN
 medium/thick models

CDN-loaded models use a custom TF.js layer that wasn't registered,
causing "Unknown layer: MultiplyBeta" error on Balanced/Quality modes.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 src/workers/upscaler.worker.ts | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/src/workers/upscaler.worker.ts b/src/workers/upscaler.worker.ts
index dd1263d8..14c431b7 100644
--- a/src/workers/upscaler.worker.ts
+++ b/src/workers/upscaler.worker.ts
@@ -1,4 +1,27 @@
 import Upscaler from "upscaler";
+import * as tf from "@tensorflow/tfjs";
+
+// Register custom layer used by ESRGAN medium/thick models (loaded from CDN)
+// Without this, TF.js throws "Unknown layer: MultiplyBeta" when loading these models
+class MultiplyBeta extends tf.layers.Layer {
+  static className = "MultiplyBeta";
+  private beta: number;
+
+  constructor(config: Record<string, unknown> = {}) {
+    super(config);
+    this.beta = (config.beta as number) ?? 0.2;
+  }
+
+  call(inputs: tf.Tensor | tf.Tensor[]): tf.Tensor {
+    const input = Array.isArray(inputs) ? inputs[0] : inputs;
+    return tf.mul(input, tf.scalar(this.beta));
+  }
+
+  getConfig() {
+    return { ...super.getConfig(), beta: this.beta };
+  }
+}
+tf.serialization.registerClass(MultiplyBeta);
 
 type ModelType = "slim" | "medium" | "thick";
 type ScaleType = "2x" | "3x" | "4x";

From d92b91ae0ca2a3b5df4abc6fc00b4b1b02188eda Mon Sep 17 00:00:00 2001
From: Haorui Jiang <143785706+HiramJiang@users.noreply.github.com>
Date: Wed, 25 Mar 2026 13:46:46 -0400
Subject: [PATCH 2/2] fix: register PixelShuffle custom layers for ESRGAN thick
 models

Thick model also uses PixelShuffle2x/3x/4x layers beyond MultiplyBeta.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 src/workers/upscaler.worker.ts | 32 ++++++++++++++++++++++++++++++++
 1 file changed, 32 insertions(+)

diff --git a/src/workers/upscaler.worker.ts b/src/workers/upscaler.worker.ts
index 14c431b7..1eafa58d 100644
--- a/src/workers/upscaler.worker.ts
+++ b/src/workers/upscaler.worker.ts
@@ -23,6 +23,38 @@ class MultiplyBeta extends tf.layers.Layer {
 }
 tf.serialization.registerClass(MultiplyBeta);
 
+// PixelShuffle layer used by ESRGAN thick models — does depth-to-space rearrangement
+function createPixelShuffleClass(scale: number) {
+  class PixelShuffle extends tf.layers.Layer {
+    static className = `PixelShuffle${scale}x`;
+    private scale: number;
+
+    constructor(config: Record<string, unknown> = {}) {
+      super(config);
+      this.scale = scale;
+    }
+
+    computeOutputShape(inputShape: Array<number | null>): Array<number | null> {
+      return [inputShape[0], inputShape[1], inputShape[2], 3];
+    }
+
+    call(inputs: tf.Tensor | tf.Tensor[]): tf.Tensor {
+      const input = Array.isArray(inputs) ? inputs[0] : inputs;
+      return tf.depthToSpace(input as tf.Tensor4D, this.scale, "NHWC");
+    }
+
+    getConfig() {
+      return { ...super.getConfig(), scale: this.scale };
+    }
+  }
+  return PixelShuffle;
+}
+
+// Register PixelShuffle for all supported scales
+[2, 3, 4].forEach((s) => {
+  tf.serialization.registerClass(createPixelShuffleClass(s));
+});
+
 type ModelType = "slim" | "medium" | "thick";
 type ScaleType = "2x" | "3x" | "4x";