Extract encoder class

Author: a-szymanska

packages/react-native-executorch/common/rnexecutorch/models/text_to_image/Encoder.cpp

@@ -0,0 +1,43 @@
+#include "Encoder.h"
+
+#include <cmath>
+#include <random>
+#include <span>
+
+#include <rnexecutorch/models/text_to_image/Constants.h>
+
+namespace rnexecutorch::models::text_to_image {
+
+Encoder::Encoder(const std::string &tokenizerSource,
+                 const std::string &encoderSource,
+                 std::shared_ptr<react::CallInvoker> callInvoker)
+    : callInvoker(callInvoker),
+      encoder(std::make_unique<embeddings::TextEmbeddings>(
+          encoderSource, tokenizerSource, callInvoker)) {}
+
+std::vector<float> Encoder::generate(std::string input) {
+  std::shared_ptr<OwningArrayBuffer> embeddingsText = encoder->generate(input);
+  std::shared_ptr<OwningArrayBuffer> embeddingsUncond =
+      encoder->generate(std::string(constants::kBosToken));
+
+  size_t embeddingsSize = embeddingsText->size() / sizeof(float);
+  auto *embeddingsTextPtr = reinterpret_cast<float *>(embeddingsText->data());
+  auto *embeddingsUncondPtr =
+      reinterpret_cast<float *>(embeddingsUncond->data());
+
+  std::vector<float> embeddingsConcat;
+  embeddingsConcat.reserve(embeddingsSize * 2);
+  embeddingsConcat.insert(embeddingsConcat.end(), embeddingsUncondPtr,
+                          embeddingsUncondPtr + embeddingsSize);
+  embeddingsConcat.insert(embeddingsConcat.end(), embeddingsTextPtr,
+                          embeddingsTextPtr + embeddingsSize);
+  return embeddingsConcat;
+}
+
+size_t Encoder::getMemoryLowerBound() const noexcept {
+  return encoder->getMemoryLowerBound();
+}
+
+void Encoder::unload() noexcept { encoder.reset(nullptr); }
+
+} // namespace rnexecutorch::models::text_to_image

packages/react-native-executorch/common/rnexecutorch/models/text_to_image/Encoder.h

@@ -0,0 +1,36 @@
+#pragma once
+
+#include <memory>
+#include <string>
+#include <vector>
+
+#include <ReactCommon/CallInvoker.h>
+#include <jsi/jsi.h>
+
+#include <rnexecutorch/jsi/OwningArrayBuffer.h>
+#include <rnexecutorch/metaprogramming/ConstructorHelpers.h>
+
+#include <rnexecutorch/models/embeddings/text/TextEmbeddings.h>
+
+namespace rnexecutorch {
+namespace models::text_to_image {
+using namespace facebook;
+
+class Encoder final {
+public:
+  explicit Encoder(const std::string &tokenizerSource,
+                   const std::string &encoderSource,
+                   std::shared_ptr<react::CallInvoker> callInvoker);
+  std::vector<float> generate(std::string input);
+  size_t getMemoryLowerBound() const noexcept;
+  void unload() noexcept;
+
+private:
+  std::shared_ptr<react::CallInvoker> callInvoker;
+  std::unique_ptr<embeddings::TextEmbeddings> encoder;
+};
+} // namespace models::text_to_image
+
+REGISTER_CONSTRUCTOR(models::text_to_image::Encoder, std::string, std::string,
+                     std::shared_ptr<react::CallInvoker>);
+} // namespace rnexecutorch

packages/react-native-executorch/common/rnexecutorch/models/text_to_image/TextToImage.cpp

@@ -5,6 +5,7 @@
 #include <span>
 
 #include <rnexecutorch/Log.h>
