| 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') |
|
|
