VoiceFocus / sdk.py
mariesig
update to aic-sdk 2.0
abae3b9
Raw History Blame
2.02 kB
from platform import processor
import numpy as np
from dotenv import load_dotenv
import aic_sdk as aic
import librosa
import soundfile as sf
load_dotenv()
class SDKParams:
def __init__(self, enhancement_level: float, sdk_key: str, model_id: str = "quail-vf-1.1-l-16khz"):
self.enhancement_level = enhancement_level
self.sdk_key = sdk_key
self.model_path= aic.Model.download(model_id, "./models")
def process_file_sdk(input_path: str, output_path: str, sdk_params: SDKParams):
# Load audio
audio, sr = librosa.load(input_path, sr=48000, mono=True)
audio = audio.reshape(1, -1) # Convert to planar format
output = np.zeros_like(audio)
load_dotenv()
model = aic.Model.from_file(sdk_params.model_path)
optimal_frames = model.get_optimal_num_frames(48000)
config = aic.ProcessorConfig(sample_rate=sr, num_channels=1, num_frames= optimal_frames,allow_variable_frames=False )
processor = aic.Processor(model, sdk_params.sdk_key, config)
proc_ctx = processor.get_processor_context()
#proc_ctx.set_parameter(aic.ProcessorParameter.EnhancementLevel, sdk_params.enhancement_level)
#proc_ctx.set_parameter(aic.ProcessorParameter.VoiceGain, 1.5)
#proc_ctx.set_parameter(aic.ProcessorParameter.Bypass, 0.0)
chunk_size = optimal_frames
for i in range(0, audio.shape[1], chunk_size):
chunk = audio[:, i : i + chunk_size]
# Pad last chunk if needed
if chunk.shape[1] < chunk_size:
last_chunk_size = chunk.shape[1]
padded = np.zeros((1, chunk_size), dtype=audio.dtype)
padded[:, : chunk.shape[1]] = chunk
chunk = padded
enhanced_chunk = processor.process(chunk)
output[:, i : i + last_chunk_size] = enhanced_chunk[:, :last_chunk_size]
break
enhanced_chunk = processor.process(chunk)
output[:, i : i + chunk_size] = enhanced_chunk[:, :chunk_size]
# Save result
sf.write(output_path, output.T, sr)