Spaces:

benhong
/

blt-patch-view

Sleeping

blt-patch-view / scripts /precompute_patch_view.py

Add app

98c0f96 about 1 year ago

10.6 kB

#!/usr/bin/env python3
"""
Precompute entropy and perplexity for test-10 audio samples.
This script loads the entropy model once and processes all 10 audio files,
saving tokens, entropies, and perplexities to JSON files for the Gradio app.
Usage:
    uv run python apps/patch_view/scripts/precompute_patch_view.py
"""

import sys
import torch
import numpy as np
import json
import math
from pathlib import Path
import torchaudio

# Add project root and WavTokenizer to path
# __file__ is in apps/patch_view/scripts/, so climb three levels to project root
PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent
sys.path.insert(0, str(PROJECT_ROOT))
sys.path.insert(0, str(PROJECT_ROOT / "WavTokenizer"))

from encoder.utils import convert_audio
from decoder.pretrained import WavTokenizer
from bytelatent.transformer import LMTransformer, LMTransformerArgs
from torch.distributed.checkpoint.state_dict import get_model_state_dict, set_model_state_dict
import torch.distributed.checkpoint as dcp


def load_entropy_model(checkpoint_dir, device):
    """Load entropy model from checkpoint (step 50000)."""
    print(f"Loading entropy model from {checkpoint_dir}...")

    params_path = Path(checkpoint_dir) / "params.json"
    with open(params_path) as f:
        params = json.load(f)

    model_args = LMTransformerArgs(**params["entropy_model"])
    model = LMTransformer(model_args)

    # Load checkpoint
    state_dict = {"model": get_model_state_dict(model)}
    dcp.load(state_dict, checkpoint_id=checkpoint_dir)
    set_model_state_dict(model, state_dict["model"])

    model = model.to(device)
    model.eval()
    num_params = sum(p.numel() for p in model.parameters())
    print(f"  ✓ Entropy model loaded ({num_params:,} parameters)")

    return model


def load_wavtokenizer(config_path, model_path, device):
    """Load WavTokenizer for audio encoding."""
    wavtokenizer = WavTokenizer.from_pretrained0802(config_path, model_path)
    wavtokenizer = wavtokenizer.to(device)
    wavtokenizer.eval()
    return wavtokenizer


def encode_audio(audio_path, wavtokenizer, device):
    """
    Load and encode audio file to tokens.
    Returns:
        tokens: List of token IDs
        duration: Duration in seconds
        original_wav: Original audio (numpy array)
        reconstructed_wav: WavTokenizer reconstructed audio (numpy array)
    """
    # Load audio using torchaudio
    wav, sr = torchaudio.load(audio_path)

    # Convert to 24kHz mono
    wav = convert_audio(wav, sr, 24000, 1)

    # Normalize
    wav = wav / (torch.abs(wav).max() + 1e-8)

    # Keep original for comparison
    original_wav = wav.squeeze().cpu().numpy()

    # Encode to tokens
    bandwidth_id = torch.tensor([0]).to(device)
    with torch.no_grad():
        _, discrete_codes = wavtokenizer.encode_infer(wav.to(device), bandwidth_id=bandwidth_id)

    # discrete_codes: [n_q, B, seq_len] - extract first codebook, first batch
    tokens = discrete_codes[0, 0].cpu()

    # Decode tokens to reconstruct audio
    with torch.no_grad():
        tokens_for_decode = tokens.unsqueeze(0).unsqueeze(0).to(device)  # [1, 1, seq_len]
        features = wavtokenizer.codes_to_features(tokens_for_decode)
        reconstructed = wavtokenizer.decode(features, bandwidth_id=bandwidth_id)
        reconstructed_wav = reconstructed.squeeze().cpu().numpy()

    tokens_list = tokens.tolist()
    duration = len(tokens_list) / 40.0  # 40 tokens/sec

    return tokens_list, duration, original_wav, reconstructed_wav


