| import hashlib |
| import json |
| import time |
| from pathlib import Path |
| import torch |
| import torch.nn.functional as F |
| from safetensors import safe_open |
| 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) |
|
|
|
|
| expected_hash = json.loads((root / 'tensor_sha256.json').read_text()) |
| with safe_open(root / 'model.safetensors', framework='pt') as f: |
| for k in f.keys(): |
| t = f.get_tensor(k) |
| assert hashlib.sha256(t.contiguous().view(torch.uint8).numpy().tobytes()).hexdigest() == expected_hash[k] |
| check('649_source_tensor_payloads', len(expected_hash) == 649) |
| m = load_model(root) |
| check('native_strict_load', len(m.state_dict()) == 649) |
| check('official_scale_schedule', m.scale_rq_layers == [1, 2, 4, 6, 8, 10, 12, 14, 16]) |
| with torch.inference_mode(): |
| x = torch.linspace(-1, 1, 3*256*256).reshape(1, 3, 256, 256) |
| q, losses, ids = m.encode(x) |
| check('multiscale_token_geometry', tuple(ids.shape) == (1, sum(s*s for s in m.scale_rq_layers)), shape=list(ids.shape)) |
| check('ID_bounds', ids.min().item() >= 0 and ids.max().item() < 32768) |
| canonical = 0. |
| cursor = 0 |
| for scale in m.scale_rq_layers: |
| n = scale*scale |
| vectors = m.quantize.get_codebook_entry(ids[:, cursor:cursor+n]) |
| vectors = vectors.reshape(1, scale, scale, m.code_dim).permute(0, 3, 1, 2) |
| canonical = canonical + F.interpolate(vectors, size=(16,16), mode='bicubic') |
| cursor += n |
| check('residual_lookup_bound', torch.allclose(q, canonical, rtol=2e-5, atol=2e-5), max_abs=(q-canonical).abs().max().item()) |
| semantic, rgb = m.decode(q) |
| check('semantic_output', tuple(semantic.shape) == (1, 256, 512) and torch.isfinite(semantic).all().item(), shape=list(semantic.shape)) |
| check('RGB_output', rgb.shape == x.shape and torch.isfinite(rgb).all().item(), shape=list(rgb.shape), mse=(rgb-x).square().mean().item()) |
| by_ids = m.decode_code(ids) |
| by_lookup = m.decode(canonical) |
| check('semantic_decode_code_lookup_exact', torch.equal(by_ids[0], by_lookup[0])) |
| check('RGB_decode_code_lookup_exact', torch.equal(by_ids[1], by_lookup[1])) |
| 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 256px native codec fixtures, not image quality benchmarks or complete VQA/T2I pipelines.', 'Includes the trained CLIP visual branch, semantic and pixel codebooks, semantic decoder and enhanced RGB decoder. No CLIP text model or external pretrained download required.', 'Native straight-through arithmetic can differ slightly from direct lookup; canonical lookup and decode-code paths are compared exactly.']} |
| (root / 'validation.json').write_text(json.dumps(report, indent=2)+'\n') |
| print('TOKENFLOW VALIDATED',len(checks),flush=True) |
|
|