+#include <rnexecutorch/models/text_to_image/Constants.h>
 
 namespace rnexecutorch::models::text_to_image {
 
@@ -20,8 +21,8 @@ TextToImage::TextToImage(const std::string &tokenizerSource,
       scheduler(std::make_unique<Scheduler>(
           schedulerBetaStart, schedulerBetaEnd, schedulerNumTrainTimesteps,
           schedulerStepsOffset, callInvoker)),
-      encoder(std::make_unique<embeddings::TextEmbeddings>(
-          encoderSource, tokenizerSource, callInvoker)),
+      encoder(std::make_unique<Encoder>(tokenizerSource, encoderSource,
+                                        callInvoker)),
       unet(std::make_unique<UNet>(unetSource, imageSize, numChannels,
                                   callInvoker)),
       decoder(std::make_unique<Decoder>(decoderSource, imageSize, numChannels,
@@ -30,21 +31,7 @@ TextToImage::TextToImage(const std::string &tokenizerSource,
 std::shared_ptr<OwningArrayBuffer>
 TextToImage::generate(std::string input, size_t numInferenceSteps,
                       std::shared_ptr<jsi::Function> callback) {
-  std::shared_ptr<OwningArrayBuffer> embeddingsText = encoder->generate(input);
-  std::shared_ptr<OwningArrayBuffer> embeddingsUncond =
-      encoder->generate(std::string(constants::kBosToken));
-
-  size_t embeddingsSize = embeddingsText->size() / sizeof(float);
-  auto *embeddingsTextPtr = reinterpret_cast<float *>(embeddingsText->data());
-  auto *embeddingsUncondPtr =
-      reinterpret_cast<float *>(embeddingsUncond->data());
-
-  std::vector<float> embeddingsConcat;
-  embeddingsConcat.reserve(embeddingsSize * 2);
-  embeddingsConcat.insert(embeddingsConcat.end(), embeddingsUncondPtr,
-                          embeddingsUncondPtr + embeddingsSize);
-  embeddingsConcat.insert(embeddingsConcat.end(), embeddingsTextPtr,
-                          embeddingsTextPtr + embeddingsSize);
+  std::vector<float> embeddings = encoder->generate(input);
 
   constexpr int32_t latentDownsample = 8;
   int32_t latentsSize = std::floor(modelImageSize / latentDownsample);
@@ -73,7 +60,7 @@ TextToImage::generate(std::string input, size_t numInferenceSteps,
     log(LOG_LEVEL::Debug, "Step:", t, "/", numInferenceSteps);
 
     std::vector<float> noisePred =
-        unet->generate(latents, timesteps[t], embeddingsConcat);
+        unet->generate(latents, timesteps[t], embeddings);
 
     size_t noiseSize = noisePred.size() / 2;
     std::span<const float> noisePredSpan{noisePred};
@@ -127,9 +114,9 @@ size_t TextToImage::getMemoryLowerBound() const noexcept {
 }
 
 void TextToImage::unload() noexcept {
-  encoder.reset(nullptr);
-  unet.reset(nullptr);
-  decoder.reset(nullptr);
+  encoder->unload();
+  unet->unload();
+  decoder->unload();
 }
 
 } // namespace rnexecutorch::models::text_to_image

packages/react-native-executorch/common/rnexecutorch/models/text_to_image/TextToImage.h

@@ -10,9 +10,8 @@
 #include <rnexecutorch/jsi/OwningArrayBuffer.h>
 #include <rnexecutorch/metaprogramming/ConstructorHelpers.h>
 
-#include <rnexecutorch/models/embeddings/text/TextEmbeddings.h>
-#include <rnexecutorch/models/text_to_image/Constants.h>
 #include <rnexecutorch/models/text_to_image/Decoder.h>
+#include <rnexecutorch/models/text_to_image/Encoder.h>
 #include <rnexecutorch/models/text_to_image/Scheduler.h>
 #include <rnexecutorch/models/text_to_image/UNet.h>
 
@@ -50,7 +49,7 @@ class TextToImage final {
 
   std::shared_ptr<react::CallInvoker> callInvoker;
   std::unique_ptr<Scheduler> scheduler;
-  std::unique_ptr<embeddings::TextEmbeddings> encoder;
+  std::unique_ptr<Encoder> encoder;
   std::unique_ptr<UNet> unet;
   std::unique_ptr<Decoder> decoder;
 };