File size: 4,126 Bytes
d2a1a41 | 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 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | """CPU-only release validation, including relocatable frozen codec round trips."""
from pathlib import Path
import argparse
import gc
import hashlib
import json
import os
import sys
ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(ROOT))
def sha(path):
value = hashlib.sha256()
with Path(path).open('rb') as stream:
for block in iter(lambda: stream.read(1024 * 1024), b''):
value.update(block)
return value.hexdigest()
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--codecs', action='store_true', help='Also load and test both pinned frozen codecs on CPU')
args = parser.parse_args()
# This command is explicitly CPU-only, including optional backend probes.
os.environ['CUDA_VISIBLE_DEVICES'] = ''
os.environ.setdefault('TORCHINDUCTOR_TRITON_DISABLE_DEVICE_DETECTION', '1')
# Full-precision codecs do not use TorchAO. Some optional TorchAO builds
# query CUDA during import even when devices are hidden; exclude it here.
sys.modules['torchao'] = None
import numpy as np
import torch
import quasar
torch.set_num_threads(2)
assert Path(quasar.__file__).resolve().is_relative_to(ROOT)
release = json.loads((ROOT / 'release.json').read_text())
manifest = json.loads((ROOT / 'files.sha256.json').read_text())
for name, item in manifest.items():
file = ROOT / name
assert file.is_file() and file.stat().st_size == item['bytes'] and sha(file) == item['sha256'], name
from quasar.training.checkpoint import load_quasar
from quasar.tokenizers.text import TextTokenizer
for folder, expected in release['weight_hashes'].items():
path = ROOT / folder
assert sha(path / 'model.safetensors') == expected
model, vocab, meta = load_quasar(path, 'cpu')
assert model.num_params() == 90247794
assert all(bool(torch.isfinite(p).all()) for p in model.parameters())
model.eval()
tokenizer = TextTokenizer.load(path / 'tokenizers/text_tokenizer.json')
tokens = [vocab.special('bos'), vocab.special('task:chat_text'), vocab.special('input'), vocab.special('text')]
tokens += tokenizer.encode('Hello.') + [vocab.special('/text'), vocab.special('output'), vocab.special('text')]
with torch.inference_mode():
logits = model(torch.tensor([tokens]))
assert torch.isfinite(logits).all() and logits.shape[-1] == 54144
audio = meta['data']['preparation']['audio_codec']
assert (path / audio['path'] / 'model.safetensors').resolve() == (ROOT / 'codecs/audio/model.safetensors').resolve()
vq = json.loads((path / 'tokenizers/vq/vq_config.json').read_text())
assert (path / 'tokenizers/vq' / vq['base_path']).resolve() == (ROOT / 'codecs/vision').resolve()
print(folder, 'finite CPU forward / original weight hash / parameter count / codec paths passed', flush=True)
del model, logits
gc.collect()
if args.codecs:
from quasar.inference.extended import ExtendedGenerator
generator = ExtendedGenerator(ROOT, 'cpu')
pixels = np.zeros((1, 64, 64, 3), dtype=np.uint8)
pixels[:, 20:40, 20:40, 1] = 255
codes = generator.vision().encode(pixels)
reconstructed = generator.vision().decode(codes)
assert codes.shape == (1, 8, 8) and reconstructed.shape == pixels.shape
assert np.isfinite(reconstructed).all()
print('Pinned visual codec CPU encode/decode passed', flush=True)
del generator.codec
generator.codec = None
gc.collect()
codec = generator.audio()
waveform = (.05 * np.sin(np.arange(2880) * 2 * np.pi * 440 / 24000)).astype(np.float32)
packet = codec.encode(waveform, 24000)
reconstructed = codec.decode(packet)
assert reconstructed.shape == waveform.shape and np.isfinite(reconstructed).all()
print('Pinned aligned audio codec CPU encode/decode passed', flush=True)
print('Release verified on CPU; no optimizer or GPU job started.', flush=True)
if __name__ == '__main__':
main()
|