add enhancement_level support to AICFilter.

# Conflicts:
#	src/pipecat/audio/filters/aic_filter.py
This commit is contained in:
Gökmen Görgen
2026-03-09 12:44:27 +01:00
parent db22bf0f75
commit df64f3f943
2 changed files with 163 additions and 61 deletions

View File

@@ -23,6 +23,7 @@ from typing import List, Optional, Tuple
import numpy as np import numpy as np
from aic_sdk import ( from aic_sdk import (
Model, Model,
ParameterOutOfRangeError,
ProcessorAsync, ProcessorAsync,
ProcessorConfig, ProcessorConfig,
ProcessorParameter, ProcessorParameter,
@@ -220,6 +221,7 @@ class AICFilter(BaseAudioFilter):
model_id: Optional[str] = None, model_id: Optional[str] = None,
model_path: Optional[Path] = None, model_path: Optional[Path] = None,
model_download_dir: Optional[Path] = None, model_download_dir: Optional[Path] = None,
enhancement_level: Optional[float] = None,
) -> None: ) -> None:
"""Initialize the AIC filter. """Initialize the AIC filter.
@@ -231,9 +233,12 @@ class AICFilter(BaseAudioFilter):
model_id is ignored and no download occurs. model_id is ignored and no download occurs.
model_download_dir: Directory for downloading models as a Path object. model_download_dir: Directory for downloading models as a Path object.
Defaults to a cache directory in user's home folder. Defaults to a cache directory in user's home folder.
enhancement_level: Optional overall enhancement strength (0.0..1.0).
If None, the model default is used.
Raises: Raises:
ValueError: If neither model_id nor model_path is provided. ValueError: If neither model_id nor model_path is provided, or if
enhancement_level is out of range.
""" """
# Set SDK ID for telemetry identification (6 = pipecat) # Set SDK ID for telemetry identification (6 = pipecat)
set_sdk_id(6) set_sdk_id(6)
@@ -244,14 +249,18 @@ class AICFilter(BaseAudioFilter):
"See https://artifacts.ai-coustics.io/ for available models." "See https://artifacts.ai-coustics.io/ for available models."
) )
if enhancement_level is not None and not 0.0 <= enhancement_level <= 1.0:
raise ValueError("'enhancement_level' must be between 0.0 and 1.0.")
self._license_key = license_key self._license_key = license_key
self._model_id = model_id self._model_id = model_id
self._model_path = model_path self._model_path = model_path
self._model_download_dir = model_download_dir or ( self._model_download_dir = model_download_dir or (
Path.home() / ".cache" / "pipecat" / "aic-models" Path.home() / ".cache" / "pipecat" / "aic-models"
) )
self._enhancement_level = enhancement_level
self._is_filter_enabled = True
self._bypass = False
self._sample_rate = 0 self._sample_rate = 0
self._aic_ready = False self._aic_ready = False
self._frames_per_block = 0 self._frames_per_block = 0
@@ -325,6 +334,20 @@ class AICFilter(BaseAudioFilter):
sensitivity=sensitivity, sensitivity=sensitivity,
) )
def _apply_enhancement_level(self):
"""Apply enhancement_level if configured and supported by the active model."""
if self._enhancement_level is None or self._processor_ctx is None:
return
level = float(self._enhancement_level if self._is_filter_enabled else 0.0)
try:
self._processor_ctx.set_parameter(ProcessorParameter.EnhancementLevel, level)
except ParameterOutOfRangeError as e:
logger.warning("AIC enhancement_level was rejected as out-of-range; ignoring it.")
logger.debug(f"AIC EnhancementLevel set_parameter out-of-range: {e}")
self._enhancement_level = None
async def start(self, sample_rate: int): async def start(self, sample_rate: int):
"""Initialize the filter with the transport's sample rate. """Initialize the filter with the transport's sample rate.
@@ -373,14 +396,19 @@ class AICFilter(BaseAudioFilter):
self._processor_ctx = self._processor.get_processor_context() self._processor_ctx = self._processor.get_processor_context()
self._vad_ctx = self._processor.get_vad_context() self._vad_ctx = self._processor.get_vad_context()
# Apply initial parameters # Apply initial enhancement settings (if configured)
self._processor_ctx.set_parameter(ProcessorParameter.Bypass, 1.0 if self._bypass else 0.0) self._apply_enhancement_level()
# Log processor information # Log processor information
logger.debug(f"ai-coustics filter started:") logger.debug(f"ai-coustics filter started:")
logger.debug(f" Model ID: {self._model.get_id()}") logger.debug(f" Model ID: {self._model.get_id()}")
logger.debug(f" Sample rate: {self._sample_rate} Hz") logger.debug(f" Sample rate: {self._sample_rate} Hz")
logger.debug(f" Frames per chunk: {self._frames_per_block}") logger.debug(f" Frames per chunk: {self._frames_per_block}")
if self._enhancement_level is not None:
level = self._enhancement_level if self._is_filter_enabled else 0.0
logger.debug(f" Enhancement strength: {int(level * 100)}%")
else:
logger.debug(" Enhancement level not configured; using the model's default behavior.")
logger.debug(f" Optimal sample rate: {self._model.get_optimal_sample_rate()} Hz") logger.debug(f" Optimal sample rate: {self._model.get_optimal_sample_rate()} Hz")
logger.debug( logger.debug(
f" Optimal number of frames for {self._sample_rate} Hz: " f" Optimal number of frames for {self._sample_rate} Hz: "
@@ -422,14 +450,9 @@ class AICFilter(BaseAudioFilter):
None None
""" """
if isinstance(frame, FilterEnableFrame): if isinstance(frame, FilterEnableFrame):
self._bypass = not frame.enable self._is_filter_enabled = frame.enable
if self._processor_ctx is not None: if self._processor_ctx is not None:
try: self._apply_enhancement_level()
self._processor_ctx.set_parameter(
ProcessorParameter.Bypass, 1.0 if self._bypass else 0.0
)
except Exception as e: # noqa: BLE001
logger.error(f"AIC set_parameter failed: {e}")
async def filter(self, audio: bytes) -> bytes: async def filter(self, audio: bytes) -> bytes:
"""Apply AIC enhancement to audio data. """Apply AIC enhancement to audio data.

View File

@@ -8,16 +8,19 @@ import asyncio
import time import time
import unittest import unittest
from pathlib import Path from pathlib import Path
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch from unittest.mock import AsyncMock, MagicMock, patch
import numpy as np import numpy as np
# Check if aic_sdk is available # Check if aic_sdk is available
aic_sdk: Any
try: try:
import aic_sdk import aic_sdk
HAS_AIC_SDK = True HAS_AIC_SDK = True
except ImportError: except ImportError:
aic_sdk = None
HAS_AIC_SDK = False HAS_AIC_SDK = False
# Module path for patching # Module path for patching
@@ -67,6 +70,22 @@ class MockProcessorContext:
self.reset_called = True self.reset_called = True
class UnsupportedEnhancementProcessorContext(MockProcessorContext):
"""Processor context mock that rejects EnhancementLevel updates."""
def __init__(self, enhancement_parameter, error_type):
super().__init__()
self._enhancement_parameter = enhancement_parameter
self._error_type = error_type
self.enhancement_attempts = 0
def set_parameter(self, param, value):
if param == self._enhancement_parameter:
self.enhancement_attempts += 1
raise self._error_type("EnhancementLevel out of range")
super().set_parameter(param, value)
class MockVadContext: class MockVadContext:
"""A lightweight mock for AIC VadContext.""" """A lightweight mock for AIC VadContext."""
@@ -157,7 +176,8 @@ class TestAICFilter(unittest.IsolatedAsyncioTestCase):
self.assertEqual(filter_instance._license_key, "test-key") self.assertEqual(filter_instance._license_key, "test-key")
self.assertEqual(filter_instance._model_id, "test-model") self.assertEqual(filter_instance._model_id, "test-model")
self.assertIsNone(filter_instance._model_path) self.assertIsNone(filter_instance._model_path)
self.assertFalse(filter_instance._bypass) self.assertIsNone(filter_instance._enhancement_level)
self.assertTrue(filter_instance._is_filter_enabled)
async def test_initialization_with_model_path(self): async def test_initialization_with_model_path(self):
"""Test filter initialization with model_path.""" """Test filter initialization with model_path."""
@@ -174,6 +194,29 @@ class TestAICFilter(unittest.IsolatedAsyncioTestCase):
self.assertEqual(filter_instance._model_download_dir, download_dir) self.assertEqual(filter_instance._model_download_dir, download_dir)
async def test_initialization_with_valid_enhancement_level(self):
"""Test filter initialization with a valid enhancement_level."""
filter_instance = self._create_filter_with_mocks(enhancement_level=0.75)
self.assertEqual(filter_instance._enhancement_level, 0.75)
async def test_initialization_with_none_enhancement_level(self):
"""Test filter initialization with enhancement_level set to None."""
filter_instance = self._create_filter_with_mocks(enhancement_level=None)
self.assertIsNone(filter_instance._enhancement_level)
async def test_initialization_invalid_enhancement_level_raises(self):
"""Test initialization rejects enhancement_level outside 0.0..1.0."""
with patch(f"{AIC_FILTER_MODULE}.set_sdk_id"):
with self.assertRaises(ValueError) as context:
self.AICFilter(
license_key="test-key",
model_id="test-model",
enhancement_level=1.5,
)
self.assertIn("enhancement_level", str(context.exception))
async def test_start_with_model_path(self): async def test_start_with_model_path(self):
"""Test starting filter with a local model path.""" """Test starting filter with a local model path."""
model_path = Path("/tmp/test.aicmodel") model_path = Path("/tmp/test.aicmodel")
@@ -241,19 +284,61 @@ class TestAICFilter(unittest.IsolatedAsyncioTestCase):
self.assertIsNotNone(filter_instance._processor_ctx) self.assertIsNotNone(filter_instance._processor_ctx)
self.assertIsNotNone(filter_instance._vad_ctx) self.assertIsNotNone(filter_instance._vad_ctx)
async def test_start_applies_initial_bypass_parameter(self): async def test_start_does_not_apply_bypass_parameter(self):
"""Test that start applies bypass parameter.""" """Test that start does not set SDK Bypass parameter."""
filter_instance = self._create_filter_with_mocks() filter_instance = self._create_filter_with_mocks()
await self._start_filter_with_mocks(filter_instance) await self._start_filter_with_mocks(filter_instance)
# Check that bypass was set to 0.0 (enabled)
bypass_params = [ bypass_params = [
(p, v) (p, v)
for p, v in self.mock_processor.processor_ctx.parameters_set for p, v in self.mock_processor.processor_ctx.parameters_set
if p == aic_sdk.ProcessorParameter.Bypass if p == aic_sdk.ProcessorParameter.Bypass
] ]
self.assertTrue(len(bypass_params) > 0) self.assertEqual(bypass_params, [])
self.assertEqual(bypass_params[-1][1], 0.0)
async def test_start_applies_enhancement_level_when_supported(self):
"""Test that start applies enhancement_level when supported by model."""
filter_instance = self._create_filter_with_mocks(enhancement_level=0.65)
await self._start_filter_with_mocks(filter_instance)
enhancement_params = [
(p, v)
for p, v in self.mock_processor.processor_ctx.parameters_set
if p == aic_sdk.ProcessorParameter.EnhancementLevel
]
self.assertTrue(len(enhancement_params) > 0)
self.assertEqual(enhancement_params[-1][1], 0.65)
async def test_start_ignores_enhancement_level_when_unsupported(self):
"""Test unsupported enhancement_level logs warning and keeps filter ready."""
filter_instance = self._create_filter_with_mocks(enhancement_level=0.65)
with patch(f"{AIC_FILTER_MODULE}.ParameterOutOfRangeError", ValueError):
unsupported_ctx = UnsupportedEnhancementProcessorContext(
aic_sdk.ProcessorParameter.EnhancementLevel, ValueError
)
self.mock_processor.processor_ctx = unsupported_ctx
await self._start_filter_with_mocks(filter_instance)
self.assertTrue(filter_instance._aic_ready)
self.assertIsNone(filter_instance._enhancement_level)
self.assertEqual(unsupported_ctx.enhancement_attempts, 1)
async def test_start_does_not_set_enhancement_level_when_none(self):
"""Test start does not attempt enhancement_level when not configured."""
filter_instance = self._create_filter_with_mocks(enhancement_level=None)
with patch(f"{AIC_FILTER_MODULE}.logger.debug") as mock_debug:
await self._start_filter_with_mocks(filter_instance)
enhancement_params = [
(p, v)
for p, v in self.mock_processor.processor_ctx.parameters_set
if p == aic_sdk.ProcessorParameter.EnhancementLevel
]
self.assertEqual(enhancement_params, [])
self.assertTrue(
any("default behavior" in str(call.args[0]) for call in mock_debug.call_args_list)
)
async def test_stop_cleans_up_resources(self): async def test_stop_cleans_up_resources(self):
"""Test that stop properly cleans up resources and releases model reference.""" """Test that stop properly cleans up resources and releases model reference."""
@@ -420,57 +505,52 @@ class TestAICFilter(unittest.IsolatedAsyncioTestCase):
self.assertIsNone(filter_instance._processor) self.assertIsNone(filter_instance._processor)
self.assertFalse(filter_instance._aic_ready) self.assertFalse(filter_instance._aic_ready)
async def test_start_parameter_fixed_error_logged(self): async def test_start_skips_unsupported_enhancement_level_after_first_attempt(self):
"""Test start() when set_parameter raises ParameterFixedError: logged, no raise.""" """Test unsupported enhancement_level is attempted once and then skipped."""
filter_instance = self._create_filter_with_mocks() filter_instance = self._create_filter_with_mocks(enhancement_level=0.9)
self.mock_processor.processor_ctx.set_parameter = MagicMock(
side_effect=aic_sdk.ParameterFixedError("fixed")
)
with ( with patch(f"{AIC_FILTER_MODULE}.ParameterOutOfRangeError", ValueError):
patch(f"{AIC_FILTER_MODULE}.AICModelManager") as mock_manager_cls, unsupported_ctx = UnsupportedEnhancementProcessorContext(
patch(f"{AIC_FILTER_MODULE}.ProcessorConfig") as mock_config_cls, aic_sdk.ProcessorParameter.EnhancementLevel, ValueError
patch(f"{AIC_FILTER_MODULE}.ProcessorAsync", return_value=self.mock_processor), )
): self.mock_processor.processor_ctx = unsupported_ctx
mock_manager_cls.acquire = AsyncMock(return_value=(self.mock_model, "test-key"))
mock_config_cls.optimal.return_value = MagicMock()
await filter_instance.start(16000) await self._start_filter_with_mocks(filter_instance)
await filter_instance.stop()
await self._start_filter_with_mocks(filter_instance)
self.assertTrue(filter_instance._aic_ready) self.assertEqual(unsupported_ctx.enhancement_attempts, 1)
async def test_process_frame_set_parameter_exception_logged(self): async def test_process_frame_disable_sets_enhancement_to_zero(self):
"""Test process_frame when set_parameter raises: exception logged, no raise.""" """Test disable frame sets enhancement level to 0.0."""
filter_instance = self._create_filter_with_mocks() filter_instance = self._create_filter_with_mocks(enhancement_level=0.7)
await self._start_filter_with_mocks(filter_instance) await self._start_filter_with_mocks(filter_instance)
filter_instance._processor_ctx.set_parameter = MagicMock(
side_effect=ValueError("param error")
)
await filter_instance.process_frame(self.FilterEnableFrame(enable=False))
enhancement_params = [
(p, v)
for p, v in self.mock_processor.processor_ctx.parameters_set
if p == aic_sdk.ProcessorParameter.EnhancementLevel
]
self.assertFalse(filter_instance._is_filter_enabled)
self.assertEqual(enhancement_params[-1][1], 0.0)
async def test_process_frame_enable_restores_configured_enhancement(self):
"""Test enable frame restores configured enhancement level."""
filter_instance = self._create_filter_with_mocks(enhancement_level=0.7)
await self._start_filter_with_mocks(filter_instance)
await filter_instance.process_frame(self.FilterEnableFrame(enable=False))
await filter_instance.process_frame(self.FilterEnableFrame(enable=True)) await filter_instance.process_frame(self.FilterEnableFrame(enable=True))
self.assertFalse(filter_instance._bypass) enhancement_params = [
(p, v)
async def test_process_frame_enable(self): for p, v in self.mock_processor.processor_ctx.parameters_set
"""Test processing FilterEnableFrame to enable filtering.""" if p == aic_sdk.ProcessorParameter.EnhancementLevel
filter_instance = self._create_filter_with_mocks() ]
await self._start_filter_with_mocks(filter_instance) self.assertTrue(filter_instance._is_filter_enabled)
filter_instance._bypass = True self.assertEqual(enhancement_params[-1][1], 0.7)
enable_frame = self.FilterEnableFrame(enable=True)
await filter_instance.process_frame(enable_frame)
self.assertFalse(filter_instance._bypass)
async def test_process_frame_disable(self):
"""Test processing FilterEnableFrame to disable filtering."""
filter_instance = self._create_filter_with_mocks()
await self._start_filter_with_mocks(filter_instance)
disable_frame = self.FilterEnableFrame(enable=False)
await filter_instance.process_frame(disable_frame)
self.assertTrue(filter_instance._bypass)
async def test_filter_when_not_ready(self): async def test_filter_when_not_ready(self):
"""Test that filter returns audio unchanged when not ready.""" """Test that filter returns audio unchanged when not ready."""
@@ -678,7 +758,6 @@ class TestAICFilter(unittest.IsolatedAsyncioTestCase):
BufferError: Existing exports of data: object cannot be re-sized BufferError: Existing exports of data: object cannot be re-sized
This is the exact error path reported in production (line 414).
The fix removes the memoryview by snapshotting data into immutable The fix removes the memoryview by snapshotting data into immutable
bytes before any await. bytes before any await.
""" """