react-native-executorch/lib/module/hooks/computer_vision/useTextToImage.js

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

"use strict";

import { useCallback, useEffect, useState } from 'react';
import { ETError, getError } from '../../Error';
import { TextToImageModule } from '../../modules/computer_vision/TextToImageModule';
export const useTextToImage = ({
  model,
  inferenceCallback,
  preventLoad = false
}) => {
  const [isReady, setIsReady] = useState(false);
  const [isGenerating, setIsGenerating] = useState(false);
  const [downloadProgress, setDownloadProgress] = useState(0);
  const [error, setError] = useState(null);
  const [module] = useState(() => new TextToImageModule(inferenceCallback));
  useEffect(() => {
    if (preventLoad) return;
    (async () => {
      setDownloadProgress(0);
      setError(null);
      try {
        setIsReady(false);
        await module.load(model, setDownloadProgress);
        setIsReady(true);
      } catch (err) {
        setError(err.message);
      }
    })();
    return () => {
      module.delete();
    };
  }, [module, model, preventLoad]);
  const generate = async (input, imageSize, numSteps, seed) => {
    if (!isReady) throw new Error(getError(ETError.ModuleNotLoaded));
    if (isGenerating) throw new Error(getError(ETError.ModelGenerating));
    try {
      setIsGenerating(true);
      return await module.forward(input, imageSize, numSteps, seed);
    } finally {
      setIsGenerating(false);
    }
  };
  const interrupt = useCallback(() => {
    if (isGenerating) {
      module.interrupt();
    }
  }, [module, isGenerating]);
  return {
    isReady,
    isGenerating,
    downloadProgress,
    error,
    generate,
    interrupt
  };
};
//# sourceMappingURL=useTextToImage.js.map