Manage input inside UNet

Author: a-szymanska

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

@@ -14,7 +14,7 @@ Decoder::Decoder(const std::string &modelSource, int32_t modelImageSize,
     : BaseModel(modelSource, callInvoker), modelImageSize(modelImageSize),
       numChannels(numChannels) {}
 
-std::vector<float> Decoder::generate(std::vector<float> &input) {
+std::vector<float> Decoder::generate(std::vector<float> &input) const {
   constexpr int32_t latentDownsample = 8;
   const int32_t latentsImageSize =
       std::floor(modelImageSize / latentDownsample);

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

@@ -17,7 +17,7 @@ class Decoder final : public BaseModel {
   explicit Decoder(const std::string &modelSource, int32_t modelImageSize,
                    int32_t numChannels,
                    std::shared_ptr<react::CallInvoker> callInvoker);
-  std::vector<float> generate(std::vector<float> &input);
+  std::vector<float> generate(std::vector<float> &input) const;
 
 private:
   int32_t modelImageSize;

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

@@ -47,9 +47,9 @@ TextToImage::generate(std::string input, size_t numInferenceSteps,
                           embeddingsTextPtr + embeddingsSize);
 
   constexpr int32_t latentDownsample = 8;
-  int32_t latentsWidth = std::floor(modelImageSize / latentDownsample);
-  int32_t latentsSize = numChannels * latentsWidth * latentsWidth;
-  std::vector<float> latents(latentsSize);
+  int32_t latentsSize = std::floor(modelImageSize / latentDownsample);
+  int32_t latentsImageSize = numChannels * latentsSize * latentsSize;
+  std::vector<float> latents(latentsImageSize);
   std::random_device rd;
   std::mt19937 gen(rd());
   std::normal_distribution<float> dist(0.0, 1.0);
@@ -71,13 +71,9 @@ TextToImage::generate(std::string input, size_t numInferenceSteps,
       return postprocess({});
     }
     log(LOG_LEVEL::Debug, "Step:", t, "/", numInferenceSteps);
-    std::vector<float> latentsConcat;
-    latentsConcat.reserve(2 * latentsSize);
-    latentsConcat.insert(latentsConcat.end(), latents.begin(), latents.end());
-    latentsConcat.insert(latentsConcat.end(), latents.begin(), latents.end());
 
     std::vector<float> noisePred =
-        unet->generate(latentsConcat, timesteps[t], embeddingsConcat);
+        unet->generate(latents, timesteps[t], embeddingsConcat);
 
     size_t noiseSize = noisePred.size() / 2;
     std::span<const float> noisePredSpan{noisePred};

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

@@ -8,26 +8,29 @@ using namespace executorch::extension;
 
 UNet::UNet(const std::string &modelSource, int32_t modelImageSize,
            int32_t numChannels, std::shared_ptr<react::CallInvoker> callInvoker)
-    : BaseModel(modelSource, callInvoker), modelImageSize(modelImageSize),
-      numChannels(numChannels) {}
+    : BaseModel(modelSource, callInvoker), numChannels(numChannels) {
+  constexpr int32_t latentDownsample = 8;
+  latentsSize = std::floor(modelImageSize / latentDownsample);
+}
 
 std::vector<float> UNet::generate(std::vector<float> &latents, int32_t timestep,
-                                  std::vector<float> &embeddings) {
-  constexpr int32_t latentDownsample = 8;
-  const int32_t latentsImageSize =
-      std::floor(modelImageSize / latentDownsample);
-  std::vector<int32_t> latentsShape = {2, numChannels, latentsImageSize,
-                                       latentsImageSize};
+                                  std::vector<float> &embeddings) const {
+  std::vector<float> latentsConcat;
+  latentsConcat.reserve(2 * latentsSize);
+  latentsConcat.insert(latentsConcat.end(), latents.begin(), latents.end());
+  latentsConcat.insert(latentsConcat.end(), latents.begin(), latents.end());
+
+  std::vector<int32_t> latentsShape = {2, numChannels, latentsSize,
+                                       latentsSize};
   std::vector<int32_t> timestepShape = {1};
   std::vector<int32_t> embeddingsShape = {2, 77, 768};
 
-  // TODO change after reexporting the model
   std::vector<int64_t> timestepData = {static_cast<int64_t>(timestep)};
   auto timestepTensor =
       make_tensor_ptr(timestepShape, timestepData.data(), ScalarType::Long);
 
   auto latentsTensor =
-      make_tensor_ptr(latentsShape, latents.data(), ScalarType::Float);
+      make_tensor_ptr(latentsShape, latentsConcat.data(), ScalarType::Float);
   auto embeddingsTensor =
       make_tensor_ptr(embeddingsShape, embeddings.data(), ScalarType::Float);
 

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

@@ -18,11 +18,11 @@ class UNet final : public BaseModel {
                 int32_t numChannels,
                 std::shared_ptr<react::CallInvoker> callInvoker);
   std::vector<float> generate(std::vector<float> &latents, int32_t timestep,
-                              std::vector<float> &embeddings);
+                              std::vector<float> &embeddings) const;
 
 private:
-  int32_t modelImageSize;
   int32_t numChannels;
+  int32_t latentsSize;
 };
 } // namespace models::text_to_image