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
try:
    # torch.hub handles cloning the repo internally
    codec = torch.hub.load(
        repo_or_dir="lucadellalib/focalcodec", 
        model="focalcodec", 
        config=MODEL_CONFIG, 
        force_reload=False # Use cached version after first load
    )
    codec.eval().requires_grad_(False) # Set to evaluation mode

    if torch.cuda.is_available():
        codec.cuda()
except Exception as e:
    print(f"Error loading model via torch.hub: {e}")
    codec = 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.
    """
    if codec is None:
        return (16000, None), None

    sr, wav_numpy = audio_input
    
    # Convert numpy to torch tensor and ensure float32
    sig = torch.tensor(wav_numpy, dtype=torch.float32).unsqueeze(0)
    
    # Resample input audio to the sample rate required by the codec (16kHz)
    if sr != codec.sample_rate_input:
        resampler = torchaudio.transforms.Resample(orig_freq=sr, new_freq=codec.sample_rate_input)
        sig = resampler(sig)
    
    # Ensure mono channel if needed
    if sig.shape[0] > 1:
        sig = sig[0, :].unsqueeze(0)

    if torch.cuda.is_available():
        sig = sig.cuda()

    # --- Process (Encode and Decode) ---
    with torch.no_grad():
        # 1. Encode signal to discrete tokens (the compressed data)
        toks = codec.sig_to_toks(sig) 
        
        # 2. Decode tokens back into a waveform
        rec_sig = codec.toks_to_sig(toks)

    # --- Save the compressed tokens to a temporary .fc file ---
    temp_dir = tempfile.mkdtemp()
    fc_file_path = os.path.join(temp_dir, "compressed_tokens.fc")
    # Save the tokens tensor
    torch.save(toks, fc_file_path)
    
    print(f"Tokens saved to {fc_file_path}")

    # Move audio back to CPU for Gradio output and formatting
    # Note: Codec output is already at sample_rate_input (16kHz)
    decoded_wav_output = rec_sig.cpu().numpy().squeeze()
    
    return (codec.sample_rate_output, decoded_wav_output), fc_file_path

# --- Gradio Interface (Use the same Blocks interface as before) ---
with gr.Blocks() as iface:
    gr.Markdown(f"## FocalCodec at 160 bps ({MODEL_CONFIG.split('/')[-1]})")
    gr.Markdown("Test the lowest bitrate neural speech codec! Optimized ONLY for speech. Upload your audio or record your voice.")

    with gr.Row():
        audio_input = gr.Audio(sources=["microphone", "upload"], type="numpy", label="Input Audio (Speech Only Recommended)")
        
        with gr.Column():
            audio_output = gr.Audio(type="numpy", label="Decoded Output Audio (160 bps)")
            file_output = gr.File(label="Download Compressed Tokens (*.fc file)", file_count="single", file_types=[".fc"])

    process_button = gr.Button("Process Audio", variant="primary")
    process_button.click(
        fn=encode_decode_focal,
        inputs=[audio_input],
        outputs=[audio_output, file_output]
    )

if __name__ == "__main__":
    iface.launch()