oidn-web 0.3.5 → 0.5.0

package/lib/modelSpec.js

@@ -0,0 +1,270 @@
+function conv(id, input) {
+    return {
+        op: 'conv2d',
+        id,
+        input,
+        weight: `${id}.weight`,
+        bias: `${id}.bias`,
+        activation: 'relu',
+        padding: 'same'
+    };
+}
+function pool(id, input) {
+    return {
+        op: 'maxPool2d',
+        id,
+        input,
+        size: 2,
+        stride: 2,
+        padding: 'same'
+    };
+}
+function upsample(id, input) {
+    return {
+        op: 'upsample2d',
+        id,
+        input,
+        scale: 2,
+        mode: 'nearest'
+    };
+}
+function concat(id, first, second) {
+    return {
+        op: 'concat',
+        id,
+        inputs: [first, second],
+        axis: 'channels'
+    };
+}
+export const OIDN_UNET_SMALL_SPEC = {
+    schemaVersion: 1,
+    id: 'oidn-unet-small-v1',
+    family: 'oidn-unet-small',
+    input: 'input',
+    output: 'dec_conv0',
+    receptiveField: 174,
+    nodes: [
+        conv('enc_conv0', 'input'),
+        conv('enc_conv1', 'enc_conv0'),
+        pool('pool1', 'enc_conv1'),
+        conv('enc_conv2', 'pool1'),
+        pool('pool2', 'enc_conv2'),
+        conv('enc_conv3', 'pool2'),
+        pool('pool3', 'enc_conv3'),
+        conv('enc_conv4', 'pool3'),
+        pool('pool4', 'enc_conv4'),
+        conv('enc_conv5a', 'pool4'),
+        conv('enc_conv5b', 'enc_conv5a'),
+        upsample('up4', 'enc_conv5b'),
+        concat('concat4', 'up4', 'pool3'),
+        conv('dec_conv4a', 'concat4'),
+        conv('dec_conv4b', 'dec_conv4a'),
+        upsample('up3', 'dec_conv4b'),
+        concat('concat3', 'up3', 'pool2'),
+        conv('dec_conv3a', 'concat3'),
+        conv('dec_conv3b', 'dec_conv3a'),
+        upsample('up2', 'dec_conv3b'),
+        concat('concat2', 'up2', 'pool1'),
+        conv('dec_conv2a', 'concat2'),
+        conv('dec_conv2b', 'dec_conv2a'),
+        upsample('up1', 'dec_conv2b'),
+        concat('concat1', 'up1', 'input'),
+        conv('dec_conv1a', 'concat1'),
+        conv('dec_conv1b', 'dec_conv1a'),
+        conv('dec_conv0', 'dec_conv1b')
+    ]
+};
+export const OIDN_UNET_LARGE_SPEC = {
+    schemaVersion: 1,
+    id: 'oidn-unet-large-v1',
+    family: 'oidn-unet-large',
+    input: 'input',
+    output: 'dec_conv1c',
+    receptiveField: 202,
+    nodes: [
+        conv('enc_conv1a', 'input'),
+        conv('enc_conv1b', 'enc_conv1a'),
+        pool('pool1', 'enc_conv1b'),
+        conv('enc_conv2a', 'pool1'),
+        conv('enc_conv2b', 'enc_conv2a'),
+        pool('pool2', 'enc_conv2b'),
+        conv('enc_conv3a', 'pool2'),
+        conv('enc_conv3b', 'enc_conv3a'),
+        pool('pool3', 'enc_conv3b'),
+        conv('enc_conv4a', 'pool3'),
+        conv('enc_conv4b', 'enc_conv4a'),
+        pool('pool4', 'enc_conv4b'),
+        conv('enc_conv5a', 'pool4'),
+        conv('enc_conv5b', 'enc_conv5a'),
+        upsample('up4', 'enc_conv5b'),
+        concat('concat4', 'up4', 'pool3'),
+        conv('dec_conv4a', 'concat4'),
+        conv('dec_conv4b', 'dec_conv4a'),
+        upsample('up3', 'dec_conv4b'),
+        concat('concat3', 'up3', 'pool2'),
+        conv('dec_conv3a', 'concat3'),
+        conv('dec_conv3b', 'dec_conv3a'),
+        upsample('up2', 'dec_conv3b'),
+        concat('concat2', 'up2', 'pool1'),
+        conv('dec_conv2a', 'concat2'),
+        conv('dec_conv2b', 'dec_conv2a'),
+        upsample('up1', 'dec_conv2b'),
+        concat('concat1', 'up1', 'input'),
+        conv('dec_conv1a', 'concat1'),
+        conv('dec_conv1b', 'dec_conv1a'),
+        conv('dec_conv1c', 'dec_conv1b')
+    ]
+};
+const BUILTIN_MODEL_SPECS = [
+    OIDN_UNET_SMALL_SPEC,
+    OIDN_UNET_LARGE_SPEC
+];
+function expectedTensorNames(spec) {
+    const names = new Set();
+    for (const node of spec.nodes) {
+        if (node.op === 'conv2d') {
+            names.add(node.weight);
+            names.add(node.bias);
+        }
+    }
+    return names;
+}
+function tensorByteSize(tensor) {
+    return tensor.desc.getByteSize();
+}
+function describeNames(names) {
+    return [...names].sort().join(', ');
+}
+export function detectUNetModelSpec(tensors, specs = BUILTIN_MODEL_SPECS) {
+    const matches = specs.filter((spec) => {
+        const expected = expectedTensorNames(spec);
+        if ([...expected].some((name) => !tensors.has(name)))
+            return false;
+        return (spec.allowAdditionalTensors === true ||
+            [...tensors.keys()].every((name) => expected.has(name)));
+    });
+    if (matches.length === 1)
+        return matches[0];
+    if (matches.length > 1) {
+        throw new Error(`Ambiguous OIDN model topology: ${matches.map((spec) => spec.id).join(', ')}`);
+    }
+    throw new Error(`Unsupported OIDN model topology. TZA tensors: ${describeNames(tensors.keys())}`);
+}
+function requireTensor(tensors, name, modelId) {
+    const tensor = tensors.get(name);
+    if (!tensor) {
+        throw new Error(`Model ${modelId} is missing tensor ${name}`);
+    }
+    if (tensor.data.byteLength !== tensorByteSize(tensor)) {
+        throw new Error(`Tensor ${name} has ${tensor.data.byteLength} bytes, expected ${tensorByteSize(tensor)}`);
+    }
+    return tensor;
+}
+/** Validate tensor layout, shapes, dtypes, graph order, and channel flow. */
+export function validateUNetModel(tensors, spec = detectUNetModelSpec(tensors)) {
+    if (spec.schemaVersion !== 1) {
+        throw new Error(`Unsupported model descriptor schema ${spec.schemaVersion}`);
+    }
+    const expected = expectedTensorNames(spec);
+    if (!spec.allowAdditionalTensors) {
+        const additional = [...tensors.keys()].filter((name) => !expected.has(name));
+        if (additional.length > 0) {
+            throw new Error(`Model ${spec.id} has unexpected tensors: ${describeNames(additional)}`);
+        }
+    }
+    const channelsByValue = new Map();
+    const convTensors = new Map();
+    const convChannels = new Map();
+    const producedValues = new Set([spec.input]);
+    let inputChannels;
+    let commonDataType;
+    const getChannels = (value, nodeId) => {
+        const channels = channelsByValue.get(value);
+        if (channels === undefined) {
+            throw new Error(`Model ${spec.id} node ${nodeId} reads unknown or forward value ${value}`);
+        }
+        return channels;
+    };
+    for (const node of spec.nodes) {
+        if (producedValues.has(node.id)) {
+            throw new Error(`Model ${spec.id} produces duplicate value ${node.id}`);
+        }
+        if (node.op === 'conv2d') {
+            const weight = requireTensor(tensors, node.weight, spec.id);
+            const bias = requireTensor(tensors, node.bias, spec.id);
+            const dims = weight.desc.dims;
+            if (weight.desc.layout !== 'oihw' || dims.length !== 4) {
+                throw new Error(`Tensor ${node.weight} must use OIHW layout`);
+            }
+            if (dims[2] !== 3 || dims[3] !== 3) {
+                throw new Error(`Tensor ${node.weight} must use a 3x3 kernel`);
+            }
+            if (bias.desc.layout !== 'x' || bias.desc.dims.length !== 1) {
+                throw new Error(`Tensor ${node.bias} must be a one-dimensional bias`);
+            }
+            if (bias.desc.dims[0] !== dims[0]) {
+                throw new Error(`Tensor ${node.bias} has ${bias.desc.dims[0]} channels, expected ${dims[0]}`);
+            }
+            if (weight.desc.dataType !== bias.desc.dataType) {
+                throw new Error(`Weight and bias dtype differ for ${node.id}`);
+            }
+            if (commonDataType && commonDataType !== weight.desc.dataType) {
+                throw new Error(`Mixed tensor dtypes are not supported by model ${spec.id}`);
+            }
+            commonDataType = weight.desc.dataType;
+            if (node.input === spec.input && inputChannels === undefined) {
+                inputChannels = dims[1];
+                channelsByValue.set(spec.input, inputChannels);
+            }
+            const actualInputChannels = getChannels(node.input, node.id);
+            if (actualInputChannels !== dims[1]) {
+                throw new Error(`Tensor ${node.weight} expects ${dims[1]} input channels, ` +
+                    `but ${node.input} provides ${actualInputChannels}`);
+            }
+            channelsByValue.set(node.id, dims[0]);
+            convTensors.set(node.id, {
+                weight,
+                bias,
+                inputChannels: dims[1],
+                outputChannels: dims[0],
+                kernelHeight: dims[2],
+                kernelWidth: dims[3]
+            });
+            convChannels.set(node.id, {
+                inputChannels: dims[1],
+                outputChannels: dims[0]
+            });
+        }
+        else if (node.op === 'concat') {
+            if (node.inputs.length < 2) {
+                throw new Error(`Concat ${node.id} requires at least two inputs`);
+            }
+            const channels = node.inputs.reduce((sum, value) => sum + getChannels(value, node.id), 0);
+            channelsByValue.set(node.id, channels);
+        }
+        else {
+            channelsByValue.set(node.id, getChannels(node.input, node.id));
+        }
+        producedValues.add(node.id);
+    }
+    if (inputChannels === undefined || commonDataType === undefined) {
+        throw new Error(`Model ${spec.id} has no convolution reading its input`);
+    }
+    const outputChannels = channelsByValue.get(spec.output);
+    if (outputChannels === undefined) {
+        throw new Error(`Model ${spec.id} output ${spec.output} is not produced`);
+    }
+    if (outputChannels !== 3) {
+        throw new Error(`Model ${spec.id} must produce 3 channels, got ${outputChannels}`);
+    }
+    return {
+        spec,
+        inputChannels,
+        outputChannels,
+        tensorDataType: commonDataType,
+        channelsByValue,
+        convChannels,
+        convTensors
+    };
+}
+//# sourceMappingURL=modelSpec.js.map

