File size: 4,019 Bytes
08620f3 | 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 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | import os, sys, json, time, hashlib
from pathlib import Path
os.environ['HF_HUB_OFFLINE']='1'
os.environ['TRANSFORMERS_OFFLINE']='1'
os.environ['TOKENIZERS_PARALLELISM']='false'
import numpy as np
import mlx.core as mx
from scipy.io import wavfile
from scipy.signal import stft
from mlx_audio.tts.utils import load_model
ROOT=Path(os.environ.get('RESULTS_DIR', '.'))
BASE=Path(os.environ['MODEL_ROOT'])
variant=sys.argv[1]
paths={
'original':Path(os.environ['ORIGINAL_MODEL']),
'affine6':Path(os.environ['AFFINE6_MODEL']),
}
path=paths.get(variant,BASE/variant)
mx.set_memory_limit(8*2**30);mx.set_cache_limit(128*2**20)
model=load_model(str(path));model.eval();mx.eval(model.parameters());mx.synchronize()
assert model.tokenizer is not None and model._conds is not None
texts=[
'This is the oMLX Chatterbox Turbo checkpoint. It was quantized from the original fp16 model.',
'Hello mate, the training checkpoint is saved, and this is Chatterbox speaking locally on the Mac.',
'Please check the appointment on Thursday at three thirty. Bring the blue folder, your keys, and a bottle of water.',
]
def log_mel(audio, sr):
_, _, z=stft(audio,fs=sr,nperseg=1024,noverlap=768,boundary='zeros')
power=np.abs(z)**2
mel=lambda hz:2595*np.log10(1+hz/700)
hz=lambda m:700*(10**(m/2595)-1)
edges=hz(np.linspace(mel(40),mel(12000),82))
freqs=np.linspace(0,sr/2,power.shape[0]);bank=np.zeros((80,power.shape[0]))
for j in range(80):
bank[j]=np.maximum(0,np.minimum((freqs-edges[j])/(edges[j+1]-edges[j]),(edges[j+2]-freqs)/(edges[j+2]-edges[j+1])))
return bank@power
rows=[]
for i,text in enumerate(texts):
tokens_file=ROOT/f'original-speech-tokens-{i}.npy'
if variant=='original':
capture={};original=model.s3gen.inference
def traced(*args,**kwargs):
capture['tokens']=np.array(kwargs['speech_tokens'])
return original(*args,**kwargs)
model.s3gen.inference=traced
mx.random.seed(20261009);np.random.seed(20261009)
for result in model.generate(text,temperature=.8,top_p=.95,top_k=1000,repetition_penalty=1.2,split_pattern=None,max_tokens=400):mx.eval(result.audio)
model.s3gen.inference=original
np.save(tokens_file,capture['tokens'])
tokens=np.load(tokens_file)
mx.random.seed(9182);mx.synchronize();start=time.perf_counter()
audio,_=model.s3gen.inference(speech_tokens=mx.array(tokens),ref_dict=model._conds.gen,n_cfm_timesteps=2)
mx.eval(audio);mx.synchronize();elapsed=time.perf_counter()-start
audio=np.array(audio).reshape(-1).astype(np.float32)
assert audio.size and np.isfinite(audio).all()
np.save(ROOT/f'{variant}-decoder-{i}.npy',audio)
wavfile.write(ROOT/f'{variant}-decoder-{i}.wav',24000,(np.clip(audio,-1,1)*32767).astype(np.int16))
reference=np.load(ROOT/f'original-decoder-{i}.npy')
assert audio.shape==reference.shape,(audio.shape,reference.shape)
a=log_mel(reference,24000);b=log_mel(audio,24000)
floor=max(float(a.max())*1e-6,1e-12)
active=(a>=max(float(a.max())*1e-4,1e-12))
delta=10*np.log10(np.maximum(a,floor))-10*np.log10(np.maximum(b,floor))
rms_delta=float(np.sqrt(np.mean(delta[active]**2)))
row={'text':text,'speech_tokens':int(tokens.size),'audio_seconds':audio.size/24000,'decoder_seconds':elapsed,'log_mel_rms_difference_db':rms_delta,'peak_amplitude':float(np.abs(audio).max()),'clipped_fraction':float(np.mean(np.abs(audio)>=1))}
rows.append(row);print('DECODER',variant,json.dumps(row),flush=True)
report={'variant':variant,'method':'Identical original speech tokens, shared bundled conditioning, identical decoder RNG seed. RMS log-mel difference over reference bins above -40 dB, floor -60 dB. Lower means spectral output closer to original; not a subjective quality score.','decoder_seed':9182,'clips':rows,'mean_log_mel_rms_difference_db':float(np.mean([r['log_mel_rms_difference_db'] for r in rows]))}
(ROOT/f'{variant}-decoder-fidelity.json').write_text(json.dumps(report,indent=2)+'\n')
|