File size: 3,067 Bytes
b9a3411 | 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 | 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)
|