@pilio/gemini-watermark-remover 1.0.19 → 1.0.21

package/src/core/allenkFdncnnOnnxExport.js

@@ -0,0 +1,337 @@
+import { halfToFloat } from './allenkFdncnnNcnnModel.js';
+
+const ONNX_IR_VERSION = 8;
+const ONNX_OPSET_VERSION = 13;
+const ONNX_TENSOR_FLOAT = 1;
+const ONNX_ATTR_INTS = 7;
+
+const WIRE_VARINT = 0;
+const WIRE_LENGTH_DELIMITED = 2;
+
+function normalizeBytes(bytes) {
+    if (bytes instanceof Uint8Array) return bytes;
+    return new Uint8Array(bytes);
+}
+
+function concatBytes(parts) {
+    const total = parts.reduce((sum, part) => sum + part.length, 0);
+    const out = new Uint8Array(total);
+    let offset = 0;
+    for (const part of parts) {
+        out.set(part, offset);
+        offset += part.length;
+    }
+    return out;
+}
+
+function encodeVarint(value) {
+    let next = BigInt(value);
+    if (next < 0n) {
+        throw new Error(`Cannot encode negative varint: ${value}`);
+    }
+
+    const bytes = [];
+    while (next >= 0x80n) {
+        bytes.push(Number((next & 0x7fn) | 0x80n));
+        next >>= 7n;
+    }
+    bytes.push(Number(next));
+    return Uint8Array.from(bytes);
+}
+
+function fieldKey(fieldNumber, wireType) {
+    return encodeVarint((BigInt(fieldNumber) << 3n) | BigInt(wireType));
+}
+
+function varintField(fieldNumber, value) {
+    return concatBytes([fieldKey(fieldNumber, WIRE_VARINT), encodeVarint(value)]);
+}
+
+function int64Field(fieldNumber, value) {
+    return varintField(fieldNumber, BigInt(value));
+}
+
+function bytesField(fieldNumber, bytesLike) {
+    const bytes = normalizeBytes(bytesLike);
+    return concatBytes([
+        fieldKey(fieldNumber, WIRE_LENGTH_DELIMITED),
+        encodeVarint(bytes.length),
+        bytes
+    ]);
+}
+
+function stringField(fieldNumber, value) {
+    return bytesField(fieldNumber, new TextEncoder().encode(String(value)));
+}
+
+function messageField(fieldNumber, message) {
+    return bytesField(fieldNumber, message);
+}
+
+function message(parts) {
+    return concatBytes(parts.filter(Boolean));
+}
+
+function float32RawData(values) {
+    const out = new Uint8Array(values.length * 4);
+    const view = new DataView(out.buffer);
+    for (let i = 0; i < values.length; i++) {
+        view.setFloat32(i * 4, values[i], true);
+    }
+    return out;
+}
+
+function readFloat32LE(bytes, offset) {
+    const buffer = normalizeBytes(bytes);
+    return new DataView(buffer.buffer, buffer.byteOffset + offset, 4).getFloat32(0, true);
+}
+
+function readUInt16LE(bytes, offset) {
+    const buffer = normalizeBytes(bytes);
+    return new DataView(buffer.buffer, buffer.byteOffset + offset, 2).getUint16(0, true);
+}
+
+function createTensorProto({ name, dims, rawData }) {
+    return message([
+        ...dims.map((dim) => int64Field(1, dim)),
+        varintField(2, ONNX_TENSOR_FLOAT),
+        stringField(8, name),
+        bytesField(9, rawData)
+    ]);
+}
+
+function createShapeProto(dims) {
+    return message(dims.map((dim) => messageField(1, message([int64Field(1, dim)]))));
+}
+
+function createValueInfoProto({ name, dims }) {
+    const tensorType = message([
+        varintField(1, ONNX_TENSOR_FLOAT),
+        messageField(2, createShapeProto(dims))
+    ]);
+    const typeProto = message([messageField(1, tensorType)]);
+
+    return message([
+        stringField(1, name),
+        messageField(2, typeProto)
+    ]);
+}
+
+function createIntsAttribute(name, values) {
+    return message([
+        stringField(1, name),
+        ...values.map((value) => int64Field(8, value)),
+        varintField(20, ONNX_ATTR_INTS)
+    ]);
+}
+
+function createNodeProto({ name, opType, inputs, outputs, attributes = [] }) {
+    return message([
+        ...inputs.map((input) => stringField(1, input)),
+        ...outputs.map((output) => stringField(2, output)),
+        stringField(3, name),
+        stringField(4, opType),
+        ...attributes.map((attribute) => messageField(5, attribute))
+    ]);
+}
+
+function decodeSegmentWeightsToFloat32(bin, segment) {
+    const values = new Array(segment.weightCount);
+    for (let i = 0; i < segment.weightCount; i++) {
+        values[i] = halfToFloat(readUInt16LE(bin, segment.weightOffset + i * 2));
+    }
+    return float32RawData(values);
+}
+
+function decodeSegmentBiasToFloat32(bin, segment) {
+    const values = new Array(segment.biasCount);
+    for (let i = 0; i < segment.biasCount; i++) {
+        values[i] = readFloat32LE(bin, segment.biasOffset + i * 4);
+    }
+    return float32RawData(values);
+}
+
+function createInitializers({ bin, segments }) {
+    const initializers = [];
+
+    for (let i = 0; i < segments.length; i++) {
+        const segment = segments[i];
+        initializers.push({
+            name: `conv${i + 1}.weight`,
+            tensor: createTensorProto({
+                name: `conv${i + 1}.weight`,
+                dims: [
+                    segment.outputChannels,
+                    segment.inputChannels,
+                    segment.kernelH,
+                    segment.kernelW
+                ],
+                rawData: decodeSegmentWeightsToFloat32(bin, segment)
+            })
+        });
+        initializers.push({
+            name: `conv${i + 1}.bias`,
+            tensor: createTensorProto({
+                name: `conv${i + 1}.bias`,
+                dims: [segment.biasCount],
+                rawData: decodeSegmentBiasToFloat32(bin, segment)
+            })
+        });
+    }
+
+    return initializers;
+}
+
+function createNodes({ segments, inputName, outputName }) {
+    const nodes = [];
+    let previous = inputName;
+
+    for (let i = 0; i < segments.length; i++) {
+        const segment = segments[i];
+        const index = i + 1;
+        const isLast = i === segments.length - 1;
+        const convOutput = isLast ? outputName : `conv${index}.out`;
+        const reluOutput = `relu${index}.out`;
+
+        nodes.push(createNodeProto({
+            name: `conv${index}`,
+            opType: 'Conv',
+            inputs: [
+                previous,
+                `conv${index}.weight`,
+                `conv${index}.bias`
+            ],
+            outputs: [convOutput],
+            attributes: [
+                createIntsAttribute('kernel_shape', [segment.kernelH, segment.kernelW]),
+                createIntsAttribute('pads', [segment.padH, segment.padW, segment.padH, segment.padW]),
+                createIntsAttribute('strides', [segment.strideH, segment.strideW])
+            ]
+        }));
+
+        if (!isLast && segment.activationType === 1) {
+            nodes.push(createNodeProto({
+                name: `relu${index}`,
+                opType: 'Relu',
+                inputs: [convOutput],
+                outputs: [reluOutput]
+            }));
+            previous = reluOutput;
+        } else {
+            previous = convOutput;
+        }
+    }
+
+    return nodes;
+}
+
+function createGraphProto({
+    bin,
+    name,
+    segments,
+    inputName,
+    outputName,
+    roiSize
+}) {
+    const initializers = createInitializers({ bin, segments });
+    const nodes = createNodes({ segments, inputName, outputName });
+    const inputShape = [1, segments[0].inputChannels, roiSize, roiSize];
+    const outputShape = [1, segments[segments.length - 1].outputChannels, roiSize, roiSize];
+
+    return {
+        nodeCount: nodes.length,
+        initializerCount: initializers.length,
+        bytes: message([
+            ...nodes.map((node) => messageField(1, node)),
+            stringField(2, name),
+            ...initializers.map((initializer) => messageField(5, initializer.tensor)),
+            messageField(11, createValueInfoProto({ name: inputName, dims: inputShape })),
+            messageField(12, createValueInfoProto({ name: outputName, dims: outputShape }))
+        ])
+    };
+}
+
+function createOpsetImport(version = ONNX_OPSET_VERSION) {
+    return message([int64Field(2, version)]);
+}
+
+function createModelProto({
+    graph,
+    irVersion = ONNX_IR_VERSION,
+    opsetVersion = ONNX_OPSET_VERSION,
+    producerName = 'gemini-watermark-remover',
+    modelVersion = 1
+}) {
+    return message([
+        int64Field(1, irVersion),
+        stringField(2, producerName),
+        int64Field(5, modelVersion),
+        messageField(7, graph),
+        messageField(8, createOpsetImport(opsetVersion))
+    ]);
+}
+
+function exportAllenkFdncnnOnnx({
+    bin,
+    weightLayout,
+    roiSize = 72,
+    inputName = 'fdncnn_input',
+    outputName = 'fdncnn_output',
+    graphName = 'allenk_fdncnn_color',
+    opsetVersion = ONNX_OPSET_VERSION,
+    irVersion = ONNX_IR_VERSION
+} = {}) {
+    const segments = weightLayout?.segments || [];
+    if (!segments.length) {
+        throw new Error('allenk FDnCNN ONNX export requires weightLayout.segments');
+    }
+    if (!Number.isInteger(roiSize) || roiSize <= 0) {
+        throw new Error(`Invalid ONNX ROI size: ${roiSize}`);
+    }
+
+    const graph = createGraphProto({
+        bin: normalizeBytes(bin),
+        name: graphName,
+        segments,
+        inputName,
+        outputName,
+        roiSize
+    });
+    const model = createModelProto({
+        graph: graph.bytes,
+        irVersion,
+        opsetVersion
+    });
+
+    return {
+        bytes: model,
+        metadata: {
+            format: 'onnx',
+            tensorDataType: 'float32-raw-data',
+            irVersion,
+            opsetVersion,
+            graphName,
+            inputName,
+            outputName,
+            inputShape: [1, segments[0].inputChannels, roiSize, roiSize],
+            outputShape: [1, segments[segments.length - 1].outputChannels, roiSize, roiSize],
+            nodeCount: graph.nodeCount,
+            convolutionNodeCount: segments.length,
+            reluNodeCount: segments.filter((segment, index) => index < segments.length - 1 && segment.activationType === 1).length,
+            initializerCount: graph.initializerCount,
+            roiSize
+        }
+    };
+}
+
+export {
+    ONNX_IR_VERSION,
+    ONNX_OPSET_VERSION,
+    concatBytes,
+    createIntsAttribute,
+    createNodeProto,
+    createTensorProto,
+    encodeVarint,
+    exportAllenkFdncnnOnnx,
+    varintField
+};

