xref: /aosp_15_r20/external/executorch/backends/apple/coreml/runtime/delegate/ETCoreMLModelLoader.mm (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a)
1//
2// ETCoreMLModelLoader.mm
3//
4// Copyright © 2024 Apple Inc. All rights reserved.
5//
6// Please refer to the license found in the LICENSE file in the root directory of the source tree.
7
8#import <ETCoreMLAsset.h>
9#import <ETCoreMLAssetManager.h>
10#import <ETCoreMLDefaultModelExecutor.h>
11#import <ETCoreMLLogging.h>
12#import <ETCoreMLModel.h>
13#import <ETCoreMLModelLoader.h>
14#import <asset.h>
15#import <model_metadata.h>
16
17using namespace executorchcoreml;
18
19namespace {
20    NSOrderedSet<NSString *> *get_ordered_set(const std::vector<std::string>& values) {
21        NSMutableOrderedSet<NSString *> *result = [NSMutableOrderedSet orderedSetWithCapacity:values.size()];
22        for (const auto& value : values) {
23            [result addObject:@(value.c_str())];
24        }
25
26        return result;
27    }
28
29    ETCoreMLModel * _Nullable get_model_from_asset(ETCoreMLAsset *asset,
30                                                   MLModelConfiguration *configuration,
31                                                   const executorchcoreml::ModelMetadata& metadata,
32                                                   NSError * __autoreleasing *error) {
33        NSOrderedSet<NSString *> *orderedInputNames = ::get_ordered_set(metadata.input_names);
34        NSOrderedSet<NSString *> *orderedOutputNames = ::get_ordered_set(metadata.output_names);
35        ETCoreMLModel *model = [[ETCoreMLModel alloc] initWithAsset:asset
36                                                      configuration:configuration
37                                                  orderedInputNames:orderedInputNames
38                                                 orderedOutputNames:orderedOutputNames
39                                                              error:error];
40        return model;
41    }
42} // namespace
43
44@implementation ETCoreMLModelLoader
45
46+ (nullable ETCoreMLModel *)loadModelWithContentsOfURL:(NSURL *)compiledModelURL
47                                         configuration:(MLModelConfiguration *)configuration
48                                              metadata:(const executorchcoreml::ModelMetadata&)metadata
49                                          assetManager:(ETCoreMLAssetManager *)assetManager
50                                                 error:(NSError * __autoreleasing *)error {
51    NSError *localError = nil;
52    NSString *identifier = @(metadata.identifier.c_str());
53    ETCoreMLAsset *asset = nil;
54    if ([assetManager hasAssetWithIdentifier:identifier error:&localError]) {
55        asset = [assetManager assetWithIdentifier:identifier error:&localError];
56    } else {
57        asset = [assetManager storeAssetAtURL:compiledModelURL withIdentifier:identifier error:&localError];
58    }
59
60    ETCoreMLModel *model = (asset != nil) ? get_model_from_asset(asset, configuration, metadata, &localError) : nil;
61    if (model) {
62        return model;
63    }
64
65    if (localError) {
66        ETCoreMLLogError(localError,
67                         "%@: Failed to load model from compiled asset with identifier = %@",
68                         NSStringFromClass(ETCoreMLModelLoader.class),
69                         identifier);
70    }
71
72    // If store failed then we will load the model from compiledURL.
73    auto backingAsset = Asset::make(compiledModelURL, identifier, assetManager.fileManager, error);
74    if (!backingAsset) {
75        return nil;
76    }
77
78    asset = [[ETCoreMLAsset alloc] initWithBackingAsset:backingAsset.value()];
79    return ::get_model_from_asset(asset, configuration, metadata, error);
80}
81
82@end
83