| """WST-Graph TorchScript loader + preprocessing (HF kwokho1/wst-graph @ 080000c1685451c7759f43d2ae7d0f08477cfdf8, MIT). |
| |
| There is no Python model code to vendor: the authors release only traced TorchScript programs (waveform -> logits, |
| Kymatio WST front-end + GraphAASIST back-end, all parameters inside). The program was reviewed before first use |
| (read from its zip `code/` dir and from `model.code`/`model.graph` of every submodule): only aten:: / prim:: nodes, |
| classes limited to the model's own modules (__torch__.models.*, kymatio ScatteringTorch1D, torch.nn modules), pickled |
| globals only torch storages / _rebuild_tensor_v2 / OrderedDict. `load_ts` re-checks the node kinds on every load. |
| |
| Preprocessing = the model card's quick start: first 64000 samples, right zero-pad to 64000, input_length = retained |
| sample count (int64). Class 0 = bona fide, 1 = spoof. The `_cuda` file was traced on cuda:0 (device constants are |
| baked in), so it runs on CUDA only. |
| """ |
| import os |
|
|
| import numpy as np |
| import torch |
|
|
| TARGET_SAMPLES = 64000 |
|
|
|
|
| def prep_np(wav, max_len=TARGET_SAMPLES): |
| """-> (float32 [max_len] right zero-padded, input_length). Empty clip -> zeros, length 1 (the card raises).""" |
| wav = np.asarray(wav, dtype=np.float32).reshape(-1) |
| n = min(wav.shape[0], max_len) |
| out = np.zeros(max_len, dtype=np.float32) |
| out[:n] = wav[:n] |
| return out, max(n, 1) |
|
|
|
|
| def check_graph(module): |
| """Assert the TorchScript program only uses aten::/prim:: ops (no custom/IO ops).""" |
| bad = set() |
|
|
| def walk(block): |
| for node in block.nodes(): |
| k = node.kind() |
| if not k.startswith(("aten::", "prim::")): |
| bad.add(k) |
| for b in node.blocks(): |
| walk(b) |
|
|
| for _, sub in module.named_modules(): |
| try: |
| g = sub.graph |
| except (AttributeError, RuntimeError): |
| continue |
| walk(g) |
| assert not bad, f"unexpected TorchScript ops: {sorted(bad)}" |
|
|
|
|
| def load_ts(path, device="cuda"): |
| m = torch.jit.load(path, map_location=device).eval() |
| check_graph(m) |
| return m |
|
|
