novasr-candle

NovaSR Candle Port

A Rust port of NovaSR - a lightning-fast audio super-resolution model - using Hugging Face Candle.

Overview

NovaSR is a tiny 52KB model that upsamples 16kHz audio to crystal-clear 48kHz at speeds exceeding 3600x realtime. This project ports the original PyTorch implementation to Rust using the Candle deep learning framework.

Features

• Tiny Model: Only ~13,000 parameters (52KB)
• Blazing Fast: 3600x realtime inference on GPU
• High Quality: 16kHz → 48kHz upsampling
• Pure Rust: No Python dependencies for inference
• WASM Compatible: Can run in the browser

Architecture

Input (16kHz) → ConvPre → Interpolate (3x) → AMPBlock0 → ConvPost → Tanh → Output (48kHz)

Key Components

1. SnakeBeta Activation: Novel activation function for periodic signals

   x + (1 - cos(2 * α * x)) / (2 * β + ε)
   

2. AMPBlock0: Residual block with dilated convolutions and SnakeBeta activation

3. Generator: Main upsampling network with pre/post convolutions

Project Structure

├── src/ │ ├── lib.rs # Core library │ └── main.rs # CLI tool ├── models/ # Safetensors weights ├── scripts/ # Utility scripts (conversion, parity) ├── data/ # Audio samples ├── Cargo.toml └── README.md

Installation

Prerequisites

• Rust 1.70+
• For weight conversion: Python 3.8+ with PyTorch and safetensors

Build

# Clone the repository
git clone <repo-url>
cd novasr-candle


# Build the library and CLI
cargo build --release

# The binary will be at target/release/novasr-cli

Usage

Convert Weights

First, convert the PyTorch weights to Candle format:

uv run --project NovaSR python scripts/convert_weights.py \
  --input novasr_model.pth \
  --output models/novasr_v1.safetensors

CLI Usage

# Upsample an audio file using a local model
./target/release/novasr-cli input.wav output.wav models/novasr_v1.safetensors

# Upsample using a model from Hugging Face Hub
./target/release/novasr-cli input.wav output.wav babybirdprd/novasr-candle

Library Usage

use candle_core::{Device, DType};
use candle_nn::VarBuilder;
use novasr_candle::{load_model, upsample_audio};

fn main() -> anyhow::Result<()> {
    let device = Device::Cpu;
    
    // Load from HF Hub
    let model = novasr_candle::from_hf("babybirdprd/novasr-candle", "main", &device)?;
    
    // OR load local
    // let vb = unsafe { VarBuilder::from_mmaped_safetensors("model.safetensors", DType::F32, &device)? };
    // let model = load_model(vb)?;

    
    // Process audio
    let input = candle_core::Tensor::from_vec(
        audio_samples,
        (1, 1, sample_count),
        &device,
    )?;
    
    let output = upsample_audio(&model, &input)?;
    
    Ok(())
}

Implementation Details

SnakeBeta Activation

The SnakeBeta activation is a learnable periodic activation function:

pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
    let a = self.alpha.exp()?;
    let b = self.beta.exp()?;
    
    // x + (1 - cos(2 * alpha * x)) / (2 * beta + epsilon)
    let two_a_x = (x.broadcast_mul(&a)? * 2.0)?;
    let cos_term = two_a_x.cos()?;
    let one_minus_cos = (Tensor::ones_like(&cos_term)? - cos_term)?;
    let inv_2b = ((b * 2.0)? + 1e-9)?.recip()?;
    
    x.add(&one_minus_cos.broadcast_mul(&inv_2b)?)
}

Generator

The Generator performs the main upsampling:

1. Pre-convolution: 7x1 conv to expand channels
2. Interpolation: 3x linear upsampling
3. Residual blocks: AMPBlock0 with SnakeBeta activations
4. Post-convolution: 7x1 conv to output channel
5. Tanh: Output normalization

Model Specifications

Comparison with Original

Web Demo

A web demo showcasing the model architecture and Rust implementation is available at:

https://vkf4cascsh44i.ok.kimi.link

Future Work

• WASM bindings for browser inference
• Quantization support (INT8)
• GPU acceleration (CUDA)
• Streaming inference for real-time processing
• Batch processing support

Credits

• Original NovaSR by Yatharth Sharma
• Candle by Hugging Face

License

Apache 2.0 - Same as the original NovaSR project