code/scripts/make_weights_upload.py

#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Stage SceneSeg's re-hosted weights for ``changh95/sceneseg-p150`` under ``weights/`` (PLAN.md D10, user decision).

SceneSeg's weights are not on the Hub upstream (``AutowareFoundation/SceneSeg`` holds only its card); VisionPilot links
Google Drive files. This script writes ``<bundle>/weights_upload/`` (git-ignored, never part of ``code/``)::

    weights_upload/UPLOAD.md                 what the orchestrator uploads and pins (not uploaded itself)
    weights_upload/weights/SceneSeg_FP32.onnx                                   the deployed graph, byte for byte
    weights_upload/weights/sceneseg_iter_140215_epoch_4_step_15999.safetensors  the .pth state_dict (bit-identical)
    weights_upload/weights/LICENSE           upstream Apache-2.0 licence (autowarefoundation/vision_pilot)
    weights_upload/weights/NOTICE            attribution and the changes made here
    weights_upload/weights/provenance.json   Drive ids, sha256 and sizes of the originals, conversion and checks
    weights_upload/weights/README.md         the weights section's own card (the repo root README is the bundle card)

Data only: the ``.pth`` is read with ``torch.load(weights_only=True)``, the ONNX through the onnx protobuf parser;
nothing downloaded is executed. Run it in the research venv (torch, safetensors, onnx; CPU)::

    OMP_NUM_THREADS=4 tools/research-venv/bin/python -I bundles/sceneseg-p150/code/scripts/make_weights_upload.py
"""
from __future__ import annotations

import argparse
import datetime as dt
import hashlib
import json
import os
import shutil
import sys
from pathlib import Path

BUNDLE = Path(__file__).resolve().parents[2]
WS = BUNDLE.parents[1]
sys.path.insert(0, str(BUNDLE / "code"))

from tt_sceneseg.reference import config  # noqa: E402

ASSETS = WS / "assets" / "sceneseg"
ONNX_SRC = ASSETS / "gdrive_SceneSeg_FP32_onnx" / config.ONNX_FILE
PTH_SRC = ASSETS / "gdrive_SceneSeg_pth" / config.PTH_FILE
LICENSE_SRC = WS / "upstream" / "autoware.privately-owned-vehicles" / "LICENSE"
LICENSE_SHA256 = "43070e2d4e532684de521b885f385d0841030efa2b1a20bafb76133a5e1379c1"   # Apache-2.0 text @ 2bae0db
UPSTREAM_HEAD = "2bae0db15dd7b3823b91f8415dd3920756dc0804"                              # vision_pilot, 2026-09-30


def sha256(path: Path) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for block in iter(lambda: f.read(1 << 20), b""):
            h.update(block)
    return h.hexdigest()


def drive_url(file_id: str) -> str:
    return f"https://drive.google.com/file/d/{file_id}"


NOTICE = """SceneSeg weights (Autoware VisionPilot), re-hosted by changh95/sceneseg-p150

Copyright The Autoware Foundation and the contributors of autowarefoundation/vision_pilot (formerly
autowarefoundation/autoware.privately-owned-vehicles). Licensed under the Apache License, Version 2.0 (see LICENSE).
The upstream README states that the complete Vision Pilot codebase, "including AI model weights", is available under
the Apache 2.0 licence (README.md lines 41-42 @ {head}).

Original files (linked from Models/model_library/SceneSeg/README.md @ {code_sha}):
  {onnx}  Google Drive {onnx_id}  sha256 {onnx_sha}
  {pth}  Google Drive {pth_id}  sha256 {pth_sha}

Changes made here:
  - {onnx} is redistributed unchanged (byte for byte).
  - {pth} (a PyTorch state_dict) was converted to
    {st}: the same 412 tensors, bit-identical, stored in the
    safetensors format instead of a pickle (torch.load(weights_only=True); no code executed).

The EfficientNet-B0 backbone of SceneSeg was initialised from torchvision's IMAGENET1K_V1 weights (BSD-3-Clause,
https://github.com/pytorch/vision). The upstream model card lists the training data (ACDC, MUSES, IDDAW, Mapillary
Vistas and comma10K; BDD100K held out); check those datasets' terms before commercial use. The Autoware DISCLAIMER
(autowarefoundation/vision_pilot DISCLAIMER.md) applies.
"""


README = """# SceneSeg weights (re-hosted)

