"""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():
    # Match the pinned pipeline's numerical settings, including convolution TF32.
    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()