app.py

import torch
import torchaudio
import gradio as gr
import os
import tempfile
import numpy as np

# Define the model ID for the 0.16 kbps codec config
MODEL_CONFIG = "lucadellalib/focalcodec_12_5hz"

# Load the model globally using torch.hub
codec = None
try:
    print("Loading FocalCodec model...")
    codec = torch.hub.load(
        repo_or_dir="lucadellalib/focalcodec",
        model="focalcodec",
        config=MODEL_CONFIG,
        force_reload=False,
        trust_repo=True
    )
    codec.eval()
    for param in codec.parameters():
        param.requires_grad = False
    
    if torch.cuda.is_available():
        codec = codec.cuda()
        print("Model loaded successfully on GPU!")
    else:
        print("Model loaded successfully on CPU!")
    
except Exception as e:
    print(f"ERROR loading model via torch.hub: {e}")
    print("\nTrying alternative installation method...")
    try:
        import subprocess
        subprocess.check_call(["pip", "install", "focalcodec@git+https://github.com/lucadellalib/focalcodec.git@main"])
        import focalcodec
        codec = focalcodec.FocalCodec.from_pretrained(MODEL_CONFIG)
        codec.eval()
        for param in codec.parameters():
            param.requires_grad = False
        if torch.cuda.is_available():
            codec = codec.cuda()
        print("Model loaded via pip installation!")
    except Exception as e2:
        print(f"ERROR with alternative method: {e2}")
        codec = None


def save_tokens_raw(toks, fc_file_path):
    """Save tokens as raw binary with NO header - pure tokens only"""
    
    toks_cpu = toks.cpu().numpy().flatten()
    max_token = int(toks_cpu.max())
    
    print(f"\n=== Saving Raw Tokens ===")
    print(f"Token shape: {toks.shape}")
    print(f"Token range: 0 to {max_token}")
    print(f"Num tokens: {len(toks_cpu)}")
    
    # Determine bits needed
    if max_token <= 1:
        bits_needed = 1
    elif max_token <= 3:
        bits_needed = 2
    elif max_token <= 7:
        bits_needed = 3
    elif max_token <= 15:
        bits_needed = 4
    elif max_token <= 31:
        bits_needed = 5
    elif max_token <= 63:
        bits_needed = 6
    elif max_token <= 127:
        bits_needed = 7
    elif max_token <= 255:
        bits_needed = 8
    elif max_token <= 511:
        bits_needed = 9
    elif max_token <= 1023:
        bits_needed = 10
    elif max_token <= 2047:
        bits_needed = 11
    elif max_token <= 4095:
        bits_needed = 12
    elif max_token <= 8191:
        bits_needed = 13
    elif max_token <= 16383:
        bits_needed = 14
    elif max_token <= 32767:
        bits_needed = 15
    else:
        bits_needed = 16
    
    print(f"Bits per token: {bits_needed}")
    
    # Create bit array
    bit_array = []
    for tok in toks_cpu:
        bits = format(int(tok), f'0{bits_needed}b')
        bit_array.extend([int(b) for b in bits])
    
    # Pad to byte boundary
    while len(bit_array) % 8 != 0:
        bit_array.append(0)
    
    # Pack into bytes
    packed_bits = np.packbits(np.array(bit_array, dtype=np.uint8))
    
    # Write ONLY the packed data (no header!)
    with open(fc_file_path, 'wb') as f:
        f.write(packed_bits.tobytes())
    
    file_size = os.path.getsize(fc_file_path)
    
    print(f"File size: {file_size} bytes (pure data, no header)")
    print(f"========================\n")
    
    return file_size, bits_needed, len(toks_cpu), toks.shape