The weights of Autoware VisionPilot's **SceneSeg** network, re-hosted unchanged for
[changh95/sceneseg-p150](https://huggingface.co/changh95/sceneseg-p150) (SceneSeg on a Tenstorrent Blackhole p150).
Upstream publishes them only as Google Drive links (the [AutowareFoundation/SceneSeg](https://huggingface.co/AutowareFoundation/SceneSeg)
card holds no weight files), so they live here, next to the port that loads them. Licence: Apache-2.0 (`LICENSE`,
`NOTICE`).

| file | bytes | sha256 | what |
|---|---:|---|---|
{rows}

`provenance.json` records the Google Drive ids, the upstream commits that link them, the sha256 of every original,
the conversion and its checks. The port loads `SceneSeg_FP32.onnx` (the graph VisionPilot deploys, BatchNorm folded at
export) by consuming node; the safetensors file is the training checkpoint the ONNX was exported from (identical
weights: per-tap PCC >= 0.9999999986 between the two, research/sceneseg/SPEC.md section 6.2), kept as the source of
truth.
"""


UPLOAD = """# Upload of SceneSeg's weights (orchestrator)

PLAN.md D10 (user decision): the weights go INSIDE the bundle repo `changh95/sceneseg-p150`, under `weights/`, as a
commit of their own BEFORE the bundle is packaged; the bundle then pins that commit.

1. Upload exactly the directory `weights/` of this folder to the repo root (paths `weights/<file>`), e.g.
   `HF_TOKEN=... hf upload changh95/sceneseg-p150 weights_upload/weights weights --repo-type model`
   (never through `bin/devrun`; the token only in the environment).
2. Pin the commit it created: `research/sceneseg/bundle_values.json` `AW_SHA` (40 hex), then
   `python research/packaging/scripts/instantiate_bundle.py --values research/sceneseg/bundle_values.json --fill
   bundles/sceneseg-p150` fills `api.SceneSeg.DEFAULT_REVISION`, `tt-model.yaml` `weights.revision`, `serve.env`
   `TT_WEIGHTS_REVISION` and the `verify:` line.
3. Check: `huggingface_hub.snapshot_download("changh95/sceneseg-p150", revision=<sha>, allow_patterns=[{patterns}])`
   gives files with the sha256 below.

`tt-model push` deletes only stale `code/**` and `image/**` files of the repo (tools/tt-model-manager
`src/tt_kernel/hub.py:164-190`), so `weights/` survives the bundle pushes (confirm on the first push: PLAN.md D10).

| file | bytes | sha256 |
|---|---:|---|
{rows}
"""


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("--out", type=Path, default=BUNDLE / "weights_upload")
    a = ap.parse_args()
    import onnx
    import safetensors
    import torch
    from safetensors.torch import load_file, save_file

    torch.set_num_threads(min(4, os.cpu_count() or 1))
    for path, want in ((ONNX_SRC, config.ONNX_SHA256), (PTH_SRC, config.PTH_SHA256), (LICENSE_SRC, LICENSE_SHA256)):
        got = sha256(path)
        if got != want:
            raise SystemExit(f"{path}: sha256 {got} != {want}")
    out = a.out / config.WEIGHTS_SUBDIR
    out.mkdir(parents=True, exist_ok=True)

    # 1. the deployed ONNX, byte for byte
    onnx_dst = out / config.ONNX_FILE
    shutil.copyfile(ONNX_SRC, onnx_dst)
    model = onnx.load(str(onnx_dst), load_external_data=False)
    n_init = len(model.graph.initializer)

    # 2. the .pth state_dict -> safetensors (tensors bit-identical)
    sd = torch.load(str(PTH_SRC), map_location="cpu", weights_only=True)
    tensors = {k: v.detach().contiguous().clone() for k, v in sd.items()}
    meta = {"source": f"{config.PTH_FILE} (Google Drive {config.PTH_DRIVE_ID})", "source_sha256": config.PTH_SHA256,
            "license": "Apache-2.0", "converted_by": "changh95/sceneseg-p150 code/scripts/make_weights_upload.py"}
    st_dst = out / config.SAFETENSORS_FILE
    # ONE metadata key: safetensors keeps the metadata in a hash map, so several keys would be written in a per-process
    # order and the file's sha256 would change from one conversion to the next (the tensor data never does)
    save_file(tensors, str(st_dst), metadata={"provenance": json.dumps(meta, sort_keys=True)})
    st_dst.chmod(0o644)
    back = load_file(str(st_dst))
    if set(back) != set(tensors) or any(not torch.equal(back[k], tensors[k]) or back[k].dtype != tensors[k].dtype
                                        for k in tensors):
        raise SystemExit("safetensors round trip is not bit-identical")
    n_elem = sum(int(v.numel()) for v in tensors.values())
    n_trainable = sum(int(v.numel()) for k, v in tensors.items()
                      if not k.endswith(("running_mean", "running_var", "num_batches_tracked")))

    # 3. the BN fold of the export, re-done from the state_dict (fp64, one rounding) against the ONNX initializers
    from tt_sceneseg.reference.weights import load_state_dict_weights, load_weights

    w_onnx, w_sd = load_weights(onnx_dst), load_state_dict_weights(st_dst)
    import numpy as np

    rel = max(float(np.abs(w_onnx[k].weight - w_sd[k].weight).max() / max(np.abs(w_sd[k].weight).max(), 1e-30))
              for k in w_onnx)
    exact = sum(np.array_equal(w_onnx[k].weight, w_sd[k].weight) and np.array_equal(w_onnx[k].bias, w_sd[k].bias)
                for k in w_onnx if w_onnx[k].spec.bn is None)
    unfolded = sum(1 for k in w_onnx if w_onnx[k].spec.bn is None)

    # 4. licence, notice
    shutil.copyfile(LICENSE_SRC, out / "LICENSE")
    (out / "NOTICE").write_text(NOTICE.format(head=UPSTREAM_HEAD, code_sha=config.POV_CODE_SHA, onnx=config.ONNX_FILE,
                                              onnx_id=config.ONNX_DRIVE_ID, onnx_sha=config.ONNX_SHA256,
                                              pth=config.PTH_FILE, pth_id=config.PTH_DRIVE_ID,
                                              pth_sha=config.PTH_SHA256, st=config.SAFETENSORS_FILE))

    # 5. provenance
    files = {p.name: {"bytes": p.stat().st_size, "sha256": sha256(p)}
             for p in (onnx_dst, st_dst, out / "LICENSE", out / "NOTICE")}
    prov = {
        "model": "SceneSeg (Autoware VisionPilot): 3-class scene segmentation, EfficientNet-B0 + SceneContext + "
                 "UNet-style decoder, input 1x3x320x640, output 1x3x320x640 logits (0 background, 1 foreground, "
                 "2 road)",
        "rehosted_in": {"repo": config.WEIGHTS_REPO, "path": f"{config.WEIGHTS_SUBDIR}/",
                        "decision": "research/PLAN.md D10 (user decision 2026-10-07): weights inside the bundle repo, "
                                    "no separate mirror"},
        "license": {"spdx": "Apache-2.0",
                    "sources": [f"{config.POV_REPO} LICENSE (sha256 {LICENSE_SHA256}) @ {UPSTREAM_HEAD}",
                                f"{config.POV_REPO} README.md:41-42 @ {UPSTREAM_HEAD}: 'the complete codebase ... "
                                "including AI model weights ... Apache 2.0 licence'",
                                f"{config.UPSTREAM_CARD} model card front matter (license: apache-2.0) @ "
                                f"{config.UPSTREAM_CARD_REVISION}"]},
        "originals": [
            {"file": config.ONNX_FILE, "bytes": config.ONNX_BYTES, "sha256": config.ONNX_SHA256,
             "google_drive_id": config.ONNX_DRIVE_ID, "url": drive_url(config.ONNX_DRIVE_ID),
             "http_last_modified": "2025-04-17",
             "linked_from": [f"{config.POV_REPO}@{config.POV_CODE_SHA}:Models/model_library/SceneSeg/README.md:20",
                             f"{config.POV_REPO}@{config.POV_DEPLOY_SHA}:VisionPilot/middleware_recipes/README.md:"
                             "107-109 (gdown)"],
             "role": "the deployed graph (autoseg.yaml model_path; TensorRT FP16 engine built from it); BatchNorm "
                     "folded at export (eps 1e-5), opset 14, IR 7, producer pytorch 2.4.1"},
            {"file": config.PTH_FILE, "bytes": config.PTH_BYTES, "sha256": config.PTH_SHA256,
             "google_drive_id": config.PTH_DRIVE_ID, "url": drive_url(config.PTH_DRIVE_ID),
             "http_last_modified": "2024-10-20",
             "linked_from": [f"{config.POV_REPO}@{config.POV_CODE_SHA}:Models/model_library/SceneSeg/README.md:18"],
             "role": "the training checkpoint the ONNX was exported from (a plain state_dict, torchvision "
                     "efficientnet_b0 key layout)"},
        ],
        "files": files,
        "conversion": {
            "safetensors": {"from": config.PTH_FILE, "method": "torch.load(weights_only=True) -> "
                                                               "safetensors.torch.save_file (tensors unchanged)",
                            "tensors": len(tensors), "elements": n_elem, "trainable_elements": n_trainable,
                            "round_trip_bit_identical": True, "metadata": meta},
            "onnx": {"from": config.ONNX_FILE, "method": "byte copy", "initializers": n_init},
            "tools": {"python": sys.version.split()[0], "torch": torch.__version__,
                      "safetensors": safetensors.__version__, "onnx": onnx.__version__,
                      "script": "code/scripts/make_weights_upload.py"},
            "date": dt.datetime.now(dt.timezone.utc).isoformat(timespec="seconds"),
        },
        "checks": {
            "onnx_vs_state_dict_bn_folded_fp64": {"layers": len(w_onnx), "max_relative_weight_difference": rel,
                                                  "unfolded_layers_bit_identical": f"{exact}/{unfolded}"},
            "reference_vs_onnxruntime": "research/sceneseg/SPEC.md 6.2: torch reference on the .pth vs ORT on the "
                                        "ONNX, "
                                        "every tap PCC >= 0.9999999986, logits max |diff| 2.3e-4",
        },
    }
    (out / "provenance.json").write_text(json.dumps(prov, indent=1) + "\n")
    files["provenance.json"] = {"bytes": (out / "provenance.json").stat().st_size,
                                "sha256": sha256(out / "provenance.json")}
    what = {config.ONNX_FILE: "the deployed graph (VisionPilot `autoseg.yaml`), unchanged",
            config.SAFETENSORS_FILE: f"`{config.PTH_FILE}` (the training state_dict) as safetensors, bit-identical",
            "LICENSE": "Apache-2.0 (autowarefoundation/vision_pilot)", "NOTICE": "attribution and changes",
            "provenance.json": "sources, sha256, conversion, checks"}
    rows = "\n".join(f"| `{n}` | {f['bytes']:,} | `{f['sha256']}` | {what[n]} |" for n, f in files.items())
    (out / "README.md").write_text(README.format(rows=rows))
    up_rows = "\n".join(f"| `weights/{n}` | {f['bytes']:,} | `{f['sha256']}` |" for n, f in files.items())
    readme = out / "README.md"
    up_rows += f"\n| `weights/README.md` | {readme.stat().st_size:,} | `{sha256(readme)}` |"
    (a.out / "UPLOAD.md").write_text(UPLOAD.format(rows=up_rows,
                                                   patterns=", ".join(f'"{p}"' for p in config.ALLOW_PATTERNS)))
    print(json.dumps({"out": str(a.out), "files": files, "tensors": len(tensors), "onnx_vs_pth_max_rel": rel,
                      "unfolded_bit_identical": f"{exact}/{unfolded}"}, indent=1))
    return 0


if __name__ == "__main__":
    sys.exit(main())