From 8c4b971c145df29f2aa870e7cf5b981f9ed6b709 Mon Sep 17 00:00:00 2001
From: MediaPipe Team <mediapipe-team@google.com>
Date: Tue, 22 Aug 2023 11:32:30 -0700
Subject: [PATCH] Add an API in model_task_graph to create or use cached model
 resources.

PiperOrigin-RevId: 559174528
---
 mediapipe/tasks/cc/core/model_task_graph.cc | 15 +++++++++++++++
 mediapipe/tasks/cc/core/model_task_graph.h  | 14 ++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/mediapipe/tasks/cc/core/model_task_graph.cc b/mediapipe/tasks/cc/core/model_task_graph.cc
index 46cc088f2..225fad418 100644
--- a/mediapipe/tasks/cc/core/model_task_graph.cc
+++ b/mediapipe/tasks/cc/core/model_task_graph.cc
@@ -186,6 +186,21 @@ absl::StatusOr<const ModelResources*> ModelTaskGraph::CreateModelResources(
   return model_resources_cache_service.GetObject().GetModelResources(tag);
 }
 
+absl::StatusOr<const ModelResources*> ModelTaskGraph::GetOrCreateModelResources(
+    SubgraphContext* sc, std::unique_ptr<proto::ExternalFile> external_file,
+    std::string tag_suffix) {
+  auto model_resources_cache_service = sc->Service(kModelResourcesCacheService);
+  if (model_resources_cache_service.IsAvailable()) {
+    std::string tag =
+        absl::StrCat(CreateModelResourcesTag(sc->OriginalNode()), tag_suffix);
+    if (model_resources_cache_service.GetObject().Exists(tag)) {
+      return model_resources_cache_service.GetObject().GetModelResources(tag);
+    }
+  }
+  return ModelTaskGraph::CreateModelResources(sc, std::move(external_file),
+                                              tag_suffix);
+}
+
 absl::StatusOr<const ModelAssetBundleResources*>
 ModelTaskGraph::CreateModelAssetBundleResources(
     SubgraphContext* sc, std::unique_ptr<proto::ExternalFile> external_file,
diff --git a/mediapipe/tasks/cc/core/model_task_graph.h b/mediapipe/tasks/cc/core/model_task_graph.h
index 10634d3d0..38367da8f 100644
--- a/mediapipe/tasks/cc/core/model_task_graph.h
+++ b/mediapipe/tasks/cc/core/model_task_graph.h
@@ -87,6 +87,20 @@ class ModelTaskGraph : public Subgraph {
       SubgraphContext* sc, std::unique_ptr<proto::ExternalFile> external_file,
       std::string tag_suffix = "");
 
+  template <typename Options>
+  absl::StatusOr<const ModelResources*> GetOrCreateModelResources(
+      SubgraphContext* sc, std::string tag_suffix = "") {
+    auto external_file = std::make_unique<proto::ExternalFile>();
+    external_file->Swap(sc->MutableOptions<Options>()
+                            ->mutable_base_options()
+                            ->mutable_model_asset());
+    return GetOrCreateModelResources(sc, std::move(external_file), tag_suffix);
+  }
+
+  absl::StatusOr<const ModelResources*> GetOrCreateModelResources(
+      SubgraphContext* sc, std::unique_ptr<proto::ExternalFile> external_file,
+      std::string tag_suffix = "");
+
   // If the model resources graph service is available, creates a model asset
   // bundle resources object from the subgraph context, and caches the created
   // model asset bundle resources into the model resources graph service on