Skip to content

Instantly share code, notes, and snippets.

@ebezzam

ebezzam/dac_layer_by_layer_debugging.py

Last active July 23, 2025 16:29
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save ebezzam/bb315efa7a416db6336a6b2a2d424ffa to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save ebezzam/bb315efa7a416db6336a6b2a2d424ffa to your computer and use it in GitHub Desktop.
Download ZIP
DAC - Transformers Integration Tests
Raw
dac_layer_by_layer_debugging.py
"""
Debugging layer-by-layer differences in DAC model conversion from original to Hugging Face format.
--------------------------------------------------------------------------------------------------
Setup:
```
# after setting up virtual environment and transformers
uv pip install descript-audio-codec==1.0.0
```
Using DAC model, as per their documentation:
https://github.com/descriptinc/descript-audio-codec?tab=readme-ov-file#programmatic-usage
"""
import dac
from audiotools import AudioSignal
from datasets import load_dataset, Audio
import torch
import numpy as np
from transformers import AutoProcessor, DacModel
# model configuration based on sampling rate
model_config = {
16000: {
"model_name": "dac_16khz",
"dac_model_type": "16khz",
},
24000: {
"model_name": "dac_24khz",
"dac_model_type": "24khz",
},
44100: {
"model_name": "dac_44khz",
"dac_model_type": "44khz",
}
}
def normalize(arr):
norm = np.linalg.norm(arr)
normalized_arr = arr / norm
return normalized_arr
def compute_rmse(arr1, arr2):
arr1_np = arr1.cpu().numpy().squeeze()
arr2_np = arr2.cpu().numpy().squeeze()
max_length = min(arr1.shape[-1], arr2.shape[-1])
arr1_np = arr1_np[..., :max_length]
arr2_np = arr2_np[..., :max_length]
arr1_normalized = normalize(arr1_np)
arr2_normalized = normalize(arr2_np)
return np.sqrt(((arr1_normalized - arr2_normalized) ** 2).mean())
def debug_encoder_error_propagation():
"""Track how small weight differences amplify through all encoder layers."""
print("\n=== ENCODER ERROR PROPAGATION ANALYSIS ===")
with torch.no_grad():
# Use identical inputs for both models
hf_x = x_hf.clone()
orig_x = x_hf.clone()
errors = [] # Track errors at each layer
layer_names = []
print(f"{'Layer':<20} {'Max Error':<15} {'Mean Error':<15} {'Error Growth':<15}")
print("-" * 70)
# Initial input (should be identical)
max_err = torch.abs(hf_x - orig_x).max().item()
mean_err = torch.abs(hf_x - orig_x).mean().item()
errors.append(max_err)
layer_names.append("Input")
print(f"{'Input':<20} {max_err:<15.2e} {mean_err:<15.2e} {'1.0x':<15}")
# Layer 0: First Conv
hf_x = model_hf.encoder.conv1(hf_x)
orig_x = model.encoder.block[0](orig_x)
max_err = torch.abs(hf_x - orig_x).max().item()
mean_err = torch.abs(hf_x - orig_x).mean().item()
growth = max_err / errors[-1] if errors[-1] > 0 else float('inf')
errors.append(max_err)
layer_names.append("Conv1")
print(f"{'Conv1':<20} {max_err:<15.2e} {mean_err:<15.2e} {f'{growth:.1f}x':<15}")
# Encoder Blocks
for block_idx in range(len(model_hf.encoder.block)):
# Apply entire block
hf_x = model_hf.encoder.block[block_idx](hf_x)
orig_x = model.encoder.block[block_idx + 1](orig_x)
max_err = torch.abs(hf_x - orig_x).max().item()
mean_err = torch.abs(hf_x - orig_x).mean().item()
growth = max_err / errors[-1] if errors[-1] > 0 else float('inf')
errors.append(max_err)
layer_name = f"Block{block_idx}"
layer_names.append(layer_name)
print(f"{layer_name:<20} {max_err:<15.2e} {mean_err:<15.2e} {f'{growth:.1f}x':<15}")
# Detailed analysis within this block
debug_encoder_block_detailed(block_idx, errors[-2]) # Pass previous error for growth calc
# Final Snake
hf_x = model_hf.encoder.snake1(hf_x)
orig_x = model.encoder.block[5](orig_x)
max_err = torch.abs(hf_x - orig_x).max().item()
mean_err = torch.abs(hf_x - orig_x).mean().item()
growth = max_err / errors[-1] if errors[-1] > 0 else float('inf')
errors.append(max_err)
layer_names.append("Snake1")
print(f"{'Snake1':<20} {max_err:<15.2e} {mean_err:<15.2e} {f'{growth:.1f}x':<15}")
# Final Conv
hf_x = model_hf.encoder.conv2(hf_x)
orig_x = model.encoder.block[6](orig_x)
max_err = torch.abs(hf_x - orig_x).max().item()
mean_err = torch.abs(hf_x - orig_x).mean().item()
growth = max_err / errors[-1] if errors[-1] > 0 else float('inf')
errors.append(max_err)
layer_names.append("Conv2")
print(f"{'Conv2':<20} {max_err:<15.2e} {mean_err:<15.2e} {f'{growth:.1f}x':<15}")
# Summary statistics
print("\n=== ERROR PROPAGATION SUMMARY ===")
total_growth = errors[-1] / errors[1] if errors[1] > 0 else float('inf') # Skip input (always 0)
print(f"Initial weight error: {errors[1]:.2e}")
print(f"Final encoder error: {errors[-1]:.2e}")
print(f"Total error amplification: {total_growth:.0f}x")
# Identify biggest amplifiers
print(f"\nTop 3 error amplifiers:")
growths = []
for i in range(1, len(errors)):
if errors[i-1] > 0:
growth = errors[i] / errors[i-1]
growths.append((layer_names[i], growth))
growths.sort(key=lambda x: x[1], reverse=True)
for i, (layer, growth) in enumerate(growths[:3]):
print(f" {i+1}. {layer}: {growth:.1f}x amplification")
return errors, layer_names
def debug_encoder_block_detailed(block_idx, prev_error):
"""Detailed analysis within a specific encoder block."""
print(f"\n --- Block {block_idx} Internal Analysis ---")
with torch.no_grad():
# Reset to block input
if block_idx == 0:
# After conv1
hf_x = model_hf.encoder.conv1(x_hf.clone())
orig_x = model.encoder.block[0](x_hf.clone())
else:
# Need to run through previous layers to get to this block's input
hf_x = x_hf.clone()
orig_x = x_hf.clone()
# Apply conv1
hf_x = model_hf.encoder.conv1(hf_x)
orig_x = model.encoder.block[0](orig_x)
# Apply previous blocks
for i in range(block_idx):
hf_x = model_hf.encoder.block[i](hf_x)
orig_x = model.encoder.block[i + 1](orig_x)
block_input_error = torch.abs(hf_x - orig_x).max().item()
# Get the blocks
hf_block = model_hf.encoder.block[block_idx]
orig_block = model.encoder.block[block_idx + 1]
# Try to go through sublayers if possible
try:
# For HF block, try to access sublayers
if hasattr(hf_block, 'res_unit1'):
# Apply each residual unit
for res_idx in range(1, 4): # res_unit1, res_unit2, res_unit3
res_unit_name = f"res_unit{res_idx}"
if hasattr(hf_block, res_unit_name):
hf_res_unit = getattr(hf_block, res_unit_name)
# Apply HF residual unit
hf_x_res = hf_res_unit(hf_x)
# For original, try to access corresponding sublayer
if hasattr(orig_block, 'block') and len(orig_block.block) > res_idx - 1:
orig_res_unit = orig_block.block[res_idx - 1]
orig_x_res = orig_res_unit(orig_x)
res_error = torch.abs(hf_x_res - orig_x_res).max().item()
growth = res_error / block_input_error if block_input_error > 0 else float('inf')
print(f" {res_unit_name}: {res_error:.2e} ({growth:.1f}x)")
# Update for next iteration (residual connection)
hf_x = hf_x + hf_x_res # Residual connection
orig_x = orig_x + orig_x_res
else:
print(f" {res_unit_name}: Structure mismatch")
# Final layers in block
if hasattr(hf_block, 'snake1') and hasattr(hf_block, 'conv1'):
hf_x = hf_block.snake1(hf_x)
hf_x = hf_block.conv1(hf_x)
# For original block, try final layers
if hasattr(orig_block, 'block') and len(orig_block.block) >= 5:
orig_x = orig_block.block[3](orig_x) # Snake
orig_x = orig_block.block[4](orig_x) # Conv
final_error = torch.abs(hf_x - orig_x).max().item()
growth = final_error / block_input_error if block_input_error > 0 else float('inf')
print(f" Final layers: {final_error:.2e} ({growth:.1f}x)")
except Exception as e:
print(f" Detailed analysis failed: {e}")
def debug_weight_differences_by_layer():
"""Compare actual weights layer by layer to identify problematic conversions."""
print("\n=== WEIGHT DIFFERENCES BY LAYER ===")
# Compare conv1 weights
hf_conv1_w = model_hf.encoder.conv1.weight
orig_conv1_w = model.encoder.block[0].weight
conv1_diff = torch.abs(hf_conv1_w - orig_conv1_w).max().item()
print(f"Conv1 weight max diff: {conv1_diff:.2e}")
# Compare encoder blocks
for block_idx in range(len(model_hf.encoder.block)):
print(f"\nBlock {block_idx} weight differences:")
hf_block = model_hf.encoder.block[block_idx]
# Compare block conv1 weight
if hasattr(hf_block, 'conv1'):
hf_block_conv_w = hf_block.conv1.weight
# Try to find corresponding weight in original
try:
orig_block = model.encoder.block[block_idx + 1]
if hasattr(orig_block, 'block') and len(orig_block.block) > 4:
orig_block_conv_w = orig_block.block[4].weight
block_conv_diff = torch.abs(hf_block_conv_w - orig_block_conv_w).max().item()
print(f" Block conv weight diff: {block_conv_diff:.2e}")
except:
print(f" Block conv: Could not compare")
# Compare residual unit weights
for res_idx in range(1, 4):
res_unit_name = f"res_unit{res_idx}"
if hasattr(hf_block, res_unit_name):
hf_res_unit = getattr(hf_block, res_unit_name)
# Compare conv1 and conv2 in residual unit
for conv_name in ['conv1', 'conv2']:
if hasattr(hf_res_unit, conv_name):
hf_conv_w = getattr(hf_res_unit, conv_name).weight
# Try to find corresponding original weight
try:
orig_res_unit = orig_block.block[res_idx - 1]
if hasattr(orig_res_unit, 'block'):
conv_idx = 1 if conv_name == 'conv1' else 3
if len(orig_res_unit.block) > conv_idx:
orig_conv_w = orig_res_unit.block[conv_idx].weight
conv_diff = torch.abs(hf_conv_w - orig_conv_w).max().item()
print(f" {res_unit_name}.{conv_name} diff: {conv_diff:.2e}")
except:
print(f" {res_unit_name}.{conv_name}: Could not compare")
# Compare final layers
hf_snake1_alpha = model_hf.encoder.snake1.alpha
orig_snake1_alpha = model.encoder.block[5].alpha
snake_diff = torch.abs(hf_snake1_alpha - orig_snake1_alpha).max().item()
print(f"\nSnake1 alpha diff: {snake_diff:.2e}")
hf_conv2_w = model_hf.encoder.conv2.weight
orig_conv2_w = model.encoder.block[6].weight
conv2_diff = torch.abs(hf_conv2_w - orig_conv2_w).max().item()
print(f"Conv2 weight diff: {conv2_diff:.2e}")
def debug_quantizer_error_propagation():
"""Track how errors propagate through all quantizer layers."""
print("\n=== QUANTIZER ERROR PROPAGATION ANALYSIS ===")
with torch.no_grad():
# Use HF encoder output as identical input for both quantizers
hf_encoder_out = model_hf.encoder(x_hf)
# Use the SAME encoder output for both quantizers
hf_x = hf_encoder_out.clone()
orig_x = hf_encoder_out.clone()
errors = []
layer_names = []
print(f"{'Layer':<25} {'Max Error':<15} {'Mean Error':<15} {'Error Growth':<15}")
print("-" * 75)
# Initial input (should be identical)
max_err = torch.abs(hf_x - orig_x).max().item()
mean_err = torch.abs(hf_x - orig_x).mean().item()
errors.append(max_err)
layer_names.append("Input")
print(f"{'Input':<25} {max_err:<15.2e} {mean_err:<15.2e} {'1.0x':<15}")
# Go through each quantizer layer
for q_idx in range(len(model_hf.quantizer.quantizers)):
print(f"\n--- Quantizer {q_idx} ---")
# Get quantizer layers
hf_quantizer = model_hf.quantizer.quantizers[q_idx]
orig_quantizer = model.quantizer.quantizers[q_idx]
# Apply in_proj
hf_x_proj = hf_quantizer.in_proj(hf_x)
orig_x_proj = orig_quantizer.in_proj(orig_x)
max_err = torch.abs(hf_x_proj - orig_x_proj).max().item()
mean_err = torch.abs(hf_x_proj - orig_x_proj).mean().item()
growth = max_err / errors[-1] if errors[-1] > 0 else float('inf')
errors.append(max_err)
layer_names.append(f"Q{q_idx}_in_proj")
print(f"{'Q' + str(q_idx) + '_in_proj':<25} {max_err:<15.2e} {mean_err:<15.2e} {f'{growth:.1f}x':<15}")
# Apply quantization step by step
hf_x_quant, hf_indices = hf_quantizer.decode_latents(hf_x_proj)
orig_x_quant, orig_indices = orig_quantizer.decode_latents(orig_x_proj)
max_err = torch.abs(hf_x_quant - orig_x_quant).max().item()
mean_err = torch.abs(hf_x_quant - orig_x_quant).mean().item()
growth = max_err / errors[-1] if errors[-1] > 0 else float('inf')
errors.append(max_err)
layer_names.append(f"Q{q_idx}_decode")
print(f"{'Q' + str(q_idx) + '_decode':<25} {max_err:<15.2e} {mean_err:<15.2e} {f'{growth:.1f}x':<15}")
# Check indices differences
indices_diff = torch.abs(hf_indices.float() - orig_indices.float()).max().item()
print(f"{'Q' + str(q_idx) + '_indices':<25} {indices_diff:<15.0f} {'N/A':<15} {'N/A':<15}")
# Apply out_proj
hf_x_out = hf_quantizer.out_proj(hf_x_quant)
orig_x_out = orig_quantizer.out_proj(orig_x_quant)
max_err = torch.abs(hf_x_out - orig_x_out).max().item()
mean_err = torch.abs(hf_x_out - orig_x_out).mean().item()
growth = max_err / errors[-2] if errors[-2] > 0 else float('inf') # Compare to decode step
errors.append(max_err)
layer_names.append(f"Q{q_idx}_out_proj")
print(f"{'Q' + str(q_idx) + '_out_proj':<25} {max_err:<15.2e} {mean_err:<15.2e} {f'{growth:.1f}x':<15}")
# Update for next quantizer (residual quantization)
hf_x = hf_x - hf_x_out # Residual for next quantizer
orig_x = orig_x - orig_x_out
# Show residual for next quantizer
residual_err = torch.abs(hf_x - orig_x).max().item()
print(f"{'Q' + str(q_idx) + '_residual':<25} {residual_err:<15.2e} {'N/A':<15} {'N/A':<15}")
print(f"\n=== QUANTIZER ERROR SUMMARY ===")
print(f"Initial input error: {errors[0]:.2e}")
print(f"Final quantizer error: {errors[-1]:.2e}")
if errors[0] > 0:
total_growth = errors[-1] / errors[0]
print(f"Total error amplification: {total_growth:.0f}x")
# Find biggest error jumps
print(f"\nTop 3 error amplifiers:")
growths = []
for i in range(1, len(errors)):
if errors[i-1] > 0:
growth = errors[i] / errors[i-1]
growths.append((layer_names[i], growth))
growths.sort(key=lambda x: x[1], reverse=True)
for i, (layer, growth) in enumerate(growths[:3]):
print(f" {i+1}. {layer}: {growth:.1f}x amplification")
def debug_quantizer_weight_differences():
"""Compare quantizer weights layer by layer."""
print("\n=== QUANTIZER WEIGHT DIFFERENCES ===")
for q_idx in range(len(model_hf.quantizer.quantizers)):
print(f"\nQuantizer {q_idx} weight differences:")
hf_quantizer = model_hf.quantizer.quantizers[q_idx]
orig_quantizer = model.quantizer.quantizers[q_idx]
# Compare in_proj weights
hf_in_weight = hf_quantizer.in_proj.weight
orig_in_weight = orig_quantizer.in_proj.weight
in_weight_diff = torch.abs(hf_in_weight - orig_in_weight).max().item()
print(f" in_proj weight diff: {in_weight_diff:.2e}")
# Compare in_proj bias
if hf_quantizer.in_proj.bias is not None and orig_quantizer.in_proj.bias is not None:
hf_in_bias = hf_quantizer.in_proj.bias
orig_in_bias = orig_quantizer.in_proj.bias
in_bias_diff = torch.abs(hf_in_bias - orig_in_bias).max().item()
print(f" in_proj bias diff: {in_bias_diff:.2e}")
# Compare codebook weights
hf_codebook = hf_quantizer.codebook.weight
orig_codebook = orig_quantizer.codebook.weight
codebook_diff = torch.abs(hf_codebook - orig_codebook).max().item()
print(f" codebook weight diff: {codebook_diff:.2e}")
# Compare out_proj weights
hf_out_weight = hf_quantizer.out_proj.weight
orig_out_weight = orig_quantizer.out_proj.weight
out_weight_diff = torch.abs(hf_out_weight - orig_out_weight).max().item()
print(f" out_proj weight diff: {out_weight_diff:.2e}")
# Compare out_proj bias
if hf_quantizer.out_proj.bias is not None and orig_quantizer.out_proj.bias is not None:
hf_out_bias = hf_quantizer.out_proj.bias
orig_out_bias = orig_quantizer.out_proj.bias
out_bias_diff = torch.abs(hf_out_bias - orig_out_bias).max().item()
print(f" out_proj bias diff: {out_bias_diff:.2e}")
def debug_quantizer_codebook_analysis():
"""Analyze codebook differences in detail."""
print("\n=== QUANTIZER CODEBOOK ANALYSIS ===")
for q_idx in range(len(model_hf.quantizer.quantizers)):
print(f"\nQuantizer {q_idx} codebook analysis:")
hf_quantizer = model_hf.quantizer.quantizers[q_idx]
orig_quantizer = model.quantizer.quantizers[q_idx]
hf_codebook = hf_quantizer.codebook.weight
orig_codebook = orig_quantizer.codebook.weight
print(f" Codebook shape: {hf_codebook.shape}")
print(f" Max difference: {torch.abs(hf_codebook - orig_codebook).max().item():.2e}")
print(f" Mean difference: {torch.abs(hf_codebook - orig_codebook).mean().item():.2e}")
print(f" Exactly equal: {torch.equal(hf_codebook, orig_codebook)}")
# Check if any codebook entries are swapped or reordered
if not torch.equal(hf_codebook, orig_codebook):
print(f" Checking for reordering...")
# Find closest matches for first few entries
for i in range(min(5, hf_codebook.shape[0])):
hf_entry = hf_codebook[i]
distances = torch.norm(orig_codebook - hf_entry.unsqueeze(0), dim=1)
closest_idx = distances.argmin().item()
closest_dist = distances[closest_idx].item()
print(f" Entry {i}: closest match at index {closest_idx}, distance {closest_dist:.2e}")
def debug_quantizer_step_by_step(q_idx=0):
"""Debug a specific quantizer step by step."""
print(f"\n=== QUANTIZER {q_idx} STEP-BY-STEP DEBUG ===")
with torch.no_grad():
# Use HF encoder output as identical input
hf_encoder_out = model_hf.encoder(x_hf)
# Get input for this specific quantizer
hf_x = hf_encoder_out.clone()
orig_x = hf_encoder_out.clone()
# Apply previous quantizers to get to the desired quantizer
for i in range(q_idx):
hf_quantizer_i = model_hf.quantizer.quantizers[i]
orig_quantizer_i = model.quantizer.quantizers[i]
hf_quant_i, _, _, _, _ = hf_quantizer_i(hf_x)
orig_quant_i, _, _, _, _ = orig_quantizer_i(orig_x)
hf_x = hf_x - hf_quant_i
orig_x = orig_x - orig_quant_i
# Get the specific quantizer
hf_quantizer = model_hf.quantizer.quantizers[q_idx]
orig_quantizer = model.quantizer.quantizers[q_idx]
# Step 1: Input to quantizer
print(f"Step 1 - Input to quantizer {q_idx}:")
print(f" Shape: {hf_x.shape}")
print(f" Max diff: {torch.abs(hf_x - orig_x).max().item():.2e}")
# Step 2: Apply in_proj
hf_x_proj = hf_quantizer.in_proj(hf_x)
orig_x_proj = orig_quantizer.in_proj(orig_x)
print(f"Step 2 - After in_proj:")
print(f" Shape: {hf_x_proj.shape}")
print(f" Max diff: {torch.abs(hf_x_proj - orig_x_proj).max().item():.2e}")
# Step 3: Apply decode_latents
hf_x_quant, hf_indices = hf_quantizer.decode_latents(hf_x_proj)
orig_x_quant, orig_indices = orig_quantizer.decode_latents(orig_x_proj)
print(f"Step 3 - After decode_latents:")
print(f" Quantized shape: {hf_x_quant.shape}")
print(f" Indices shape: {hf_indices.shape}")
print(f" Quantized max diff: {torch.abs(hf_x_quant - orig_x_quant).max().item():.2e}")
print(f" Indices max diff: {torch.abs(hf_indices.float() - orig_indices.float()).max().item():.0f}")
# Check if indices are identical
if torch.equal(hf_indices, orig_indices):
print(f" ✅ Indices are identical!")
else:
print(f" ❌ Indices differ!")
# Show some example differences
diff_mask = (hf_indices != orig_indices)
if diff_mask.any():
print(f" First few differing indices:")
diff_positions = diff_mask.nonzero()[:5]
for pos in diff_positions:
pos_tuple = tuple(pos.tolist())
print(f" Position {pos_tuple}: HF={hf_indices[pos_tuple]}, Orig={orig_indices[pos_tuple]}")
# Step 4: Apply out_proj
hf_x_out = hf_quantizer.out_proj(hf_x_quant)
orig_x_out = orig_quantizer.out_proj(orig_x_quant)
print(f"Step 4 - After out_proj:")
print(f" Shape: {hf_x_out.shape}")
print(f" Max diff: {torch.abs(hf_x_out - orig_x_out).max().item():.2e}")
# Step 5: Full forward pass for comparison
print(f"Step 5 - Full forward pass:")
hf_full_out, hf_commitment, hf_codebook, hf_indices_full, hf_proj_latents = hf_quantizer(hf_x)
orig_full_out, orig_commitment, orig_codebook, orig_indices_full, orig_proj_latents = orig_quantizer(orig_x)
print(f" Full forward max diff: {torch.abs(hf_full_out - orig_full_out).max().item():.2e}")
print(f" Commitment loss diff: {torch.abs(hf_commitment - orig_commitment).max().item():.2e}")
print(f" Codebook loss diff: {torch.abs(hf_codebook - orig_codebook).max().item():.2e}")
def debug_quantizer_full_pipeline():
"""Debug the full quantizer pipeline with identical inputs."""
print("\n=== QUANTIZER FULL PIPELINE DEBUG ===")
with torch.no_grad():
# Use HF encoder output as identical input
hf_encoder_out = model_hf.encoder(x_hf)
# Run full quantizer pipeline
hf_quant_out = model_hf.quantizer(hf_encoder_out)
orig_quant_out = model.quantizer(hf_encoder_out, n_quantizers=None)
print(f"Full pipeline comparison:")
print(f" HF quantized shape: {hf_quant_out[0].shape}")
print(f" Original quantized shape: {orig_quant_out[0].shape}")
print(f" Quantized max diff: {torch.abs(hf_quant_out[0] - orig_quant_out[0]).max().item():.2e}")
print(f" HF codes shape: {hf_quant_out[1].shape}")
print(f" Original codes shape: {orig_quant_out[1].shape}")
print(f" Codes max diff: {torch.abs(hf_quant_out[1].float() - orig_quant_out[1].float()).max().item():.0f}")
print(f" HF latents shape: {hf_quant_out[2].shape}")
print(f" Original latents shape: {orig_quant_out[2].shape}")
print(f" Latents max diff: {torch.abs(hf_quant_out[2] - orig_quant_out[2]).max().item():.2e}")
print(f" HF commitment loss: {hf_quant_out[3].mean().item():.2e}")
print(f" Original commitment loss: {orig_quant_out[3].mean().item():.2e}")
print(f" Commitment loss diff: {torch.abs(hf_quant_out[3] - orig_quant_out[3]).max().item():.2e}")
print(f" HF codebook loss: {hf_quant_out[4].mean().item():.2e}")
print(f" Original codebook loss: {orig_quant_out[4].mean().item():.2e}")
print(f" Codebook loss diff: {torch.abs(hf_quant_out[4] - orig_quant_out[4]).max().item():.2e}")
def debug_decoder_error_propagation():
"""Track how errors propagate through all decoder layers."""
print("\n=== DECODER ERROR PROPAGATION ANALYSIS ===")
with torch.no_grad():
# Use identical quantized inputs for both decoders
hf_encoder_out = model_hf.encoder(x_hf)
hf_quantizer_out = model_hf.quantizer(hf_encoder_out)
quantized_input = hf_quantizer_out[0]
# Use the SAME quantized input for both decoders
hf_x = quantized_input.clone()
orig_x = quantized_input.clone()
errors = []
layer_names = []
print(f"{'Layer':<20} {'Max Error':<15} {'Mean Error':<15} {'Error Growth':<15}")
print("-" * 70)
# Initial input (should be identical)
max_err = torch.abs(hf_x - orig_x).max().item()
mean_err = torch.abs(hf_x - orig_x).mean().item()
errors.append(max_err)
layer_names.append("Input")
print(f"{'Input':<20} {max_err:<15.2e} {mean_err:<15.2e} {'1.0x':<15}")
# Layer 0: First Conv
hf_x = model_hf.decoder.conv1(hf_x)
orig_x = model.decoder.model[0](orig_x)
max_err = torch.abs(hf_x - orig_x).max().item()
mean_err = torch.abs(hf_x - orig_x).mean().item()
growth = max_err / errors[-1] if errors[-1] > 0 else float('inf')
errors.append(max_err)
layer_names.append("Conv1")
print(f"{'Conv1':<20} {max_err:<15.2e} {mean_err:<15.2e} {f'{growth:.1f}x':<15}")
# Decoder Blocks (layers 1-4 in original)
for block_idx in range(len(model_hf.decoder.block)):
hf_x = model_hf.decoder.block[block_idx](hf_x)
orig_x = model.decoder.model[block_idx + 1](orig_x) # layers 1,2,3,4
max_err = torch.abs(hf_x - orig_x).max().item()
mean_err = torch.abs(hf_x - orig_x).mean().item()
growth = max_err / errors[-1] if errors[-1] > 0 else float('inf')
errors.append(max_err)
layer_name = f"Block{block_idx}"
layer_names.append(layer_name)
print(f"{layer_name:<20} {max_err:<15.2e} {mean_err:<15.2e} {f'{growth:.1f}x':<15}")
# Final Snake (layer 5)
hf_x = model_hf.decoder.snake1(hf_x)
orig_x = model.decoder.model[5](orig_x)
max_err = torch.abs(hf_x - orig_x).max().item()
mean_err = torch.abs(hf_x - orig_x).mean().item()
growth = max_err / errors[-1] if errors[-1] > 0 else float('inf')
errors.append(max_err)
layer_names.append("Snake1")
print(f"{'Snake1':<20} {max_err:<15.2e} {mean_err:<15.2e} {f'{growth:.1f}x':<15}")
# Final Conv (layer 6)
hf_x = model_hf.decoder.conv2(hf_x)
orig_x = model.decoder.model[6](orig_x)
max_err = torch.abs(hf_x - orig_x).max().item()
mean_err = torch.abs(hf_x - orig_x).mean().item()
growth = max_err / errors[-1] if errors[-1] > 0 else float('inf')
errors.append(max_err)
layer_names.append("Conv2")
print(f"{'Conv2':<20} {max_err:<15.2e} {mean_err:<15.2e} {f'{growth:.1f}x':<15}")
# Final Tanh (layer 7) - HF model might not have this explicitly
if len(model.decoder.model) > 7:
orig_x = model.decoder.model[7](orig_x) # Tanh activation
max_err = torch.abs(hf_x - orig_x).max().item()
mean_err = torch.abs(hf_x - orig_x).mean().item()
growth = max_err / errors[-1] if errors[-1] > 0 else float('inf')
errors.append(max_err)
layer_names.append("Tanh")
print(f"{'Tanh':<20} {max_err:<15.2e} {mean_err:<15.2e} {f'{growth:.1f}x':<15}")
# Summary statistics
print("\n=== DECODER ERROR PROPAGATION SUMMARY ===")
total_growth = errors[-1] / errors[1] if errors[1] > 0 else float('inf')
print(f"Initial weight error: {errors[1]:.2e}")
print(f"Final decoder error: {errors[-1]:.2e}")
print(f"Total error amplification: {total_growth:.0f}x")
return errors, layer_names
def debug_decoder_block_detailed(block_idx, prev_error, quantized_input):
"""Detailed analysis within a specific decoder block."""
print(f"\n --- Decoder Block {block_idx} Internal Analysis ---")
with torch.no_grad():
# Get to this block's input
hf_x = quantized_input.clone()
orig_x = quantized_input.clone()
# Apply conv1
hf_x = model_hf.decoder.conv1(hf_x)
orig_x = model.decoder.model[0](orig_x)
# Apply previous blocks
for i in range(block_idx):
hf_x = model_hf.decoder.block[i](hf_x)
orig_x = model.decoder.model[i + 1](orig_x)
block_input_error = torch.abs(hf_x - orig_x).max().item()
# Get the blocks
hf_block = model_hf.decoder.block[block_idx]
orig_block = model.decoder.model[block_idx + 1]
# Try to go through sublayers if possible
try:
# For HF block, try to access sublayers
if hasattr(hf_block, 'res_unit1'):
# Apply each residual unit
for res_idx in range(1, 4): # res_unit1, res_unit2, res_unit3
res_unit_name = f"res_unit{res_idx}"
if hasattr(hf_block, res_unit_name):
hf_res_unit = getattr(hf_block, res_unit_name)
# Apply HF residual unit
hf_x_res = hf_res_unit(hf_x)
# For original, try to access corresponding sublayer
if hasattr(orig_block, 'block') and len(orig_block.block) > res_idx - 1:
orig_res_unit = orig_block.block[res_idx - 1]
orig_x_res = orig_res_unit(orig_x)
res_error = torch.abs(hf_x_res - orig_x_res).max().item()
growth = res_error / block_input_error if block_input_error > 0 else float('inf')
print(f" {res_unit_name}: {res_error:.2e} ({growth:.1f}x)")
# Update for next iteration (residual connection)
hf_x = hf_x + hf_x_res # Residual connection
orig_x = orig_x + orig_x_res
else:
print(f" {res_unit_name}: Structure mismatch")
# Final layers in block (Snake + Conv + Upsample)
if hasattr(hf_block, 'snake1'):
hf_x = hf_block.snake1(hf_x)
if hasattr(orig_block, 'block') and len(orig_block.block) > 3:
orig_x = orig_block.block[3](orig_x) # Snake
snake_error = torch.abs(hf_x - orig_x).max().item()
growth = snake_error / block_input_error if block_input_error > 0 else float('inf')
print(f" Snake: {snake_error:.2e} ({growth:.1f}x)")
if hasattr(hf_block, 'conv_t1'):
hf_x = hf_block.conv_t1(hf_x)
if hasattr(orig_block, 'block') and len(orig_block.block) > 4:
orig_x = orig_block.block[4](orig_x) # ConvTranspose
conv_error = torch.abs(hf_x - orig_x).max().item()
growth = conv_error / block_input_error if block_input_error > 0 else float('inf')
print(f" ConvTranspose: {conv_error:.2e} ({growth:.1f}x)")
except Exception as e:
print(f" Detailed analysis failed: {e}")
def debug_decoder_weight_differences():
"""Compare decoder weights layer by layer."""
print("\n=== DECODER WEIGHT DIFFERENCES ===")
# Debug: Show original decoder structure
print("Original decoder structure:")
for i, layer in enumerate(model.decoder.model):
print(f" Layer {i}: {type(layer).__name__}")
if hasattr(layer, 'weight'):
print(f" Weight shape: {layer.weight.shape}")
# Compare conv1 weights
try:
hf_conv1_w = model_hf.decoder.conv1.weight
orig_conv1_w = model.decoder.model[0].weight # First layer
conv1_diff = torch.abs(hf_conv1_w - orig_conv1_w).max().item()
print(f"\nConv1 weight max diff: {conv1_diff:.2e}")
if model_hf.decoder.conv1.bias is not None and model.decoder.model[0].bias is not None:
hf_conv1_b = model_hf.decoder.conv1.bias
orig_conv1_b = model.decoder.model[0].bias
conv1_bias_diff = torch.abs(hf_conv1_b - orig_conv1_b).max().item()
print(f"Conv1 bias max diff: {conv1_bias_diff:.2e}")
except Exception as e:
print(f"Conv1 comparison failed: {e}")
# Compare other layers by going through Sequential
print("\nOther layer comparisons:")
for i, orig_layer in enumerate(model.decoder.model[1:], 1):
try:
print(f"Layer {i}: {type(orig_layer).__name__}")
if hasattr(orig_layer, 'weight'):
print(f" Weight shape: {orig_layer.weight.shape}")
# Fix: Use correct layer indices
if i == 5: # Snake layer (layer 5, not second-to-last)
if hasattr(model_hf.decoder, 'snake1') and hasattr(orig_layer, 'alpha'):
hf_param = model_hf.decoder.snake1.alpha
orig_param = orig_layer.alpha
param_diff = torch.abs(hf_param - orig_param).max().item()
print(f" Snake alpha diff: {param_diff:.2e}")
elif i == 6: # Final conv (layer 6, not last)
if hasattr(model_hf.decoder, 'conv2'):
hf_param = model_hf.decoder.conv2.weight
orig_param = orig_layer.weight
param_diff = torch.abs(hf_param - orig_param).max().item()
print(f" Final conv weight diff: {param_diff:.2e}")
if model_hf.decoder.conv2.bias is not None and orig_layer.bias is not None:
hf_bias = model_hf.decoder.conv2.bias
orig_bias = orig_layer.bias
bias_diff = torch.abs(hf_bias - orig_bias).max().item()
print(f" Final conv bias diff: {bias_diff:.2e}")
except Exception as e:
print(f"Layer {i} comparison failed: {e}")
def debug_decoder_step_by_step():
"""Debug decoder step by step with identical quantized inputs."""
print("\n=== DECODER STEP-BY-STEP DEBUG ===")
with torch.no_grad():
# Use identical quantized inputs
hf_encoder_out = model_hf.encoder(x_hf)
hf_quantizer_out = model_hf.quantizer(hf_encoder_out)
quantized_input = hf_quantizer_out[0]
print(f"Step 1 - Quantized Input:")
print(f" Shape: {quantized_input.shape}")
print(f" Value range: [{quantized_input.min().item():.3f}, {quantized_input.max().item():.3f}]")
# Step 2: Apply first conv
hf_x = model_hf.decoder.conv1(quantized_input)
orig_x = model.decoder.model[0](quantized_input)
print(f"Step 2 - After first conv:")
print(f" HF shape: {hf_x.shape}")
print(f" Original shape: {orig_x.shape}")
print(f" Max diff: {torch.abs(hf_x - orig_x).max().item():.2e}")
print(f" Mean diff: {torch.abs(hf_x - orig_x).mean().item():.2e}")
# Step 3: Apply layers step by step
orig_layer_idx = 1
# Apply HF decoder blocks
for block_idx in range(len(model_hf.decoder.block)):
hf_x = model_hf.decoder.block[block_idx](hf_x)
# Apply corresponding original layer
if orig_layer_idx < len(model.decoder.model):
orig_x = model.decoder.model[orig_layer_idx](orig_x)
orig_layer_idx += 1
print(f"Step 3.{block_idx} - After decoder block {block_idx}:")
print(f" HF shape: {hf_x.shape}")
print(f" Original shape: {orig_x.shape}")
print(f" Max diff: {torch.abs(hf_x - orig_x).max().item():.2e}")
print(f" Mean diff: {torch.abs(hf_x - orig_x).mean().item():.2e}")
# Step 4: Apply final snake
hf_x = model_hf.decoder.snake1(hf_x)
if orig_layer_idx < len(model.decoder.model):
orig_x = model.decoder.model[orig_layer_idx](orig_x)
orig_layer_idx += 1
print(f"Step 4 - After final snake:")
print(f" HF shape: {hf_x.shape}")
print(f" Original shape: {orig_x.shape}")
print(f" Max diff: {torch.abs(hf_x - orig_x).max().item():.2e}")
print(f" Mean diff: {torch.abs(hf_x - orig_x).mean().item():.2e}")
# Step 5: Apply final conv
hf_output = model_hf.decoder.conv2(hf_x)
if orig_layer_idx < len(model.decoder.model):
orig_output = model.decoder.model[orig_layer_idx](orig_x)
print(f"Step 5 - Final decoder output:")
print(f" HF shape: {hf_output.shape}")
print(f" Original shape: {orig_output.shape}")
print(f" Max diff: {torch.abs(hf_output - orig_output).max().item():.2e}")
print(f" Mean diff: {torch.abs(hf_output - orig_output).mean().item():.2e}")
# Step 6: Full pipeline comparison
print(f"Step 6 - Full decoder pipeline:")
hf_full_output = model_hf.decoder(quantized_input)
orig_full_output = model.decoder(quantized_input)
print(f" Full pipeline max diff: {torch.abs(hf_full_output - orig_full_output).max().item():.2e}")
print(f" Full pipeline mean diff: {torch.abs(hf_full_output - orig_full_output).mean().item():.2e}")
def debug_decoder_full_pipeline():
"""Debug the full decoder pipeline with identical quantized inputs."""
print("\n=== DECODER FULL PIPELINE DEBUG ===")
with torch.no_grad():
# Use identical quantized inputs from HF model
hf_encoder_out = model_hf.encoder(x_hf)
hf_quantizer_out = model_hf.quantizer(hf_encoder_out)
quantized_input = hf_quantizer_out[0]
print(f"Quantized input shape: {quantized_input.shape}")
print(f"Quantized input range: [{quantized_input.min().item():.3f}, {quantized_input.max().item():.3f}]")
# Run full decoder pipeline
hf_decoded = model_hf.decoder(quantized_input)
orig_decoded = model.decoder(quantized_input)
print(f"HF decoded shape: {hf_decoded.shape}")
print(f"Original decoded shape: {orig_decoded.shape}")
print(f"Max difference: {torch.abs(hf_decoded - orig_decoded).max().item():.2e}")
print(f"Mean difference: {torch.abs(hf_decoded - orig_decoded).mean().item():.2e}")
print(f"Relative error: {(torch.abs(hf_decoded - orig_decoded).max() / torch.abs(orig_decoded).max()).item():.2e}")
# Check output ranges
print(f"HF output range: [{hf_decoded.min().item():.3f}, {hf_decoded.max().item():.3f}]")
print(f"Original output range: [{orig_decoded.min().item():.3f}, {orig_decoded.max().item():.3f}]")
# Test with HF's decode method
print(f"\nUsing HF decode method:")
hf_decode_output = model_hf.decode(quantized_input)["audio_values"]
print(f"HF decode output shape: {hf_decode_output.shape}")
print(f"HF decode vs direct decoder max diff: {torch.abs(hf_decode_output - hf_decoded).max().item():.2e}")
# Test reconstruction quality
print(f"\nReconstruction quality analysis:")
# Compute SNR
signal_power = torch.mean(orig_decoded ** 2)
noise_power = torch.mean((hf_decoded - orig_decoded) ** 2)
snr_db = 10 * torch.log10(signal_power / noise_power)
print(f"Signal-to-Noise Ratio: {snr_db:.2f} dB")
# Compute relative error
relative_error = torch.abs(hf_decoded - orig_decoded).mean() / torch.abs(orig_decoded).mean()
print(f"Relative error: {relative_error:.2e}")
def debug_full_codec_pipeline():
"""Debug the complete encode-decode pipeline."""
print("\n=== FULL CODEC PIPELINE DEBUG ===")
with torch.no_grad():
# Full HF pipeline
print("Running full HF pipeline...")
hf_encoded = model_hf.encode(x_hf)
hf_decoded = model_hf.decode(hf_encoded.quantized_representation)["audio_values"]
# Full original pipeline
print("Running full original pipeline...")
orig_encoded = model.encode(x)
orig_decoded = model.decode(orig_encoded[0])
print(f"Input shape: {x_hf.shape}")
print(f"HF encoded shape: {hf_encoded.quantized_representation.shape}")
print(f"Original encoded shape: {orig_encoded[0].shape}")
print(f"HF decoded shape: {hf_decoded.shape}")
print(f"Original decoded shape: {orig_decoded.shape}")
# Compare final outputs
# Trim to same length
min_length = min(hf_decoded.shape[-1], orig_decoded.shape[-1])
hf_trimmed = hf_decoded[..., :min_length]
orig_trimmed = orig_decoded[..., :min_length]
print(f"Final codec comparison (trimmed to {min_length} samples):")
print(f" Max difference: {torch.abs(hf_trimmed - orig_trimmed).max().item():.2e}")
print(f" Mean difference: {torch.abs(hf_trimmed - orig_trimmed).mean().item():.2e}")
print(f" Relative error: {(torch.abs(hf_trimmed - orig_trimmed).max() / torch.abs(orig_trimmed).max()).item():.2e}")
# Compute RMSE
rmse = compute_rmse(hf_trimmed, orig_trimmed)
print(f" RMSE: {rmse:.6f}")
# Compute SNR
signal_power = torch.mean(orig_trimmed ** 2)
noise_power = torch.mean((hf_trimmed - orig_trimmed) ** 2)
snr_db = 10 * torch.log10(signal_power / noise_power)
print(f" SNR: {snr_db:.2f} dB")
#### Main execution block
torch_device = "cuda" if torch.cuda.is_available() else "cpu"
for sampling_rate in model_config:
print(f"\nTesting DAC model for sampling rate: {sampling_rate} Hz")
hf_model_name = model_config[sampling_rate]["model_name"]
dac_model_type = model_config[sampling_rate]["dac_model_type"]
#### 1) LOAD MODEL
# -- Hugging Face
model_id = f"descript/{hf_model_name}"
model_hf = DacModel.from_pretrained(model_id).to(torch_device).eval()
processor = AutoProcessor.from_pretrained(model_id)
# -- Original
model_path = dac.utils.download(model_type=dac_model_type)
model = dac.DAC.load(model_path).eval()
model.to(torch_device)
#### 2) PREPARE AUDIO DATA
# get audio data
librispeech_dummy = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
librispeech_dummy = librispeech_dummy.cast_column("audio", Audio(sampling_rate=processor.sampling_rate))
audio_array = librispeech_dummy[0]['audio']['array']
# -- Hugging Face
inputs = processor(
raw_audio=audio_array,
return_tensors="pt",
sampling_rate=sampling_rate,
).to(torch_device)
x_hf = inputs["input_values"]
# -- Original
signal = AudioSignal(audio_array, sample_rate=sampling_rate)
signal.to(model.device)
x = model.preprocess(signal.audio_data, signal.sample_rate)
# -- compare
torch.testing.assert_close(x_hf, x, rtol=1e-6, atol=1e-6)
### 3) layer-by-layer debugging
with torch.no_grad():
# --- DEBUG ENCODER
print("\n" + "="*50)
print("ENCODER DEBUGGING")
print("="*50)
debug_encoder_error_propagation()
debug_weight_differences_by_layer()
# -- DEBUG QUANTIZER (with same input)
print("\n" + "="*50)
print("QUANTIZER DEBUGGING")
print("="*50)
debug_quantizer_error_propagation()
debug_quantizer_weight_differences()
debug_quantizer_codebook_analysis()
debug_quantizer_step_by_step(q_idx=4) # Debug quantizer 4 specifically since that's where differences start
debug_quantizer_full_pipeline()
# -- DEBUG DECODER (with same input)
print("\n" + "="*50)
print("DECODER DEBUGGING")
print("="*50)
debug_decoder_error_propagation()
debug_decoder_weight_differences()
debug_decoder_step_by_step()
debug_decoder_full_pipeline()
# -- FULL CODEC
debug_full_codec_pipeline()
print("\n" + "="*50)
Raw
test_dac.py
"""
Script for computing expected outputs of DAC model.
Run the following to run and store outputs in a TXT file:
```bash
python scripts/test_dac.py > dac_single.txt
```
----------------------------------------------
Setup:
```
# after setting up virtual environment and transformers
uv pip install descript-audio-codec==1.0.0
```
Using DAC model, as per their documentation:
https://github.com/descriptinc/descript-audio-codec?tab=readme-ov-file#programmatic-usage
"""
import dac
from audiotools import AudioSignal
from datasets import load_dataset, Audio
import torch
import numpy as np
from transformers import AutoProcessor, DacModel
N_ELEM_PRINT_CODES = 15
N_ELEM_PRINT_LATENTS = None
N_ELEM_PRINT_DEC = 50
torch.set_printoptions(threshold=100000)
# model configuration based on sampling rate
model_config = {
16000: {
"model_name": "dac_16khz",
"dac_model_type": "16khz",
},
24000: {
"model_name": "dac_24khz",
"dac_model_type": "24khz",
},
44100: {
"model_name": "dac_44khz",
"dac_model_type": "44khz",
}
}
def normalize(arr):
norm = np.linalg.norm(arr)
normalized_arr = arr / norm
return normalized_arr
def compute_rmse(arr1, arr2):
arr1_np = arr1.cpu().numpy().squeeze()
arr2_np = arr2.cpu().numpy().squeeze()
max_length = min(arr1.shape[-1], arr2.shape[-1])
arr1_np = arr1_np[..., :max_length]
arr2_np = arr2_np[..., :max_length]
arr1_normalized = normalize(arr1_np)
arr2_normalized = normalize(arr2_np)
return np.sqrt(((arr1_normalized - arr2_normalized) ** 2).mean())
### Main execution
torch_device = "cuda" if torch.cuda.is_available() else "cpu"
for sampling_rate in model_config:
print(f"\nTesting DAC model for sampling rate: {sampling_rate} Hz")
hf_model_name = model_config[sampling_rate]["model_name"]
dac_model_type = model_config[sampling_rate]["dac_model_type"]
#### 1) LOAD MODEL
# -- Hugging Face
model_id = f"descript/{hf_model_name}"
model_hf = DacModel.from_pretrained(model_id).to(torch_device).eval()
processor = AutoProcessor.from_pretrained(model_id)
# -- Original
model_path = dac.utils.download(model_type=dac_model_type)
model = dac.DAC.load(model_path).eval()
model.to(torch_device)
#### 2) PREPARE AUDIO DATA
# get audio data
librispeech_dummy = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
librispeech_dummy = librispeech_dummy.cast_column("audio", Audio(sampling_rate=processor.sampling_rate))
audio_array = librispeech_dummy[0]['audio']['array']
# -- Hugging Face
inputs = processor(
raw_audio=audio_array,
return_tensors="pt",
sampling_rate=sampling_rate,
).to(torch_device)
x_hf = inputs["input_values"]
# -- Original
signal = AudioSignal(audio_array, sample_rate=sampling_rate)
signal.to(model.device)
x = model.preprocess(signal.audio_data, signal.sample_rate)
# -- compare
torch.testing.assert_close(x_hf, x, rtol=1e-6, atol=1e-6)
print("Expected input shape : ", x.shape)
### 3) TESTS
with torch.no_grad():
print("=== ENCODER COMPARISON ===")
# -- Hugging Face
outputs_hf = model_hf.encode(x_hf)
# -- DAC
z, codes, latents, commitment_loss, codebook_loss = model.encode(x)
expected_loss = model_hf.config.commitment_loss_weight * commitment_loss + model_hf.config.codebook_loss_weight * codebook_loss
"""
we expect high error due to:
- precision errors of weight normalization
- errors accumulate through layers
- downsampling operations being more sensitive to precision
"""
print(f"Expected encoder loss: {expected_loss}")
torch.testing.assert_close(expected_loss, outputs_hf[0].squeeze(), rtol=1e-3, atol=1e-3)
# Below fail due to precision errors, but final codec quality is similar
# print(f"Expected encoder quantized_representation: {z[..., :N_ELEM_PRINT_LATENTS]}")
# torch.testing.assert_close(z, outputs_hf[1], rtol=1e-3, atol=1e-3)
# print(f"Expected encoder codes: {codes[..., :N_ELEM_PRINT_CODES]}")
# torch.testing.assert_close(codes, outputs_hf[2], rtol=1e-3, atol=1e-3)
# print(f"Expected encoder latents: {latents[..., :N_ELEM_PRINT_LATENTS]}")
# torch.testing.assert_close(latents, outputs_hf[3], rtol=1e-3, atol=1e-3)
print("=== QUANTIZER COMPARISON ===")
output_quant_hf = model_hf.quantizer(outputs_hf[1])
output_quant_dac = model.quantizer(outputs_hf[1])
"""
We expect less error for the discrete codes since discretization "resets" precision errors.
Continuous representations and latents have high error, similar to reasons above for encoder.
"""
# # -- quantized continuous representation (large to print)
if N_ELEM_PRINT_LATENTS is not None:
print(f"Expected quantizer continuous representation: {output_quant_dac[0][..., :N_ELEM_PRINT_LATENTS]}")
torch.testing.assert_close(output_quant_hf[0], output_quant_dac[0], rtol=1e-3, atol=1e-3)
# -- latents (large to print)
if N_ELEM_PRINT_LATENTS is not None:
print(f"Expected quantizer latent: {output_quant_dac[2][..., :N_ELEM_PRINT_LATENTS]}")
torch.testing.assert_close(output_quant_hf[2], output_quant_dac[2], rtol=1e-3, atol=1e-3)
# -- codes
print(f"Expected codes: {output_quant_dac[1][..., :N_ELEM_PRINT_CODES]}")
torch.testing.assert_close(output_quant_hf[1], output_quant_dac[1], rtol=1e-6, atol=1e-6)
# -- commitment loss
print(f"Expected quantizer commitment loss: {output_quant_dac[3]}")
torch.testing.assert_close(output_quant_hf[3][0], output_quant_dac[3], rtol=1e-6, atol=1e-6)
# -- codebook loss
print(f"Expected quantizer codebook loss: {output_quant_dac[4]}")
torch.testing.assert_close(output_quant_hf[4][0], output_quant_dac[4], rtol=1e-6, atol=1e-6)
print("=== DECODER COMPARISON ===")
# compare decoders with same input
"""
we again expect high error due to:
- precision errors of weight normalization
- errors accumulate through layers
"""
hf_decoded = model_hf.decode(outputs_hf[1])["audio_values"]
dac_decoded = model.decode(outputs_hf[1])[0]
print(f"Expected DAC decoded output: {dac_decoded[..., :N_ELEM_PRINT_DEC]}")
torch.testing.assert_close(dac_decoded, hf_decoded, rtol=1e-3, atol=1e-3)
# codec lossiness from original implementation
print("=== CODEC ERROR ===")
x_dac = model.decode(model.encode(x)[0]).squeeze()
expected_rmse = compute_rmse(x_dac, x)
print(f"Expected codec error: {expected_rmse}")
x_dac_hf = model_hf.decode(outputs_hf[1])["audio_values"].squeeze()
hf_rmse = compute_rmse(x_dac_hf, x_hf)
torch.testing.assert_close(hf_rmse, expected_rmse, rtol=1e-5, atol=1e-5)
# make sure forward and decode gives same result
_, quantized_representation, _, _ = outputs_hf.to_tuple()
input_values_dec = model_hf.decode(quantized_representation)[0]
input_values_enc_dec = model_hf(inputs["input_values"])[1]
torch.testing.assert_close(input_values_dec, input_values_enc_dec, rtol=1e-6, atol=1e-6)
print("\n" + "="*50)
Raw
test_dac_batch.py
"""
Script for computing expected outputs of DAC model,
for batched processing.
Run the following to run and store outputs in a TXT file:
```bash
python scripts/test_dac_batch.py > dac_batch.txt
```
---------------------------------------------------
Setup:
```
# after setting up virtual environment and transformers
uv pip install descript-audio-codec==1.0.0
```
Using DAC model, as per their documentation:
https://github.com/descriptinc/descript-audio-codec?tab=readme-ov-file#programmatic-usage
"""
import dac
from audiotools import AudioSignal
from datasets import load_dataset, Audio
import torch
import numpy as np
from transformers import AutoProcessor, DacModel
N_ELEM_PRINT_CODES = 15
N_ELEM_PRINT_LATENTS = None
N_ELEM_PRINT_DEC = 50
torch.set_printoptions(threshold=100000)
# model configuration based on sampling rate
model_config = {
16000: {
"model_name": "dac_16khz",
"dac_model_type": "16khz",
},
24000: {
"model_name": "dac_24khz",
"dac_model_type": "24khz",
},
44100: {
"model_name": "dac_44khz",
"dac_model_type": "44khz",
}
}
def normalize(arr):
norm = np.linalg.norm(arr)
normalized_arr = arr / norm
return normalized_arr
def compute_rmse(arr1, arr2):
arr1_np = arr1.cpu().numpy().squeeze()
arr2_np = arr2.cpu().numpy().squeeze()
max_length = min(arr1.shape[-1], arr2.shape[-1])
arr1_np = arr1_np[..., :max_length]
arr2_np = arr2_np[..., :max_length]
arr1_normalized = normalize(arr1_np)
arr2_normalized = normalize(arr2_np)
return np.sqrt(((arr1_normalized - arr2_normalized) ** 2).mean())
### Main execution
torch_device = "cuda" if torch.cuda.is_available() else "cpu"
for sampling_rate in model_config:
print(f"\nTesting DAC model for sampling rate: {sampling_rate} Hz")
hf_model_name = model_config[sampling_rate]["model_name"]
dac_model_type = model_config[sampling_rate]["dac_model_type"]
#### 1) LOAD MODEL
# -- Hugging Face
model_id = f"descript/{hf_model_name}"
model_hf = DacModel.from_pretrained(model_id).to(torch_device).eval()
processor = AutoProcessor.from_pretrained(model_id)
# -- Original
model_path = dac.utils.download(model_type=dac_model_type)
model = dac.DAC.load(model_path).eval()
model.to(torch_device)
#### 2) PREPARE AUDIO DATA
# get audio data
librispeech_dummy = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
librispeech_dummy = librispeech_dummy.cast_column("audio", Audio(sampling_rate=processor.sampling_rate))
audio_samples = [np.array([audio_sample["array"]])[0] for audio_sample in librispeech_dummy[-2:]["audio"]]
# -- Hugging Face
inputs = processor(
raw_audio=audio_samples,
return_tensors="pt",
sampling_rate=sampling_rate,
).to(torch_device)
x_hf = inputs["input_values"]
# -- Original
# create AudioSignal objects for each audio sample
audio_signal = [AudioSignal(audio_sample, sample_rate=sampling_rate) for audio_sample in audio_samples]
signal = AudioSignal.batch(audio_signal, pad_signals=True)
signal.to(model.device)
x = model.preprocess(signal.audio_data, signal.sample_rate)
# -- compare
torch.testing.assert_close(x_hf, x, rtol=1e-6, atol=1e-6)
print("Expected input shape : ", x.shape)
### 3) TESTS
with torch.no_grad():
print("=== ENCODER COMPARISON ===")
# -- Hugging Face
outputs_hf = model_hf.encode(x_hf)
# -- DAC
z, codes, latents, commitment_loss, codebook_loss = model.encode(x)
expected_loss = model_hf.config.commitment_loss_weight * commitment_loss + model_hf.config.codebook_loss_weight * codebook_loss
"""
we expect high error due to:
- precision errors of weight normalization
- errors accumulate through layers
- downsampling operations being more sensitive to precision
"""
print(f"Expected encoder loss: {expected_loss.item()}")
torch.testing.assert_close(expected_loss.item(), outputs_hf[0].mean().item(), rtol=1e-3, atol=1e-3)
# Below fail due to precision errors, but final codec quality is similar
# print(f"Expected encoder quantized_representation: {z[..., :N_ELEM_PRINT_LATENTS]}")
# torch.testing.assert_close(z, outputs_hf[1], rtol=1e-3, atol=1e-3)
# print(f"Expected encoder codes: {codes[..., :N_ELEM_PRINT_CODES]}")
# torch.testing.assert_close(codes, outputs_hf[2], rtol=1e-3, atol=1e-3)
# print(f"Expected encoder latents: {latents[..., :N_ELEM_PRINT_LATENTS]}")
# torch.testing.assert_close(latents, outputs_hf[3], rtol=1e-3, atol=1e-3)
print("=== QUANTIZER COMPARISON ===")
output_quant_hf = model_hf.quantizer(outputs_hf[1])
output_quant_dac = model.quantizer(outputs_hf[1])
"""
We expect less error for the discrete codes since discretization "resets" precision errors.
Continuous representations and latents have high error, similar to reasons above for encoder.
"""
# # -- quantized continuous representation (large to print)
if N_ELEM_PRINT_LATENTS is not None:
print(f"Expected quantizer continuous representation: {output_quant_dac[0][..., :N_ELEM_PRINT_LATENTS]}")
torch.testing.assert_close(output_quant_hf[0], output_quant_dac[0], rtol=1e-3, atol=1e-3)
# -- latents (large to print)
if N_ELEM_PRINT_LATENTS is not None:
print(f"Expected quantizer latent: {output_quant_dac[2][..., :N_ELEM_PRINT_LATENTS]}")
torch.testing.assert_close(output_quant_hf[2], output_quant_dac[2], rtol=1e-3, atol=1e-3)
# -- codes
print(f"Expected codes: {output_quant_dac[1][..., :N_ELEM_PRINT_CODES]}")
torch.testing.assert_close(output_quant_hf[1], output_quant_dac[1], rtol=1e-6, atol=1e-6)
# -- commitment loss
print(f"Expected quantizer commitment loss: {output_quant_dac[3].item()}")
torch.testing.assert_close(output_quant_hf[3][0].mean().item(), output_quant_dac[3].item(), rtol=1e-6, atol=1e-6)
# -- codebook loss
print(f"Expected quantizer codebook loss: {output_quant_dac[4].item()}")
torch.testing.assert_close(output_quant_hf[4][0].mean().item(), output_quant_dac[4].item(), rtol=1e-6, atol=1e-6)
print("=== DECODER COMPARISON ===")
# compare decoders with same input
"""
we again expect high error due to:
- precision errors of weight normalization
- errors accumulate through layers
"""
hf_decoded = model_hf.decode(outputs_hf[1])["audio_values"]
dac_decoded = model.decode(outputs_hf[1]).squeeze()
print(f"Expected DAC decoded output: {dac_decoded[..., :N_ELEM_PRINT_DEC]}")
torch.testing.assert_close(dac_decoded, hf_decoded, rtol=1e-3, atol=1e-3)
# codec lossiness from original implementation
print("=== CODEC ERROR ===")
x_dac = model.decode(model.encode(x)[0])
expected_rmse = compute_rmse(x_dac, x)
print(f"Expected codec error: {expected_rmse}")
x_dac_hf = model_hf.decode(outputs_hf[1])["audio_values"]
hf_rmse = compute_rmse(x_dac_hf, x_hf)
torch.testing.assert_close(expected_rmse, hf_rmse, rtol=1e-6, atol=1e-6)
# make sure forward and decode gives same result
_, quantized_representation, _, _ = outputs_hf.to_tuple()
input_values_dec = model_hf.decode(quantized_representation)[0]
input_values_enc_dec = model_hf(inputs["input_values"])[1]
torch.testing.assert_close(input_values_dec, input_values_enc_dec, rtol=1e-6, atol=1e-6)
print("\n" + "="*50)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment