diff --git a/.agents/skills/add-api-tests/SKILL.md b/.agents/skills/add-api-tests/SKILL.md index f0f699bcdd..48c332b301 100644 --- a/.agents/skills/add-api-tests/SKILL.md +++ b/.agents/skills/add-api-tests/SKILL.md @@ -80,7 +80,11 @@ Key helpers: 5. **Options** — every default in `modelOpts`, and every per-call override. 6. **Disposal** — `dispose()` leaves `fakeJsi.liveTensors()` at 0 and `fakeJsi.liveModels()` empty, and repeated calls do not accumulate scratch - tensors. + tensors. The same has to hold when construction *fails*: allocate through a + `createResourceScope()` and wrap the factory body in `try`/`catch` so a + schema mismatch releases the model instead of stranding it. Add the factory + to `__tests__/tasks/constructionFailure.test.ts`, which drives every one of + them. 7. **Sync/async parity** — `runTaskWorklet(x)` equals `await runTask(x)`. Only the *weights* are out of scope, not the pipeline that runs on them. Before @@ -144,5 +148,8 @@ When adding or changing code under `src/`, verify that: is intentional (a removal or rename is a breaking change). - [ ] Any new fake behaviour in `__tests__/support/` is faithful where fidelity changes an assertion, and its simplifications are commented. -- [ ] No test was made to pass by calling `allowNativeLeaks()` without an - explanation. +- [ ] A new `create` allocates through `createResourceScope()` and is + listed in `__tests__/tasks/constructionFailure.test.ts`. +- [ ] No test was made to pass by calling `allowNativeLeaks()`. Nothing in the + suite needs it today, so reach for it only when a leak is genuinely the + point of the test, and say why. diff --git a/.agents/skills/add-task-pipeline/SKILL.md b/.agents/skills/add-task-pipeline/SKILL.md index 6eefdabd72..35c372f7af 100644 --- a/.agents/skills/add-task-pipeline/SKILL.md +++ b/.agents/skills/add-task-pipeline/SKILL.md @@ -27,15 +27,28 @@ When implementing task constructors like `create` (e.g. `createClassifier` const [tReshape, tUint8] = tensors; ``` -2. **Immediate `dispose()` Definition**: - - Right after allocating the static tensors, define the `dispose` function immediately. This makes it instantly visible and verifiable that all native memory will be cleaned up: +2. **Allocate Through a Resource Scope**: + - Open a `createResourceScope()` as the first statement of the constructor, take its `dispose` as the pipeline's, and wrap the whole body in `try`/`catch`. `track` every native resource as it is created: ```typescript - const dispose = () => { - tensors.forEach((t) => t.dispose()); - preprocessor.dispose(); - model.dispose(); - }; + const scope = createResourceScope(); + const dispose = scope.dispose; + + try { + const model = scope.track(await wrapAsync(loadModel, runtime)(modelPath)); + const { dims } = validateSpec(model.schema, { ... }); // may throw + + const tensors = [tensor('float32', shapeA), tensor('float32', shapeB)] as const; + tensors.forEach(scope.track); + const preprocessor = scope.track(createImagePreprocessor(modelOpts, inpShape)); + + return { runTask, runTaskWorklet, dispose }; + } catch (error) { + dispose(); + throw error; + } ``` + - This is not style. A caller whose `create` throws never receives a `dispose`, so anything already allocated would be stranded in native memory for the life of the process, and `useModel` re-runs the factory on every config change. The scope makes success and failure share one teardown path. + - Track resources **as they land**, not afterwards. For parallel loads that means `Promise.all([load(a).then(scope.track), load(b).then(scope.track)])`, so one rejecting does not strand the other. 3. **Dynamic Tensors & `try/finally` Pattern**: - If you must allocate dynamically sized tensors during inference execution (e.g. resizing an output tensor to match the input image dimensions), you must wrap the execution inside a `try {} finally {}` block. @@ -51,7 +64,7 @@ When implementing task constructors like `create` (e.g. `createClassifier` 4. **Pure Helper Functions**: - Write all auxiliary/helper logic as pure, worklet-compatible functions **outside** the `create` constructor. Any helper functions invoked inside the worklet executor thread must contain the `'worklet';` directive. - - **Push Back Hard on Inner Helpers:** You must push back hard against any request to add internal closures or nested functions inside `create` (other than `dispose` and the worklet executor itself). Keep the constructor scope flat to avoid scope leak and dependency chain bugs. + - **Push Back Hard on Inner Helpers:** You must push back hard against any request to add internal closures or nested functions inside `create` (other than the worklet executor itself). Keep the constructor scope flat to avoid scope leak and dependency chain bugs. 5. **PTE Model Export & Optimizations**: - **Shift Heavy Ops to PyTorch**: Push complex tensor reshaping, data normalization, activations (e.g. `softmax`), or bounding box decoding into the PyTorch model itself so they execute on native backends (e.g., XNNPACK or CoreML). diff --git a/packages/react-native-executorch/__tests__/README.md b/packages/react-native-executorch/__tests__/README.md index 3fcc6e9f9c..5879ed002d 100644 --- a/packages/react-native-executorch/__tests__/README.md +++ b/packages/react-native-executorch/__tests__/README.md @@ -65,7 +65,10 @@ for free. Wrap construction in `tracked()` and the harness disposes it at the end of the test — which also keeps a failing assertion from cascading into a second, -misleading leak error. A test that means to leak calls `allowNativeLeaks()`. +misleading leak error. A test that means to leak calls `allowNativeLeaks()`; +nothing in the suite currently does, including the construction-failure cases, +because a factory that throws part-way now releases what it had allocated +(`src/core/lifetime.ts`). ## Layout diff --git a/packages/react-native-executorch/__tests__/api/__snapshots__/apiSurface.test.ts.snap b/packages/react-native-executorch/__tests__/api/__snapshots__/apiSurface.test.ts.snap index 7318448638..a78c99b254 100644 --- a/packages/react-native-executorch/__tests__/api/__snapshots__/apiSurface.test.ts.snap +++ b/packages/react-native-executorch/__tests__/api/__snapshots__/apiSurface.test.ts.snap @@ -139,6 +139,7 @@ exports[`public API surface matches the recorded export list 1`] = ` "createObjectDetector", "createPaddleOcr", "createPrivacyFilter", + "createResourceScope", "createSdxsTextToImage", "createSemanticSegmenter", "createStyleTransfer", diff --git a/packages/react-native-executorch/__tests__/hooks/taskHooks.test.ts b/packages/react-native-executorch/__tests__/hooks/taskHooks.test.ts index beae69f524..972d265617 100644 --- a/packages/react-native-executorch/__tests__/hooks/taskHooks.test.ts +++ b/packages/react-native-executorch/__tests__/hooks/taskHooks.test.ts @@ -13,7 +13,6 @@ import { deferred, fakeNet } from '../support/blobUtilMock'; import { cachePathFor } from '../support/cachePath'; import { fakeJsi } from '../support/fakeJsi'; import { STRETCH_PREPROCESSING, exported, imageBuffer, writesOutputs } from '../support/fixtures'; -import { allowNativeLeaks } from '../support/setup'; const MODEL_URL = 'https://huggingface.co/software-mansion/model/resolve/v1/model.pte'; const TOKENIZER_URL = 'https://huggingface.co/software-mansion/model/resolve/v1/tokenizer.json'; @@ -121,11 +120,11 @@ describe('useClassifier', () => { expect(result.current.error?.message).toMatch(/labels length \(2\)/); expect(result.current.isReady).toBe(false); - // The failed construction abandons the native model it had already loaded, - // and the hook has no handle to release — this is the app-level shape of - // the leak recorded in `tasks/constructionFailure.test.ts`. - expect(fakeJsi.liveModels()).toEqual([cachePathFor(MODEL_URL)]); - allowNativeLeaks(); + // The hook never receives a `dispose` for a pipeline that failed to build, + // so the factory has to have released the model itself. `useModel` re-runs + // the factory on every config change, which is what makes this the + // app-level shape of `tasks/constructionFailure.test.ts`. + expect(fakeJsi.liveModels()).toEqual([]); }); it('loads nothing while preventLoad is set', async () => { diff --git a/packages/react-native-executorch/__tests__/support/lifetime.ts b/packages/react-native-executorch/__tests__/support/lifetime.ts index acabb15c74..b189427d8d 100644 --- a/packages/react-native-executorch/__tests__/support/lifetime.ts +++ b/packages/react-native-executorch/__tests__/support/lifetime.ts @@ -9,9 +9,9 @@ * Disposal is idempotent, so tests that assert on `dispose()` explicitly can * still call it themselves. */ -type Disposable = { dispose: () => void }; +type NativeResource = { dispose: () => void }; -const created: Disposable[] = []; +const created: NativeResource[] = []; /** * Registers `instance` for disposal at the end of the current test. @@ -19,7 +19,7 @@ const created: Disposable[] = []; * @param instance The pipeline to track. * @returns The same instance. */ -export function tracked(instance: T): T { +export function tracked(instance: T): T { created.push(instance); return instance; } diff --git a/packages/react-native-executorch/__tests__/tasks/classification.test.ts b/packages/react-native-executorch/__tests__/tasks/classification.test.ts index 3ac20e0a4b..20f48bc510 100644 --- a/packages/react-native-executorch/__tests__/tasks/classification.test.ts +++ b/packages/react-native-executorch/__tests__/tasks/classification.test.ts @@ -3,7 +3,6 @@ import { createClassifier } from '../../src/extensions/cv/tasks/classification'; import { fakeJsi } from '../support/fakeJsi'; import { tracked } from '../support/lifetime'; import { STRETCH_PREPROCESSING, exported, imageBuffer, writesOutputs } from '../support/fixtures'; -import { allowNativeLeaks } from '../support/setup'; const MODEL_PATH = '/models/classifier.pte'; const LABELS = ['cat', 'dog', 'bird'] as const; @@ -46,7 +45,6 @@ describe('createClassifier — model acceptance', () => { }); await expect(createClassifier(config())).rejects.toThrow(/doesn't match any of the provided/); - allowNativeLeaks(); // see `tasks/constructionFailure.test.ts` }); it('rejects a labels array that does not match the output dimension', async () => { @@ -54,7 +52,6 @@ describe('createClassifier — model acceptance', () => { await expect(createClassifier(config(['cat', 'dog']))).rejects.toThrow( /labels length \(2\) must match model output dimension \(3\)/ ); - allowNativeLeaks(); // see `tasks/constructionFailure.test.ts` }); it('surfaces a load failure', async () => { diff --git a/packages/react-native-executorch/__tests__/tasks/constructionFailure.test.ts b/packages/react-native-executorch/__tests__/tasks/constructionFailure.test.ts index 4d5d738e24..bb4b07cb84 100644 --- a/packages/react-native-executorch/__tests__/tasks/constructionFailure.test.ts +++ b/packages/react-native-executorch/__tests__/tasks/constructionFailure.test.ts @@ -1,32 +1,47 @@ /** * What a `create` factory leaves behind when it throws. * - * Every factory follows the same shape: load the model (and, for some tasks, a - * tokenizer), validate its schema, pre-allocate the execution tensors, and only - * then hand back a `dispose`. If validation throws, the caller never receives a - * `dispose` — so anything already allocated is unreachable from JavaScript and - * stays alive in native memory for the process's lifetime. + * Every factory follows the same shape: load a model (and, depending on the + * task, a tokenizer, a phonemizer, an LLM runner or a whole nested pipeline), + * validate the schema, pre-allocate the execution tensors, and only then hand + * back a `dispose`. A failure anywhere after the first allocation means the + * caller never receives that `dispose`, so whatever was already allocated is + * unreachable from JavaScript — and native memory is not garbage collected, so + * it stays alive for the rest of the process. * - * That is what these tests record. They are written as assertions on the - * current* behavior rather than on the desired behavior, so the day a - * factory starts cleaning up after itself they fail loudly and can be flipped, - * instead of quietly passing either way. + * The exposure is not theoretical: `useModel` re-runs its factory whenever the + * config changes, so an app pointed at a mismatched model would leak one + * resource set per attempt. * - * The exposure is real: `useModel` re-runs its factory whenever the config - * changes, so an app pointed at a mismatched model leaks one native model per - * attempt. + * So every factory releases what it allocated before rethrowing, and this suite + * holds them to it. The assertions are per-resource-kind rather than left to + * the setup file's global leak check, so a failure says which factory and which + * kind of handle, not just that something leaked. */ import { f32, method } from '../../src/core/schema'; import { createClassifier } from '../../src/extensions/cv/tasks/classification'; import { createImageEmbedder } from '../../src/extensions/cv/tasks/imageEmbedding'; +import { createInstanceSegmenter } from '../../src/extensions/cv/tasks/instanceSegmentation'; +import { createKeypointDetector } from '../../src/extensions/cv/tasks/keypointDetection'; import { createObjectDetector } from '../../src/extensions/cv/tasks/objectDetection'; +import { createPaddleOcr } from '../../src/extensions/cv/tasks/paddleOcr'; +import { createSdxsTextToImage } from '../../src/extensions/cv/tasks/sdxsTextToImage'; import { createSemanticSegmenter } from '../../src/extensions/cv/tasks/semanticSegmentation'; import { createStyleTransfer } from '../../src/extensions/cv/tasks/styleTransfer'; +import { createPrivacyFilter } from '../../src/extensions/nlp/tasks/privacyFilter'; +import { createTextEmbedder } from '../../src/extensions/nlp/tasks/textEmbedding'; +import { createFsmnVoiceActivityDetector } from '../../src/extensions/speech/tasks/fsmnVoiceActivityDetection'; +import { createKokoroTextToSpeech } from '../../src/extensions/speech/tasks/kokoroTextToSpeech'; +import { createSupertonicTextToSpeech } from '../../src/extensions/speech/tasks/supertonicTextToSpeech'; +import { createWhisperSpeechToText } from '../../src/extensions/speech/tasks/whisperSpeechToText'; +import { fakeFs } from '../support/blobUtilMock'; import { fakeJsi } from '../support/fakeJsi'; import { STRETCH_PREPROCESSING, exported } from '../support/fixtures'; -import { allowNativeLeaks } from '../support/setup'; const MODEL_PATH = '/models/mismatched.pte'; +const TOKENIZER_PATH = '/models/tokenizer.json'; +const CHARSET_PATH = '/models/charset.json'; +const VOICE_PATH = '/models/voice.bin'; /** A schema no task pipeline declares: two inputs, three outputs, wrong ranks. */ const MISMATCHED = exported(method('forward', [f32(9), f32(9)], [f32(9), f32(9), f32(9)])); @@ -37,63 +52,149 @@ const CV_OPTS = { outInterpolation: 'linear', outNormalizeOpts: { alpha: 255, beta: 0 }, labels: ['a'], + landmarks: ['nose'], boxFormat: 'xyxy', defaultIouThreshold: 0.5, + defaultMaskThreshold: 0.5, defaultConfidenceThreshold: 0.5, } as const; -const factories = { - classifier: () => createClassifier({ modelPath: MODEL_PATH, modelOpts: CV_OPTS }), - objectDetector: () => createObjectDetector({ modelPath: MODEL_PATH, modelOpts: CV_OPTS }), - semanticSegmenter: () => createSemanticSegmenter({ modelPath: MODEL_PATH, modelOpts: CV_OPTS }), - styleTransfer: () => createStyleTransfer({ modelPath: MODEL_PATH, modelOpts: CV_OPTS }), - imageEmbedder: () => createImageEmbedder({ modelPath: MODEL_PATH, modelOpts: CV_OPTS }), +const cv = { modelPath: MODEL_PATH, modelOpts: CV_OPTS } as const; +const tokenized = { modelPath: MODEL_PATH, tokenizerPath: TOKENIZER_PATH } as const; +const vadModel = { + modelPath: MODEL_PATH, + defaultOptions: { + speechThreshold: 0.5, + minSpeechDurationMs: 100, + minSilenceDurationMs: 100, + speechPadMs: 30, + mergeGapMs: 100, + }, +} as const; + +/** + * Every factory that allocates something before it can fail, with whatever it + * needs registered so the failure lands on schema validation rather than on the + * load itself. + * + * `createTokenizer` is absent on purpose: it loads a tokenizer and returns, with + * nothing in between that can throw, so it has no window to leak through. + */ +const factories: Record Promise> = { + classifier: () => createClassifier(cv), + objectDetector: () => createObjectDetector(cv), + semanticSegmenter: () => createSemanticSegmenter(cv), + styleTransfer: () => createStyleTransfer(cv), + imageEmbedder: () => createImageEmbedder(cv), + instanceSegmenter: () => createInstanceSegmenter(cv), + keypointDetector: () => createKeypointDetector(cv), + paddleOcr: () => + createPaddleOcr({ + modelPath: MODEL_PATH, + charsetPath: CHARSET_PATH, + modelOpts: { defaultConfidenceThreshold: 0.5 }, + }), + sdxsTextToImage: () => createSdxsTextToImage(tokenized), + textEmbedder: () => createTextEmbedder(tokenized), + privacyFilter: () => + createPrivacyFilter({ + ...tokenized, + modelOpts: { labelNames: ['O', 'B-p', 'I-p', 'E-p', 'S-p'], padTokenId: 0 }, + }), + whisperSpeechToText: () => + createWhisperSpeechToText({ + ...tokenized, + vadModel, + supportedLanguages: ['en'], + } as never), + fsmnVoiceActivityDetector: () => createFsmnVoiceActivityDetector(vadModel), + supertonicTextToSpeech: () => + createSupertonicTextToSpeech({ + name: 'supertonic', + modelPaths: { + durationPredictor: MODEL_PATH, + vectorEstimator: MODEL_PATH, + textEncoder: MODEL_PATH, + vocoder: MODEL_PATH, + }, + voices: { v: VOICE_PATH }, + } as never), + kokoroTextToSpeech: () => + createKokoroTextToSpeech({ + name: 'kokoro', + modelPaths: { durationPredictor: MODEL_PATH, synthesizer: MODEL_PATH }, + phonemizer: { lang: 'en-us' }, + voices: { v: VOICE_PATH }, + } as never), }; +/** Everything the fake still considers allocated, as one comparable object. */ +const stillAllocated = () => ({ + models: fakeJsi.liveModels(), + tokenizers: fakeJsi.liveTokenizers(), + runners: fakeJsi.liveRunners(), + phonemizers: fakeJsi.livePhonemizers(), + tensors: fakeJsi.liveTensorDescriptions(), +}); + +const NOTHING = { models: [], tokenizers: [], runners: [], phonemizers: [], tensors: [] }; + describe('create — schema validation failure', () => { beforeEach(() => { fakeJsi.registerModel(MODEL_PATH, { schema: MISMATCHED }); + fakeJsi.registerTokenizer(TOKENIZER_PATH, { tokens: ['a', 'b'] }); + fakeFs.write(CHARSET_PATH, JSON.stringify(['a', 'b', 'c'])); + fakeFs.write(VOICE_PATH, new Uint8Array(1024)); + }); + + it.each(Object.entries(factories))('create%s rejects rather than resolving', async (_n, make) => { + await expect(make()).rejects.toThrow(); }); it.each(Object.entries(factories))( - 'create%s rejects with a message naming every variant it tried', - async (_name, factory) => { - await expect(factory()).rejects.toThrow(/doesn't match any of the provided variants/); - allowNativeLeaks(); + 'create%s releases everything it had allocated', + async (_name, make) => { + await expect(make()).rejects.toThrow(); + + expect(stillAllocated()).toEqual(NOTHING); } ); - it.each(Object.entries(factories))( - 'create%s abandons the loaded native model (known leak)', - async (_name, factory) => { - await expect(factory()).rejects.toThrow(); - - expect(fakeJsi.liveModels()).toEqual([MODEL_PATH]); - allowNativeLeaks(); + it.each(Object.entries(factories).filter(([name]) => name !== 'kokoroTextToSpeech'))( + 'create%s names every variant it tried', + async (_name, make) => { + // A caller should learn what is wrong from the message. Kokoro validates + // two sub-models and reports whichever failed first, so its wording is + // pinned in its own suite instead. + await expect(make()).rejects.toThrow(/doesn't match any of the provided variants/); } ); - it('leaks one model per failed attempt, the way a re-rendering hook would', async () => { + it('stays clean across repeated attempts, the way a re-rendering hook would', async () => { + // `useModel` re-runs its factory on every config change. Before the + // factories cleaned up, this leaked one model per attempt. for (const path of ['/a.pte', '/b.pte', '/c.pte']) { fakeJsi.registerModel(path, { schema: MISMATCHED }); - await expect(createClassifier({ modelPath: path, modelOpts: CV_OPTS })).rejects.toThrow(); + await expect(createClassifier({ ...cv, modelPath: path })).rejects.toThrow(); } - expect(fakeJsi.liveModels()).toEqual(['/a.pte', '/b.pte', '/c.pte']); - allowNativeLeaks(); + expect(stillAllocated()).toEqual(NOTHING); }); }); describe('create — load failure', () => { it.each(Object.entries(factories))( 'create%s leaves nothing allocated when the model itself cannot be loaded', - async (_name, factory) => { - // Nothing registered at MODEL_PATH, so `loadModel` throws before any - // allocation happens — the one failure path that is already clean. - await expect(factory()).rejects.toThrow(/mismatched.pte/); + async (_name, make) => { + // Nothing registered, so the very first load throws. The factory has + // allocated nothing yet, but for the multi-resource pipelines a sibling + // load may already have succeeded. + fakeFs.write(CHARSET_PATH, JSON.stringify(['a', 'b', 'c'])); + fakeFs.write(VOICE_PATH, new Uint8Array(1024)); + + await expect(make()).rejects.toThrow(); - expect(fakeJsi.liveModels()).toEqual([]); - expect(fakeJsi.liveTensors()).toBe(0); + expect(stillAllocated()).toEqual(NOTHING); } ); }); diff --git a/packages/react-native-executorch/__tests__/tasks/embedding.test.ts b/packages/react-native-executorch/__tests__/tasks/embedding.test.ts index 90c46db8eb..ade276639a 100644 --- a/packages/react-native-executorch/__tests__/tasks/embedding.test.ts +++ b/packages/react-native-executorch/__tests__/tasks/embedding.test.ts @@ -5,7 +5,6 @@ import { fakeJsi } from '../support/fakeJsi'; import { tracked } from '../support/lifetime'; import type { FakeTensor } from '../support/fakeTensor'; import { STRETCH_PREPROCESSING, exported, imageBuffer, writesOutputs } from '../support/fixtures'; -import { allowNativeLeaks } from '../support/setup'; const MODEL_PATH = '/models/embedder.pte'; const TOKENIZER_PATH = '/models/tokenizer.json'; @@ -107,7 +106,6 @@ describe('createTextEmbedder', () => { await expect(createTextEmbedder(config)).rejects.toThrow( /doesn't match any of the provided variants/ ); - allowNativeLeaks(); // see `tasks/constructionFailure.test.ts` }); it('feeds tokens at their exact length with an all-ones attention mask', async () => { diff --git a/packages/react-native-executorch/__tests__/tasks/kokoroTextToSpeech.test.ts b/packages/react-native-executorch/__tests__/tasks/kokoroTextToSpeech.test.ts index 635b2b5e62..00297560c2 100644 --- a/packages/react-native-executorch/__tests__/tasks/kokoroTextToSpeech.test.ts +++ b/packages/react-native-executorch/__tests__/tasks/kokoroTextToSpeech.test.ts @@ -20,7 +20,6 @@ import { fakePhonemizer } from '../support/fakeOps'; import { fakeFs } from '../support/blobUtilMock'; import { exported } from '../support/fixtures'; import { tracked } from '../support/lifetime'; -import { allowNativeLeaks } from '../support/setup'; const PREDICTOR_PATH = '/models/duration_predictor.pte'; const SYNTHESIZER_PATH = '/models/synthesizer.pte'; @@ -216,7 +215,6 @@ describe('createKokoroTextToSpeech — the model contract', () => { fakeFs.remove(VOICE_PATH); await expect(createKokoroTextToSpeech(config)).rejects.toThrow(/ENOENT/); - allowNativeLeaks(); // see `tasks/constructionFailure.test.ts` }); it('releases both models, the phonemizer and its tensors on dispose', async () => { diff --git a/packages/react-native-executorch/__tests__/tasks/llmChatSession.test.ts b/packages/react-native-executorch/__tests__/tasks/llmChatSession.test.ts index 919ace54af..e6003dc6d0 100644 --- a/packages/react-native-executorch/__tests__/tasks/llmChatSession.test.ts +++ b/packages/react-native-executorch/__tests__/tasks/llmChatSession.test.ts @@ -17,7 +17,6 @@ import type { ToolCall, ToolParserResult } from '../../src/extensions/llm/utils/ import { fakeJsi } from '../support/fakeJsi'; import { fakeFs } from '../support/blobUtilMock'; import { tracked } from '../support/lifetime'; -import { allowNativeLeaks } from '../support/setup'; const MODEL_PATH = '/models/llm.pte'; const TOKENIZER_PATH = '/models/tokenizer.json'; @@ -405,6 +404,20 @@ describe('createLLMChatSession — failure', () => { writeTokenizerConfig(); await expect(createLLMChatSession(config)).rejects.toThrow(/no runner registered/); - allowNativeLeaks(); // see `tasks/constructionFailure.test.ts` + expect(fakeJsi.liveRunners()).toEqual([]); + }); + + it('releases the runner when the initial prefill fails', async () => { + // The one window where construction allocates and can still throw: the + // runner exists, and rendering the initial messages blows up. The caller + // never gets a `dispose`, so the factory has to release it itself. + // eslint-disable-next-line camelcase + writeRawTokenizerConfig({ chat_template: '{{ not_a_filter() }}', eos_token: EOS }); + + await expect( + createLLMChatSession(config, { initialMessages: [{ role: 'user', content: 'hi' }] }) + ).rejects.toThrow(); + + expect(fakeJsi.liveRunners()).toEqual([]); }); }); diff --git a/packages/react-native-executorch/__tests__/tasks/objectDetection.test.ts b/packages/react-native-executorch/__tests__/tasks/objectDetection.test.ts index ee711b213b..ba5fbfbfd4 100644 --- a/packages/react-native-executorch/__tests__/tasks/objectDetection.test.ts +++ b/packages/react-native-executorch/__tests__/tasks/objectDetection.test.ts @@ -3,7 +3,6 @@ import { createObjectDetector } from '../../src/extensions/cv/tasks/objectDetect import { fakeJsi } from '../support/fakeJsi'; import { tracked } from '../support/lifetime'; import { STRETCH_PREPROCESSING, exported, imageBuffer, writesOutputs } from '../support/fixtures'; -import { allowNativeLeaks } from '../support/setup'; const MODEL_PATH = '/models/detector.pte'; const LABELS = ['person', 'car', 'dog'] as const; @@ -182,7 +181,6 @@ describe('createObjectDetector — model acceptance', () => { }); await expect(createObjectDetector(config)).rejects.toThrow(/Output count mismatch/); - allowNativeLeaks(); // see `tasks/constructionFailure.test.ts` }); it('requires boxes, scores and classes to agree on the candidate count', async () => { @@ -193,7 +191,6 @@ describe('createObjectDetector — model acceptance', () => { }); await expect(createObjectDetector(config)).rejects.toThrow(/inconsistent bindings/); - allowNativeLeaks(); // see `tasks/constructionFailure.test.ts` }); }); diff --git a/packages/react-native-executorch/__tests__/tasks/paddleOcr.test.ts b/packages/react-native-executorch/__tests__/tasks/paddleOcr.test.ts index 35bbf88e20..b797a111a2 100644 --- a/packages/react-native-executorch/__tests__/tasks/paddleOcr.test.ts +++ b/packages/react-native-executorch/__tests__/tasks/paddleOcr.test.ts @@ -19,7 +19,6 @@ import { fakeJsi } from '../support/fakeJsi'; import { fakeFs } from '../support/blobUtilMock'; import { exported } from '../support/fixtures'; import { tracked } from '../support/lifetime'; -import { allowNativeLeaks } from '../support/setup'; const MODEL_PATH = '/models/paddle_ocr.pte'; const CHARSET_PATH = '/models/charset.json'; @@ -143,7 +142,6 @@ describe('createPaddleOcr — the model contract', () => { }); await expect(createPaddleOcr(config)).rejects.toThrow(/recognize/); - allowNativeLeaks(); // see `tasks/constructionFailure.test.ts` }); it('rejects a recognizer whose width is not 8x its CTC time axis', async () => { @@ -161,21 +159,18 @@ describe('createPaddleOcr — the model contract', () => { }); await expect(createPaddleOcr(config)).rejects.toThrow(/constraint/i); - allowNativeLeaks(); // see `tasks/constructionFailure.test.ts` }); it('rejects a charset that does not cover the recognizer vocabulary', async () => { fakeFs.write(CHARSET_PATH, JSON.stringify([...CHARSET, 'd'])); await expect(createPaddleOcr(config)).rejects.toThrow(/charset size/); - allowNativeLeaks(); // see `tasks/constructionFailure.test.ts` }); it('surfaces a missing charset file', async () => { fakeFs.remove(CHARSET_PATH); await expect(createPaddleOcr(config)).rejects.toThrow(/ENOENT/); - allowNativeLeaks(); // see `tasks/constructionFailure.test.ts` }); it('releases the model on dispose', async () => { diff --git a/packages/react-native-executorch/__tests__/tasks/privacyFilter.test.ts b/packages/react-native-executorch/__tests__/tasks/privacyFilter.test.ts index 561a6d83c6..30ffbe0390 100644 --- a/packages/react-native-executorch/__tests__/tasks/privacyFilter.test.ts +++ b/packages/react-native-executorch/__tests__/tasks/privacyFilter.test.ts @@ -19,7 +19,6 @@ import { createPrivacyFilter } from '../../src/extensions/nlp/tasks/privacyFilte import { fakeJsi, type FakeExecute } from '../support/fakeJsi'; import { exported } from '../support/fixtures'; import { tracked } from '../support/lifetime'; -import { allowNativeLeaks } from '../support/setup'; const MODEL_PATH = '/models/privacy-filter.pte'; const TOKENIZER_PATH = '/models/tokenizer.json'; @@ -126,14 +125,12 @@ describe('createPrivacyFilter — the label space', () => { await expect(createPrivacyFilter(config)).rejects.toThrow( /output #0 Tensor dim #2: Constant dimension mismatch/ ); - allowNativeLeaks(); // see `tasks/constructionFailure.test.ts` }); it('rejects a window too short to hold a span and its context', async () => { fakeJsi.registerModel(MODEL_PATH, { schema: staticSchema(1) }); await expect(createPrivacyFilter(config)).rejects.toThrow(/at least 2 tokens/); - allowNativeLeaks(); // see `tasks/constructionFailure.test.ts` }); }); diff --git a/packages/react-native-executorch/__tests__/tasks/remainingTasks.test.ts b/packages/react-native-executorch/__tests__/tasks/remainingTasks.test.ts index 3f2713418f..7840ce60b0 100644 --- a/packages/react-native-executorch/__tests__/tasks/remainingTasks.test.ts +++ b/packages/react-native-executorch/__tests__/tasks/remainingTasks.test.ts @@ -22,7 +22,6 @@ import { createFsmnVoiceActivityDetector } from '../../src/extensions/speech/tas import { createWhisperSpeechToText } from '../../src/extensions/speech/tasks/whisperSpeechToText'; import { fakeJsi } from '../support/fakeJsi'; import { STRETCH_PREPROCESSING, exported } from '../support/fixtures'; -import { allowNativeLeaks } from '../support/setup'; const MODEL_PATH = '/models/task.pte'; const TOKENIZER_PATH = '/models/tokenizer.json'; @@ -67,7 +66,6 @@ describe('createKeypointDetector', () => { }); await expect(createKeypointDetector(config)).rejects.toThrow(/Constant dimension mismatch/); - allowNativeLeaks(); // see `tasks/constructionFailure.test.ts` }); it('has no unbatched variant', async () => { @@ -78,7 +76,6 @@ describe('createKeypointDetector', () => { }); await expect(createKeypointDetector(config)).rejects.toThrow(/Rank mismatch/); - allowNativeLeaks(); // see `tasks/constructionFailure.test.ts` }); }); @@ -138,7 +135,6 @@ describe('createInstanceSegmenter', () => { }); await expect(createInstanceSegmenter(config)).rejects.toThrow(/inconsistent bindings/); - allowNativeLeaks(); // see `tasks/constructionFailure.test.ts` }); }); @@ -182,7 +178,6 @@ describe('createFsmnVoiceActivityDetector', () => { await expect(createFsmnVoiceActivityDetector(VAD_CONFIG)).rejects.toThrow( /Cannot match symbolic 'dynamic' with concrete 'constant'/ ); - allowNativeLeaks(); // see `tasks/constructionFailure.test.ts` }); it('requires the input and output frame dimensions to share a domain', async () => { @@ -195,7 +190,6 @@ describe('createFsmnVoiceActivityDetector', () => { await expect(createFsmnVoiceActivityDetector(VAD_CONFIG)).rejects.toThrow( /inconsistent bindings/ ); - allowNativeLeaks(); // see `tasks/constructionFailure.test.ts` }); it('returns no segments for a waveform shorter than one analysis frame', async () => { @@ -267,7 +261,6 @@ describe('createWhisperSpeechToText', () => { await expect(createWhisperSpeechToText(config)).rejects.toThrow( /Method 'decode' not found in exported model spec/ ); - allowNativeLeaks(); // see `tasks/constructionFailure.test.ts` }); it('requires encode and decode to agree on the encoder state shape', async () => { @@ -280,7 +273,6 @@ describe('createWhisperSpeechToText', () => { }); await expect(createWhisperSpeechToText(config)).rejects.toThrow(/inconsistent bindings/); - allowNativeLeaks(); // see `tasks/constructionFailure.test.ts` }); it('fails when the tokenizer has no end-of-text token', async () => { @@ -288,7 +280,6 @@ describe('createWhisperSpeechToText', () => { fakeJsi.registerTokenizer(TOKENIZER_PATH, { tokens: ['hello'] }); await expect(createWhisperSpeechToText(config)).rejects.toThrow(/<\|endoftext\|>/); - allowNativeLeaks(); // see `tasks/constructionFailure.test.ts` }); }); @@ -328,7 +319,6 @@ describe('createSdxsTextToImage', () => { fakeJsi.registerTokenizer(TOKENIZER_PATH, { tokens: ['a'] }); await expect(createSdxsTextToImage(config)).rejects.toThrow(/Method 'denoise' not found/); - allowNativeLeaks(); // see `tasks/constructionFailure.test.ts` }); it('rejects a text encoder with a different hidden size', async () => { @@ -342,6 +332,5 @@ describe('createSdxsTextToImage', () => { fakeJsi.registerTokenizer(TOKENIZER_PATH, { tokens: ['a'] }); await expect(createSdxsTextToImage(config)).rejects.toThrow(/Constant dimension mismatch/); - allowNativeLeaks(); // see `tasks/constructionFailure.test.ts` }); }); diff --git a/packages/react-native-executorch/__tests__/tasks/semanticSegmentation.test.ts b/packages/react-native-executorch/__tests__/tasks/semanticSegmentation.test.ts index 5e066e2e41..4e907abea4 100644 --- a/packages/react-native-executorch/__tests__/tasks/semanticSegmentation.test.ts +++ b/packages/react-native-executorch/__tests__/tasks/semanticSegmentation.test.ts @@ -3,7 +3,6 @@ import { createSemanticSegmenter } from '../../src/extensions/cv/tasks/semanticS import { fakeJsi } from '../support/fakeJsi'; import { tracked } from '../support/lifetime'; import { STRETCH_PREPROCESSING, exported, imageBuffer, writesOutputs } from '../support/fixtures'; -import { allowNativeLeaks } from '../support/setup'; const MODEL_PATH = '/models/segmenter.pte'; const LABELS = ['background', 'person', 'cat'] as const; @@ -54,7 +53,6 @@ describe('createSemanticSegmenter — multi-class models', () => { await expect( createSemanticSegmenter({ ...config, modelOpts: { ...options, labels: ['only-one'] } }) ).rejects.toThrow(/Model outputs 3 classes, but 1 labels were provided/); - allowNativeLeaks(); // see `tasks/constructionFailure.test.ts` }); it('returns an RGBA mask at the input image resolution', async () => { diff --git a/packages/react-native-executorch/__tests__/tasks/styleTransfer.test.ts b/packages/react-native-executorch/__tests__/tasks/styleTransfer.test.ts index 606f6612ed..310cf1317b 100644 --- a/packages/react-native-executorch/__tests__/tasks/styleTransfer.test.ts +++ b/packages/react-native-executorch/__tests__/tasks/styleTransfer.test.ts @@ -8,7 +8,6 @@ import { exported, imageBuffer, } from '../support/fixtures'; -import { allowNativeLeaks } from '../support/setup'; const MODEL_PATH = '/models/style.pte'; const SIZE = 4; @@ -77,7 +76,6 @@ describe('createStyleTransfer', () => { }); await expect(createStyleTransfer(config)).rejects.toThrow(/inconsistent bindings/); - allowNativeLeaks(); // see `tasks/constructionFailure.test.ts` }); it('produces the same result synchronously and asynchronously', async () => { diff --git a/packages/react-native-executorch/src/core/lifetime.ts b/packages/react-native-executorch/src/core/lifetime.ts new file mode 100644 index 0000000000..b0ed1f11a5 --- /dev/null +++ b/packages/react-native-executorch/src/core/lifetime.ts @@ -0,0 +1,80 @@ +/** + * Construction-time ownership of native resources. + * + * A factory that builds a pipeline allocates as it goes: it loads models, + * validates their schemas, pre-allocates execution tensors, and only at the end + * hands the caller a `dispose`. Anything that throws in between leaves the + * caller with no reference to what was already allocated, and native memory is + * not garbage collected, so it stays alive for the rest of the process. + * + * A scope closes that window: every resource is tracked as it is created, and + * one `dispose` releases whatever exists at the moment it is called. The same + * function is the pipeline's own `dispose`, so there is a single teardown path + * rather than one for failure and one for success. + */ + +/** + * Anything holding native memory that has to be released explicitly. + * @category Core / Types + */ +export type NativeResource = { dispose: () => void }; + +/** + * A set of native resources with a single teardown. + * @category Core / Types + */ +export type ResourceScope = { + /** + * Takes ownership of a resource and returns it unchanged, so it can wrap an + * allocation in place. + */ + readonly track: (resource: R) => R; + + /** + * Releases every tracked resource, most recently allocated first. Safe to + * call more than once: a second call has nothing left to release. + */ + readonly dispose: () => void; +}; + +/** + * Creates a {@link ResourceScope} for semi-automatic lifetime management of + * native resources. + * + * Wrap the allocating body in `try`/`catch`, `track` each resource as it is + * created, and return `scope.dispose` as the resulting object's `dispose`. A + * failure part-way through then releases everything allocated so far, and a + * successful build hands the caller the same teardown path. + * @category Core / Functions + * @returns A scope that tracks resources and releases them on `dispose`. + * @example + * ```typescript + * const scope = createResourceScope(); + * const dispose = scope.dispose; + * try { + * const model = scope.track(await wrapAsync(loadModel, runtime)(modelPath)); + * const { dims } = validateSpec(model.schema, { ... }); // may throw + * const tensors = [scope.track(tensor('float32', shape))]; + * return { run, dispose }; + * } catch (error) { + * dispose(); + * throw error; + * } + * ``` + */ +export function createResourceScope(): ResourceScope { + const allocated: NativeResource[] = []; + + return { + track: (resource: R): R => { + allocated.push(resource); + return resource; + }, + // Reverse order, so a resource is released before whatever it was built + // from. `splice` empties the list as it goes, which is what makes a second + // call a no-op rather than a double dispose. + dispose: (): void => { + for (const resource of allocated.splice(0).reverse()) resource.dispose(); + }, + }; +} diff --git a/packages/react-native-executorch/src/extensions/cv/tasks/classification.ts b/packages/react-native-executorch/src/extensions/cv/tasks/classification.ts index bb4c059c3e..adec88a345 100644 --- a/packages/react-native-executorch/src/extensions/cv/tasks/classification.ts +++ b/packages/react-native-executorch/src/extensions/cv/tasks/classification.ts @@ -14,6 +14,7 @@ import { softmax } from '../../math'; import type { ImageBuffer } from '../image'; import { createImagePreprocessor, type ImagePreprocessorOptions } from '../utils/imagePreprocessor'; import { RnExecuTorchError } from '../../../core/error'; +import { createResourceScope } from '../../../core/lifetime'; /** * Options for configuring an image classifier preprocessor and label @@ -120,69 +121,76 @@ export async function createClassifier( config: ClassifierModel, runtime?: WorkletRuntime ): Promise> { - const { modelPath, modelOpts } = config; - const model = await wrapAsync(loadModel, runtime)(modelPath); - - const { variant, dims } = validateSpec(model.schema, { - batched: method( - 'forward', // prettier-ignore - [f32(1, 3, 'H', 'W')], - [f32(1, 'N')] - ), - unbatched: method( - 'forward', // prettier-ignore - [f32(3, 'H', 'W')], - [f32('N')] - ), - }); - - const [N, H, W] = dims.constant('N', 'H', 'W'); - const inpShape = { batched: [1, 3, H, W], unbatched: [3, H, W] }[variant]; - const outShape = { batched: [1, N], unbatched: [N] }[variant]; - - if (modelOpts.labels.length !== N) { - throw RnExecuTorchError( - 'INVALID_ARGUMENT', - `Classifier labels length (${modelOpts.labels.length}) must match model output dimension (${N}).` - ); - } - - const tensors = [ - tensor('float32', outShape), // prettier-ignore - tensor('float32', outShape), - ] as const; - - const [tLogits, tProbas] = tensors; - const preprocessor = createImagePreprocessor(modelOpts, inpShape); - - const dispose = () => { - preprocessor.dispose(); - tensors.forEach((t) => t.dispose()); - model.dispose(); - }; - - const classifyWorklet = (input: ImageBuffer, options?: ClassifyOptions): Classification[] => { - 'worklet'; - if (options?.topk !== undefined && options.topk < 0) { + const scope = createResourceScope(); + const dispose = scope.dispose; + + try { + const { modelPath, modelOpts } = config; + const model = scope.track(await wrapAsync(loadModel, runtime)(modelPath)); + + const { variant, dims } = validateSpec(model.schema, { + batched: method( + 'forward', // prettier-ignore + [f32(1, 3, 'H', 'W')], + [f32(1, 'N')] + ), + unbatched: method( + 'forward', // prettier-ignore + [f32(3, 'H', 'W')], + [f32('N')] + ), + }); + + const [N, H, W] = dims.constant('N', 'H', 'W'); + const inpShape = { batched: [1, 3, H, W], unbatched: [3, H, W] }[variant]; + const outShape = { batched: [1, N], unbatched: [N] }[variant]; + + if (modelOpts.labels.length !== N) { throw RnExecuTorchError( 'INVALID_ARGUMENT', - `Classifier topk option must be non-negative, got ${options.topk}` + `Classifier labels length (${modelOpts.labels.length}) must match model output dimension (${N}).` ); } - const tInput = preprocessor.process(input); - model.execute('forward', [tInput], [tLogits]); - const probas = tLogits - .through(softmax, tProbas) // prettier-ignore - .getData(new Float32Array(tProbas.numel)); - - return Array.from(probas) - .map((confidence, index) => ({ confidence, label: modelOpts.labels[index]! })) - .sort((a, b) => b.confidence - a.confidence) - .slice(0, options?.topk); - }; - - const classify = wrapAsync(classifyWorklet, runtime); - - return { classify, classifyWorklet, dispose }; + const tensors = [ + tensor('float32', outShape), // prettier-ignore + tensor('float32', outShape), + ] as const; + + tensors.forEach(scope.track); + + const [tLogits, tProbas] = tensors; + const preprocessor = scope.track(createImagePreprocessor(modelOpts, inpShape)); + + const classifyWorklet = ( + input: ImageBuffer, + options?: ClassifyOptions + ): Classification[] => { + 'worklet'; + if (options?.topk !== undefined && options.topk < 0) { + throw RnExecuTorchError( + 'INVALID_ARGUMENT', + `Classifier topk option must be non-negative, got ${options.topk}` + ); + } + const tInput = preprocessor.process(input); + model.execute('forward', [tInput], [tLogits]); + + const probas = tLogits + .through(softmax, tProbas) // prettier-ignore + .getData(new Float32Array(tProbas.numel)); + + return Array.from(probas) + .map((confidence, index) => ({ confidence, label: modelOpts.labels[index]! })) + .sort((a, b) => b.confidence - a.confidence) + .slice(0, options?.topk); + }; + + const classify = wrapAsync(classifyWorklet, runtime); + + return { classify, classifyWorklet, dispose }; + } catch (error) { + dispose(); + throw error; + } } diff --git a/packages/react-native-executorch/src/extensions/cv/tasks/imageEmbedding.ts b/packages/react-native-executorch/src/extensions/cv/tasks/imageEmbedding.ts index 868e3bdada..816af7e879 100644 --- a/packages/react-native-executorch/src/extensions/cv/tasks/imageEmbedding.ts +++ b/packages/react-native-executorch/src/extensions/cv/tasks/imageEmbedding.ts @@ -8,6 +8,7 @@ import { tensor } from '../../../core/tensor'; import { loadModel } from '../../../core/model'; import { validateSpec, method, f32 } from '../../../core/schema'; import { wrapAsync } from '../../../core/runtime'; +import { createResourceScope } from '../../../core/lifetime'; import type { ImageBuffer } from '../image'; import { createImagePreprocessor, type ImagePreprocessorOptions } from '../utils/imagePreprocessor'; @@ -75,44 +76,48 @@ export async function createImageEmbedder( config: ImageEmbedderModel, runtime?: WorkletRuntime ): Promise { - const { modelPath, modelOpts } = config; - const model = await wrapAsync(loadModel, runtime)(modelPath); - - const { variant, dims } = validateSpec(model.schema, { - batched: method( - 'forward', // prettier-ignore - [f32(1, 3, 'H', 'W')], - [f32(1, 'D')] - ), - unbatched: method( - 'forward', // prettier-ignore - [f32(3, 'H', 'W')], - [f32('D')] - ), - }); - - const [D, H, W] = dims.constant('D', 'H', 'W'); - const inpShape = { batched: [1, 3, H, W], unbatched: [3, H, W] }[variant]; - const outShape = { batched: [1, D], unbatched: [D] }[variant]; - - const tensors = [tensor('float32', outShape)] as const; - const [tEmbedding] = tensors; - const preprocessor = createImagePreprocessor(modelOpts, inpShape); - - const dispose = () => { - preprocessor.dispose(); - tensors.forEach((t) => t.dispose()); - model.dispose(); - }; - - const embedWorklet = (input: ImageBuffer): Float32Array => { - 'worklet'; - const tInput = preprocessor.process(input); - model.execute('forward', [tInput], [tEmbedding]); - return tEmbedding.getData(new Float32Array(tEmbedding.numel)); - }; - - const embed = wrapAsync(embedWorklet, runtime); - - return { embed, embedWorklet, dispose }; + const scope = createResourceScope(); + const dispose = scope.dispose; + + try { + const { modelPath, modelOpts } = config; + const model = scope.track(await wrapAsync(loadModel, runtime)(modelPath)); + + const { variant, dims } = validateSpec(model.schema, { + batched: method( + 'forward', // prettier-ignore + [f32(1, 3, 'H', 'W')], + [f32(1, 'D')] + ), + unbatched: method( + 'forward', // prettier-ignore + [f32(3, 'H', 'W')], + [f32('D')] + ), + }); + + const [D, H, W] = dims.constant('D', 'H', 'W'); + const inpShape = { batched: [1, 3, H, W], unbatched: [3, H, W] }[variant]; + const outShape = { batched: [1, D], unbatched: [D] }[variant]; + + const tensors = [tensor('float32', outShape)] as const; + + tensors.forEach(scope.track); + const [tEmbedding] = tensors; + const preprocessor = scope.track(createImagePreprocessor(modelOpts, inpShape)); + + const embedWorklet = (input: ImageBuffer): Float32Array => { + 'worklet'; + const tInput = preprocessor.process(input); + model.execute('forward', [tInput], [tEmbedding]); + return tEmbedding.getData(new Float32Array(tEmbedding.numel)); + }; + + const embed = wrapAsync(embedWorklet, runtime); + + return { embed, embedWorklet, dispose }; + } catch (error) { + dispose(); + throw error; + } } diff --git a/packages/react-native-executorch/src/extensions/cv/tasks/instanceSegmentation.ts b/packages/react-native-executorch/src/extensions/cv/tasks/instanceSegmentation.ts index 1390ab185c..72dc8f8abb 100644 --- a/packages/react-native-executorch/src/extensions/cv/tasks/instanceSegmentation.ts +++ b/packages/react-native-executorch/src/extensions/cv/tasks/instanceSegmentation.ts @@ -23,6 +23,7 @@ import { type BoxFormat, } from '../ops/box'; import { RnExecuTorchError } from '../../../core/error'; +import { createResourceScope } from '../../../core/lifetime'; /** * Options for configuring an instance segmenter preprocessor, label @@ -166,134 +167,138 @@ export async function createInstanceSegmenter( config: InstanceSegmenterModel, runtime?: WorkletRuntime ): Promise> { - const { modelPath, modelOpts } = config; - const model = await wrapAsync(loadModel, runtime)(modelPath); - - const { variant, dims } = validateSpec(model.schema, { - batched: method( - 'forward', - [f32(1, 3, 'H', 'W')], - [f32('N', 4), f32('N'), f32('N'), f32('N', 'MH', 'MW')] - ), - unbatched: method( - 'forward', - [f32(3, 'H', 'W')], - [f32('N', 4), f32('N'), f32('N'), f32('N', 'MH', 'MW')] - ), - }); - - const [N, H, W, maskH, maskW] = dims.constant('N', 'H', 'W', 'MH', 'MW'); - const inpShape = { batched: [1, 3, H, W], unbatched: [3, H, W] }[variant]; - const outShape = { boxes: [N, 4], scores: [N], classes: [N], masks: [N, maskH, maskW] }; - - const tensors = [ - tensor('float32', outShape.boxes), - tensor('float32', outShape.scores), - tensor('float32', outShape.classes), - tensor('float32', outShape.masks), - tensor('float32', [maskH, maskW, 1]), - ] as const; - - const [tBoxes, tScores, tClasses, tAllMasks, tMask] = tensors; - - const preprocessor = createImagePreprocessor(modelOpts, inpShape); - - const dispose = () => { - preprocessor.dispose(); - tensors.forEach((t) => t.dispose()); - model.dispose(); - }; - - const segmentInstancesWorklet = ( - input: ImageBuffer, - options?: { confidenceThreshold?: number; iouThreshold?: number; maskThreshold?: number } - ): InstanceSegmentationResult[] => { - 'worklet'; - const tInput = preprocessor.process(input); - model.execute('forward', [tInput], [tBoxes, tScores, tClasses, tAllMasks]); - - const iouThreshold = options?.iouThreshold ?? modelOpts.defaultIouThreshold; - const maskThreshold = options?.maskThreshold ?? modelOpts.defaultMaskThreshold; - const confidenceThreshold = - options?.confidenceThreshold ?? modelOpts.defaultConfidenceThreshold; - - const eps = 1e-7; - const clampedMaskThreshold = Math.max(eps, Math.min(1 - eps, maskThreshold)); - const logitMaskThreshold = Math.log(clampedMaskThreshold / (1 - clampedMaskThreshold)); - - const indices = nms(tBoxes, tScores, { - boxFormat: modelOpts.boxFormat, - iouThreshold, - confidenceThreshold, - nmsType: 'standard', + const scope = createResourceScope(); + const dispose = scope.dispose; + + try { + const { modelPath, modelOpts } = config; + const model = scope.track(await wrapAsync(loadModel, runtime)(modelPath)); + + const { variant, dims } = validateSpec(model.schema, { + batched: method( + 'forward', + [f32(1, 3, 'H', 'W')], + [f32('N', 4), f32('N'), f32('N'), f32('N', 'MH', 'MW')] + ), + unbatched: method( + 'forward', + [f32(3, 'H', 'W')], + [f32('N', 4), f32('N'), f32('N'), f32('N', 'MH', 'MW')] + ), }); - const boxes = tBoxes.getData(new Float32Array(tBoxes.numel)); - const scores = tScores.getData(new Float32Array(tScores.numel)); - const classes = tClasses.getData(new Float32Array(tClasses.numel)); + const [N, H, W, maskH, maskW] = dims.constant('N', 'H', 'W', 'MH', 'MW'); + const inpShape = { batched: [1, 3, H, W], unbatched: [3, H, W] }[variant]; + const outShape = { boxes: [N, 4], scores: [N], classes: [N], masks: [N, maskH, maskW] }; - const auxTensors = [ - tensor('float32', [input.height, input.width, 1]), - tensor('float32', [input.height, input.width, 1]), - tensor('float32', [input.height, input.width, 1]), - tensor('uint8', [input.height, input.width, 1]), + const tensors = [ + tensor('float32', outShape.boxes), + tensor('float32', outShape.scores), + tensor('float32', outShape.classes), + tensor('float32', outShape.masks), + tensor('float32', [maskH, maskW, 1]), ] as const; - const [tResize, tThreshold, tCrop, tUint8] = auxTensors; - - const results: InstanceSegmentationResult[] = []; - - try { - for (const idx of indices) { - const confidence = scores[idx]!; - const classIdx = Math.round(classes[idx]!); - const label = modelOpts.labels[classIdx]; - - if (label === undefined) { - throw RnExecuTorchError( - 'INVALID_ARGUMENT', - `InstanceSegmenter: Predicted class index ${classIdx} is ` + - `out of bounds for labels array of size ${modelOpts.labels.length}.` - ); + tensors.forEach(scope.track); + + const [tBoxes, tScores, tClasses, tAllMasks, tMask] = tensors; + + const preprocessor = scope.track(createImagePreprocessor(modelOpts, inpShape)); + + const segmentInstancesWorklet = ( + input: ImageBuffer, + options?: { confidenceThreshold?: number; iouThreshold?: number; maskThreshold?: number } + ): InstanceSegmentationResult[] => { + 'worklet'; + const tInput = preprocessor.process(input); + model.execute('forward', [tInput], [tBoxes, tScores, tClasses, tAllMasks]); + + const iouThreshold = options?.iouThreshold ?? modelOpts.defaultIouThreshold; + const maskThreshold = options?.maskThreshold ?? modelOpts.defaultMaskThreshold; + const confidenceThreshold = + options?.confidenceThreshold ?? modelOpts.defaultConfidenceThreshold; + + const eps = 1e-7; + const clampedMaskThreshold = Math.max(eps, Math.min(1 - eps, maskThreshold)); + const logitMaskThreshold = Math.log(clampedMaskThreshold / (1 - clampedMaskThreshold)); + + const indices = nms(tBoxes, tScores, { + boxFormat: modelOpts.boxFormat, + iouThreshold, + confidenceThreshold, + nmsType: 'standard', + }); + + const boxes = tBoxes.getData(new Float32Array(tBoxes.numel)); + const scores = tScores.getData(new Float32Array(tScores.numel)); + const classes = tClasses.getData(new Float32Array(tClasses.numel)); + + const auxTensors = [ + tensor('float32', [input.height, input.width, 1]), + tensor('float32', [input.height, input.width, 1]), + tensor('float32', [input.height, input.width, 1]), + tensor('uint8', [input.height, input.width, 1]), + ] as const; + + const [tResize, tThreshold, tCrop, tUint8] = auxTensors; + + const results: InstanceSegmentationResult[] = []; + + try { + for (const idx of indices) { + const confidence = scores[idx]!; + const classIdx = Math.round(classes[idx]!); + const label = modelOpts.labels[classIdx]; + + if (label === undefined) { + throw RnExecuTorchError( + 'INVALID_ARGUMENT', + `InstanceSegmenter: Predicted class index ${classIdx} is ` + + `out of bounds for labels array of size ${modelOpts.labels.length}.` + ); + } + + const a = boxes[idx * 4]!; + const b = boxes[idx * 4 + 1]!; + const c = boxes[idx * 4 + 2]!; + const d = boxes[idx * 4 + 3]!; + + const box = scaleBox(decodeBox([a, b, c, d], modelOpts.boxFormat), { + from: { width: W, height: H }, + to: { width: input.width, height: input.height }, + resizeMode: 'stretch', + }); + + const maskData = tAllMasks + .copyTo(tMask, { offset: idx * maskH * maskW, length: maskH * maskW }) + .through(resize, tResize, { mode: 'stretch', interpolation: 'linear' }) + .through(threshold, tThreshold, logitMaskThreshold) + .through(restrictToBox, tCrop, box) + .through(normalize, tUint8, { alpha: 255.0 }) + .getData(new Uint8Array(tUint8.numel)); + + const mask = { + data: maskData, + width: input.width, + height: input.height, + format: 'gray' as const, + layout: 'hwc' as const, + }; + + results.push({ box, mask, confidence, label }); } - - const a = boxes[idx * 4]!; - const b = boxes[idx * 4 + 1]!; - const c = boxes[idx * 4 + 2]!; - const d = boxes[idx * 4 + 3]!; - - const box = scaleBox(decodeBox([a, b, c, d], modelOpts.boxFormat), { - from: { width: W, height: H }, - to: { width: input.width, height: input.height }, - resizeMode: 'stretch', - }); - - const maskData = tAllMasks - .copyTo(tMask, { offset: idx * maskH * maskW, length: maskH * maskW }) - .through(resize, tResize, { mode: 'stretch', interpolation: 'linear' }) - .through(threshold, tThreshold, logitMaskThreshold) - .through(restrictToBox, tCrop, box) - .through(normalize, tUint8, { alpha: 255.0 }) - .getData(new Uint8Array(tUint8.numel)); - - const mask = { - data: maskData, - width: input.width, - height: input.height, - format: 'gray' as const, - layout: 'hwc' as const, - }; - - results.push({ box, mask, confidence, label }); + } finally { + auxTensors.forEach((t) => t.dispose()); } - } finally { - auxTensors.forEach((t) => t.dispose()); - } - return results; - }; + return results; + }; - const segmentInstances = wrapAsync(segmentInstancesWorklet, runtime); + const segmentInstances = wrapAsync(segmentInstancesWorklet, runtime); - return { segmentInstances, segmentInstancesWorklet, dispose }; + return { segmentInstances, segmentInstancesWorklet, dispose }; + } catch (error) { + dispose(); + throw error; + } } diff --git a/packages/react-native-executorch/src/extensions/cv/tasks/keypointDetection.ts b/packages/react-native-executorch/src/extensions/cv/tasks/keypointDetection.ts index 7b168625a4..6f1465e12a 100644 --- a/packages/react-native-executorch/src/extensions/cv/tasks/keypointDetection.ts +++ b/packages/react-native-executorch/src/extensions/cv/tasks/keypointDetection.ts @@ -9,6 +9,7 @@ import { tensor, type Tensor } from '../../../core/tensor'; import { loadModel } from '../../../core/model'; import { validateSpec, method, f32 } from '../../../core/schema'; import { wrapAsync } from '../../../core/runtime'; +import { createResourceScope } from '../../../core/lifetime'; import type { ImageBuffer } from '../image'; import { createImagePreprocessor, type ImagePreprocessorOptions } from '../utils/imagePreprocessor'; @@ -222,59 +223,63 @@ export async function createKeypointDetector, runtime?: WorkletRuntime ): Promise> { - const { modelPath, modelOpts } = config; - const { landmarks } = modelOpts; - const model = await wrapAsync(loadModel, runtime)(modelPath); - - const { dims } = validateSpec(model.schema, { - default: method( - 'forward', - [f32(1, 3, 'H', 'W')], - [f32('N', 4), f32('N'), f32('N', landmarks.length, 3)] - ), - }); - - const [N, targetH, targetW] = dims.constant('N', 'H', 'W'); - const inpShape = [1, 3, targetH, targetW]; - const outShape = { boxes: [N, 4], scores: [N], keypoints: [N, landmarks.length, 3] }; - - const tensors = [ - tensor('float32', outShape.boxes), - tensor('float32', outShape.scores), - tensor('float32', outShape.keypoints), - ] as const; - - const [tBoxes, tScores, tKeypoints] = tensors; - const preprocessor = createImagePreprocessor(modelOpts, inpShape); - - const dispose = () => { - preprocessor.dispose(); - tensors.forEach((t) => t.dispose()); - model.dispose(); - }; - - const detectKeypointsWorklet = ( - input: ImageBuffer, - options?: { confidenceThreshold?: number; iouThreshold?: number } - ): KeypointDetection[] => { - 'worklet'; - const tInput = preprocessor.process(input); - model.execute('forward', [tInput], [tBoxes, tScores, tKeypoints]); - - const iouThreshold = options?.iouThreshold ?? modelOpts.defaultIouThreshold; - const confidenceThreshold = - options?.confidenceThreshold ?? modelOpts.defaultConfidenceThreshold; - - return postprocess(tBoxes, tScores, tKeypoints, { - ...modelOpts, - iouThreshold, - confidenceThreshold, - from: { width: targetW, height: targetH }, - to: { width: input.width, height: input.height }, + const scope = createResourceScope(); + const dispose = scope.dispose; + + try { + const { modelPath, modelOpts } = config; + const { landmarks } = modelOpts; + const model = scope.track(await wrapAsync(loadModel, runtime)(modelPath)); + + const { dims } = validateSpec(model.schema, { + default: method( + 'forward', + [f32(1, 3, 'H', 'W')], + [f32('N', 4), f32('N'), f32('N', landmarks.length, 3)] + ), }); - }; - const detectKeypoints = wrapAsync(detectKeypointsWorklet, runtime); + const [N, targetH, targetW] = dims.constant('N', 'H', 'W'); + const inpShape = [1, 3, targetH, targetW]; + const outShape = { boxes: [N, 4], scores: [N], keypoints: [N, landmarks.length, 3] }; + + const tensors = [ + tensor('float32', outShape.boxes), + tensor('float32', outShape.scores), + tensor('float32', outShape.keypoints), + ] as const; + + tensors.forEach(scope.track); + + const [tBoxes, tScores, tKeypoints] = tensors; + const preprocessor = scope.track(createImagePreprocessor(modelOpts, inpShape)); + + const detectKeypointsWorklet = ( + input: ImageBuffer, + options?: { confidenceThreshold?: number; iouThreshold?: number } + ): KeypointDetection[] => { + 'worklet'; + const tInput = preprocessor.process(input); + model.execute('forward', [tInput], [tBoxes, tScores, tKeypoints]); + + const iouThreshold = options?.iouThreshold ?? modelOpts.defaultIouThreshold; + const confidenceThreshold = + options?.confidenceThreshold ?? modelOpts.defaultConfidenceThreshold; + + return postprocess(tBoxes, tScores, tKeypoints, { + ...modelOpts, + iouThreshold, + confidenceThreshold, + from: { width: targetW, height: targetH }, + to: { width: input.width, height: input.height }, + }); + }; + + const detectKeypoints = wrapAsync(detectKeypointsWorklet, runtime); - return { detectKeypoints, detectKeypointsWorklet, dispose }; + return { detectKeypoints, detectKeypointsWorklet, dispose }; + } catch (error) { + dispose(); + throw error; + } } diff --git a/packages/react-native-executorch/src/extensions/cv/tasks/objectDetection.ts b/packages/react-native-executorch/src/extensions/cv/tasks/objectDetection.ts index 833890527a..882aad0cd3 100644 --- a/packages/react-native-executorch/src/extensions/cv/tasks/objectDetection.ts +++ b/packages/react-native-executorch/src/extensions/cv/tasks/objectDetection.ts @@ -15,6 +15,7 @@ import type { ImageBuffer } from '../image'; import { createImagePreprocessor, type ImagePreprocessorOptions } from '../utils/imagePreprocessor'; import { nms, scaleBox, decodeBox, type BoundingBox, type BoxFormat } from '../ops/box'; import { RnExecuTorchError } from '../../../core/error'; +import { createResourceScope } from '../../../core/lifetime'; /** * Options for configuring an object detector preprocessor, label vocabulary, @@ -141,100 +142,104 @@ export async function createObjectDetector( config: ObjectDetectorModel, runtime?: WorkletRuntime ): Promise> { - const { modelPath, modelOpts } = config; - const model = await wrapAsync(loadModel, runtime)(modelPath); - - const { variant, dims } = validateSpec(model.schema, { - batched: method( - 'forward', // prettier-ignore - [f32(1, 3, 'H', 'W')], - [f32('N', 4), f32('N'), f32('N')] - ), - unbatched: method( - 'forward', // prettier-ignore - [f32(3, 'H', 'W')], - [f32('N', 4), f32('N'), f32('N')] - ), - }); - - const [N, H, W] = dims.constant('N', 'H', 'W'); - const inpShape = { batched: [1, 3, H, W], unbatched: [3, H, W] }[variant]; - const outShape = { boxes: [N, 4], scores: [N], classes: [N] }; - - const tensors = [ - tensor('float32', outShape.boxes), - tensor('float32', outShape.scores), - tensor('float32', outShape.classes), - ] as const; - - const [tBoxes, tScores, tClasses] = tensors; - const preprocessor = createImagePreprocessor(modelOpts, inpShape); - - const { boxFormat } = modelOpts; - - const dispose = () => { - preprocessor.dispose(); - tensors.forEach((t) => t.dispose()); - model.dispose(); - }; - - const detectObjectsWorklet = ( - input: ImageBuffer, - options?: { confidenceThreshold?: number; iouThreshold?: number } - ): ObjectDetection[] => { - 'worklet'; - const tInput = preprocessor.process(input); - model.execute('forward', [tInput], [tBoxes, tScores, tClasses]); - - const boxes = tBoxes.getData(new Float32Array(tBoxes.numel)); - const scores = tScores.getData(new Float32Array(tScores.numel)); - const classes = tClasses.getData(new Float32Array(tClasses.numel)); - - const iouThreshold = options?.iouThreshold ?? modelOpts.defaultIouThreshold; - const confidenceThreshold = - options?.confidenceThreshold ?? modelOpts.defaultConfidenceThreshold; - - const results: ObjectDetection[] = []; - const indices = nms(tBoxes, tScores, { - boxFormat, - iouThreshold, - confidenceThreshold, - nmsType: 'standard', + const scope = createResourceScope(); + const dispose = scope.dispose; + + try { + const { modelPath, modelOpts } = config; + const model = scope.track(await wrapAsync(loadModel, runtime)(modelPath)); + + const { variant, dims } = validateSpec(model.schema, { + batched: method( + 'forward', // prettier-ignore + [f32(1, 3, 'H', 'W')], + [f32('N', 4), f32('N'), f32('N')] + ), + unbatched: method( + 'forward', // prettier-ignore + [f32(3, 'H', 'W')], + [f32('N', 4), f32('N'), f32('N')] + ), }); - for (const index of indices) { - const confidence = scores[index]!; - const classIdx = Math.round(classes[index]!); - const label = modelOpts.labels[classIdx]; - - if (label === undefined) { - throw RnExecuTorchError( - 'INVALID_ARGUMENT', - `ObjectDetector: Predicted class index ${classIdx} is out of bounds for` + - `labels array of size ${modelOpts.labels.length}.` - ); - } - - const a = boxes[index * 4]!; - const b = boxes[index * 4 + 1]!; - const c = boxes[index * 4 + 2]!; - const d = boxes[index * 4 + 3]!; - - results.push({ - label, - confidence, - box: scaleBox(decodeBox([a, b, c, d], boxFormat), { - from: { width: W, height: H }, - to: { width: input.width, height: input.height }, - ...modelOpts, - }), + const [N, H, W] = dims.constant('N', 'H', 'W'); + const inpShape = { batched: [1, 3, H, W], unbatched: [3, H, W] }[variant]; + const outShape = { boxes: [N, 4], scores: [N], classes: [N] }; + + const tensors = [ + tensor('float32', outShape.boxes), + tensor('float32', outShape.scores), + tensor('float32', outShape.classes), + ] as const; + + tensors.forEach(scope.track); + + const [tBoxes, tScores, tClasses] = tensors; + const preprocessor = scope.track(createImagePreprocessor(modelOpts, inpShape)); + + const { boxFormat } = modelOpts; + + const detectObjectsWorklet = ( + input: ImageBuffer, + options?: { confidenceThreshold?: number; iouThreshold?: number } + ): ObjectDetection[] => { + 'worklet'; + const tInput = preprocessor.process(input); + model.execute('forward', [tInput], [tBoxes, tScores, tClasses]); + + const boxes = tBoxes.getData(new Float32Array(tBoxes.numel)); + const scores = tScores.getData(new Float32Array(tScores.numel)); + const classes = tClasses.getData(new Float32Array(tClasses.numel)); + + const iouThreshold = options?.iouThreshold ?? modelOpts.defaultIouThreshold; + const confidenceThreshold = + options?.confidenceThreshold ?? modelOpts.defaultConfidenceThreshold; + + const results: ObjectDetection[] = []; + const indices = nms(tBoxes, tScores, { + boxFormat, + iouThreshold, + confidenceThreshold, + nmsType: 'standard', }); - } - return results; - }; + for (const index of indices) { + const confidence = scores[index]!; + const classIdx = Math.round(classes[index]!); + const label = modelOpts.labels[classIdx]; + + if (label === undefined) { + throw RnExecuTorchError( + 'INVALID_ARGUMENT', + `ObjectDetector: Predicted class index ${classIdx} is out of bounds for` + + `labels array of size ${modelOpts.labels.length}.` + ); + } + + const a = boxes[index * 4]!; + const b = boxes[index * 4 + 1]!; + const c = boxes[index * 4 + 2]!; + const d = boxes[index * 4 + 3]!; + + results.push({ + label, + confidence, + box: scaleBox(decodeBox([a, b, c, d], boxFormat), { + from: { width: W, height: H }, + to: { width: input.width, height: input.height }, + ...modelOpts, + }), + }); + } + + return results; + }; - const detectObjects = wrapAsync(detectObjectsWorklet, runtime); + const detectObjects = wrapAsync(detectObjectsWorklet, runtime); - return { detectObjects, detectObjectsWorklet, dispose }; + return { detectObjects, detectObjectsWorklet, dispose }; + } catch (error) { + dispose(); + throw error; + } } diff --git a/packages/react-native-executorch/src/extensions/cv/tasks/paddleOcr.ts b/packages/react-native-executorch/src/extensions/cv/tasks/paddleOcr.ts index 6d5f4bf7cf..154b6b72f8 100644 --- a/packages/react-native-executorch/src/extensions/cv/tasks/paddleOcr.ts +++ b/packages/react-native-executorch/src/extensions/cv/tasks/paddleOcr.ts @@ -10,6 +10,7 @@ import { loadModel } from '../../../core/model'; import { RnExecuTorchError } from '../../../core/error'; import { wrapAsync } from '../../../core/runtime'; import { tensor, type Tensor } from '../../../core/tensor'; +import { createResourceScope } from '../../../core/lifetime'; import { validateSpec, method, @@ -400,160 +401,174 @@ export async function createPaddleOcr( config: PaddleOcrModel, runtime?: WorkletRuntime ): Promise { - const { modelPath, charsetPath, modelOpts } = config; - const model = await wrapAsync(loadModel, runtime)(modelPath); - - // Both methods are `DynamicDim` on every backend — a range on XNNPACK and - // Vulkan, an enum on CoreML, which has no RangeDim — so one variant covers - // all three. - const { dims } = validateSpec(model.schema, { - ppOcrV6: { - ...method( - 'detect', - [f32(1, 3, DynamicDim('detH'), DynamicDim('detW'))], - [f32(1, 1, DynamicDim('detOutH'), DynamicDim('detOutW'))] - ), - ...method( - 'recognize', - [f32(1, 3, 'recH', DynamicDim('recW'))], - [f32(1, DynamicDim('recT'), 'vocab')], - [ - constraint.linear( - { paramSide: 'input', tensorIdx: 0, dimIdx: 3 }, - { paramSide: 'output', tensorIdx: 0, dimIdx: 1 }, - SVTR_CTC_STRIDE - ), - ] - ), - }, - }); - - const [detHDim, detWDim, recWDim] = dims.dynamic('detH', 'detW', 'recW'); - const [recH, vocabSize] = dims.constant('recH', 'vocab'); - const [detHMax, detWMax, recWMax] = [dimMax(detHDim), dimMax(detWDim), dimMax(recWDim)]; - - // Kept off the JS bundle on purpose: the table is ~128 KB, and an app that - // never runs OCR should not carry it. CTC lookup: index 0 is the blank, then - // the model's characters. - const charsetString = await RNBlobUtil.fs.readFile(charsetPath, 'utf8'); - const charset: readonly string[] = ['[blank]', ...JSON.parse(charsetString)]; - if (charset.length !== vocabSize) { - throw RnExecuTorchError( - 'SCHEMA_MISMATCH', - `createPaddleOcr: charset size (${charset.length}, incl. blank) must match the ` + - `recognizer output vocab (${vocabSize}).` - ); - } - - const dispose = () => model.dispose(); + const scope = createResourceScope(); + const dispose = scope.dispose; - const recognizeCharactersWorklet = ( - input: ImageBuffer, - options?: RecognizeCharactersOptions - ): OcrDetection[] => { - 'worklet'; - const confidenceThreshold = - options?.confidenceThreshold ?? modelOpts.defaultConfidenceThreshold; - - const [H, W] = [input.height, input.width]; - const numChannels = FORMAT_CHANNELS[input.format]; - const colorCode = FORMAT_CONVERSION[input.format].rgb; - - // Scale the page down to fit the detector before snapping, so a wide page - // keeps its aspect ratio instead of paying for padding on the short side. - const scale = Math.min(1, detHMax / H, detWMax / W); - const [detH, detW] = [ - snap(Math.round(H * scale), detHDim), - snap(Math.round(W * scale), detWDim), - ]; - - // 1. Detect the quads holding text. - let quads: Quad[]; - const detPreprocessor = createImagePreprocessor(DETECTOR_PREPROCESSOR_OPTS, [1, 3, detH, detW]); - // DBNet emits one full-resolution probability map, [1, 1, H, W]. - const tDetProbas = tensor('float32', [1, 1, detH, detW]); - try { - const tDetInput = detPreprocessor.process(input); - model.execute('detect', [tDetInput], [tDetProbas]); - - quads = extractDbnetTextQuads(tDetProbas, DBNET_DECODE_OPTS) - .map((q) => orderQuad(scaleQuad(q, { from: { width: detW, height: detH }, to: input }))) - .filter((q) => Object.values(quadSize(q)).every((s) => s >= MIN_RECOGNIZABLE_SIDE)); - } finally { - tDetProbas.dispose(); - detPreprocessor.dispose(); + try { + const { modelPath, charsetPath, modelOpts } = config; + const model = scope.track(await wrapAsync(loadModel, runtime)(modelPath)); + + // Both methods are `DynamicDim` on every backend — a range on XNNPACK and + // Vulkan, an enum on CoreML, which has no RangeDim — so one variant covers + // all three. + const { dims } = validateSpec(model.schema, { + ppOcrV6: { + ...method( + 'detect', + [f32(1, 3, DynamicDim('detH'), DynamicDim('detW'))], + [f32(1, 1, DynamicDim('detOutH'), DynamicDim('detOutW'))] + ), + ...method( + 'recognize', + [f32(1, 3, 'recH', DynamicDim('recW'))], + [f32(1, DynamicDim('recT'), 'vocab')], + [ + constraint.linear( + { paramSide: 'input', tensorIdx: 0, dimIdx: 3 }, + { paramSide: 'output', tensorIdx: 0, dimIdx: 1 }, + SVTR_CTC_STRIDE + ), + ] + ), + }, + }); + + const [detHDim, detWDim, recWDim] = dims.dynamic('detH', 'detW', 'recW'); + const [recH, vocabSize] = dims.constant('recH', 'vocab'); + const [detHMax, detWMax, recWMax] = [dimMax(detHDim), dimMax(detWDim), dimMax(recWDim)]; + + // Kept off the JS bundle on purpose: the table is ~128 KB, and an app that + // never runs OCR should not carry it. CTC lookup: index 0 is the blank, then + // the model's characters. + const charsetString = await RNBlobUtil.fs.readFile(charsetPath, 'utf8'); + const charset: readonly string[] = ['[blank]', ...JSON.parse(charsetString)]; + if (charset.length !== vocabSize) { + throw RnExecuTorchError( + 'SCHEMA_MISMATCH', + `createPaddleOcr: charset size (${charset.length}, incl. blank) must match the ` + + `recognizer output vocab (${vocabSize}).` + ); } - // 2. Read the text inside each quad. - const detections: OcrDetection[] = []; - const tImage = tensor('uint8', [H, W, numChannels], input.data); - try { - for (const quad of quads) { - const size = quadSize(quad); - const aspectW = Math.max(1, Math.round((recH * size.width) / size.height)); - - // Known limitation: segments abut with no overlap, so a glyph straddling - // a cut can be mangled at the seam. Acceptable for now, since a line this - // wide is the rare case. - const splits = - aspectW > recWMax * WIDE_SQUISH_TOLERANCE - ? splitWideQuad(quad, Math.ceil(aspectW / recWMax)) - : [quad]; - - let text = ''; - let weightedConf = 0; - - for (const splitQuad of splits) { - const splitSize = quadSize(splitQuad); - const splitAspectW = Math.max(1, Math.round((recH * splitSize.width) / splitSize.height)); - const recW = snap(splitAspectW, recWDim); - - const auxTensors = [ - tensor('uint8', [recH, recW, numChannels]), // tQuad - tensor('uint8', [recH, recW, 3]), // tColor - tensor('uint8', [3, recH, recW]), // tChanFirst - tensor('float32', [3, recH, recW]), // tNorm - tensor('float32', [1, 3, recH, recW]), // tRecInput - // Every width the domain admits is a multiple of SVTR_CTC_STRIDE, so - // the dynamically-sized probs output can be pre-allocated exactly. - tensor('float32', [1, recW / SVTR_CTC_STRIDE, vocabSize]), // tRecProbas - ] as const; - const [tQuad, tColor, tChanFirst, tNorm, tRecInput, tRecProbas] = auxTensors; - - try { - tImage - .through(rectifyQuad, tQuad, splitQuad, { - contentWidth: Math.min(recWMax, splitAspectW), - padValue: RECOGNIZER_PREPROCESSOR_OPTS.padValue, - }) - .throughIf(colorCode !== null, cvtColor, tColor, colorCode!) - .through(toChannelsFirst, tChanFirst) - .through(normalize, tNorm, RECOGNIZER_PREPROCESSOR_OPTS.normalizeOpts) - .copyTo(tRecInput); - - model.execute('recognize', [tRecInput], [tRecProbas]); - - const split = greedyCtcDecode(tRecProbas, charset); - text += split.text; - weightedConf += split.conf * split.text.length; - } finally { - auxTensors.forEach((t) => t.dispose()); + const recognizeCharactersWorklet = ( + input: ImageBuffer, + options?: RecognizeCharactersOptions + ): OcrDetection[] => { + 'worklet'; + const confidenceThreshold = + options?.confidenceThreshold ?? modelOpts.defaultConfidenceThreshold; + + const [H, W] = [input.height, input.width]; + const numChannels = FORMAT_CHANNELS[input.format]; + const colorCode = FORMAT_CONVERSION[input.format].rgb; + + // Scale the page down to fit the detector before snapping, so a wide page + // keeps its aspect ratio instead of paying for padding on the short side. + const scale = Math.min(1, detHMax / H, detWMax / W); + const [detH, detW] = [ + snap(Math.round(H * scale), detHDim), + snap(Math.round(W * scale), detWDim), + ]; + + // 1. Detect the quads holding text. + let quads: Quad[]; + const detPreprocessor = createImagePreprocessor(DETECTOR_PREPROCESSOR_OPTS, [ + 1, + 3, + detH, + detW, + ]); + // DBNet emits one full-resolution probability map, [1, 1, H, W]. + const tDetProbas = tensor('float32', [1, 1, detH, detW]); + try { + const tDetInput = detPreprocessor.process(input); + model.execute('detect', [tDetInput], [tDetProbas]); + + quads = extractDbnetTextQuads(tDetProbas, DBNET_DECODE_OPTS) + .map((q) => orderQuad(scaleQuad(q, { from: { width: detW, height: detH }, to: input }))) + .filter((q) => Object.values(quadSize(q)).every((s) => s >= MIN_RECOGNIZABLE_SIDE)); + } finally { + tDetProbas.dispose(); + detPreprocessor.dispose(); + } + + // 2. Read the text inside each quad. + const detections: OcrDetection[] = []; + const tImage = tensor('uint8', [H, W, numChannels], input.data); + try { + for (const quad of quads) { + const size = quadSize(quad); + const aspectW = Math.max(1, Math.round((recH * size.width) / size.height)); + + // Known limitation: segments abut with no overlap, so a glyph straddling + // a cut can be mangled at the seam. Acceptable for now, since a line this + // wide is the rare case. + const splits = + aspectW > recWMax * WIDE_SQUISH_TOLERANCE + ? splitWideQuad(quad, Math.ceil(aspectW / recWMax)) + : [quad]; + + let text = ''; + let weightedConf = 0; + + for (const splitQuad of splits) { + const splitSize = quadSize(splitQuad); + const splitAspectW = Math.max( + 1, + Math.round((recH * splitSize.width) / splitSize.height) + ); + const recW = snap(splitAspectW, recWDim); + + const auxTensors = [ + tensor('uint8', [recH, recW, numChannels]), // tQuad + tensor('uint8', [recH, recW, 3]), // tColor + tensor('uint8', [3, recH, recW]), // tChanFirst + tensor('float32', [3, recH, recW]), // tNorm + tensor('float32', [1, 3, recH, recW]), // tRecInput + // Every width the domain admits is a multiple of SVTR_CTC_STRIDE, so + // the dynamically-sized probs output can be pre-allocated exactly. + tensor('float32', [1, recW / SVTR_CTC_STRIDE, vocabSize]), // tRecProbas + ] as const; + const [tQuad, tColor, tChanFirst, tNorm, tRecInput, tRecProbas] = auxTensors; + + try { + tImage + .through(rectifyQuad, tQuad, splitQuad, { + contentWidth: Math.min(recWMax, splitAspectW), + padValue: RECOGNIZER_PREPROCESSOR_OPTS.padValue, + }) + .throughIf(colorCode !== null, cvtColor, tColor, colorCode!) + .through(toChannelsFirst, tChanFirst) + .through(normalize, tNorm, RECOGNIZER_PREPROCESSOR_OPTS.normalizeOpts) + .copyTo(tRecInput); + + model.execute('recognize', [tRecInput], [tRecProbas]); + + const split = greedyCtcDecode(tRecProbas, charset); + text += split.text; + weightedConf += split.conf * split.text.length; + } finally { + auxTensors.forEach((t) => t.dispose()); + } } - } - const confidence = text.length === 0 ? 0 : weightedConf / text.length; - if (text.length > 0 && confidence >= confidenceThreshold) { - detections.push({ text, confidence, quad }); + const confidence = text.length === 0 ? 0 : weightedConf / text.length; + if (text.length > 0 && confidence >= confidenceThreshold) { + detections.push({ text, confidence, quad }); + } } + } finally { + tImage.dispose(); } - } finally { - tImage.dispose(); - } - return orderByReadingOrder(detections); - }; + return orderByReadingOrder(detections); + }; - const recognizeCharacters = wrapAsync(recognizeCharactersWorklet, runtime); + const recognizeCharacters = wrapAsync(recognizeCharactersWorklet, runtime); - return { recognizeCharacters, recognizeCharactersWorklet, dispose }; + return { recognizeCharacters, recognizeCharactersWorklet, dispose }; + } catch (error) { + dispose(); + throw error; + } } diff --git a/packages/react-native-executorch/src/extensions/cv/tasks/sdxsTextToImage.ts b/packages/react-native-executorch/src/extensions/cv/tasks/sdxsTextToImage.ts index 37512c737d..5dcc136d5a 100644 --- a/packages/react-native-executorch/src/extensions/cv/tasks/sdxsTextToImage.ts +++ b/packages/react-native-executorch/src/extensions/cv/tasks/sdxsTextToImage.ts @@ -9,6 +9,7 @@ import { tensor } from '../../../core/tensor'; import { loadModel } from '../../../core/model'; import { validateSpec, method, i64, f32 } from '../../../core/schema'; import { wrapAsync } from '../../../core/runtime'; +import { createResourceScope } from '../../../core/lifetime'; import { randomNormal } from '../../math'; import { loadTokenizer } from '../../nlp/tokenizer'; @@ -94,90 +95,94 @@ export async function createSdxsTextToImage( config: SdxsTextToImageModel, runtime?: WorkletRuntime ): Promise { - const { modelPath, tokenizerPath } = config; - const model = await wrapAsync(loadModel, runtime)(modelPath); - const tokenizer = await wrapAsync(loadTokenizer, runtime)(tokenizerPath); - - validateSpec(model.schema, { - default: { - ...method( - 'encode', // prettier-ignore - [i64(1, CLIP_MAX_TOKENS)], - [f32(1, CLIP_MAX_TOKENS, CLIP_HIDDEN_SIZE)] - ), - ...method( - 'denoise', - [f32(...LATENT_SHAPE), i64(1), f32(1, CLIP_MAX_TOKENS, CLIP_HIDDEN_SIZE)], - [f32(...LATENT_SHAPE)] - ), - ...method( - 'decode', // prettier-ignore - [f32(...LATENT_SHAPE)], - [f32(1, 3, IMAGE_SIZE, IMAGE_SIZE)] - ), - }, - }); - - const tensors = [ - tensor('int64', [1, CLIP_MAX_TOKENS]), - tensor('float32', [1, CLIP_MAX_TOKENS, CLIP_HIDDEN_SIZE]), - tensor('int64', [1]), - tensor('float32', LATENT_SHAPE), - tensor('float32', LATENT_SHAPE), - tensor('float32', [1, 3, IMAGE_SIZE, IMAGE_SIZE]), - tensor('float32', [3, IMAGE_SIZE, IMAGE_SIZE]), - tensor('uint8', [3, IMAGE_SIZE, IMAGE_SIZE]), - tensor('uint8', [IMAGE_SIZE, IMAGE_SIZE, 3]), - tensor('uint8', [IMAGE_SIZE, IMAGE_SIZE, 4]), - ] as const; - - // prettier-ignore - const [ - tTokens, tEmbeddings, tTimestep, tLatents, tNoisePred, - tDecoded, tReshape, tUint8, tChanLast, tRgba - ] = tensors; - - const dispose = () => { - tensors.forEach((t) => t.dispose()); - tokenizer.dispose(); - model.dispose(); - }; - - const generateWorklet = (prompt: string, seed?: number): ImageBuffer => { - 'worklet'; - - const ids = tokenizer.encode(prompt); - const tokens = new BigInt64Array(CLIP_MAX_TOKENS); - for (let i = 0; i < CLIP_MAX_TOKENS; i++) { - tokens[i] = BigInt(i < ids.length ? ids[i]! : CLIP_PAD_TOKEN_ID); - } - tTokens.setData(tokens); - model.execute('encode', [tTokens], [tEmbeddings]); - - tLatents.setData(randomNormal(tLatents.numel, { std: INIT_NOISE_SIGMA, seed })); - tTimestep.setData(new BigInt64Array([BigInt(TIMESTEP)])); - model.execute('denoise', [tLatents, tTimestep, tEmbeddings], [tNoisePred]); - - const latents = tLatents.getData(new Float32Array(tLatents.numel)); - const modelOutput = tNoisePred.getData(new Float32Array(tNoisePred.numel)); - for (let i = 0; i < latents.length; i++) { - latents[i] = SAMPLE_COEFF * latents[i]! + NOISE_COEFF * modelOutput[i]!; - } - tLatents.setData(latents); - - model.execute('decode', [tLatents], [tDecoded]); - - const data = tDecoded - .copyTo(tReshape) - .through(normalize, tUint8, { alpha: 255.0 }) - .through(toChannelsLast, tChanLast) - .through(cvtColor, tRgba, 'RGB2RGBA') - .getData(new Uint8Array(IMAGE_SIZE * IMAGE_SIZE * 4)); - - return { data, width: IMAGE_SIZE, height: IMAGE_SIZE, format: 'rgba', layout: 'hwc' }; - }; - - const generate = wrapAsync(generateWorklet, runtime); - - return { generate, generateWorklet, dispose }; + const scope = createResourceScope(); + const dispose = scope.dispose; + + try { + const { modelPath, tokenizerPath } = config; + const model = scope.track(await wrapAsync(loadModel, runtime)(modelPath)); + const tokenizer = scope.track(await wrapAsync(loadTokenizer, runtime)(tokenizerPath)); + + validateSpec(model.schema, { + default: { + ...method( + 'encode', // prettier-ignore + [i64(1, CLIP_MAX_TOKENS)], + [f32(1, CLIP_MAX_TOKENS, CLIP_HIDDEN_SIZE)] + ), + ...method( + 'denoise', + [f32(...LATENT_SHAPE), i64(1), f32(1, CLIP_MAX_TOKENS, CLIP_HIDDEN_SIZE)], + [f32(...LATENT_SHAPE)] + ), + ...method( + 'decode', // prettier-ignore + [f32(...LATENT_SHAPE)], + [f32(1, 3, IMAGE_SIZE, IMAGE_SIZE)] + ), + }, + }); + + const tensors = [ + tensor('int64', [1, CLIP_MAX_TOKENS]), + tensor('float32', [1, CLIP_MAX_TOKENS, CLIP_HIDDEN_SIZE]), + tensor('int64', [1]), + tensor('float32', LATENT_SHAPE), + tensor('float32', LATENT_SHAPE), + tensor('float32', [1, 3, IMAGE_SIZE, IMAGE_SIZE]), + tensor('float32', [3, IMAGE_SIZE, IMAGE_SIZE]), + tensor('uint8', [3, IMAGE_SIZE, IMAGE_SIZE]), + tensor('uint8', [IMAGE_SIZE, IMAGE_SIZE, 3]), + tensor('uint8', [IMAGE_SIZE, IMAGE_SIZE, 4]), + ] as const; + + tensors.forEach(scope.track); + + // prettier-ignore + const [ + tTokens, tEmbeddings, tTimestep, tLatents, tNoisePred, + tDecoded, tReshape, tUint8, tChanLast, tRgba + ] = tensors; + + const generateWorklet = (prompt: string, seed?: number): ImageBuffer => { + 'worklet'; + + const ids = tokenizer.encode(prompt); + const tokens = new BigInt64Array(CLIP_MAX_TOKENS); + for (let i = 0; i < CLIP_MAX_TOKENS; i++) { + tokens[i] = BigInt(i < ids.length ? ids[i]! : CLIP_PAD_TOKEN_ID); + } + tTokens.setData(tokens); + model.execute('encode', [tTokens], [tEmbeddings]); + + tLatents.setData(randomNormal(tLatents.numel, { std: INIT_NOISE_SIGMA, seed })); + tTimestep.setData(new BigInt64Array([BigInt(TIMESTEP)])); + model.execute('denoise', [tLatents, tTimestep, tEmbeddings], [tNoisePred]); + + const latents = tLatents.getData(new Float32Array(tLatents.numel)); + const modelOutput = tNoisePred.getData(new Float32Array(tNoisePred.numel)); + for (let i = 0; i < latents.length; i++) { + latents[i] = SAMPLE_COEFF * latents[i]! + NOISE_COEFF * modelOutput[i]!; + } + tLatents.setData(latents); + + model.execute('decode', [tLatents], [tDecoded]); + + const data = tDecoded + .copyTo(tReshape) + .through(normalize, tUint8, { alpha: 255.0 }) + .through(toChannelsLast, tChanLast) + .through(cvtColor, tRgba, 'RGB2RGBA') + .getData(new Uint8Array(IMAGE_SIZE * IMAGE_SIZE * 4)); + + return { data, width: IMAGE_SIZE, height: IMAGE_SIZE, format: 'rgba', layout: 'hwc' }; + }; + + const generate = wrapAsync(generateWorklet, runtime); + + return { generate, generateWorklet, dispose }; + } catch (error) { + dispose(); + throw error; + } } diff --git a/packages/react-native-executorch/src/extensions/cv/tasks/semanticSegmentation.ts b/packages/react-native-executorch/src/extensions/cv/tasks/semanticSegmentation.ts index 7cf25928b6..56493c621d 100644 --- a/packages/react-native-executorch/src/extensions/cv/tasks/semanticSegmentation.ts +++ b/packages/react-native-executorch/src/extensions/cv/tasks/semanticSegmentation.ts @@ -22,6 +22,7 @@ import { } from '../ops/image'; import { sigmoid, argmax } from '../../math'; import { RnExecuTorchError } from '../../../core/error'; +import { createResourceScope } from '../../../core/lifetime'; /** * Options for configuring a semantic segmenter preprocessor and label @@ -149,111 +150,115 @@ export async function createSemanticSegmenter( config: SemanticSegmenterModel, runtime?: WorkletRuntime ): Promise> { - const { modelPath, modelOpts } = config; - const model = await wrapAsync(loadModel, runtime)(modelPath); + const scope = createResourceScope(); + const dispose = scope.dispose; - const { variant, dims } = validateSpec(model.schema, { - batched: method( - 'forward', // prettier-ignore - [f32(1, 3, 'H', 'W')], - [f32(1, 'K', 'H', 'W')] - ), - unbatched: method( - 'forward', // prettier-ignore - [f32(3, 'H', 'W')], - [f32('K', 'H', 'W')] - ), - }); + try { + const { modelPath, modelOpts } = config; + const model = scope.track(await wrapAsync(loadModel, runtime)(modelPath)); - const [nClasses, H, W] = dims.constant('K', 'H', 'W'); - const inpShape = { batched: [1, 3, H, W], unbatched: [3, H, W] }[variant]; - const outShape = { batched: [1, nClasses, H, W], unbatched: [nClasses, H, W] }[variant]; + const { variant, dims } = validateSpec(model.schema, { + batched: method( + 'forward', // prettier-ignore + [f32(1, 3, 'H', 'W')], + [f32(1, 'K', 'H', 'W')] + ), + unbatched: method( + 'forward', // prettier-ignore + [f32(3, 'H', 'W')], + [f32('K', 'H', 'W')] + ), + }); - // Generate highly distinct, high-contrast colors, see: - // https://martin.ankerl.com/2009/12/09/how-to-create-random-colors-programmatically/ - const defaultColormap = modelOpts.labels.map((_, i) => { - if (i === 0) return [0, 0, 0, 0] as const; - return [...hslToRgb((i * 137.5) % 360, 95, 50), 255] as const; - }); + const [nClasses, H, W] = dims.constant('K', 'H', 'W'); + const inpShape = { batched: [1, 3, H, W], unbatched: [3, H, W] }[variant]; + const outShape = { batched: [1, nClasses, H, W], unbatched: [nClasses, H, W] }[variant]; - if (nClasses > 1 && modelOpts.labels.length !== nClasses) { - throw RnExecuTorchError( - 'INVALID_ARGUMENT', - `Model outputs ${nClasses} classes, but ${modelOpts.labels.length} labels were provided in the configuration.` - ); - } + // Generate highly distinct, high-contrast colors, see: + // https://martin.ankerl.com/2009/12/09/how-to-create-random-colors-programmatically/ + const defaultColormap = modelOpts.labels.map((_, i) => { + if (i === 0) return [0, 0, 0, 0] as const; + return [...hslToRgb((i * 137.5) % 360, 95, 50), 255] as const; + }); - const tensors = [ - tensor('float32', outShape), - tensor('float32', [nClasses, H, W]), - tensor('float32', [nClasses, H, W]), - tensor('float32', [H, W, nClasses]), - tensor(nClasses > 1 ? 'int32' : 'uint8', [H, W, 1]), - tensor('uint8', [H, W, 4]), - ] as const; + if (nClasses > 1 && modelOpts.labels.length !== nClasses) { + throw RnExecuTorchError( + 'INVALID_ARGUMENT', + `Model outputs ${nClasses} classes, but ${modelOpts.labels.length} labels were provided in the configuration.` + ); + } - const [tOutput, tReshape, tSigmoid, tChanLast, tMask, tRgba] = tensors; - const preprocessor = createImagePreprocessor(modelOpts, inpShape); + const tensors = [ + tensor('float32', outShape), + tensor('float32', [nClasses, H, W]), + tensor('float32', [nClasses, H, W]), + tensor('float32', [H, W, nClasses]), + tensor(nClasses > 1 ? 'int32' : 'uint8', [H, W, 1]), + tensor('uint8', [H, W, 4]), + ] as const; - const dispose = () => { - tensors.forEach((t) => t.dispose()); - preprocessor.dispose(); - model.dispose(); - }; + tensors.forEach(scope.track); - const segmentWorklet = ( - input: ImageBuffer, - colormap?: Partial> - ): SemanticSegmentationResult => { - 'worklet'; - const tInput = preprocessor.process(input); - model.execute('forward', [tInput], [tOutput]); + const [tOutput, tReshape, tSigmoid, tChanLast, tMask, tRgba] = tensors; + const preprocessor = scope.track(createImagePreprocessor(modelOpts, inpShape)); - let returnColormap: ColorMap | undefined; - if (nClasses > 1) { - if (colormap) { - returnColormap = Object.fromEntries( - modelOpts.labels.map((l) => [l, colormap[l] ?? [0, 0, 0, 0]]) - ) as ColorMap; - } else { - returnColormap = Object.fromEntries( - modelOpts.labels.map((l, i) => [l, defaultColormap[i]!]) - ) as ColorMap; - } + const segmentWorklet = ( + input: ImageBuffer, + colormap?: Partial> + ): SemanticSegmentationResult => { + 'worklet'; + const tInput = preprocessor.process(input); + model.execute('forward', [tInput], [tOutput]); - const colormapData = modelOpts.labels.map((l) => returnColormap![l]); + let returnColormap: ColorMap | undefined; + if (nClasses > 1) { + if (colormap) { + returnColormap = Object.fromEntries( + modelOpts.labels.map((l) => [l, colormap[l] ?? [0, 0, 0, 0]]) + ) as ColorMap; + } else { + returnColormap = Object.fromEntries( + modelOpts.labels.map((l, i) => [l, defaultColormap[i]!]) + ) as ColorMap; + } - tOutput - .copyTo(tReshape) - .through(toChannelsLast, tChanLast) - .through(argmax, tMask, -1) - .through(applyColormap, tRgba, colormapData); - } else { - tOutput - .copyTo(tReshape) - .through(sigmoid, tSigmoid) - .through(toChannelsLast, tChanLast) - .through(normalize, tMask, { alpha: 255.0 }) - .through(cvtColor, tRgba, 'GRAY2RGBA'); - } + const colormapData = modelOpts.labels.map((l) => returnColormap![l]); - const data = new Uint8Array(input.width * input.height * 4); - const tResize = tensor('uint8', [input.height, input.width, 4]); - try { - tRgba - .through(resize, tResize, { mode: 'stretch', interpolation: modelOpts.outInterpolation }) - .getData(data); - } finally { - tResize.dispose(); - } + tOutput + .copyTo(tReshape) + .through(toChannelsLast, tChanLast) + .through(argmax, tMask, -1) + .through(applyColormap, tRgba, colormapData); + } else { + tOutput + .copyTo(tReshape) + .through(sigmoid, tSigmoid) + .through(toChannelsLast, tChanLast) + .through(normalize, tMask, { alpha: 255.0 }) + .through(cvtColor, tRgba, 'GRAY2RGBA'); + } - return { - buffer: { data, width: input.width, height: input.height, format: 'rgba', layout: 'hwc' }, - colormap: returnColormap, + const data = new Uint8Array(input.width * input.height * 4); + const tResize = tensor('uint8', [input.height, input.width, 4]); + try { + tRgba + .through(resize, tResize, { mode: 'stretch', interpolation: modelOpts.outInterpolation }) + .getData(data); + } finally { + tResize.dispose(); + } + + return { + buffer: { data, width: input.width, height: input.height, format: 'rgba', layout: 'hwc' }, + colormap: returnColormap, + }; }; - }; - const segment = wrapAsync(segmentWorklet, runtime); + const segment = wrapAsync(segmentWorklet, runtime); - return { segment, segmentWorklet, dispose }; + return { segment, segmentWorklet, dispose }; + } catch (error) { + dispose(); + throw error; + } } diff --git a/packages/react-native-executorch/src/extensions/cv/tasks/styleTransfer.ts b/packages/react-native-executorch/src/extensions/cv/tasks/styleTransfer.ts index 2232c0eff6..ce26eac675 100644 --- a/packages/react-native-executorch/src/extensions/cv/tasks/styleTransfer.ts +++ b/packages/react-native-executorch/src/extensions/cv/tasks/styleTransfer.ts @@ -9,6 +9,7 @@ import { tensor } from '../../../core/tensor'; import { loadModel } from '../../../core/model'; import { validateSpec, method, f32 } from '../../../core/schema'; import { wrapAsync } from '../../../core/runtime'; +import { createResourceScope } from '../../../core/lifetime'; import type { ImageBuffer } from '../image'; import { createImagePreprocessor, type ImagePreprocessorOptions } from '../utils/imagePreprocessor'; @@ -93,66 +94,70 @@ export async function createStyleTransfer( config: StyleTransferModel, runtime?: WorkletRuntime ): Promise { - const { modelPath, modelOpts } = config; - const model = await wrapAsync(loadModel, runtime)(modelPath); - - const { variant, dims } = validateSpec(model.schema, { - batched: method( - 'forward', // prettier-ignore - [f32(1, 3, 'H', 'W')], - [f32(1, 3, 'H', 'W')] - ), - unbatched: method( - 'forward', // prettier-ignore - [f32(3, 'H', 'W')], - [f32(3, 'H', 'W')] - ), - }); - - const [H, W] = dims.constant('H', 'W'); - const inpShape = { batched: [1, 3, H, W], unbatched: [3, H, W] }[variant]; - const outShape = inpShape; - - const tensors = [ - tensor('float32', outShape), - tensor('float32', [3, H, W]), - tensor('uint8', [3, H, W]), - tensor('uint8', [H, W, 3]), - tensor('uint8', [H, W, 4]), - ] as const; - - const [tOutput, tReshape, tUint8, tChanLast, tRgba] = tensors; - const preprocessor = createImagePreprocessor(modelOpts, inpShape); - - const dispose = () => { - tensors.forEach((t) => t.dispose()); - preprocessor.dispose(); - model.dispose(); - }; - - const transferStyleWorklet = (input: ImageBuffer): ImageBuffer => { - 'worklet'; - const tInput = preprocessor.process(input); - model.execute('forward', [tInput], [tOutput]); - - const data = new Uint8Array(input.width * input.height * 4); - const tResize = tensor('uint8', [input.height, input.width, 4]); - try { - tOutput - .copyTo(tReshape) - .through(normalize, tUint8, modelOpts.outNormalizeOpts) - .through(toChannelsLast, tChanLast) - .through(cvtColor, tRgba, 'RGB2RGBA') - .through(resize, tResize, { mode: 'stretch', interpolation: modelOpts.outInterpolation }) - .getData(data); - } finally { - tResize.dispose(); - } - - return { data, width: input.width, height: input.height, format: 'rgba', layout: 'hwc' }; - }; - - const transferStyle = wrapAsync(transferStyleWorklet, runtime); - - return { transferStyle, transferStyleWorklet, dispose }; + const scope = createResourceScope(); + const dispose = scope.dispose; + + try { + const { modelPath, modelOpts } = config; + const model = scope.track(await wrapAsync(loadModel, runtime)(modelPath)); + + const { variant, dims } = validateSpec(model.schema, { + batched: method( + 'forward', // prettier-ignore + [f32(1, 3, 'H', 'W')], + [f32(1, 3, 'H', 'W')] + ), + unbatched: method( + 'forward', // prettier-ignore + [f32(3, 'H', 'W')], + [f32(3, 'H', 'W')] + ), + }); + + const [H, W] = dims.constant('H', 'W'); + const inpShape = { batched: [1, 3, H, W], unbatched: [3, H, W] }[variant]; + const outShape = inpShape; + + const tensors = [ + tensor('float32', outShape), + tensor('float32', [3, H, W]), + tensor('uint8', [3, H, W]), + tensor('uint8', [H, W, 3]), + tensor('uint8', [H, W, 4]), + ] as const; + + tensors.forEach(scope.track); + + const [tOutput, tReshape, tUint8, tChanLast, tRgba] = tensors; + const preprocessor = scope.track(createImagePreprocessor(modelOpts, inpShape)); + + const transferStyleWorklet = (input: ImageBuffer): ImageBuffer => { + 'worklet'; + const tInput = preprocessor.process(input); + model.execute('forward', [tInput], [tOutput]); + + const data = new Uint8Array(input.width * input.height * 4); + const tResize = tensor('uint8', [input.height, input.width, 4]); + try { + tOutput + .copyTo(tReshape) + .through(normalize, tUint8, modelOpts.outNormalizeOpts) + .through(toChannelsLast, tChanLast) + .through(cvtColor, tRgba, 'RGB2RGBA') + .through(resize, tResize, { mode: 'stretch', interpolation: modelOpts.outInterpolation }) + .getData(data); + } finally { + tResize.dispose(); + } + + return { data, width: input.width, height: input.height, format: 'rgba', layout: 'hwc' }; + }; + + const transferStyle = wrapAsync(transferStyleWorklet, runtime); + + return { transferStyle, transferStyleWorklet, dispose }; + } catch (error) { + dispose(); + throw error; + } } diff --git a/packages/react-native-executorch/src/extensions/llm/tasks/llmChatSession.ts b/packages/react-native-executorch/src/extensions/llm/tasks/llmChatSession.ts index a75ee7704b..b16413bb66 100644 --- a/packages/react-native-executorch/src/extensions/llm/tasks/llmChatSession.ts +++ b/packages/react-native-executorch/src/extensions/llm/tasks/llmChatSession.ts @@ -8,6 +8,7 @@ import { scheduleOnRN, type WorkletRuntime } from 'react-native-worklets'; import RNBlobUtil from 'react-native-blob-util'; import { wrapAsync } from '../../../core/runtime'; +import { createResourceScope } from '../../../core/lifetime'; import { createLLMRunner, type LLMRunner, @@ -177,171 +178,177 @@ export async function createLLMChatSession( options: LLMChatSessionOptions = {}, runtime?: WorkletRuntime ): Promise { - const { - generationConfig: defaultGenerationConfig, - initialMessages = [], - stopRegex, - toolOpts, - resetOnTurn = false, - } = options; - - const { modelPath, tokenizerPath, tokenizerConfigPath, modalities, preprocessorConfig } = config; - const { tools, parseToolCalls, maxToolTurns = DEFAULT_MAX_TURNS } = toolOpts ?? {}; - - // Read and parse tokenizer_config.json - const tokenizerConfigStr = await RNBlobUtil.fs.readFile(tokenizerConfigPath, 'utf8'); - const tokenizerConfig = parseTokenizerConfig(JSON.parse(tokenizerConfigStr)); - const { chatTemplate, eosToken } = tokenizerConfig; - - // Prepare chat preprocessor - const chatPreprocessorConfig = { chatTemplate, modalities, preprocessorConfig, tools }; - const chatPreprocessor = createChatPreprocessor(chatPreprocessorConfig); - - // Prepare runner - const runner = await wrapAsync(createLLMRunner, runtime)(modelPath, tokenizerPath, modalities); - const prefill = wrapAsync(runner.prefill, runtime); - - const history: ChatMessage[] = []; - - // Tracks the number of messages in `history` whose tokens and closing - // delimiters have been permanently prefilled and committed into the runner's KV cache. - let committed = 0; - - // Prefill initial messages if provided - if (initialMessages.length > 0) { - history.push(...initialMessages); - const prompt = chatPreprocessor.process(history, history.length, { addGenPrompt: false }); - await prefill(prompt); - chatPreprocessor.clear(); - committed = history.length; - } + const scope = createResourceScope(); + const dispose = scope.dispose; + + try { + const { + generationConfig: defaultGenerationConfig, + initialMessages = [], + stopRegex, + toolOpts, + resetOnTurn = false, + } = options; + + const { modelPath, tokenizerPath, tokenizerConfigPath, modalities, preprocessorConfig } = + config; + const { tools, parseToolCalls, maxToolTurns = DEFAULT_MAX_TURNS } = toolOpts ?? {}; + + // Read and parse tokenizer_config.json + const tokenizerConfigStr = await RNBlobUtil.fs.readFile(tokenizerConfigPath, 'utf8'); + const tokenizerConfig = parseTokenizerConfig(JSON.parse(tokenizerConfigStr)); + const { chatTemplate, eosToken } = tokenizerConfig; + + // Prepare chat preprocessor + const chatPreprocessorConfig = { chatTemplate, modalities, preprocessorConfig, tools }; + const chatPreprocessor = scope.track(createChatPreprocessor(chatPreprocessorConfig)); + + // Prepare runner + const runner = scope.track( + await wrapAsync(createLLMRunner, runtime)(modelPath, tokenizerPath, modalities) + ); + const prefill = wrapAsync(runner.prefill, runtime); + + const history: ChatMessage[] = []; + + // Tracks the number of messages in `history` whose tokens and closing + // delimiters have been permanently prefilled and committed into the runner's KV cache. + let committed = 0; + + // Prefill initial messages if provided + if (initialMessages.length > 0) { + history.push(...initialMessages); + const prompt = chatPreprocessor.process(history, history.length, { addGenPrompt: false }); + await prefill(prompt); + chatPreprocessor.clear(); + committed = history.length; + } - const dispose = () => { - runner.dispose(); - chatPreprocessor.dispose(); - }; + const stop = () => runner.stop(); - const stop = () => runner.stop(); + const generateChatTurn = wrapAsync(generateChatTurnWorklet, runtime); - const generateChatTurn = wrapAsync(generateChatTurnWorklet, runtime); + const sendMessage = async ( + message: ChatMessageContent, + onToken?: (token: string) => void, + genConfig?: LLMGenerationConfig + ): Promise => { + const turnGenConfig = { ...defaultGenerationConfig, ...genConfig }; + const generationOpts = { genConfig: turnGenConfig, eosToken, stopRegex, onToken }; - const sendMessage = async ( - message: ChatMessageContent, - onToken?: (token: string) => void, - genConfig?: LLMGenerationConfig - ): Promise => { - const turnGenConfig = { ...defaultGenerationConfig, ...genConfig }; - const generationOpts = { genConfig: turnGenConfig, eosToken, stopRegex, onToken }; + const initialCommitted = committed; + const initialPos = runner.getKVCacheState().pos; - const initialCommitted = committed; - const initialPos = runner.getKVCacheState().pos; + const turnStartIdx = history.length; + const generationStatsList: LLMGenerationStats[] = []; - const turnStartIdx = history.length; - const generationStatsList: LLMGenerationStats[] = []; + history.push({ role: 'user', content: message }); - history.push({ role: 'user', content: message }); + if (resetOnTurn) { + runner.reset(); + committed = 0; + } - if (resetOnTurn) { - runner.reset(); - committed = 0; - } + try { + let prefillStartMs = Date.now(); - try { - let prefillStartMs = Date.now(); + // Prefill newly committed messages up to current user message without generation prompt + const toCommit = history.length - committed; + const userPrompt = chatPreprocessor.process(history, toCommit, { addGenPrompt: false }); + await prefill(userPrompt); + chatPreprocessor.clear(); - // Prefill newly committed messages up to current user message without generation prompt - const toCommit = history.length - committed; - const userPrompt = chatPreprocessor.process(history, toCommit, { addGenPrompt: false }); - await prefill(userPrompt); - chatPreprocessor.clear(); + // Record exact position at the end of the user message (before assistant generation header) + const posAtEndOfUser = runner.getKVCacheState().pos; + committed = history.length; - // Record exact position at the end of the user message (before assistant generation header) - const posAtEndOfUser = runner.getKVCacheState().pos; - committed = history.length; + let finishReason: 'stop' | 'maxToolTurns' = 'maxToolTurns'; - let finishReason: 'stop' | 'maxToolTurns' = 'maxToolTurns'; + for (let currentTurn = 0; currentTurn < maxToolTurns; ++currentTurn) { + const uncommitted = history.length - committed; + const prompt = chatPreprocessor.process(history, uncommitted, { addGenPrompt: true }); - for (let currentTurn = 0; currentTurn < maxToolTurns; ++currentTurn) { - const uncommitted = history.length - committed; - const prompt = chatPreprocessor.process(history, uncommitted, { addGenPrompt: true }); + const prefillDurationMs = Date.now() - prefillStartMs; - const prefillDurationMs = Date.now() - prefillStartMs; + const { response, stats } = await generateChatTurn(runner, prompt, generationOpts); + chatPreprocessor.clear(); + generationStatsList.push({ ...stats, prefillDurationMs }); - const { response, stats } = await generateChatTurn(runner, prompt, generationOpts); - chatPreprocessor.clear(); - generationStatsList.push({ ...stats, prefillDurationMs }); + // Always rewind KV cache back to posAtEndOfUser so next turn prefills + // cleanly formatted message with tool outputs + runner.reset(posAtEndOfUser); - // Always rewind KV cache back to posAtEndOfUser so next turn prefills - // cleanly formatted message with tool outputs - runner.reset(posAtEndOfUser); + // Check for tool calls + const parsedTools = parseToolCalls?.(response); - // Check for tool calls - const parsedTools = parseToolCalls?.(response); + if (!parsedTools || parsedTools.toolCalls.length === 0) { + history.push({ role: 'assistant', content: response }); + finishReason = 'stop'; + break; + } - if (!parsedTools || parsedTools.toolCalls.length === 0) { - history.push({ role: 'assistant', content: response }); - finishReason = 'stop'; - break; - } + // Execute tool calls + history.push({ + role: 'assistant', + content: parsedTools.textContent, + toolCalls: parsedTools.toolCalls, + }); - // Execute tool calls - history.push({ - role: 'assistant', - content: parsedTools.textContent, - toolCalls: parsedTools.toolCalls, - }); - - for (const toolCall of parsedTools.toolCalls) { - const tool = tools?.find((t) => t.function.name === toolCall.function.name); - - let toolContent: ChatMessageContent; - if (!tool) { - toolContent = `Error: Tool '${toolCall.function.name}' is not recognized or not available.`; - } else { - try { - toolContent = await tool.execute(toolCall.function.arguments); - } catch (err) { - toolContent = `Error executing tool ${toolCall.function.name}: ${String(err)}`; + for (const toolCall of parsedTools.toolCalls) { + const tool = tools?.find((t) => t.function.name === toolCall.function.name); + + let toolContent: ChatMessageContent; + if (!tool) { + toolContent = `Error: Tool '${toolCall.function.name}' is not recognized or not available.`; + } else { + try { + toolContent = await tool.execute(toolCall.function.arguments); + } catch (err) { + toolContent = `Error executing tool ${toolCall.function.name}: ${String(err)}`; + } } + + history.push({ + role: 'tool', + toolCallId: toolCall.id, + name: toolCall.function.name, + content: toolContent, + }); } - history.push({ - role: 'tool', - toolCallId: toolCall.id, - name: toolCall.function.name, - content: toolContent, - }); + prefillStartMs = Date.now(); } - prefillStartMs = Date.now(); - } + // Prefill all uncommitted assistant & tool messages so KV cache contains + // full closed conversation + const uncommitted = history.length - committed; + if (uncommitted > 0) { + const prompt = chatPreprocessor.process(history, uncommitted, { addGenPrompt: false }); + await prefill(prompt); + chatPreprocessor.clear(); + committed = history.length; + } - // Prefill all uncommitted assistant & tool messages so KV cache contains - // full closed conversation - const uncommitted = history.length - committed; - if (uncommitted > 0) { - const prompt = chatPreprocessor.process(history, uncommitted, { addGenPrompt: false }); - await prefill(prompt); + return { messages: history.slice(turnStartIdx), stats: generationStatsList, finishReason }; + } catch (err) { + // Roll back history, KV cache, and active tensors to pre-turn state on failure + history.length = turnStartIdx; + committed = initialCommitted; + runner.reset(initialPos); chatPreprocessor.clear(); - committed = history.length; + throw err; } - - return { messages: history.slice(turnStartIdx), stats: generationStatsList, finishReason }; - } catch (err) { - // Roll back history, KV cache, and active tensors to pre-turn state on failure - history.length = turnStartIdx; - committed = initialCommitted; - runner.reset(initialPos); - chatPreprocessor.clear(); - throw err; - } - }; - - return { - stop, - dispose, - sendMessage, - getHistory: () => [...history], - getKVCacheState: () => runner.getKVCacheState(), - }; + }; + + return { + stop, + dispose, + sendMessage, + getHistory: () => [...history], + getKVCacheState: () => runner.getKVCacheState(), + }; + } catch (error) { + dispose(); + throw error; + } } diff --git a/packages/react-native-executorch/src/extensions/nlp/tasks/privacyFilter.ts b/packages/react-native-executorch/src/extensions/nlp/tasks/privacyFilter.ts index 7230f5581f..e9633d7f62 100644 --- a/packages/react-native-executorch/src/extensions/nlp/tasks/privacyFilter.ts +++ b/packages/react-native-executorch/src/extensions/nlp/tasks/privacyFilter.ts @@ -18,6 +18,7 @@ import { } from '../../../core/schema'; import { wrapAsync } from '../../../core/runtime'; import { RnExecuTorchError } from '../../../core/error'; +import { createResourceScope } from '../../../core/lifetime'; import { loadTokenizer } from '../tokenizer'; import { @@ -130,161 +131,162 @@ export async function createPrivacyFilter