| """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()) |
|
|
