File size: 1,518 Bytes
b1882a9 | 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 | """Strict phase-policy checkpoint loading, retaining the frozen older factories."""
import io
from pathlib import Path
import torch
from tinyasr.codec_factory_v2 import load_codec_checkpoint as load_previous
from tinyasr.codec_phase_residual_v1 import ARCHITECTURE, POLICIES, PhaseResidualCodec
def load_codec_checkpoint(path):
# A fresh stream is needed by older factories that reopen their input.
if isinstance(path, (str, Path)):
with Path(path).open('rb') as source:
raw = source.read(100_000_001)
else:
raw = path.read(100_000_001)
if not isinstance(raw, bytes) or not 0 < len(raw) <= 100_000_000:
raise ValueError('Invalid checkpoint bytes')
state = torch.load(io.BytesIO(raw), map_location='cpu', weights_only=False)
if state.get('architecture') != ARCHITECTURE:
# Factory v2 reopens fallback checkpoints; v1 must receive a fresh stream.
from tinyasr.codec_projected_warmup_v1 import ARCHITECTURE as PROJECTED
from tinyasr.codec_factory_v1 import load_codec_checkpoint as load_v1
return (load_previous if state.get('architecture') == PROJECTED else load_v1)(io.BytesIO(raw))
policy = state.get('phase_policy')
if policy not in POLICIES or state.get('config', {}).get('phase_policy', policy) != policy:
raise ValueError('Missing or inconsistent phase policy')
model = PhaseResidualCodec(policy)
model.load_state_dict(state['codec'], strict=True)
return model.eval().requires_grad_(False)
|