Add cache_support flag (#9)

Signed-off-by: Greg Clark <[email protected]>

Author: messiaen

nemo2riva/args.py

@@ -29,6 +29,12 @@ def get_args(argv):
     parser.add_argument("--onnx-opset", type=int, default=14, help="ONNX opset for model export")
     parser.add_argument("--device", default="cuda", help="Device to export for")
     parser.add_argument("--export-subnet", default=None, help="Export specified subnetwork/layer")
+    parser.add_argument(
+        "--cache-support",
+        type=lambda x: bool(strtobool(x)),
+        default=None,
+        help="cache support for models that support it"
+    )
 
     args = parser.parse_args(argv)
     return args

nemo2riva/cookbook.py

@@ -51,6 +51,10 @@ def export_model(model, cfg, args, artifacts, metadata):
     runtime = format_meta["runtime"]
     metadata.update({"runtime": runtime})
 
+    if cfg.cache_support and hasattr(model, "encoder") and hasattr(model.encoder, "export_cache_support"):
+        model.encoder.export_cache_support = True
+        logging.info("Caching support is enabled.")
+
     with tempfile.TemporaryDirectory() as tmpdir:
         export_filename = cfg.export_file
         export_file = os.path.join(tmpdir, export_filename)

nemo2riva/schema.py

@@ -31,6 +31,7 @@ class ExportConfig:
     encryption: Optional[str] = None
     autocast: bool = False
     max_dim: int = None
+    cache_support: bool = False
 
 
 @dataclass
@@ -90,6 +91,9 @@ def get_export_config(export_obj, args):
             "Format `{}` is invalid. Please pick one of the ({})".format(conf.export_format, supported_formats)
         )
 
+    if args.cache_support is not None:
+        conf.cache_support = args.cache_support
+
     return conf