package/lib/modelSpec.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"modelSpec.js","sourceRoot":"","sources":["../src/modelSpec.ts"],"names":[],"mappings":"AA2FA,SAAS,IAAI,CAAC,EAAU,EAAE,KAAa;IACrC,OAAO;QACL,EAAE,EAAE,QAAQ;QACZ,EAAE;QACF,KAAK;QACL,MAAM,EAAE,GAAG,EAAE,SAAS;QACtB,IAAI,EAAE,GAAG,EAAE,OAAO;QAClB,UAAU,EAAE,MAAM;QAClB,OAAO,EAAE,MAAM;KAChB,CAAC;AACJ,CAAC;AAED,SAAS,IAAI,CAAC,EAAU,EAAE,KAAa;IACrC,OAAO;QACL,EAAE,EAAE,WAAW;QACf,EAAE;QACF,KAAK;QACL,IAAI,EAAE,CAAC;QACP,MAAM,EAAE,CAAC;QACT,OAAO,EAAE,MAAM;KAChB,CAAC;AACJ,CAAC;AAED,SAAS,QAAQ,CAAC,EAAU,EAAE,KAAa;IACzC,OAAO;QACL,EAAE,EAAE,YAAY;QAChB,EAAE;QACF,KAAK;QACL,KAAK,EAAE,CAAC;QACR,IAAI,EAAE,SAAS;KAChB,CAAC;AACJ,CAAC;AAED,SAAS,MAAM,CACb,EAAU,EACV,KAAa,EACb,MAAc;IAEd,OAAO;QACL,EAAE,EAAE,QAAQ;QACZ,EAAE;QACF,MAAM,EAAE,CAAC,KAAK,EAAE,MAAM,CAAC;QACvB,IAAI,EAAE,UAAU;KACjB,CAAC;AACJ,CAAC;AAED,MAAM,CAAC,MAAM,oBAAoB,GAAkB;IACjD,aAAa,EAAE,CAAC;IAChB,EAAE,EAAE,oBAAoB;IACxB,MAAM,EAAE,iBAAiB;IACzB,KAAK,EAAE,OAAO;IACd,MAAM,EAAE,WAAW;IACnB,cAAc,EAAE,GAAG;IACnB,KAAK,EAAE;QACL,IAAI,CAAC,WAAW,EAAE,OAAO,CAAC;QAC1B,IAAI,CAAC,WAAW,EAAE,WAAW,CAAC;QAC9B,IAAI,CAAC,OAAO,EAAE,WAAW,CAAC;QAC1B,IAAI,CAAC,WAAW,EAAE,OAAO,CAAC;QAC1B,IAAI,CAAC,OAAO,EAAE,WAAW,CAAC;QAC1B,IAAI,CAAC,WAAW,EAAE,OAAO,CAAC;QAC1B,IAAI,CAAC,OAAO,EAAE,WAAW,CAAC;QAC1B,IAAI,CAAC,WAAW,EAAE,OAAO,CAAC;QAC1B,IAAI,CAAC,OAAO,EAAE,WAAW,CAAC;QAC1B,IAAI,CAAC,YAAY,EAAE,OAAO,CAAC;QAC3B,IAAI,CAAC,YAAY,EAAE,YAAY,CAAC;QAChC,QAAQ,CAAC,KAAK,EAAE,YAAY,CAAC;QAC7B,MAAM,CAAC,SAAS,EAAE,KAAK,EAAE,OAAO,CAAC;QACjC,IAAI,CAAC,YAAY,EAAE,SAAS,CAAC;QAC7B,IAAI,CAAC,YAAY,EAAE,YAAY,CAAC;QAChC,QAAQ,CAAC,KAAK,EAAE,YAAY,CAAC;QAC7B,MAAM,CAAC,SAAS,EAAE,KAAK,EAAE,OAAO,CAAC;QACjC,IAAI,CAAC,YAAY,EAAE,SAAS,CAAC;QAC7B,IAAI,CAAC,YAAY,EAAE,YAAY,CAAC;QAChC,QAAQ,CAAC,KAAK,EAAE,YAAY,CAAC;QAC7B,MAAM,CAAC,SAAS,EAAE,KAAK,EAAE,OAAO,CAAC;QACjC,IAAI,CAAC,YAAY,EAAE,SAAS,CAAC;QAC7B,IAAI,CAAC,YAAY,EAAE,YAAY,CAAC;QAChC,QAAQ,CAAC,KAAK,EAAE,YAAY,CAAC;QAC7B,MAAM,CAAC,SAAS,EAAE,KAAK,EAAE,OAAO,CAAC;QACjC,IAAI,CAAC,YAAY,EAAE,SAAS,CAAC;QAC7B,IAAI,CAAC,YAAY,EAAE,YAAY,CAAC;QAChC,IAAI,CAAC,WAAW,EAAE,YAAY,CAAC;KAChC;CACF,CAAC;AAEF,MAAM,CAAC,MAAM,oBAAoB,GAAkB;IACjD,aAAa,EAAE,CAAC;IAChB,EAAE,EAAE,oBAAoB;IACxB,MAAM,EAAE,iBAAiB;IACzB,KAAK,EAAE,OAAO;IACd,MAAM,EAAE,YAAY;IACpB,cAAc,EAAE,GAAG;IACnB,KAAK,EAAE;QACL,IAAI,CAAC,YAAY,EAAE,OAAO,CAAC;QAC3B,IAAI,CAAC,YAAY,EAAE,YAAY,CAAC;QAChC,IAAI,CAAC,OAAO,EAAE,YAAY,CAAC;QAC3B,IAAI,CAAC,YAAY,EAAE,OAAO,CAAC;QAC3B,IAAI,CAAC,YAAY,EAAE,YAAY,CAAC;QAChC,IAAI,CAAC,OAAO,EAAE,YAAY,CAAC;QAC3B,IAAI,CAAC,YAAY,EAAE,OAAO,CAAC;QAC3B,IAAI,CAAC,YAAY,EAAE,YAAY,CAAC;QAChC,IAAI,CAAC,OAAO,EAAE,YAAY,CAAC;QAC3B,IAAI,CAAC,YAAY,EAAE,OAAO,CAAC;QAC3B,IAAI,CAAC,YAAY,EAAE,YAAY,CAAC;QAChC,IAAI,CAAC,OAAO,EAAE,YAAY,CAAC;QAC3B,IAAI,CAAC,YAAY,EAAE,OAAO,CAAC;QAC3B,IAAI,CAAC,YAAY,EAAE,YAAY,CAAC;QAChC,QAAQ,CAAC,KAAK,EAAE,YAAY,CAAC;QAC7B,MAAM,CAAC,SAAS,EAAE,KAAK,EAAE,OAAO,CAAC;QACjC,IAAI,CAAC,YAAY,EAAE,SAAS,CAAC;QAC7B,IAAI,CAAC,YAAY,EAAE,YAAY,CAAC;QAChC,QAAQ,CAAC,KAAK,EAAE,YAAY,CAAC;QAC7B,MAAM,CAAC,SAAS,EAAE,KAAK,EAAE,OAAO,CAAC;QACjC,IAAI,CAAC,YAAY,EAAE,SAAS,CAAC;QAC7B,IAAI,CAAC,YAAY,EAAE,YAAY,CAAC;QAChC,QAAQ,CAAC,KAAK,EAAE,YAAY,CAAC;QAC7B,MAAM,CAAC,SAAS,EAAE,KAAK,EAAE,OAAO,CAAC;QACjC,IAAI,CAAC,YAAY,EAAE,SAAS,CAAC;QAC7B,IAAI,CAAC,YAAY,EAAE,YAAY,CAAC;QAChC,QAAQ,CAAC,KAAK,EAAE,YAAY,CAAC;QAC7B,MAAM,CAAC,SAAS,EAAE,KAAK,EAAE,OAAO,CAAC;QACjC,IAAI,CAAC,YAAY,EAAE,SAAS,CAAC;QAC7B,IAAI,CAAC,YAAY,EAAE,YAAY,CAAC;QAChC,IAAI,CAAC,YAAY,EAAE,YAAY,CAAC;KACjC;CACF,CAAC;AAEF,MAAM,mBAAmB,GAAG;IAC1B,oBAAoB;IACpB,oBAAoB;CACZ,CAAC;AAEX,SAAS,mBAAmB,CAAC,IAAmB;IAC9C,MAAM,KAAK,GAAG,IAAI,GAAG,EAAU,CAAC;IAChC,KAAK,MAAM,IAAI,IAAI,IAAI,CAAC,KAAK,EAAE,CAAC;QAC9B,IAAI,IAAI,CAAC,EAAE,KAAK,QAAQ,EAAE,CAAC;YACzB,KAAK,CAAC,GAAG,CAAC,IAAI,CAAC,MAAM,CAAC,CAAC;YACvB,KAAK,CAAC,GAAG,CAAC,IAAI,CAAC,IAAI,CAAC,CAAC;QACvB,CAAC;IACH,CAAC;IACD,OAAO,KAAK,CAAC;AACf,CAAC;AAED,SAAS,cAAc,CAAC,MAAkB;IACxC,OAAO,MAAM,CAAC,IAAI,CAAC,WAAW,EAAE,CAAC;AACnC,CAAC;AAED,SAAS,aAAa,CAAC,KAAuB;IAC5C,OAAO,CAAC,GAAG,KAAK,CAAC,CAAC,IAAI,EAAE,CAAC,IAAI,CAAC,IAAI,CAAC,CAAC;AACtC,CAAC;AAED,MAAM,UAAU,mBAAmB,CACjC,OAAwC,EACxC,QAAkC,mBAAmB;IAErD,MAAM,OAAO,GAAG,KAAK,CAAC,MAAM,CAAC,CAAC,IAAI,EAAE,EAAE;QACpC,MAAM,QAAQ,GAAG,mBAAmB,CAAC,IAAI,CAAC,CAAC;QAC3C,IAAI,CAAC,GAAG,QAAQ,CAAC,CAAC,IAAI,CAAC,CAAC,IAAI,EAAE,EAAE,CAAC,CAAC,OAAO,CAAC,GAAG,CAAC,IAAI,CAAC,CAAC;YAAE,OAAO,KAAK,CAAC;QACnE,OAAO,CACL,IAAI,CAAC,sBAAsB,KAAK,IAAI;YACpC,CAAC,GAAG,OAAO,CAAC,IAAI,EAAE,CAAC,CAAC,KAAK,CAAC,CAAC,IAAI,EAAE,EAAE,CAAC,QAAQ,CAAC,GAAG,CAAC,IAAI,CAAC,CAAC,CACxD,CAAC;IACJ,CAAC,CAAC,CAAC;IAEH,IAAI,OAAO,CAAC,MAAM,KAAK,CAAC;QAAE,OAAO,OAAO,CAAC,CAAC,CAAC,CAAC;IAC5C,IAAI,OAAO,CAAC,MAAM,GAAG,CAAC,EAAE,CAAC;QACvB,MAAM,IAAI,KAAK,CACb,kCAAkC,OAAO,CAAC,GAAG,CAAC,CAAC,IAAI,EAAE,EAAE,CAAC,IAAI,CAAC,EAAE,CAAC,CAAC,IAAI,CAAC,IAAI,CAAC,EAAE,CAC9E,CAAC;IACJ,CAAC;IAED,MAAM,IAAI,KAAK,CACb,iDAAiD,aAAa,CAAC,OAAO,CAAC,IAAI,EAAE,CAAC,EAAE,CACjF,CAAC;AACJ,CAAC;AAED,SAAS,aAAa,CACpB,OAAwC,EACxC,IAAY,EACZ,OAAe;IAEf,MAAM,MAAM,GAAG,OAAO,CAAC,GAAG,CAAC,IAAI,CAAC,CAAC;IACjC,IAAI,CAAC,MAAM,EAAE,CAAC;QACZ,MAAM,IAAI,KAAK,CAAC,SAAS,OAAO,sBAAsB,IAAI,EAAE,CAAC,CAAC;IAChE,CAAC;IACD,IAAI,MAAM,CAAC,IAAI,CAAC,UAAU,KAAK,cAAc,CAAC,MAAM,CAAC,EAAE,CAAC;QACtD,MAAM,IAAI,KAAK,CACb,UAAU,IAAI,QAAQ,MAAM,CAAC,IAAI,CAAC,UAAU,oBAAoB,cAAc,CAAC,MAAM,CAAC,EAAE,CACzF,CAAC;IACJ,CAAC;IACD,OAAO,MAAM,CAAC;AAChB,CAAC;AAED,6EAA6E;AAC7E,MAAM,UAAU,iBAAiB,CAC/B,OAAwC,EACxC,OAAsB,mBAAmB,CAAC,OAAO,CAAC;IAElD,IAAI,IAAI,CAAC,aAAa,KAAK,CAAC,EAAE,CAAC;QAC7B,MAAM,IAAI,KAAK,CAAC,uCAAuC,IAAI,CAAC,aAAa,EAAE,CAAC,CAAC;IAC/E,CAAC;IAED,MAAM,QAAQ,GAAG,mBAAmB,CAAC,IAAI,CAAC,CAAC;IAC3C,IAAI,CAAC,IAAI,CAAC,sBAAsB,EAAE,CAAC;QACjC,MAAM,UAAU,GAAG,CAAC,GAAG,OAAO,CAAC,IAAI,EAAE,CAAC,CAAC,MAAM,CAAC,CAAC,IAAI,EAAE,EAAE,CAAC,CAAC,QAAQ,CAAC,GAAG,CAAC,IAAI,CAAC,CAAC,CAAC;QAC7E,IAAI,UAAU,CAAC,MAAM,GAAG,CAAC,EAAE,CAAC;YAC1B,MAAM,IAAI,KAAK,CACb,SAAS,IAAI,CAAC,EAAE,4BAA4B,aAAa,CAAC,UAAU,CAAC,EAAE,CACxE,CAAC;QACJ,CAAC;IACH,CAAC;IAED,MAAM,eAAe,GAAG,IAAI,GAAG,EAAkB,CAAC;IAClD,MAAM,WAAW,GAAG,IAAI,GAAG,EAA+B,CAAC;IAC3D,MAAM,YAAY,GAAG,IAAI,GAAG,EAA6B,CAAC;IAC1D,MAAM,cAAc,GAAG,IAAI,GAAG,CAAS,CAAC,IAAI,CAAC,KAAK,CAAC,CAAC,CAAC;IACrD,IAAI,aAAiC,CAAC;IACtC,IAAI,cAA0D,CAAC;IAE/D,MAAM,WAAW,GAAG,CAAC,KAAa,EAAE,MAAc,EAAE,EAAE;QACpD,MAAM,QAAQ,GAAG,eAAe,CAAC,GAAG,CAAC,KAAK,CAAC,CAAC;QAC5C,IAAI,QAAQ,KAAK,SAAS,EAAE,CAAC;YAC3B,MAAM,IAAI,KAAK,CACb,SAAS,IAAI,CAAC,EAAE,SAAS,MAAM,mCAAmC,KAAK,EAAE,CAC1E,CAAC;QACJ,CAAC;QACD,OAAO,QAAQ,CAAC;IAClB,CAAC,CAAC;IAEF,KAAK,MAAM,IAAI,IAAI,IAAI,CAAC,KAAK,EAAE,CAAC;QAC9B,IAAI,cAAc,CAAC,GAAG,CAAC,IAAI,CAAC,EAAE,CAAC,EAAE,CAAC;YAChC,MAAM,IAAI,KAAK,CAAC,SAAS,IAAI,CAAC,EAAE,6BAA6B,IAAI,CAAC,EAAE,EAAE,CAAC,CAAC;QAC1E,CAAC;QAED,IAAI,IAAI,CAAC,EAAE,KAAK,QAAQ,EAAE,CAAC;YACzB,MAAM,MAAM,GAAG,aAAa,CAAC,OAAO,EAAE,IAAI,CAAC,MAAM,EAAE,IAAI,CAAC,EAAE,CAAC,CAAC;YAC5D,MAAM,IAAI,GAAG,aAAa,CAAC,OAAO,EAAE,IAAI,CAAC,IAAI,EAAE,IAAI,CAAC,EAAE,CAAC,CAAC;YACxD,MAAM,IAAI,GAAG,MAAM,CAAC,IAAI,CAAC,IAAI,CAAC;YAE9B,IAAI,MAAM,CAAC,IAAI,CAAC,MAAM,KAAK,MAAM,IAAI,IAAI,CAAC,MAAM,KAAK,CAAC,EAAE,CAAC;gBACvD,MAAM,IAAI,KAAK,CAAC,UAAU,IAAI,CAAC,MAAM,uBAAuB,CAAC,CAAC;YAChE,CAAC;YACD,IAAI,IAAI,CAAC,CAAC,CAAC,KAAK,CAAC,IAAI,IAAI,CAAC,CAAC,CAAC,KAAK,CAAC,EAAE,CAAC;gBACnC,MAAM,IAAI,KAAK,CAAC,UAAU,IAAI,CAAC,MAAM,wBAAwB,CAAC,CAAC;YACjE,CAAC;YACD,IAAI,IAAI,CAAC,IAAI,CAAC,MAAM,KAAK,GAAG,IAAI,IAAI,CAAC,IAAI,CAAC,IAAI,CAAC,MAAM,KAAK,CAAC,EAAE,CAAC;gBAC5D,MAAM,IAAI,KAAK,CAAC,UAAU,IAAI,CAAC,IAAI,iCAAiC,CAAC,CAAC;YACxE,CAAC;YACD,IAAI,IAAI,CAAC,IAAI,CAAC,IAAI,CAAC,CAAC,CAAC,KAAK,IAAI,CAAC,CAAC,CAAC,EAAE,CAAC;gBAClC,MAAM,IAAI,KAAK,CACb,UAAU,IAAI,CAAC,IAAI,QAAQ,IAAI,CAAC,IAAI,CAAC,IAAI,CAAC,CAAC,CAAC,uBAAuB,IAAI,CAAC,CAAC,CAAC,EAAE,CAC7E,CAAC;YACJ,CAAC;YACD,IAAI,MAAM,CAAC,IAAI,CAAC,QAAQ,KAAK,IAAI,CAAC,IAAI,CAAC,QAAQ,EAAE,CAAC;gBAChD,MAAM,IAAI,KAAK,CAAC,oCAAoC,IAAI,CAAC,EAAE,EAAE,CAAC,CAAC;YACjE,CAAC;YACD,IAAI,cAAc,IAAI,cAAc,KAAK,MAAM,CAAC,IAAI,CAAC,QAAQ,EAAE,CAAC;gBAC9D,MAAM,IAAI,KAAK,CAAC,kDAAkD,IAAI,CAAC,EAAE,EAAE,CAAC,CAAC;YAC/E,CAAC;YACD,cAAc,GAAG,MAAM,CAAC,IAAI,CAAC,QAAQ,CAAC;YAEtC,IAAI,IAAI,CAAC,KAAK,KAAK,IAAI,CAAC,KAAK,IAAI,aAAa,KAAK,SAAS,EAAE,CAAC;gBAC7D,aAAa,GAAG,IAAI,CAAC,CAAC,CAAC,CAAC;gBACxB,eAAe,CAAC,GAAG,CAAC,IAAI,CAAC,KAAK,EAAE,aAAa,CAAC,CAAC;YACjD,CAAC;YACD,MAAM,mBAAmB,GAAG,WAAW,CAAC,IAAI,CAAC,KAAK,EAAE,IAAI,CAAC,EAAE,CAAC,CAAC;YAC7D,IAAI,mBAAmB,KAAK,IAAI,CAAC,CAAC,CAAC,EAAE,CAAC;gBACpC,MAAM,IAAI,KAAK,CACb,UAAU,IAAI,CAAC,MAAM,YAAY,IAAI,CAAC,CAAC,CAAC,mBAAmB;oBACzD,OAAO,IAAI,CAAC,KAAK,aAAa,mBAAmB,EAAE,CACtD,CAAC;YACJ,CAAC;YAED,eAAe,CAAC,GAAG,CAAC,IAAI,CAAC,EAAE,EAAE,IAAI,CAAC,CAAC,CAAC,CAAC,CAAC;YACtC,WAAW,CAAC,GAAG,CAAC,IAAI,CAAC,EAAE,EAAE;gBACvB,MAAM;gBACN,IAAI;gBACJ,aAAa,EAAE,IAAI,CAAC,CAAC,CAAC;gBACtB,cAAc,EAAE,IAAI,CAAC,CAAC,CAAC;gBACvB,YAAY,EAAE,IAAI,CAAC,CAAC,CAAC;gBACrB,WAAW,EAAE,IAAI,CAAC,CAAC,CAAC;aACrB,CAAC,CAAC;YACH,YAAY,CAAC,GAAG,CAAC,IAAI,CAAC,EAAE,EAAE;gBACxB,aAAa,EAAE,IAAI,CAAC,CAAC,CAAC;gBACtB,cAAc,EAAE,IAAI,CAAC,CAAC,CAAC;aACxB,CAAC,CAAC;QACL,CAAC;aAAM,IAAI,IAAI,CAAC,EAAE,KAAK,QAAQ,EAAE,CAAC;YAChC,IAAI,IAAI,CAAC,MAAM,CAAC,MAAM,GAAG,CAAC,EAAE,CAAC;gBAC3B,MAAM,IAAI,KAAK,CAAC,UAAU,IAAI,CAAC,EAAE,+BAA+B,CAAC,CAAC;YACpE,CAAC;YACD,MAAM,QAAQ,GAAG,IAAI,CAAC,MAAM,CAAC,MAAM,CACjC,CAAC,GAAG,EAAE,KAAK,EAAE,EAAE,CAAC,GAAG,GAAG,WAAW,CAAC,KAAK,EAAE,IAAI,CAAC,EAAE,CAAC,EACjD,CAAC,CACF,CAAC;YACF,eAAe,CAAC,GAAG,CAAC,IAAI,CAAC,EAAE,EAAE,QAAQ,CAAC,CAAC;QACzC,CAAC;aAAM,CAAC;YACN,eAAe,CAAC,GAAG,CAAC,IAAI,CAAC,EAAE,EAAE,WAAW,CAAC,IAAI,CAAC,KAAK,EAAE,IAAI,CAAC,EAAE,CAAC,CAAC,CAAC;QACjE,CAAC;QAED,cAAc,CAAC,GAAG,CAAC,IAAI,CAAC,EAAE,CAAC,CAAC;IAC9B,CAAC;IAED,IAAI,aAAa,KAAK,SAAS,IAAI,cAAc,KAAK,SAAS,EAAE,CAAC;QAChE,MAAM,IAAI,KAAK,CAAC,SAAS,IAAI,CAAC,EAAE,uCAAuC,CAAC,CAAC;IAC3E,CAAC;IACD,MAAM,cAAc,GAAG,eAAe,CAAC,GAAG,CAAC,IAAI,CAAC,MAAM,CAAC,CAAC;IACxD,IAAI,cAAc,KAAK,SAAS,EAAE,CAAC;QACjC,MAAM,IAAI,KAAK,CAAC,SAAS,IAAI,CAAC,EAAE,WAAW,IAAI,CAAC,MAAM,kBAAkB,CAAC,CAAC;IAC5E,CAAC;IACD,IAAI,cAAc,KAAK,CAAC,EAAE,CAAC;QACzB,MAAM,IAAI,KAAK,CAAC,SAAS,IAAI,CAAC,EAAE,iCAAiC,cAAc,EAAE,CAAC,CAAC;IACrF,CAAC;IAED,OAAO;QACL,IAAI;QACJ,aAAa;QACb,cAAc;QACd,cAAc,EAAE,cAAc;QAC9B,eAAe;QACf,YAAY;QACZ,WAAW;KACZ,CAAC;AACJ,CAAC"}