def load_tokens_raw(fc_file_path, bits_per_token, num_tokens, original_shape):
    """Load raw tokens from headerless binary file"""
    
    print(f"\n=== Loading Raw Tokens ===")
    print(f"Expected bits/token: {bits_per_token}")
    print(f"Expected num tokens: {num_tokens}")
    print(f"Expected shape: {original_shape}")
    
    # Read all bytes
    with open(fc_file_path, 'rb') as f:
        packed_data = np.frombuffer(f.read(), dtype=np.uint8)
    
    # Unpack bits
    unpacked_bits = np.unpackbits(packed_data)
    
    # Extract exact number of bits needed
    total_bits = num_tokens * bits_per_token
    token_bits = unpacked_bits[:total_bits]
    
    # Reconstruct tokens
    tokens = []
    for i in range(num_tokens):
        start = i * bits_per_token
        end = start + bits_per_token
        token_bits_slice = token_bits[start:end]
        
        # Convert binary to integer
        token_value = 0
        for bit in token_bits_slice:
            token_value = (token_value << 1) | bit
        tokens.append(token_value)
    
    # Reshape to original shape
    tokens_array = np.array(tokens, dtype=np.int64).reshape(original_shape)
    tokens_tensor = torch.from_numpy(tokens_array)
    
    print(f"Loaded tokens: {tokens_tensor.shape}")
    print(f"Token range: {tokens_tensor.min().item()} to {tokens_tensor.max().item()}")
    print(f"==========================\n")
    
    return tokens_tensor


# Global variables to store metadata for decoding
last_encoding_metadata = {
    'bits_per_token': None,
    'num_tokens': None,
    'shape': None,
    'duration': None
}


def encode_decode_focal(audio_input):
    """
    Processes input audio through the 160 bps FocalCodec, saves the tokens,
    and returns both the decoded WAV and the path to the FC file for download.
    """
    global last_encoding_metadata
    
    if codec is None:
        return None, None, "❌ ERROR: Model failed to load. Check console for details."
    
    if audio_input is None:
        return None, None, "❌ Please provide audio input."
    
    try:
        sr, wav_numpy = audio_input
        
        print(f"\n{'='*50}")
        print(f"Processing new audio...")
        print(f"Input audio: sample_rate={sr}, shape={wav_numpy.shape}")
        
        # Handle stereo to mono conversion
        if len(wav_numpy.shape) > 1:
            if wav_numpy.shape[1] == 2:
                wav_numpy = wav_numpy.mean(axis=1)
                print("Converted stereo to mono")
            elif wav_numpy.shape[0] == 2:
                wav_numpy = wav_numpy.mean(axis=0)
                print("Converted stereo to mono (channels first)")
        
        # Ensure float32 and normalize
        wav_numpy = wav_numpy.astype(np.float32)
        if wav_numpy.max() > 1.0 or wav_numpy.min() < -1.0:
            wav_numpy = wav_numpy / 32768.0
        
        # Convert to torch tensor
        sig = torch.from_numpy(wav_numpy).unsqueeze(0)
        
        # Resample to 16kHz
        if sr != codec.sample_rate_input:
            print(f"Resampling from {sr}Hz to {codec.sample_rate_input}Hz...")
            resampler = torchaudio.transforms.Resample(
                orig_freq=sr, 
                new_freq=codec.sample_rate_input
            )
            sig = resampler(sig)
        
        print(f"Signal shape: {sig.shape}")
        
        if torch.cuda.is_available():
            sig = sig.cuda()
        
        # --- Encode and Decode ---
        with torch.no_grad():
            print("\n--- Encoding ---")
            toks = codec.sig_to_toks(sig)
            
            duration_sec = sig.shape[-1] / codec.sample_rate_input
            token_rate = toks.shape[1] / duration_sec
            
            print(f"Tokens shape: {toks.shape}")
            print(f"Token range: {toks.min().item()} to {toks.max().item()}")
            print(f"Duration: {duration_sec:.2f}s")
            print(f"Token rate: {token_rate:.2f} tokens/sec")
            
            print("\n--- Decoding ---")
            rec_sig = codec.toks_to_sig(toks)
            print(f"Reconstructed signal shape: {rec_sig.shape}")
        
        # --- Save raw tokens (no header) ---
        temp_dir = tempfile.mkdtemp()
        fc_file_path = os.path.join(temp_dir, "compressed_tokens.fc")
        
        file_size, bits_per_token, num_tokens, shape = save_tokens_raw(toks, fc_file_path)
        
        # Store metadata globally for decoding
        last_encoding_metadata = {
            'bits_per_token': bits_per_token,
            'num_tokens': num_tokens,
            'shape': shape,
            'duration': duration_sec
        }
        
        # Calculate bitrates
        bitrate = (file_size * 8) / duration_sec
        theoretical_bitrate = token_rate * bits_per_token
        
        print(f"--- Results ---")
        print(f"File bitrate: {bitrate:.1f} bps (pure data)")
        print(f"Theoretical: {theoretical_bitrate:.1f} bps")
        print(f"Target: 160 bps")
        print(f"Efficiency: {(160/bitrate)*100:.1f}% of target")
        print(f"{'='*50}\n")
        
        # Prepare output
        decoded_wav_output = rec_sig.cpu().numpy().squeeze()
        
        if len(decoded_wav_output.shape) == 0:
            decoded_wav_output = decoded_wav_output.reshape(1)
        
        metadata_info = f"\n\nℹ️ SAVE THIS: bits={bits_per_token}, tokens={num_tokens}, shape={shape}"
        status_msg = f"✅ {duration_sec:.1f}s | {file_size}B | {bitrate:.0f} bps | {bits_per_token} bits/tok{metadata_info}"
        
        return (codec.sample_rate_output, decoded_wav_output), fc_file_path, status_msg
    
    except Exception as e:
        error_msg = f"❌ Error: {str(e)}"
        print(error_msg)
        import traceback
        traceback.print_exc()
        return None, None, error_msg