package/src/core/allenkFdncnnOnnxRuntime.js

@@ -0,0 +1,138 @@
+import * as wasmOrt from 'onnxruntime-web/wasm';
+
+import {
+    buildAllenkFdncnnInput,
+    convertAllenkFdncnnOutputToRgba
+} from './allenkFdncnnDenoise.js';
+
+function getNow() {
+    return globalThis.performance?.now ? globalThis.performance.now() : Date.now();
+}
+
+function normalizeShape(shape, fallback = null) {
+    if (!Array.isArray(shape)) return fallback;
+    const normalized = shape.map((value) => Number(value));
+    return normalized.every((value) => Number.isInteger(value) && value > 0)
+        ? normalized
+        : fallback;
+}
+
+function validateImageShape(imageData, inputShape) {
+    const expectedHeight = inputShape?.[2];
+    const expectedWidth = inputShape?.[3];
+    if (!imageData?.data || imageData.width <= 0 || imageData.height <= 0) {
+        throw new Error('allenk FDnCNN ONNX runtime requires ImageData-like input');
+    }
+    if (imageData.width !== expectedWidth || imageData.height !== expectedHeight) {
+        const error = new Error(
+            `allenk FDnCNN ONNX runtime expected ${expectedWidth}x${expectedHeight}, got ${imageData.width}x${imageData.height}`
+        );
+        error.code = 'ALLENK_FDNCNN_ONNX_SHAPE_MISMATCH';
+        error.expectedWidth = expectedWidth;
+        error.expectedHeight = expectedHeight;
+        error.actualWidth = imageData.width;
+        error.actualHeight = imageData.height;
+        throw error;
+    }
+}
+
+async function createAllenkFdncnnOnnxRuntime({
+    ort = null,
+    modelBytes,
+    session = null,
+    executionProvider = 'wasm',
+    inputName = 'fdncnn_input',
+    outputName = 'fdncnn_output',
+    inputShape = [1, 4, 72, 72],
+    outputShape = [1, 3, 72, 72],
+    graphOptimizationLevel = 'all',
+    wasmPaths = null,
+    numThreads = 'auto'
+} = {}) {
+    if (executionProvider === 'webgpu' && !ort) {
+        throw new Error('allenk FDnCNN WebGPU runtime requires an injected onnxruntime-web/webgpu module');
+    }
+    const resolvedOrt = ort || wasmOrt;
+    const resolvedInputShape = normalizeShape(inputShape, [1, 4, 72, 72]);
+    const resolvedOutputShape = normalizeShape(outputShape, [1, 3, resolvedInputShape[2], resolvedInputShape[3]]);
+    if (resolvedOrt.env?.wasm) {
+        const canUseThreads = Boolean(globalThis.crossOriginIsolated && typeof SharedArrayBuffer !== 'undefined');
+        const hardwareThreads = Number(globalThis.navigator?.hardwareConcurrency) || 1;
+        const resolvedNumThreads = numThreads === 'auto'
+            ? (canUseThreads ? Math.max(1, Math.min(4, Math.ceil(hardwareThreads / 2))) : 1)
+            : Math.max(1, Math.round(Number(numThreads) || 1));
+        resolvedOrt.env.wasm.numThreads = executionProvider === 'wasm' ? resolvedNumThreads : 1;
+        resolvedOrt.env.wasm.proxy = false;
+        if (wasmPaths) {
+            resolvedOrt.env.wasm.wasmPaths = wasmPaths;
+        }
+    }
+    if (executionProvider === 'webgpu' && resolvedOrt.env?.webgpu) {
+        resolvedOrt.env.webgpu.powerPreference = 'high-performance';
+    }
+
+    const createStarted = getNow();
+    const resolvedSession = session || await resolvedOrt.InferenceSession.create(modelBytes, {
+        executionProviders: [executionProvider],
+        graphOptimizationLevel
+    });
+    const createMs = getNow() - createStarted;
+
+    return {
+        id: `allenk-fdncnn-onnx-${executionProvider}`,
+        status: 'prototype',
+        executionProvider,
+        inputName,
+        outputName,
+        inputShape: resolvedInputShape,
+        outputShape: resolvedOutputShape,
+        createMs,
+        session: resolvedSession,
+        numThreads: executionProvider === 'wasm'
+            ? resolvedOrt.env?.wasm?.numThreads ?? null
+            : null,
+        estimateMacs(width, height) {
+            // Matches the decoded allenk FDnCNN graph: 4->64, 18x 64->64, 64->3, all 3x3.
+            return width * height * ((4 * 64 * 9) + (18 * 64 * 64 * 9) + (64 * 3 * 9));
+        },
+        async execute(input) {
+            const tensor = new resolvedOrt.Tensor('float32', input, resolvedInputShape);
+            const started = getNow();
+            const outputs = await resolvedSession.run({ [inputName]: tensor });
+            const runMs = getNow() - started;
+            const outputTensor = outputs[outputName];
+            if (!outputTensor?.data) {
+                throw new Error(`allenk FDnCNN ONNX runtime did not return ${outputName}`);
+            }
+            return {
+                output: outputTensor.data,
+                outputShape: [...outputTensor.dims],
+                runtime: this.id,
+                runMs,
+                macs: this.estimateMacs(resolvedInputShape[3], resolvedInputShape[2])
+            };
+        },
+        async denoiseImageData({ imageData, sigma } = {}) {
+            validateImageShape(imageData, resolvedInputShape);
+            const input = buildAllenkFdncnnInput({ imageData, sigma });
+            const result = await this.execute(input);
+            return {
+                ...result,
+                imageData: {
+                    width: imageData.width,
+                    height: imageData.height,
+                    data: convertAllenkFdncnnOutputToRgba({
+                        output: result.output,
+                        width: imageData.width,
+                        height: imageData.height
+                    })
+                }
+            };
+        }
+    };
+}
+
+export {
+    createAllenkFdncnnOnnxRuntime,
+    validateImageShape
+};

