File size: 5,112 Bytes
4510d70 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 | """
HuggingFace Inference Endpoint Handler for SongFormer
Supports binary audio input (WAV, MP3, etc.) via base64 encoding or direct bytes
"""
import os
import sys
import io
import base64
import json
import tempfile
from typing import Dict, Any, Union
import librosa
import numpy as np
import torch
from transformers import AutoModel
class EndpointHandler:
"""
HuggingFace Inference Endpoint Handler for SongFormer model.
Accepts base64-encoded audio (WAV, MP3, FLAC, etc.)
"""
def __init__(self, path: str = ""):
"""
Initialize the handler and load the SongFormer model.
Args:
path: Path to the model directory (provided by HuggingFace)
"""
# Set up environment
self.model_path = path or os.getcwd()
os.environ["SONGFORMER_LOCAL_DIR"] = self.model_path
sys.path.insert(0, self.model_path)
# Import after setting up path
# Load the model
self.device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Loading SongFormer model on {self.device}...")
# Load model without device_map to avoid meta device initialization
# The SongFormerModel.__init__ now handles meta device detection
self.model = AutoModel.from_pretrained(
self.model_path,
trust_remote_code=True,
device_map=None,
)
self.model.to(self.device)
self.model.eval()
# Expected sampling rate for the model
self.target_sr = 24000
print("SongFormer model loaded successfully!")
def _decode_base64_audio(self, audio_b64: str) -> np.ndarray:
"""
Decode base64-encoded audio to numpy array.
Args:
audio_b64: Base64-encoded audio string
Returns:
numpy array of audio samples at 24kHz
"""
# Decode base64 string to bytes
try:
audio_bytes = base64.b64decode(audio_b64)
except Exception as e:
raise ValueError(f"Failed to decode base64 audio data: {e}")
# Load audio from bytes using librosa
# Create a file-like object from bytes
audio_io = io.BytesIO(audio_bytes)
# Load with librosa (automatically handles WAV, MP3, etc.)
audio_array, _ = librosa.load(audio_io, sr=self.target_sr, mono=True)
return audio_array
def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
"""
Process inference request with base64-encoded audio.
Expected input:
{
"inputs": "<base64-encoded-audio-data>"
}
Returns:
{
"segments": [
{
"label": "intro",
"start": 0.0,
"end": 15.2
},
...
],
"duration": 180.5,
"num_segments": 8
}
"""
try:
# Extract base64-encoded audio
audio_b64 = data.get("inputs")
if not audio_b64:
raise ValueError("Missing 'inputs' key with base64-encoded audio")
if not isinstance(audio_b64, str):
raise ValueError("Input must be a base64-encoded string")
# Decode audio
audio_array = self._decode_base64_audio(audio_b64)
# Run inference
with torch.no_grad():
result = self.model(audio_array)
# Calculate duration
duration = len(audio_array) / self.target_sr
# Format output
output = {
"segments": result,
"duration": float(duration),
"num_segments": len(result)
}
return output
except Exception as e:
# Return error in a structured format
return {
"error": str(e),
"error_type": type(e).__name__,
"segments": [],
"duration": 0.0,
"num_segments": 0
}
# For local testing
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(description="Test SongFormer handler locally")
parser.add_argument("audio_file", help="Path to audio file to test")
parser.add_argument("--model-path", default=".", help="Path to model directory")
args = parser.parse_args()
# Initialize handler
handler = EndpointHandler(args.model_path)
# Read and encode audio file
with open(args.audio_file, "rb") as f:
audio_bytes = f.read()
audio_b64 = base64.b64encode(audio_bytes).decode('utf-8')
# Test with base64 input
print("\n=== Testing with base64-encoded audio ===")
result = handler({"inputs": audio_b64})
print(json.dumps(result, indent=2))
# Test with file path directly (for comparison)
print("\n=== Testing with direct file path (not typical for endpoint) ===")
result_direct = handler.model(args.audio_file)
print(json.dumps(result_direct, indent=2))
|