chore: migrate TextToImageModule to factory pattern, add TextToImageModelName type

chmjkb · claude · chmjkb · commit ba3518307edb · 2026-03-09T17:26:05.000+01:00
- Add TextToImageModelName union type
- Add modelName to TextToImageProps.model
- TextToImageModule: private constructor, fromModelName, fromCustomModel
- useTextToImage: use factory, add model.modelName to deps

Co-Authored-By: Claude Sonnet 4.6 &lt;noreply@anthropic.com&gt;
diff --git a/packages/react-native-executorch/src/hooks/computer_vision/useTextToImage.ts b/packages/react-native-executorch/src/hooks/computer_vision/useTextToImage.ts
@@ -20,36 +20,71 @@ export const useTextToImage = ({
   const [isGenerating, setIsGenerating] = useState(false);
   const [downloadProgress, setDownloadProgress] = useState(0);
   const [error, setError] = useState<RnExecutorchError | null>(null);
-
-  const [module] = useState(() => new TextToImageModule(inferenceCallback));
+  const [moduleInstance, setModuleInstance] =
+    useState<TextToImageModule | null>(null);
 
   useEffect(() => {
     if (preventLoad) return;
 
-    (async () => {
-      setDownloadProgress(0);
-      setError(null);
-      try {
-        setIsReady(false);
-        await module.load(model, setDownloadProgress);
-        setIsReady(true);
-      } catch (err) {
-        setError(parseUnknownError(err));
+    let active = true;
+    setDownloadProgress(0);
+    setError(null);
+    setIsReady(false);
+
+    TextToImageModule.fromModelName(
+      {
+        modelName: model.modelName,
+        tokenizerSource: model.tokenizerSource,
+        schedulerSource: model.schedulerSource,
+        encoderSource: model.encoderSource,
+        unetSource: model.unetSource,
+        decoderSource: model.decoderSource,
+        inferenceCallback,
+      },
+      (p) => {
+        if (active) setDownloadProgress(p);
       }
-    })();
+    )
+      .then((mod) => {
+        if (!active) {
+          mod.delete();
+          return;
+        }
+        setModuleInstance((prev) => {
+          prev?.delete();
+          return mod;
+        });
+        setIsReady(true);
+      })
+      .catch((err) => {
+        if (active) setError(parseUnknownError(err));
+      });
 
     return () => {
-      module.delete();
+      active = false;
+      setModuleInstance((prev) => {
+        prev?.delete();
+        return null;
+      });
     };
-  }, [module, model, preventLoad]);
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [
+    model.modelName,
+    model.tokenizerSource,
+    model.schedulerSource,
+    model.encoderSource,
+    model.unetSource,
+    model.decoderSource,
+    preventLoad,
+  ]);
 
   const generate = async (
     input: string,
     imageSize?: number,
     numSteps?: number,
     seed?: number
   ): Promise<string> => {
-    if (!isReady)
+    if (!isReady || !moduleInstance)
       throw new RnExecutorchError(
         RnExecutorchErrorCode.ModuleNotLoaded,
         'The model is currently not loaded. Please load the model before calling forward().'
@@ -61,17 +96,17 @@ export const useTextToImage = ({
       );
     try {
       setIsGenerating(true);
-      return await module.forward(input, imageSize, numSteps, seed);
+      return await moduleInstance.forward(input, imageSize, numSteps, seed);
     } finally {
       setIsGenerating(false);
     }
   };
 
   const interrupt = useCallback(() => {
-    if (isGenerating) {
-      module.interrupt();
+    if (isGenerating && moduleInstance) {
+      moduleInstance.interrupt();
     }
-  }, [module, isGenerating]);
+  }, [moduleInstance, isGenerating]);
 
   return {
     isReady,
diff --git a/packages/react-native-executorch/src/modules/computer_vision/TextToImageModule.ts b/packages/react-native-executorch/src/modules/computer_vision/TextToImageModule.ts
@@ -1,5 +1,6 @@
 import { ResourceFetcher } from '../../utils/ResourceFetcher';
 import { ResourceSource } from '../../types/common';
+import { TextToImageModelName } from '../../types/tti';
 import { BaseModule } from '../BaseModule';
 
 import { PNG } from 'pngjs/browser';
@@ -15,82 +16,132 @@ import { Logger } from '../../common/Logger';
 export class TextToImageModule extends BaseModule {
   private inferenceCallback: (stepIdx: number) => void;
 
-  /**
-   * Creates a new instance of `TextToImageModule` with optional callback on inference step.
-   *
-   * @param inferenceCallback - Optional callback function that receives the current step index during inference.
-   */
-  constructor(inferenceCallback?: (stepIdx: number) => void) {
+  private constructor(inferenceCallback?: (stepIdx: number) => void) {
     super();
     this.inferenceCallback = (stepIdx: number) => {
       inferenceCallback?.(stepIdx);
     };
   }
 
   /**
-   * Loads the model from specified resources.
+   * Creates a Text to Image instance for a built-in model.
    *
-   * @param model - Object containing sources for tokenizer, scheduler, encoder, unet, and decoder.
-   * @param onDownloadProgressCallback - Optional callback to monitor download progress.
+   * @param namedSources - An object specifying the model name, pipeline sources, and optional inference callback.
+   * @param onDownloadProgress - Optional callback to monitor download progress, receiving a value between 0 and 1.
+   * @returns A Promise resolving to a `TextToImageModule` instance.
+   *
+   * @example
+   * ```ts
+   * import { TextToImageModule, BK_SDM_TINY_VPRED_512 } from 'react-native-executorch';
+   * const tti = await TextToImageModule.fromModelName(BK_SDM_TINY_VPRED_512);
+   * ```
    */
-  async load(
-    model: {
+  static async fromModelName(
+    namedSources: {
+      modelName: TextToImageModelName;
       tokenizerSource: ResourceSource;
       schedulerSource: ResourceSource;
       encoderSource: ResourceSource;
       unetSource: ResourceSource;
       decoderSource: ResourceSource;
+      inferenceCallback?: (stepIdx: number) => void;
     },
-    onDownloadProgressCallback: (progress: number) => void = () => {}
-  ): Promise<void> {
+    onDownloadProgress: (progress: number) => void = () => {}
+  ): Promise<TextToImageModule> {
+    const instance = new TextToImageModule(namedSources.inferenceCallback);
     try {
-      const results = await ResourceFetcher.fetch(
-        onDownloadProgressCallback,
-        model.tokenizerSource,
-        model.schedulerSource,
-        model.encoderSource,
-        model.unetSource,
-        model.decoderSource
-      );
-      if (!results) {
-        throw new RnExecutorchError(
-          RnExecutorchErrorCode.DownloadInterrupted,
-          'The download has been interrupted. As a result, not every file was downloaded. Please retry the download.'
-        );
-      }
-      const [tokenizerPath, schedulerPath, encoderPath, unetPath, decoderPath] =
-        results;
+      await instance.internalLoad(namedSources, onDownloadProgress);
+      return instance;
+    } catch (error) {
+      Logger.error('Load failed:', error);
+      throw parseUnknownError(error);
+    }
+  }
 
-      if (
-        !tokenizerPath ||
-        !schedulerPath ||
-        !encoderPath ||
-        !unetPath ||
-        !decoderPath
-      ) {
-        throw new RnExecutorchError(
-          RnExecutorchErrorCode.DownloadInterrupted,
-          'The download has been interrupted. As a result, not every file was downloaded. Please retry the download.'
-        );
-      }
+  /**
+   * Creates a Text to Image instance with user-provided model binaries.
+   * Use this when working with a custom-exported diffusion pipeline.
+   * Internally uses `'custom'` as the model name for telemetry.
+   *
+   * @param sources - An object containing the pipeline source paths.
+   * @param onDownloadProgress - Optional callback to monitor download progress, receiving a value between 0 and 1.
+   * @param inferenceCallback - Optional callback triggered after each diffusion step.
+   * @returns A Promise resolving to a `TextToImageModule` instance.
+   */
+  static fromCustomModel(
+    sources: {
+      tokenizerSource: ResourceSource;
+      schedulerSource: ResourceSource;
+      encoderSource: ResourceSource;
+      unetSource: ResourceSource;
+      decoderSource: ResourceSource;
+    },
+    onDownloadProgress: (progress: number) => void = () => {},
+    inferenceCallback?: (stepIdx: number) => void
+  ): Promise<TextToImageModule> {
+    return TextToImageModule.fromModelName(
+      {
+        modelName: 'custom' as TextToImageModelName,
+        ...sources,
+        inferenceCallback,
+      },
+      onDownloadProgress
+    );
+  }
 
-      const response = await fetch('file://' + schedulerPath);
-      const schedulerConfig = await response.json();
+  private async internalLoad(
+    model: {
+      tokenizerSource: ResourceSource;
+      schedulerSource: ResourceSource;
+      encoderSource: ResourceSource;
+      unetSource: ResourceSource;
+      decoderSource: ResourceSource;
+    },
+    onDownloadProgressCallback: (progress: number) => void
+  ): Promise<void> {
+    const results = await ResourceFetcher.fetch(
+      onDownloadProgressCallback,
+      model.tokenizerSource,
+      model.schedulerSource,
+      model.encoderSource,
+      model.unetSource,
+      model.decoderSource
+    );
+    if (!results) {
+      throw new RnExecutorchError(
+        RnExecutorchErrorCode.DownloadInterrupted,
+        'The download has been interrupted. As a result, not every file was downloaded. Please retry the download.'
+      );
+    }
+    const [tokenizerPath, schedulerPath, encoderPath, unetPath, decoderPath] =
+      results;
 
-      this.nativeModule = global.loadTextToImage(
-        tokenizerPath,
-        encoderPath,
-        unetPath,
-        decoderPath,
-        schedulerConfig.beta_start,
-        schedulerConfig.beta_end,
-        schedulerConfig.num_train_timesteps,
-        schedulerConfig.steps_offset
+    if (
+      !tokenizerPath ||
+      !schedulerPath ||
+      !encoderPath ||
+      !unetPath ||
+      !decoderPath
+    ) {
+      throw new RnExecutorchError(
+        RnExecutorchErrorCode.DownloadInterrupted,
+        'The download has been interrupted. As a result, not every file was downloaded. Please retry the download.'
       );
-    } catch (error) {
-      Logger.error('Load failed:', error);
-      throw parseUnknownError(error);
     }
+
+    const response = await fetch('file://' + schedulerPath);
+    const schedulerConfig = await response.json();
+
+    this.nativeModule = global.loadTextToImage(
+      tokenizerPath,
+      encoderPath,
+      unetPath,
+      decoderPath,
+      schedulerConfig.beta_start,
+      schedulerConfig.beta_end,
+      schedulerConfig.num_train_timesteps,
+      schedulerConfig.steps_offset
+    );
   }
 
   /**
diff --git a/packages/react-native-executorch/src/types/tti.ts b/packages/react-native-executorch/src/types/tti.ts
@@ -1,6 +1,15 @@
 import { RnExecutorchError } from '../errors/errorUtils';
 import { ResourceSource } from '../types/common';
 
+/**
+ * Union of all built-in Text-to-Image model names.
+ *
+ * @category Types
+ */
+export type TextToImageModelName =
+  | 'bk-sdm-tiny-vpred-512'
+  | 'bk-sdm-tiny-vpred-256';
+
 /**
  * Configuration properties for the `useTextToImage` hook.
  *
@@ -11,6 +20,11 @@ export interface TextToImageProps {
    * Object containing the required model sources for the diffusion pipeline.
    */
   model: {
+    /**
+     * The built-in model name (e.g. `'bk-sdm-tiny-vpred-512'`). Used for telemetry and hook reload triggers.
+     * Pass one of the pre-built TTI constants (e.g. `BK_SDM_TINY_VPRED_512`) to populate all required fields.
+     */
+    modelName: TextToImageModelName;
     /** Source for the text tokenizer binary/config. */
     tokenizerSource: ResourceSource;
     /** Source for the diffusion scheduler binary/config. */
