File size: 2,812 Bytes
9711758 | 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 | """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)
|