react-native-executorch/src/hooks/computer_vision/useStyleTransfer.ts

An easy way to run AI models in React Native with ExecuTorch

import { StyleTransferModule } from '../../modules/computer_vision/StyleTransferModule';
import { PixelData } from '../../types/common';
import {
  StyleTransferProps,
  StyleTransferType,
} from '../../types/styleTransfer';
import { useModuleFactory } from '../useModuleFactory';

/**
 * React hook for managing a Style Transfer model instance.
 * @category Hooks
 * @param StyleTransferProps - Configuration object containing `model` source and optional `preventLoad` flag.
 * @returns Ready to use Style Transfer model.
 */
export const useStyleTransfer = ({
  model,
  preventLoad = false,
}: StyleTransferProps): StyleTransferType => {
  const {
    error,
    isReady,
    isGenerating,
    downloadProgress,
    runForward,
    runOnFrame,
  } = useModuleFactory({
    factory: (config, onProgress) =>
      StyleTransferModule.fromModelName(config, onProgress),
    config: model,
    deps: [model.modelName, model.modelSource],
    preventLoad,
  });

  const forward = <O extends 'pixelData' | 'url' = 'pixelData'>(
    imageSource: string | PixelData,
    outputType?: O
  ) => runForward((inst) => inst.forward(imageSource, outputType));

  return {
    error,
    isReady,
    isGenerating,
    downloadProgress,
    forward,
    runOnFrame,
  } as StyleTransferType;
};