def compute_entropies(tokens, entropy_model, device, max_context=512):
    """
    Compute entropy for each token position using batched contexts.
    Args:
        tokens: List of token IDs
        entropy_model: Loaded entropy model
        max_context: Maximum context window used per forward pass
    Returns:
        numpy array of entropy values (bits) aligned with token positions
    """
    num_tokens = len(tokens)
    print(f"  Computing entropies for {num_tokens} tokens...")

    if num_tokens == 0:
        return np.array([])

    token_tensor = torch.tensor(tokens, dtype=torch.long, device=device)
    entropies = torch.zeros(num_tokens, dtype=torch.float32, device=device)

    chunk_size = max(1, max_context // 2)

    with torch.no_grad():
        processed = 0
        while processed < num_tokens:
            chunk_end = min(num_tokens, processed + chunk_size)
            context_start = max(0, chunk_end - max_context)
            context_tokens = token_tensor[context_start:chunk_end].unsqueeze(0)

            logits = entropy_model(context_tokens)
            probs = torch.softmax(logits[0], dim=-1)
            log_probs = torch.log(probs + 1e-10) / math.log(2)
            chunk_entropy = -(probs * log_probs).sum(dim=-1)

            new_count = chunk_end - processed
            entropies[processed:chunk_end] = chunk_entropy[-new_count:]

            processed = chunk_end
            print(f"    Processed {processed}/{num_tokens} tokens", end="\r")

    print()  # newline after progress carriage-return
    return entropies.cpu().numpy()


def process_audio_file(audio_path, entropy_model, wavtokenizer, device, output_dir):
    """
    Process single audio file: encode, compute entropies, save JSON and reconstructed audio.
    Args:
        audio_path: Path to FLAC file
        entropy_model: Loaded entropy model
        wavtokenizer: Loaded WavTokenizer
        device: torch device
        output_dir: Directory to save JSON
    Returns:
        Path to saved JSON file
    """
    import soundfile as sf

    file_id = audio_path.stem
    print(f"\nProcessing {file_id}...")

    # Load transcript if available
    txt_path = audio_path.with_suffix(".txt")
    transcript = txt_path.read_text().strip() if txt_path.exists() else ""

    # Encode audio and get reconstruction
    print("  Encoding audio...")
    tokens, duration, original_wav, reconstructed_wav = encode_audio(
        audio_path, wavtokenizer, device
    )
    print(f"  ✓ Encoded {len(tokens)} tokens ({duration:.2f}s)")

    # Save reconstructed audio for comparison
    reconstructed_path = output_dir / f"{file_id}_reconstructed.wav"
    sf.write(reconstructed_path, reconstructed_wav, 24000)
    print(f"  ✓ Saved reconstructed audio ({reconstructed_path.stat().st_size / 1024:.1f} KB)")

    # Compute entropies
    entropies = compute_entropies(tokens, entropy_model, device)
    print(f"  ✓ Computed entropies")

    # Compute perplexities: PPL = 2^H
    perplexities = 2 ** entropies
    print(f"  ✓ Computed perplexities")

    # Create output data
    data = {
        "file_id": file_id,
        "audio_path": str(audio_path),
        "reconstructed_audio_path": str(reconstructed_path),
        "transcript": transcript,
        "audio_duration_sec": float(duration),
        "sample_rate": 24000,
        "tokens": tokens,
        "entropies": entropies.tolist(),
        "perplexities": perplexities.tolist(),
        "num_tokens": len(tokens),
        "avg_entropy": float(np.mean(entropies)),
        "min_entropy": float(np.min(entropies)),
        "max_entropy": float(np.max(entropies)),
        "avg_perplexity": float(np.mean(perplexities)),
    }

    # Save to JSON
    output_file = output_dir / f"{file_id}.json"
    with open(output_file, 'w') as f:
        json.dump(data, f, indent=2)

    print(f"  ✓ Saved to {output_file} ({output_file.stat().st_size / 1024:.1f} KB)")

    return output_file


def main():
    # Configuration
    ENTROPY_CHECKPOINT = "archive/checkpoints/audio_peoplespeech_50k_scratch/checkpoints/0000050000"
    WAVTOKENIZER_CONFIG = "WavTokenizer/configs/wavtokenizer_smalldata_frame40_3s_nq1_code4096_dim512_kmeans200_attn.yaml"
    WAVTOKENIZER_MODEL = "WavTokenizer/WavTokenizer_small_600_24k_4096.ckpt"
    TEST_AUDIO_DIR = Path("data/test-10")
    OUTPUT_DIR = Path("data/test-10/precomputed")

    print("=" * 70)
    print("BLT Patch Visualization - Precomputation")
    print("=" * 70)
    print(f"\nEntropy Model: {ENTROPY_CHECKPOINT}")
    print(f"Audio Files: {TEST_AUDIO_DIR}")
    print(f"Output: {OUTPUT_DIR}")

    # Create output directory
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

    # Setup device
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print(f"\nUsing device: {device}")

    # Load models
    print("\n" + "=" * 70)
    print("Loading Models")
    print("=" * 70)

    print("\nLoading WavTokenizer...")
    wavtokenizer = load_wavtokenizer(WAVTOKENIZER_CONFIG, WAVTOKENIZER_MODEL, device)
    print("✓ WavTokenizer loaded")

    print("\nLoading Entropy Model...")
    entropy_model = load_entropy_model(ENTROPY_CHECKPOINT, device)
    print("✓ Entropy Model loaded")

    # Process all audio files
    print("\n" + "=" * 70)
    print("Processing Audio Files")
    print("=" * 70)

    audio_files = sorted(TEST_AUDIO_DIR.glob("*.flac"))
    print(f"\nFound {len(audio_files)} audio files")

    results = []
    for i, audio_path in enumerate(audio_files, 1):
        print(f"\n[{i}/{len(audio_files)}] {audio_path.name}")
        try:
            output_file = process_audio_file(
                audio_path, entropy_model, wavtokenizer, device, OUTPUT_DIR
            )
            results.append({
                "file_id": audio_path.stem,
                "status": "success",
                "output_file": str(output_file)
            })
        except Exception as e:
            print(f"  ✗ Error: {e}")
            results.append({
                "file_id": audio_path.stem,
                "status": "error",
                "error": str(e)
            })

    # Summary
    print("\n" + "=" * 70)
    print("Summary")
    print("=" * 70)

    successful = sum(1 for r in results if r["status"] == "success")
    failed = sum(1 for r in results if r["status"] == "error")

    print(f"\nProcessed: {len(results)} files")
    print(f"  ✓ Successful: {successful}")
    if failed > 0:
        print(f"  ✗ Failed: {failed}")

    # Calculate total size
    total_size = sum(
        (OUTPUT_DIR / f"{r['file_id']}.json").stat().st_size
        for r in results if r["status"] == "success"
    )
    print(f"\nTotal precomputed data size: {total_size / 1024:.1f} KB")
    print(f"Average per file: {total_size / successful / 1024:.1f} KB")

    print(f"\n✓ Precomputation complete!")
    print(f"  Data saved to: {OUTPUT_DIR}/")
    print(f"\nYou can now run the Gradio app:")
    print(f"  uv run python -m patch_view.app")


if __name__ == "__main__":
    main()

	#!/usr/bin/env python3
	"""
	Precompute entropy and perplexity for test-10 audio samples.

	This script loads the entropy model once and processes all 10 audio files,
	saving tokens, entropies, and perplexities to JSON files for the Gradio app.

	Usage:
	uv run python apps/patch_view/scripts/precompute_patch_view.py
	"""

	import sys
	import torch
	import numpy as np
	import json
	import math
	from pathlib import Path
	import torchaudio

	# Add project root and WavTokenizer to path
	# __file__ is in apps/patch_view/scripts/, so climb three levels to project root
	PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent
	sys.path.insert(0, str(PROJECT_ROOT))
	sys.path.insert(0, str(PROJECT_ROOT / "WavTokenizer"))

	from encoder.utils import convert_audio
	from decoder.pretrained import WavTokenizer
	from bytelatent.transformer import LMTransformer, LMTransformerArgs
	from torch.distributed.checkpoint.state_dict import get_model_state_dict, set_model_state_dict
	import torch.distributed.checkpoint as dcp


	def load_entropy_model(checkpoint_dir, device):
	"""Load entropy model from checkpoint (step 50000)."""
	print(f"Loading entropy model from {checkpoint_dir}...")

	params_path = Path(checkpoint_dir) / "params.json"
	with open(params_path) as f:
	params = json.load(f)

	model_args = LMTransformerArgs(**params["entropy_model"])
	model = LMTransformer(model_args)

	# Load checkpoint
	state_dict = {"model": get_model_state_dict(model)}
	dcp.load(state_dict, checkpoint_id=checkpoint_dir)
	set_model_state_dict(model, state_dict["model"])

	model = model.to(device)
	model.eval()
	num_params = sum(p.numel() for p in model.parameters())
	print(f" ✓ Entropy model loaded ({num_params:,} parameters)")

	return model


	def load_wavtokenizer(config_path, model_path, device):
	"""Load WavTokenizer for audio encoding."""
	wavtokenizer = WavTokenizer.from_pretrained0802(config_path, model_path)
	wavtokenizer = wavtokenizer.to(device)
	wavtokenizer.eval()
	return wavtokenizer


	def encode_audio(audio_path, wavtokenizer, device):
	"""
	Load and encode audio file to tokens.

	Returns:
	tokens: List of token IDs
	duration: Duration in seconds
	original_wav: Original audio (numpy array)
	reconstructed_wav: WavTokenizer reconstructed audio (numpy array)
	"""
	# Load audio using torchaudio
	wav, sr = torchaudio.load(audio_path)

	# Convert to 24kHz mono
	wav = convert_audio(wav, sr, 24000, 1)

	# Normalize
	wav = wav / (torch.abs(wav).max() + 1e-8)

	# Keep original for comparison
	original_wav = wav.squeeze().cpu().numpy()

	# Encode to tokens
	bandwidth_id = torch.tensor([0]).to(device)
	with torch.no_grad():
	_, discrete_codes = wavtokenizer.encode_infer(wav.to(device), bandwidth_id=bandwidth_id)

	# discrete_codes: [n_q, B, seq_len] - extract first codebook, first batch
	tokens = discrete_codes[0, 0].cpu()

	# Decode tokens to reconstruct audio
	with torch.no_grad():
	tokens_for_decode = tokens.unsqueeze(0).unsqueeze(0).to(device) # [1, 1, seq_len]
	features = wavtokenizer.codes_to_features(tokens_for_decode)
	reconstructed = wavtokenizer.decode(features, bandwidth_id=bandwidth_id)
	reconstructed_wav = reconstructed.squeeze().cpu().numpy()

	tokens_list = tokens.tolist()
	duration = len(tokens_list) / 40.0 # 40 tokens/sec

	return tokens_list, duration, original_wav, reconstructed_wav


	def compute_entropies(tokens, entropy_model, device, max_context=512):
	"""
	Compute entropy for each token position using batched contexts.

	Args:
	tokens: List of token IDs
	entropy_model: Loaded entropy model
	max_context: Maximum context window used per forward pass

	Returns:
	numpy array of entropy values (bits) aligned with token positions
	"""
	num_tokens = len(tokens)
	print(f" Computing entropies for {num_tokens} tokens...")

	if num_tokens == 0:
	return np.array([])

	token_tensor = torch.tensor(tokens, dtype=torch.long, device=device)
	entropies = torch.zeros(num_tokens, dtype=torch.float32, device=device)

	chunk_size = max(1, max_context // 2)

	with torch.no_grad():
	processed = 0
	while processed < num_tokens:
	chunk_end = min(num_tokens, processed + chunk_size)
	context_start = max(0, chunk_end - max_context)
	context_tokens = token_tensor[context_start:chunk_end].unsqueeze(0)

	logits = entropy_model(context_tokens)
	probs = torch.softmax(logits[0], dim=-1)
	log_probs = torch.log(probs + 1e-10) / math.log(2)
	chunk_entropy = -(probs * log_probs).sum(dim=-1)

	new_count = chunk_end - processed
	entropies[processed:chunk_end] = chunk_entropy[-new_count:]

	processed = chunk_end
	print(f" Processed {processed}/{num_tokens} tokens", end="\r")

	print() # newline after progress carriage-return
	return entropies.cpu().numpy()


	def process_audio_file(audio_path, entropy_model, wavtokenizer, device, output_dir):
	"""
	Process single audio file: encode, compute entropies, save JSON and reconstructed audio.

	Args:
	audio_path: Path to FLAC file
	entropy_model: Loaded entropy model
	wavtokenizer: Loaded WavTokenizer
	device: torch device
	output_dir: Directory to save JSON

	Returns:
	Path to saved JSON file
	"""
	import soundfile as sf

	file_id = audio_path.stem
	print(f"\nProcessing {file_id}...")

	# Load transcript if available
	txt_path = audio_path.with_suffix(".txt")
	transcript = txt_path.read_text().strip() if txt_path.exists() else ""

	# Encode audio and get reconstruction
	print(" Encoding audio...")
	tokens, duration, original_wav, reconstructed_wav = encode_audio(
	audio_path, wavtokenizer, device
	)
	print(f" ✓ Encoded {len(tokens)} tokens ({duration:.2f}s)")

	# Save reconstructed audio for comparison
	reconstructed_path = output_dir / f"{file_id}_reconstructed.wav"
	sf.write(reconstructed_path, reconstructed_wav, 24000)
	print(f" ✓ Saved reconstructed audio ({reconstructed_path.stat().st_size / 1024:.1f} KB)")

	# Compute entropies
	entropies = compute_entropies(tokens, entropy_model, device)
	print(f" ✓ Computed entropies")

	# Compute perplexities: PPL = 2^H
	perplexities = 2 ** entropies
	print(f" ✓ Computed perplexities")

	# Create output data
	data = {
	"file_id": file_id,
	"audio_path": str(audio_path),
	"reconstructed_audio_path": str(reconstructed_path),
	"transcript": transcript,
	"audio_duration_sec": float(duration),
	"sample_rate": 24000,
	"tokens": tokens,
	"entropies": entropies.tolist(),
	"perplexities": perplexities.tolist(),
	"num_tokens": len(tokens),
	"avg_entropy": float(np.mean(entropies)),
	"min_entropy": float(np.min(entropies)),
	"max_entropy": float(np.max(entropies)),
	"avg_perplexity": float(np.mean(perplexities)),
	}

	# Save to JSON
	output_file = output_dir / f"{file_id}.json"
	with open(output_file, 'w') as f:
	json.dump(data, f, indent=2)

	print(f" ✓ Saved to {output_file} ({output_file.stat().st_size / 1024:.1f} KB)")

	return output_file


	def main():
	# Configuration
	ENTROPY_CHECKPOINT = "archive/checkpoints/audio_peoplespeech_50k_scratch/checkpoints/0000050000"
	WAVTOKENIZER_CONFIG = "WavTokenizer/configs/wavtokenizer_smalldata_frame40_3s_nq1_code4096_dim512_kmeans200_attn.yaml"
	WAVTOKENIZER_MODEL = "WavTokenizer/WavTokenizer_small_600_24k_4096.ckpt"
	TEST_AUDIO_DIR = Path("data/test-10")
	OUTPUT_DIR = Path("data/test-10/precomputed")

	print("=" * 70)
	print("BLT Patch Visualization - Precomputation")
	print("=" * 70)
	print(f"\nEntropy Model: {ENTROPY_CHECKPOINT}")
	print(f"Audio Files: {TEST_AUDIO_DIR}")
	print(f"Output: {OUTPUT_DIR}")

	# Create output directory
	OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

	# Setup device
	device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
	print(f"\nUsing device: {device}")

	# Load models
	print("\n" + "=" * 70)
	print("Loading Models")
	print("=" * 70)

	print("\nLoading WavTokenizer...")
	wavtokenizer = load_wavtokenizer(WAVTOKENIZER_CONFIG, WAVTOKENIZER_MODEL, device)
	print("✓ WavTokenizer loaded")

	print("\nLoading Entropy Model...")
	entropy_model = load_entropy_model(ENTROPY_CHECKPOINT, device)
	print("✓ Entropy Model loaded")

	# Process all audio files
	print("\n" + "=" * 70)
	print("Processing Audio Files")
	print("=" * 70)

	audio_files = sorted(TEST_AUDIO_DIR.glob("*.flac"))
	print(f"\nFound {len(audio_files)} audio files")

	results = []
	for i, audio_path in enumerate(audio_files, 1):
	print(f"\n[{i}/{len(audio_files)}] {audio_path.name}")
	try:
	output_file = process_audio_file(
	audio_path, entropy_model, wavtokenizer, device, OUTPUT_DIR
	)
	results.append({
	"file_id": audio_path.stem,
	"status": "success",
	"output_file": str(output_file)
	})
	except Exception as e:
	print(f" ✗ Error: {e}")
	results.append({
	"file_id": audio_path.stem,
	"status": "error",
	"error": str(e)
	})

	# Summary
	print("\n" + "=" * 70)
	print("Summary")
	print("=" * 70)

	successful = sum(1 for r in results if r["status"] == "success")
	failed = sum(1 for r in results if r["status"] == "error")

	print(f"\nProcessed: {len(results)} files")
	print(f" ✓ Successful: {successful}")
	if failed > 0:
	print(f" ✗ Failed: {failed}")

	# Calculate total size
	total_size = sum(
	(OUTPUT_DIR / f"{r['file_id']}.json").stat().st_size
	for r in results if r["status"] == "success"
	)
	print(f"\nTotal precomputed data size: {total_size / 1024:.1f} KB")
	print(f"Average per file: {total_size / successful / 1024:.1f} KB")

	print(f"\n✓ Precomputation complete!")
	print(f" Data saved to: {OUTPUT_DIR}/")
	print(f"\nYou can now run the Gradio app:")
	print(f" uv run python -m patch_view.app")


	if __name__ == "__main__":
	main()