| """Full, pinned Socialtrait checkpoint with one on-disk copy of each LFS object. |
| |
| Install huggingface_hub and git-lfs, authenticate with `hf auth login`, then: |
| python clone_checkpoint.py /path/with/at-least-31-GiB-free/flux3-klein |
| Do not edit weight files in place: the checkout and LFS store use hardlinks. |
| """ |
| import argparse |
| import hashlib |
| import json |
| import os |
| from pathlib import Path |
| import shutil |
| import subprocess |
|
|
| from huggingface_hub import HfApi, get_token, hf_hub_download |
|
|
| from model_config import UPSTREAM as REPO, REVISION |
| OVERLAY = ('model_config.py', 'flux3_engine.py', 'handler.py', 'web.py', 'modal_app.py', |
| 'requirements.txt', 'prompts.json', 'README.md', 'DEPLOYMENT.txt', |
| 'verify_checkpoint.py', 'validate_deployment.py') |
|
|
|
|
| def digest(path): |
| sha = hashlib.sha256() |
| with path.open('rb') as source: |
| for chunk in iter(lambda: source.read(8 * 1024 * 1024), b''): |
| sha.update(chunk) |
| return sha.hexdigest() |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument('destination', type=Path) |
| parser.add_argument('--reserve-gib', type=float, default=5, |
| help='Free disk space to retain after download (default: 5 GiB).') |
| args = parser.parse_args() |
| destination = args.destination.expanduser().resolve() |
| token = get_token() |
| if not token: |
| raise SystemExit('Authenticate first with hf auth login; do not put a token in this command.') |
| info = HfApi(token=token).model_info(REPO, revision=REVISION, files_metadata=True) |
| files = list(info.siblings) |
| |
| identity = {'repo': REPO, 'revision': REVISION} |
| marker = destination / '.git' / 'full-checkpoint-identity.json' |
| if destination.exists() and any(destination.iterdir()): |
| if not marker.exists() or json.loads(marker.read_text()) != identity: |
| raise SystemExit('Destination is not a matching checkpoint clone. Choose an empty directory.') |
| needed = 0 |
| for item in files: |
| if item.lfs: |
| sha = item.lfs.sha256 |
| obj = destination / '.git/lfs/objects' / sha[:2] / sha[2:4] / sha |
| if not obj.exists() or obj.stat().st_size != item.size: |
| needed += item.size |
| parent = destination |
| while not parent.exists(): |
| parent = parent.parent |
| free = shutil.disk_usage(parent).free |
| reserve = max(args.reserve_gib, 0) * 2**30 |
| if free < needed + reserve: |
| raise SystemExit(f'Need {needed/2**30:.2f} GiB plus {reserve/2**30:.1f} GiB reserve; ' |
| f'only {free/2**30:.2f} GiB free. Select a larger disk.') |
| env = os.environ.copy() |
| env.update(GIT_LFS_SKIP_SMUDGE='1', GIT_TERMINAL_PROMPT='0', |
| GIT_CONFIG_COUNT='1', GIT_CONFIG_KEY_0='http.https://huggingface.co/.extraheader', |
| GIT_CONFIG_VALUE_0=f'Authorization: Bearer {token}') |
| if not marker.exists(): |
| destination.parent.mkdir(parents=True, exist_ok=True) |
| subprocess.run(['git', 'clone', '--no-checkout', f'https://huggingface.co/{REPO}', |
| str(destination)], env=env, check=True) |
| subprocess.run(['git', '-C', str(destination), 'checkout', REVISION], env=env, check=True) |
| marker.write_text(json.dumps(identity)) |
| staging = destination / '.git' / 'hf-download' |
| inventory = [] |
| for index, item in enumerate(files, 1): |
| target = destination / item.rfilename |
| if item.lfs: |
| sha = item.lfs.sha256 |
| obj = destination / '.git/lfs/objects' / sha[:2] / sha[2:4] / sha |
| obj.parent.mkdir(parents=True, exist_ok=True) |
| if not obj.exists(): |
| download = Path(hf_hub_download(REPO, item.rfilename, revision=REVISION, |
| token=token, local_dir=staging)) |
| if digest(download) != sha: |
| raise RuntimeError(f'Checksum mismatch: {item.rfilename}') |
| os.replace(download, obj) |
| elif obj.stat().st_size != item.size or digest(obj) != sha: |
| raise RuntimeError(f'Existing LFS object failed verification: {item.rfilename}') |
| target.parent.mkdir(parents=True, exist_ok=True) |
| if target.exists(): |
| target.unlink() |
| os.link(obj, target) |
| if not target.is_file() or target.stat().st_size != item.size: |
| raise RuntimeError(f'Missing or incomplete file: {item.rfilename}') |
| inventory.append({'path': item.rfilename, 'bytes': item.size, |
| 'sha256': item.lfs.sha256 if item.lfs else None}) |
| print(f'[{index}/{len(files)}] Verified {item.rfilename}', flush=True) |
| (destination / 'checkpoint_inventory.json').write_text(json.dumps( |
| {**identity, 'files': inventory}, indent=2) + '\n') |
| source = Path(__file__).resolve().parent |
| for name in OVERLAY: |
| shutil.copy2(source / name, destination / name) |
| with (destination / '.git/info/exclude').open('a') as exclude: |
| exclude.write('\n.env*\n__pycache__/\nbenchmark_outputs/\n') |
| print(f'Complete checkpoint and deployment files: {destination}') |
| print('Socialtrait BF16/Q8 weights, tokenizers, configurations, licenses and notices retained.') |
| print('No repository has been uploaded or made public.') |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