def decode_from_fc_file(fc_file, bits_per_token_input, num_tokens_input, batch_size_input, seq_length_input):
    """Decode audio from uploaded .fc file using provided metadata"""
    
    if codec is None:
        return None, "❌ Model not loaded"
    
    if fc_file is None:
        return None, "❌ Please upload a .fc file"
    
    # Try to use provided metadata, or fall back to last encoding
    try:
        bits_per_token = int(bits_per_token_input) if bits_per_token_input else last_encoding_metadata.get('bits_per_token')
        num_tokens = int(num_tokens_input) if num_tokens_input else last_encoding_metadata.get('num_tokens')
        
        if batch_size_input and seq_length_input:
            shape = (int(batch_size_input), int(seq_length_input))
        else:
            shape = last_encoding_metadata.get('shape')
        
        if not all([bits_per_token, num_tokens, shape]):
            return None, "❌ Please provide metadata (bits/token, num tokens, batch, seq_length) OR encode a file first"
        
    except Exception as e:
        return None, f"❌ Invalid metadata format: {str(e)}"
    
    try:
        print(f"\n{'='*50}")
        print(f"Decoding from file: {fc_file.name}")
        
        # Load tokens
        toks = load_tokens_raw(fc_file.name, bits_per_token, num_tokens, shape)
        
        if torch.cuda.is_available():
            toks = toks.cuda()
        
        # Decode to audio
        with torch.no_grad():
            print("Decoding tokens to audio...")
            rec_sig = codec.toks_to_sig(toks)
            print(f"Reconstructed signal shape: {rec_sig.shape}")
        
        decoded_wav = rec_sig.cpu().numpy().squeeze()
        
        # Calculate stats
        duration_sec = decoded_wav.shape[0] / codec.sample_rate_output
        file_size = os.path.getsize(fc_file.name)
        bitrate = (file_size * 8) / duration_sec
        
        print(f"Duration: {duration_sec:.2f}s")
        print(f"Bitrate: {bitrate:.1f} bps")
        print(f"{'='*50}\n")
        
        status = f"✅ Decoded! {duration_sec:.1f}s | {bitrate:.0f} bps | {bits_per_token} bits/token"
        
        return (codec.sample_rate_output, decoded_wav), status
    
    except Exception as e:
        import traceback
        traceback.print_exc()
        return None, f"❌ Error: {str(e)}"


