scripts/prepare_model_repo.py

"""Create a NEW, verified model upload directory. Never upload or modify source."""
import argparse
import hashlib
import json
import shutil
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from grounding_demo.native_runtime import inspect_contract

METADATA = [
    "config.json", "generation_config.json", "preprocessor_config.json",
    "tokenizer.json", "tokenizer_config.json", "chat_template.jinja",
    "configuration_oxygen_vl_v3.py", "configuration_kimi_k3.py",
    "modeling_oxygen_vl_v3.py", "modeling_kimi_k3.py", "processing_oxygen_vl_v3.py",
    "kimi_k3_vision_processing.py", "media_utils.py",
]
OPTIONAL = ["model.safetensors.index.json", "added_tokens.json", "special_tokens_map.json",
            "vocab.json", "merges.txt", "streammind_gate.py"]


def sha256(path):
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for block in iter(lambda: stream.read(8 * 1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("source", type=Path)
    p.add_argument("output", type=Path)
    args = p.parse_args()
    source = args.source.resolve(strict=True)
    output = args.output.resolve()
    if output.exists() or output.is_relative_to(source):
        raise SystemExit("Choose a NEW output directory outside the original model")
    contract = inspect_contract(source)
    names = METADATA + [n for n in OPTIONAL if (source / n).is_file()] + contract["weight_shards"]
    output.parent.mkdir(parents=True, exist_ok=True)
    required = sum((source / n).stat().st_size for n in names)
    if shutil.disk_usage(output.parent).free < required + 1024 ** 3:
        raise SystemExit("Insufficient free disk for independent model copy")
    output.mkdir()
    manifest = {}
    for name in names:
        original, copied = source / name, output / name
        before = original.stat()
        source_hash = sha256(original)
        shutil.copy2(original, copied)
        after = original.stat()
        if (before.st_size, before.st_mtime_ns) != (after.st_size, after.st_mtime_ns):
            raise RuntimeError(f"Source changed during copy: {name}")
        if sha256(copied) != source_hash:
            raise RuntimeError(f"Copy checksum failed: {name}")
        manifest[name] = {"bytes": copied.stat().st_size, "sha256": source_hash}
        print(f"Verified {name} {source_hash}", flush=True)
    inspect_contract(output)
    (output / "model_manifest.json").write_text(json.dumps(manifest, indent=2) + "\n")
    # No license is invented. The owner must add the approved model card/license.
    (output / "README.md").write_text(
        "# GroundingPI model upload preparation\n\n"
        "Oxygen-VL-v3 with Kimi-K3 vision backbone and GAM tokens.\n\n"
        "This is an internal preparation, not permission to publish. Before upload, "
        "the owner must supply the approved model card, license, training-data "
        "disclosure and visibility. Custom code requires review and a pinned commit.\n\n"
        "Use Transformers 5.3.0, the bundled processor, BF16 and SDPA. "
        "The grounding demo does not enable StreamMind or require mamba-ssm.\n"
    )
    print(json.dumps({"status": "PASS", "files": len(names), "bytes": required,
                      "output": str(output), "uploaded": False}))


if __name__ == "__main__":
    main()