| """Measure the existing VAE chunk-size knob independently of transformer changes.""" |
| import json,time |
| from pathlib import Path |
| import numpy as np |
| import torch |
| import soundfile as sf |
| from yue2.modeling_vae import YuE2VAE |
| from benchmark import Monitor |
|
|
| @torch.inference_mode() |
| def main(): |
| |
| torch.backends.cudnn.benchmark=False |
| torch.backends.cudnn.deterministic=True |
| torch.backends.cuda.matmul.allow_tf32=False |
| torch.backends.cudnn.allow_tf32=False |
| torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction=False |
| torch.set_float32_matmul_precision('highest') |
| model=YuE2VAE.from_pretrained('models/YuE2-Vae',decoder_only=True,device='cuda',local_files_only=True) |
| root=Path('benchmark/ru-final-w4a4-full/run-01') |
| z=torch.from_numpy(np.load(root/'latent.npy')).T.unsqueeze(0) |
| reference,sr=sf.read(root/'audio.wav',dtype='float32') |
| results=[] |
| for frames in (1024,512,256): |
| for repeat in range(3): |
| torch.cuda.empty_cache() |
| with Monitor() as monitor: |
| torch.cuda.synchronize();start=time.perf_counter() |
| audio=model.decode_tiled(z,core_frames=frames,halo_frames=16,output_device='cpu')[0].float().clamp(-1,1).T.contiguous().numpy() |
| torch.cuda.synchronize();elapsed=time.perf_counter()-start |
| diff=audio-reference |
| row=dict(pipeline_numerics=True,core_frames=frames,repeat=repeat,seconds=elapsed,nvml_peak_bytes=monitor.peak,torch_peak_allocated_bytes=torch.cuda.max_memory_allocated(),max_abs_difference=float(np.abs(diff).max()),relative_l2=float(np.linalg.norm(diff)/(np.linalg.norm(reference)+1e-12)),waveform_correlation=float(np.corrcoef(audio.ravel(),reference.ravel())[0,1])) |
| results.append(row);print(json.dumps(row),flush=True) |
| Path('reports/vae-chunk-benchmark.json').write_text(json.dumps(results,indent=2)) |
|
|
| if __name__=='__main__': main() |
|
|
