File size: 2,640 Bytes
d651476 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | """Validate the exact exported student and invoke its pinned four-call runner."""
import argparse
import hashlib
import json
import os
from pathlib import Path
import subprocess
import sys
HERE = Path(__file__).resolve().parent
def sha(path):
digest = hashlib.sha256()
with path.open('rb') as stream:
for block in iter(lambda: stream.read(16*1024*1024), b''):
digest.update(block)
return digest.hexdigest()
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--model-dir', '--model-id', default='MiniMaxAI/MiniMax-H3')
parser.add_argument('--diffusers-src', type=Path)
parser.add_argument('--prompts', type=Path, default=HERE/'prompts.example.json')
parser.add_argument('--expected-prompt-count', type=int, default=4)
parser.add_argument('--output-dir', type=Path, default=Path('outputs'))
args, rest = parser.parse_known_args()
for argument in rest:
if argument.split('=')[0] in ('--checkpoint', '--revision', '--num-inference-steps'):
parser.error('Checkpoint, revision and four-call sampler are fixed by this release')
provenance = json.loads((HERE/'provenance.json').read_text())
checkpoint = HERE/provenance['export']['file']
if checkpoint.is_symlink() or sha(checkpoint) != provenance['export']['sha256']:
raise RuntimeError('Student checkpoint checksum mismatch')
for name, digest in json.loads((HERE/'source_hashes.json').read_text()).items():
if sha(HERE/name) != digest:
raise RuntimeError('Source snapshot checksum mismatch: '+name)
paths = [str(HERE/'source_snapshot')]
if args.diffusers_src is not None:
if not (args.diffusers_src/'diffusers').is_dir():
parser.error('--diffusers-src must contain the diffusers package')
paths.append(str(args.diffusers_src.resolve()))
env = os.environ.copy()
env['PYTHONPATH'] = os.pathsep.join(paths + ([env['PYTHONPATH']] if env.get('PYTHONPATH') else []))
env.update(PYTHONDONTWRITEBYTECODE='1', PYTHONUNBUFFERED='1', TOKENIZERS_PARALLELISM='false')
env.setdefault('PYTORCH_CUDA_ALLOC_CONF', 'expandable_segments:True')
return subprocess.run([sys.executable, '-m', 'minimax_h3.infer_student_4step',
'--checkpoint', str(checkpoint), '--model-id', args.model_dir,
'--revision', provenance['inference']['base_revision'],
'--prompts', str(args.prompts), '--expected-prompt-count', str(args.expected_prompt_count),
'--output-dir', str(args.output_dir), *rest], env=env).returncode
if __name__ == '__main__':
raise SystemExit(main())
|