package/src/core/allenkFdncnnReferenceRuntime.js

@@ -0,0 +1,169 @@
+import {
+    buildAllenkFdncnnInput,
+    convertAllenkFdncnnOutputToRgba
+} from './allenkFdncnnDenoise.js';
+import { halfToFloat } from './allenkFdncnnNcnnModel.js';
+
+function clamp(value, min, max) {
+    return Math.max(min, Math.min(max, value));
+}
+
+function readUInt16LE(bytes, offset) {
+    return bytes[offset] | (bytes[offset + 1] << 8);
+}
+
+function readFloat32LE(bytes, offset) {
+    return new DataView(bytes.buffer, bytes.byteOffset + offset, 4).getFloat32(0, true);
+}
+
+function calculateSegmentMacs(segment, width, height) {
+    return width * height *
+        segment.outputChannels *
+        segment.inputChannels *
+        segment.kernelW *
+        segment.kernelH;
+}
+
+function calculateLayoutMacs(segments, width, height) {
+    return segments.reduce((sum, segment) => sum + calculateSegmentMacs(segment, width, height), 0);
+}
+
+function decodeSegmentWeights(segment, weightBin) {
+    const weights = new Float32Array(segment.weightCount);
+    for (let i = 0; i < segment.weightCount; i++) {
+        weights[i] = halfToFloat(readUInt16LE(weightBin, segment.weightOffset + i * 2));
+    }
+
+    const bias = new Float32Array(segment.biasCount);
+    for (let i = 0; i < segment.biasCount; i++) {
+        bias[i] = readFloat32LE(weightBin, segment.biasOffset + i * 4);
+    }
+
+    return { weights, bias };
+}
+
+function runSamePaddingConvolution({
+    input,
+    width,
+    height,
+    segment,
+    weights,
+    bias
+}) {
+    const outputChannels = segment.outputChannels;
+    const inputChannels = segment.inputChannels;
+    const output = new Float32Array(width * height * outputChannels);
+    const kernelW = segment.kernelW;
+    const kernelH = segment.kernelH;
+    const padX = segment.padW;
+    const padY = segment.padH;
+    const relu = segment.activationType === 1;
+    const plane = width * height;
+
+    for (let oc = 0; oc < outputChannels; oc++) {
+        const outputBase = oc * plane;
+        const biasValue = bias[oc] || 0;
+        for (let y = 0; y < height; y++) {
+            for (let x = 0; x < width; x++) {
+                let sum = biasValue;
+                for (let ic = 0; ic < inputChannels; ic++) {
+                    const inputBase = ic * plane;
+                    const weightBase = (((oc * inputChannels) + ic) * kernelH) * kernelW;
+                    for (let ky = 0; ky < kernelH; ky++) {
+                        const sy = y + ky - padY;
+                        if (sy < 0 || sy >= height) continue;
+                        for (let kx = 0; kx < kernelW; kx++) {
+                            const sx = x + kx - padX;
+                            if (sx < 0 || sx >= width) continue;
+                            sum += input[inputBase + sy * width + sx] * weights[weightBase + ky * kernelW + kx];
+                        }
+                    }
+                }
+                output[outputBase + y * width + x] = relu ? Math.max(0, sum) : sum;
+            }
+        }
+    }
+
+    return output;
+}
+
+function createAllenkFdncnnReferenceRuntime({
+    weightBin,
+    weightLayout,
+    maxMacs = 50_000_000
+} = {}) {
+    if (!weightBin || !weightLayout?.segments?.length) {
+        throw new Error('Missing allenk FDnCNN weight bin or layout');
+    }
+
+    const bin = weightBin instanceof Uint8Array ? weightBin : new Uint8Array(weightBin);
+    const decodedSegments = weightLayout.segments.map((segment) => ({
+        segment,
+        ...decodeSegmentWeights(segment, bin)
+    }));
+
+    return {
+        id: 'allenk-fdncnn-pure-js-reference',
+        status: 'debug-only',
+        maxMacs,
+        estimateMacs(width, height) {
+            return calculateLayoutMacs(weightLayout.segments, width, height);
+        },
+        execute(input, width, height) {
+            const macs = calculateLayoutMacs(weightLayout.segments, width, height);
+            if (macs > maxMacs) {
+                const error = new Error(`allenk FDnCNN pure JS reference refused ${macs} MACs; max=${maxMacs}`);
+                error.code = 'ALLENK_FDNCNN_REFERENCE_MAC_LIMIT';
+                error.macs = macs;
+                error.maxMacs = maxMacs;
+                throw error;
+            }
+
+            let current = input;
+            for (const decoded of decodedSegments) {
+                current = runSamePaddingConvolution({
+                    input: current,
+                    width,
+                    height,
+                    segment: decoded.segment,
+                    weights: decoded.weights,
+                    bias: decoded.bias
+                });
+            }
+
+            for (let i = 0; i < current.length; i++) {
+                current[i] = clamp(current[i], 0, 1);
+            }
+
+            return {
+                output: current,
+                macs,
+                runtime: this.id
+            };
+        },
+        denoiseImageData({ imageData, sigma } = {}) {
+            const input = buildAllenkFdncnnInput({ imageData, sigma });
+            const result = this.execute(input, imageData.width, imageData.height);
+            return {
+                ...result,
+                imageData: {
+                    width: imageData.width,
+                    height: imageData.height,
+                    data: convertAllenkFdncnnOutputToRgba({
+                        output: result.output,
+                        width: imageData.width,
+                        height: imageData.height
+                    })
+                }
+            };
+        }
+    };
+}
+
+export {
+    calculateLayoutMacs,
+    calculateSegmentMacs,
+    createAllenkFdncnnReferenceRuntime,
+    decodeSegmentWeights,
+    runSamePaddingConvolution
+};