diff --git a/release/ci/model_generator.py b/release/ci/model_generator.py index 80b102a6f1..651dd207de 100755 --- a/release/ci/model_generator.py +++ b/release/ci/model_generator.py @@ -123,7 +123,7 @@ def generate_chunked_model(driving_pkl: Path) -> dict: } -def create_metadata_json(models: list, output_dir: Path, custom_name=None, short_name=None, is_20hz=False, upstream_branch="unknown") -> None: +def create_metadata_json(models: list, output_dir: Path, custom_name=None, short_name=None, is_20hz=False, upstream_branch="unknown", is_big=False) -> None: bundle_json = { "short_name": short_name, "display_name": custom_name or upstream_branch, @@ -136,6 +136,7 @@ def create_metadata_json(models: list, output_dir: Path, custom_name=None, short "generation": "-1", "build_time": datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ"), "overrides": {}, + "is_big": is_big, "models": models, } @@ -170,6 +171,8 @@ if __name__ == "__main__": print(f"No driving_tinygrad.pkl found in {_output_dir}", file=sys.stderr) sys.exit(1) + is_big = _driving_pkl.name.startswith('big_') + if _pkl: new_pkl = _output_dir / f"driving_{_pkl}_tinygrad.pkl" if not new_pkl.exists(): @@ -178,4 +181,4 @@ if __name__ == "__main__": _driving_pkl = new_pkl _model_metadata = generate_chunked_model(_driving_pkl) - create_metadata_json([_model_metadata], _output_dir, args.custom_name, _short_name, args.is_20hz, args.upstream_branch) + create_metadata_json([_model_metadata], _output_dir, args.custom_name, _short_name, args.is_20hz, args.upstream_branch, is_big=is_big)