Refactor upscale function into a generator (#236)

thekevinscott · Mar 3, 2022 · fcedc68 · fcedc68
1 parent 151d9e5
commit fcedc68
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 9 deletions.
diff --git a/packages/upscalerjs/src/upscale.test.ts b/packages/upscalerjs/src/upscale.test.ts
@@ -1,11 +1,12 @@
 import { tf } from './dependencies.generated';
-import upscale, {
+import {
   predict,
   getRowsAndColumns,
   getTensorDimensions,
   getCopyOfInput,
   getProcessedPixels,
   concatTensors,
+  upscale,
   WARNING_PROGRESS_WITHOUT_PATCH_SIZE,
   WARNING_UNDEFINED_PADDING,
 } from './upscale';
@@ -1408,7 +1409,7 @@ describe('upscale', () => {
       predict: jest.fn(() => tf.ones([1, 2, 2, 3,])),
     } as unknown as tf.LayersModel;
     (mockedTensorAsBase as any).default = async() => 'foobarbaz';
-    const result = await upscale(model, img, { scale: 2, } as IModelDefinition);
+    const result = await wrapGen(upscale(model, img, { scale: 2, } as IModelDefinition));
     expect(result).toEqual('foobarbaz');
   });
 
@@ -1429,9 +1430,9 @@ describe('upscale', () => {
       predict: jest.fn(() => upscaledTensor),
     } as unknown as tf.LayersModel;
     (mockedTensorAsBase as any).default = async() => 'foobarbaz';
-    const result = await upscale(model, img, { scale: 2, } as IModelDefinition, {
+    const result = await wrapGen(upscale(model, img, { scale: 2, } as IModelDefinition, {
       output: 'tensor',
-    });
+    }));
     if (typeof result === 'string') {
       throw new Error('Unexpected string type');
     }

diff --git a/packages/upscalerjs/src/upscale.ts b/packages/upscalerjs/src/upscale.ts
@@ -347,12 +347,12 @@ export function getProcessedPixels<T extends tf.Tensor3D | tf.Tensor4D>(
 // what input is in which format
 export const getCopyOfInput = (input: GetImageAsTensorInput) => isTensor(input) ? input.clone() : input;
 
-async function upscale<P extends Progress<O, PO>, O extends ReturnType = 'src', PO extends ReturnType = undefined>(
+export async function* upscale<P extends Progress<O, PO>, O extends ReturnType = 'src', PO extends ReturnType = undefined>(
   model: tf.LayersModel,
   input: GetImageAsTensorInput,
   modelDefinition: IModelDefinition,
   options: IUpscaleOptions<P, O, PO> = {},
-): Promise<UpscaleResponse<O>> {
+): AsyncGenerator<UpscaleResponse<O>> {
   const parsedInput = getCopyOfInput(input);
   const startingPixels = await getImageAsTensor(parsedInput);
 
@@ -391,4 +391,25 @@ async function upscale<P extends Progress<O, PO>, O extends ReturnType = 'src',
   return <UpscaleResponse<O>>base64Src;
 };
 
-export default upscale;
+async function wrappedUpscale<P extends Progress<O, PO>, O extends ReturnType = 'src', PO extends ReturnType = undefined>(
+  model: tf.LayersModel,
+  input: GetImageAsTensorInput,
+  modelDefinition: IModelDefinition,
+  options: IUpscaleOptions<P, O, PO> = {},
+): Promise<UpscaleResponse<O>> {
+  const gen = upscale(
+    model,
+    input,
+    modelDefinition,
+    options,
+  );
+  let { value: upscaledPixels, done } = await gen.next();
+  while (done === false) {
+    const genResult = await gen.next();
+    upscaledPixels = genResult.value;
+    done = genResult.done;
+  }
+  return upscaledPixels;
+}
+
+export default wrappedUpscale;
diff --git a/packages/upscalerjs/src/upscaler.ts b/packages/upscalerjs/src/upscaler.ts
@@ -9,7 +9,7 @@ import {
 } from './types';
 import loadModel, { getModelDefinitions, } from './loadModel';
 import warmup from './warmup';
-import upscale from './upscale';
+import upscaleImage from './upscale';
 import type { GetImageAsTensorInput, } from './image.generated';
 
 class Upscaler {
@@ -42,7 +42,7 @@ class Upscaler {
     options: IUpscaleOptions<P, O, PO> = {},
   ) => {
     const { model, modelDefinition, } = await this._model;
-    return upscale(model, image, modelDefinition, options);
+    return upscaleImage(model, image, modelDefinition, options);
   };
 
   getModelDefinitions = async () => {