package/lib/nativeUNet.d.ts

@@ -0,0 +1,103 @@
+import type { ValidatedUNetModel } from './modelSpec';
+import { type OIDNResourceSnapshot } from './resourceTracker.js';
+export type NativeUNetPrecision = 'fp32' | 'fp16';
+export type NativeUNetPrecisionSetting = NativeUNetPrecision | 'auto';
+export type NativeUNetKernel = 'direct' | 'implicit-gemm' | 'spatial' | 'subgroup';
+export type NativeUNetKernelSetting = NativeUNetKernel | 'auto';
+export type NativeUNetGemmWorkgroup = readonly [4 | 8 | 16, 4 | 8];
+export interface NativeUNetGemmOptions {
+    /** Internal per-layer tile selection; explicit workgroup sizes remain fixed by default. */
+    tilePolicy?: 'fixed' | 'output-aligned';
+    /** Internal decoder source-branch scheduling experiment. */
+    decoderLoad?: 'per-load' | 'source-first';
+    /** Internal pooling dispatch experiment for the GEMM execution path. */
+    poolLayout?: 'spatial' | 'channels';
+    /** Internal shared-memory padding experiment; does not change storage ABI. */
+    sharedLayout?: 'linear' | 'padded' | 'padded-input' | 'padded-weights';
+    /** Internal loop scheduling experiment; each output keeps the same K order. */
+    accumulationOrder?: 'k-major' | 'row-major';
+    /** Internal final-convolution experiments; retains the original FP16 partial order. */
+    finalLayer?: 'direct' | 'shared-input' | 'shared-input-weights' | 'shared-auto';
+    /** Internal load-layout experiments; packed loads preserve the FP16 storage bits. */
+    loadMode?: 'native' | 'packed-weights' | 'packed-all';
+    /** Incremental addressing (default), or the original analytic path for A/B comparisons. */
+    addressMode?: 'analytic' | 'incremental' | 'base-offset';
+    /** K-major (default) coalesces neighboring output-block loads; output-major preserves the old ABI. */
+    weightLayout?: 'output-major' | 'k-major';
+    /** Register tile height (default 8). The K reduction grouping stays fixed. */
+    rowsPerThread?: 2 | 4 | 8;
+    /** Workgroup width/height (default [8, 8]); width also selects output channel blocks. */
+    workgroupSize?: NativeUNetGemmWorkgroup;
+}
+export interface NativeUNetOptions {
+    precision?: NativeUNetPrecisionSetting;
+    /**
+     * Convolution kernel selection. `auto` uses implicit GEMM for FP16/FP32
+     * convolutions and the direct kernel for the final output layer.
+     */
+    kernel?: NativeUNetKernelSetting;
+    /** Optional implicit-GEMM tuning; other convolution kernels ignore it. */
+    gemm?: NativeUNetGemmOptions;
+    /** Maximum number of shape-dependent activation plans retained. */
+    shapeCacheSize?: number;
+}
+export interface NativeUNetLayerTiming {
+    id: string;
+    durationMs: number;
+}
+export interface NativeUNetExecutionProfile {
+    totalMs: number;
+    layers: NativeUNetLayerTiming[];
+}
+export declare function resolveNativeUNetPrecision(device: GPUDevice, requested?: NativeUNetPrecisionSetting): NativeUNetPrecision;
+/** Native, model-driven OIDN U-Net executor. */
+export declare class NativeUNetExecutor {
+    private _device;
+    readonly precision: NativeUNetPrecision;
+    readonly kernelSetting: NativeUNetKernelSetting;
+    readonly gemm: Readonly<Required<NativeUNetGemmOptions>>;
+    readonly maxSpatialInputBlocks: number;
+    readonly subgroupsAvailable: boolean;
+    private _model;
+    private _gemmByOutputBlocks;
+    private _packedConvs;
+    private _pipelineCache;
+    private _pipelinePromises;
+    private _executionCache;
+    private _retiredExecutions;
+    private _clock;
+    private _shapeCacheSize;
+    private _profileNextExecution;
+    private _lastExecutionProfile?;
+    private _profileOperations;
+    private _resources;
+    private _disposed;
+    constructor(_device: GPUDevice, model: ValidatedUNetModel, options?: NativeUNetOptions);
+    private _pipeline;
+    private _pipelineAsync;
+    private _nodePipelineSpec;
+    private _gemmForOutput;
+    private _coalescedPool;
+    private _selectConvKernel;
+    private _nodePipeline;
+    /** Compiles all shape-independent kernels before the model reports ready. */
+    prepare(): Promise<void>;
+    private _createExecution;
+    private _execution;
+    /** Allocates shape-dependent execution resources before interactive use. */
+    prewarm(shapes: readonly {
+        width: number;
+        height: number;
+    }[]): void;
+    /** Captures per-pass GPU timestamps for the next execute call when supported. */
+    profileNextExecution(): boolean;
+    getLastExecutionProfile(): Promise<NativeUNetExecutionProfile> | undefined;
+    execute(inputBuffers: readonly GPUBuffer[], width: number, height: number): GPUBuffer;
+    /** Executes interleaved CPU image data through the native GPU runtime. */
+    executeCPU(interleavedInput: Float32Array, width: number, height: number): Promise<Float32Array>;
+    private _releaseBuffer;
+    private _releaseQuerySet;
+    private _destroyExecution;
+    getResourceInfo(): OIDNResourceSnapshot;
+    dispose(): void;
+}