app.py

import os
import gradio as gr
import numpy as np
import soundfile as sf
from semanticodec import SemantiCodec
from huggingface_hub import HfApi
import spaces
import torch
import tempfile
import io
import uuid
from pathlib import Path

# Initialize the model
def load_model():
    return SemantiCodec(token_rate=100, semantic_vocab_size=32768)  # 1.40 kbps

semanticodec = load_model()

@spaces.GPU(duration=20)
def encode_audio(audio_path):
    """Encode audio file to tokens and return them as a file"""
    try:
        tokens = semanticodec.encode(audio_path)
        # Move tokens to CPU before converting to numpy
        if isinstance(tokens, torch.Tensor):
            tokens = tokens.cpu().numpy()
        
        # Ensure tokens are in the right shape for later decoding
        if tokens.ndim == 1:
            # Reshape to match expected format [batch, seq_len, features]
            tokens = tokens.reshape(1, -1, 1)
        
        # Create a temporary file in /tmp which is writable in Spaces
        temp_dir = "/tmp"
        os.makedirs(temp_dir, exist_ok=True)
        temp_file_path = os.path.join(temp_dir, f"tokens_{uuid.uuid4()}.oterin")
        
        # Save tokens directly to file
        np.save(temp_file_path, tokens)
        
        # Verify the file exists and has content
        if not os.path.exists(temp_file_path) or os.path.getsize(temp_file_path) == 0:
            raise Exception("Failed to create token file")
        
        return temp_file_path, f"Encoded to {tokens.shape[1]} tokens"
    except Exception as e:
        return None, f"Error encoding audio: {str(e)}"

@spaces.GPU(duration=60)
def decode_tokens(token_file):
    """Decode tokens to audio"""
    # Ensure the file exists and has content
    if not token_file or not os.path.exists(token_file):
        return None, "Error: Empty or missing token file"
    
    try:
        # Load tokens from file
        tokens = np.load(token_file, allow_pickle=True)
        
        # Convert to torch tensor with proper dimensions
        if isinstance(tokens, np.ndarray):
            # Ensure tokens are in the right shape
            if tokens.ndim == 1:
                # Reshape to match expected format [batch, seq_len, features]
                tokens = tokens.reshape(1, -1, 1)
            
            # Convert to torch tensor
            tokens = torch.tensor(tokens)
        
        # Ensure tokens are on the right device
        if torch.cuda.is_available():
            tokens = tokens.cuda()
        
        # Decode the tokens
        waveform = semanticodec.decode(tokens)
        
        # Move waveform to CPU if it's a tensor
        if isinstance(waveform, torch.Tensor):
            waveform = waveform.cpu().numpy()
        
        # Create in-memory file for audio
        output_buffer = io.BytesIO()
        sf.write(output_buffer, waveform[0, 0], 32000, format='WAV')
        output_buffer.seek(0)
        
        # Verify the buffer has content
        if output_buffer.getbuffer().nbytes == 0:
            return None, "Error: Failed to generate audio"
        
        return output_buffer, f"Decoded {tokens.shape[1]} tokens to audio"
    except Exception as e:
        return None, f"Error decoding tokens: {str(e)}"

@spaces.GPU(duration=80)
def process_both(audio_path):
    """Encode and then decode the audio without saving intermediate files"""
    try:
        # Encode
        tokens = semanticodec.encode(audio_path)
        if isinstance(tokens, torch.Tensor):
            tokens = tokens.cpu().numpy()
        
        # Ensure tokens are in the right shape for decoding
        if tokens.ndim == 1:
            # Reshape to match expected format [batch, seq_len, features]
            tokens = tokens.reshape(1, -1, 1)
        
        # Convert back to tensor for decoding
        tokens_tensor = torch.tensor(tokens)
        
        # Ensure tokens are on the right device
        if torch.cuda.is_available():
            tokens_tensor = tokens_tensor.cuda()
        
        # Decode
        waveform = semanticodec.decode(tokens_tensor)
        if isinstance(waveform, torch.Tensor):
            waveform = waveform.cpu().numpy()
        
        # Create in-memory file
        output_buffer = io.BytesIO()
        sf.write(output_buffer, waveform[0, 0], 32000, format='WAV')
        output_buffer.seek(0)
        
        # Verify the buffer has content
        if output_buffer.getbuffer().nbytes == 0:
            return None, "Error: Failed to generate audio"
        
        return output_buffer, f"Encoded to {tokens.shape[1]} tokens\nDecoded {tokens.shape[1]} tokens to audio"
    except Exception as e:
        return None, f"Error processing audio: {str(e)}"

# Create Gradio interface
with gr.Blocks(title="Oterin Audio Codec") as demo:
    gr.Markdown("# Oterin Audio Codec")
    gr.Markdown("Upload an audio file to encode it to semantic tokens, decode tokens back to audio, or do both.")
    
    with gr.Tab("Encode Audio"):
        with gr.Row():
            encode_input = gr.Audio(type="filepath", label="Input Audio")
            encode_output = gr.File(label="Encoded Tokens (.oterin)", file_types=[".oterin"])
        encode_status = gr.Textbox(label="Status")
        encode_btn = gr.Button("Encode")
        encode_btn.click(encode_audio, inputs=encode_input, outputs=[encode_output, encode_status])
    
    with gr.Tab("Decode Tokens"):
        with gr.Row():
            decode_input = gr.File(label="Token File (.oterin)", file_types=[".oterin"])
            decode_output = gr.Audio(label="Decoded Audio")
        decode_status = gr.Textbox(label="Status")
        decode_btn = gr.Button("Decode")
        decode_btn.click(decode_tokens, inputs=decode_input, outputs=[decode_output, decode_status])
    
    with gr.Tab("Both (Encode & Decode)"):
        with gr.Row():
            both_input = gr.Audio(type="filepath", label="Input Audio")
            both_output = gr.Audio(label="Reconstructed Audio")
        both_status = gr.Textbox(label="Status")
        both_btn = gr.Button("Process")
        both_btn.click(process_both, inputs=both_input, outputs=[both_output, both_status])

if __name__ == "__main__":
    demo.launch(share=True)