# --- Gradio Interface ---
with gr.Blocks(title="FocalCodec 160 bps", theme=gr.themes.Soft()) as iface:
    gr.Markdown("# 🎙️ FocalCodec at 160 bps")
    gr.Markdown(f"**Neural speech codec at insanely low bitrate!** Using `{MODEL_CONFIG}`")
    gr.Markdown("⚠️ **Optimized for speech only** | 🔥 **Pure tokens, no header overhead!**")
    
    with gr.Tab("🎤 Encode Audio"):
        gr.Markdown("### Compress audio to ~160 bps (pure tokens, no header)")
        
        with gr.Row():
            audio_input = gr.Audio(
                sources=["microphone", "upload"], 
                type="numpy", 
                label="Input Audio (any format/sample rate)"
            )
            
            with gr.Column():
                audio_output = gr.Audio(
                    type="numpy", 
                    label="🔊 Decoded Output (16kHz)"
                )
                file_output = gr.File(
                    label="💾 Download Compressed .fc File (headerless)"
                )
                status_output = gr.Textbox(label="📊 Status", lines=4)
        
        encode_btn = gr.Button("🔄 Encode & Decode", variant="primary", size="lg")
        encode_btn.click(
            fn=encode_decode_focal,
            inputs=[audio_input],
            outputs=[audio_output, file_output, status_output]
        )
        
        gr.Markdown("### ⚠️ Important:")
        gr.Markdown("- The .fc file contains ONLY raw token data (no metadata/header)")
        gr.Markdown("- **Save the metadata** from the status message to decode later!")
        gr.Markdown("- You need: bits per token, number of tokens, and shape")
    
    with gr.Tab("📂 Decode from .fc File"):
        gr.Markdown("### Decode raw .fc file (requires metadata)")
        
        with gr.Row():
            with gr.Column():
                fc_input = gr.File(
                    label="Upload .fc File", 
                    file_types=[".fc"]
                )
                
                gr.Markdown("#### Metadata (required for decoding):")
                
                with gr.Row():
                    bits_input = gr.Number(
                        label="Bits per token",
                        value=13,
                        precision=0,
                        info="Usually 13 for this model"
                    )
                    tokens_input = gr.Number(
                        label="Number of tokens",
                        precision=0,
                        info="Total tokens in file"
                    )
                
                with gr.Row():
                    batch_input = gr.Number(
                        label="Batch size",
                        value=1,
                        precision=0,
                        info="Usually 1"
                    )
                    seq_input = gr.Number(
                        label="Sequence length",
                        precision=0,
                        info="Tokens per batch"
                    )
                
                gr.Markdown("💡 If you just encoded a file, leave these blank to use saved metadata")
            
            with gr.Column():
                decoded_output = gr.Audio(
                    type="numpy", 
                    label="🔊 Decoded Audio"
                )
                decode_status = gr.Textbox(label="📊 Status", lines=2)
        
        decode_btn = gr.Button("🔊 Decode Audio", variant="primary", size="lg")
        decode_btn.click(
            fn=decode_from_fc_file,
            inputs=[fc_input, bits_input, tokens_input, batch_input, seq_input],
            outputs=[decoded_output, decode_status]
        )
    
    with gr.Tab("ℹ️ About"):
        gr.Markdown("""
        ## FocalCodec - Ultra Low Bitrate Neural Audio Codec
        
        ### 🎯 Pure Token Format (No Headers!)
        
        This version saves **ONLY the compressed tokens** with no metadata overhead.
        
        **Benefits:**
        - ✅ Absolute minimum file size
        - ✅ True 160 bps (no header padding)
        - ✅ Maximum compression efficiency
        
        **Trade-off:**
        - ⚠️ You must save the metadata separately to decode
        - Required info: bits per token, number of tokens, shape
        
        ### 📊 Compression Ratios:
        | Format | Bitrate | 1-Hour File Size |
        |--------|---------|------------------|
        | Uncompressed PCM | 256 kbps | ~115 MB |
        | MP3 | 128 kbps | ~57 MB |
        | Opus | 16 kbps | ~7.2 MB |
        | **FocalCodec** | **0.16 kbps** | **~72 KB** 🔥 |
        
        ### 🔧 Technical Details:
        - **Token Rate:** ~12.5 tokens/sec
        - **Bits per Token:** 13 bits (for most speech)
        - **Bitrate:** 12.5 × 13 = 162.5 bps ≈ **160 bps**
        - **Format:** Raw bit-packed tokens (no header)
        
        ### 📝 Example Metadata:
        After encoding, you'll see:
        ```
        ℹ️ SAVE THIS: bits=13, tokens=113, shape=(1, 113)
        ```
        
        Save this to decode the file later!
        
        ### 💡 Pro Tip:
        If you're building a system, embed the metadata in a separate JSON file:
        ```json
        {
          "audio.fc": {
            "bits_per_token": 13,
            "num_tokens": 113,
            "shape": [1, 113],
            "duration": 9.04
          }
        }
        ```
        
        ---
        
        🔗 [FocalCodec GitHub](https://github.com/lucadellalib/focalcodec)
        """)

if __name__ == "__main__":
    print("\n" + "="*50)
    print("🎙️  FocalCodec 160 bps Demo (Headerless Format)")
    print("="*50 + "\n")
    iface.launch()