| """Source integrity and native image-codebook-decode tests; no language model required.""" |
| import gc |
| import hashlib |
| import json |
| import time |
| from pathlib import Path |
| import torch |
| from safetensors.torch import load_file |
| from component import load_model |
|
|
| root = Path(__file__).parent |
| torch.set_num_threads(4) |
| checks = [] |
| started = time.monotonic() |
|
|
|
|
| def check(name, ok, **details): |
| assert ok, (name, details) |
| checks.append({'name': name, 'passed': True, **details}) |
| print(name, details, flush=True) |
|
|
|
|
| with (root / 'model.safetensors').open('rb') as f: |
| file_sha = hashlib.file_digest(f, 'sha256').hexdigest() |
| check('official_LFS_SHA256', file_sha == 'a808dc617e489c37be7f51a63a0e7cb77a97a54ceb59f674255d2e7bc7b2c080') |
| m = load_model(root) |
| source = load_file(root / 'model.safetensors') |
| state = m.state_dict() |
| check('native_strict_543_entries', set(state) == set(source) and len(state) == 543) |
| check('all_original_tensor_values', all(torch.equal(source[k], state[k]) for k in source)) |
| del source, state |
| with torch.inference_mode(): |
| for height, width in [(256, 256), (128, 192)]: |
| label = f'{height}x{width}' |
| x = torch.linspace(-1, 1, 3 * height * width).reshape(1, 3, height, width) |
| q, _, info = m.encode(x) |
| ids = info[-1] |
| check(label + '_encode_geometry', tuple(q.shape) == (1, 256, height // 16, width // 16), shape=list(q.shape)) |
| check(label + '_ID_range', ids.numel() == height * width // 256 and ids.min().item() >= 0 and ids.max().item() < 131072) |
| lookup = m.quantize.get_codebook_entry(ids, shape=(1, height // 16, width // 16)) |
| check(label + '_lookup_exact', torch.equal(q, lookup)) |
| h = m.quant_conv(m.encoder(x)) |
| expected = (h.permute(0, 2, 3, 1).reshape(-1, 256) @ m.quantize.embedding.weight.T).argmax(-1) |
| check(label + '_max_dot_assignment', torch.equal(ids, expected)) |
| decoded = m.decode(q) |
| by_code = m.decode_code(ids, shape=(1, height // 16, width // 16)) |
| check(label + '_decode_code_exact', torch.equal(decoded, by_code)) |
| check(label + '_reconstruction', decoded.shape == x.shape and torch.isfinite(decoded).all().item(), mse=(decoded - x).square().mean().item()) |
| del q, lookup, h, decoded, by_code |
| gc.collect() |
| report = {'strict_load': True, 'source_tensor_bytes_equal': True, 'checks': checks, 'smoke': {'passed': True, 'device': 'cpu', 'dtype': 'float32'}, 'elapsed_seconds': time.monotonic() - started, 'limitations': ['Synthetic codec fixtures only; no visual quality benchmark or full Emu3.5 language/generation backbone evaluation.', 'This is a 2D image VQ codec. No temporal video codec behavior is asserted.']} |
| (root / 'validation.json').write_text(json.dumps(report, indent=2) + '\n') |
| print('EMU3.5 VALIDATED', len(checks), flush=True) |
|
|
