atom-audio-engine 0.1.0__py3-none-any.whl

core/pipeline.py

@@ -0,0 +1,282 @@
+"""Main pipeline orchestrator for audio-to-audio conversation."""
+
+import asyncio
+import logging
+import time
+from typing import AsyncIterator, Optional, Callable, Any
+
+from asr.base import BaseASR
+from llm.base import BaseLLM
+from tts.base import BaseTTS
+from core.types import (
+    AudioChunk,
+    TranscriptChunk,
+    ResponseChunk,
+    ConversationContext,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class Pipeline:
+    """
+    Main orchestrator for the audio-to-audio conversational pipeline.
+
+    Coordinates the flow: Audio Input → ASR → LLM → TTS → Audio Output
+
+    Example:
+        ```python
+        pipeline = Pipeline(
+            asr=WhisperASR(api_key="..."),
+            llm=AnthropicLLM(api_key="...", model="claude-sonnet-4-20250514"),
+            tts=CartesiaTTS(api_key="...", voice_id="...")
+        )
+
+        # Simple usage
+        response_audio = await pipeline.process(input_audio_bytes)
+
+        # Streaming usage
+        async for chunk in pipeline.stream(audio_stream):
+            play_audio(chunk)
+        ```
+    """
+
+    def __init__(
+        self,
+        asr: BaseASR,
+        llm: BaseLLM,
+        tts: BaseTTS,
+        system_prompt: Optional[str] = None,
+        on_transcript: Optional[Callable[[str], Any]] = None,
+        on_llm_response: Optional[Callable[[str], Any]] = None,
+        debug: bool = False,
+    ):
+        """
+        Initialize the pipeline with providers.
+
+        Args:
+            asr: Speech-to-Text provider instance
+            llm: Language model provider instance
+            tts: Text-to-Speech provider instance
+            system_prompt: System prompt for the LLM
+            on_transcript: Callback when transcript is ready
+            on_llm_response: Callback when LLM response is ready
+            debug: Enable debug logging
+        """
+        self.asr = asr
+        self.llm = llm
+        self.tts = tts
+
+        self.context = ConversationContext(system_prompt=system_prompt)
+        self.on_transcript = on_transcript
+        self.on_llm_response = on_llm_response
+        self.debug = debug
+
+        if debug:
+            logging.basicConfig(level=logging.DEBUG)
+
+        self._is_connected = False
+
+    async def connect(self):
+        """Initialize connections to all providers."""
+        if self._is_connected:
+            return
+
+        await asyncio.gather(
+            self.asr.connect(),
+            self.llm.connect(),
+            self.tts.connect(),
+        )
+        self._is_connected = True
+        logger.info("Pipeline connected to all providers")
+
+    async def disconnect(self):
+        """Close connections to all providers."""
+        if not self._is_connected:
+            return
+
+        await asyncio.gather(
+            self.asr.disconnect(),
+            self.llm.disconnect(),
+            self.tts.disconnect(),
+        )
+        self._is_connected = False
+        logger.info("Pipeline disconnected from all providers")
+
+    async def __aenter__(self):
+        """Async context manager entry."""
+        await self.connect()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Async context manager exit."""
+        await self.disconnect()
+
+    async def process(self, audio: bytes, sample_rate: int = 16000) -> bytes:
+        """
+        Process audio input and return complete audio response.
+
+        This is the simple, non-streaming interface. Use `stream()` for
+        lower latency.
+
+        Args:
+            audio: Input audio bytes (PCM format)
+            sample_rate: Sample rate of input audio
+
+        Returns:
+            Response audio bytes
+        """
+        start_time = time.time()
+
+        # Step 1: ASR - Transcribe audio to text
+        transcript = await self.asr.transcribe(audio, sample_rate)
+        asr_time = time.time() - start_time
+        logger.debug(f"ASR completed in {asr_time:.2f}s: {transcript[:50]}...")
+
+        if self.on_transcript:
+            self.on_transcript(transcript)
+
+        # Add user message to context
+        self.context.add_message("user", transcript)
+
+        # Step 2: LLM - Generate response
+        llm_start = time.time()
+        response_text = await self.llm.generate(transcript, self.context)
+        llm_time = time.time() - llm_start
+        logger.debug(f"LLM completed in {llm_time:.2f}s: {response_text[:50]}...")
+
+        if self.on_llm_response:
+            self.on_llm_response(response_text)
+
+        # Add assistant message to context
+        self.context.add_message("assistant", response_text)
+
+        # Step 3: TTS - Synthesize response to audio
+        tts_start = time.time()
+        response_audio = await self.tts.synthesize(response_text)
+        tts_time = time.time() - tts_start
+        logger.debug(f"TTS completed in {tts_time:.2f}s")
+
+        total_time = time.time() - start_time
+        logger.info(
+            f"Pipeline total: {total_time:.2f}s "
+            f"(ASR: {asr_time:.2f}s, LLM: {llm_time:.2f}s, TTS: {tts_time:.2f}s)"
+        )
+
+        return response_audio
+
+    async def stream(
+        self, audio_stream: AsyncIterator[AudioChunk]
+    ) -> AsyncIterator[AudioChunk]:
+        """
+        Process streaming audio input and yield streaming audio output.
+
+        This provides the lowest latency by:
+        1. Streaming ASR transcription
+        2. Streaming LLM generation
+        3. Streaming TTS synthesis
+
+        Args:
+            audio_stream: Async iterator of input AudioChunk objects
+
+        Yields:
+            AudioChunk objects containing response audio
+        """
+        start_time = time.time()
+
+        # Step 1: Stream audio through ASR
+        transcript_buffer = ""
+        async for transcript_chunk in self.asr.transcribe_stream(audio_stream):
+            transcript_buffer += transcript_chunk.text
+            if transcript_chunk.is_final:
+                break
+
+        if not transcript_buffer.strip():
+            return
+
+        asr_time = time.time() - start_time
+        logger.debug(f"ASR streaming completed in {asr_time:.2f}s")
+
+        if self.on_transcript:
+            self.on_transcript(transcript_buffer)
+
+        self.context.add_message("user", transcript_buffer)
+
+        # Step 2: Stream LLM response and pipe to TTS
+        llm_start = time.time()
+        response_buffer = ""
+
+        async def llm_text_stream() -> AsyncIterator[str]:
+            nonlocal response_buffer
+            async for chunk in self.llm.generate_stream(
+                transcript_buffer, self.context
+            ):
+                response_buffer += chunk.text
+                yield chunk.text
+                if chunk.is_final:
+                    break
+
+        # Step 3: Stream TTS audio as LLM generates text
+        first_audio_time = None
+        async for audio_chunk in self.tts.synthesize_stream_text(llm_text_stream()):
+            if first_audio_time is None:
+                first_audio_time = time.time() - start_time
+                logger.debug(f"Time to first audio: {first_audio_time:.2f}s")
+            yield audio_chunk
+
+        llm_time = time.time() - llm_start
+        logger.debug(f"LLM+TTS streaming completed in {llm_time:.2f}s")
+
+        if self.on_llm_response:
+            self.on_llm_response(response_buffer)
+
+        self.context.add_message("assistant", response_buffer)
+
+        total_time = time.time() - start_time
+        logger.info(f"Pipeline streaming total: {total_time:.2f}s")
+
+    async def stream_text_input(self, text: str) -> AsyncIterator[AudioChunk]:
+        """
+        Process text input (skip ASR) and yield streaming audio output.
+
+        Useful for text-based input or when ASR is handled externally.
+
+        Args:
+            text: User's text input
+
+        Yields:
+            AudioChunk objects containing response audio
+        """
+        self.context.add_message("user", text)
+
+        response_buffer = ""
+
+        async def llm_text_stream() -> AsyncIterator[str]:
+            nonlocal response_buffer
+            async for chunk in self.llm.generate_stream(text, self.context):
+                response_buffer += chunk.text
+                yield chunk.text
+
+        async for audio_chunk in self.tts.synthesize_stream_text(llm_text_stream()):
+            yield audio_chunk
+
+        self.context.add_message("assistant", response_buffer)
+
+    def reset_context(self):
+        """Clear conversation history."""
+        self.context.clear()
+        logger.debug("Conversation context cleared")
+
+    def set_system_prompt(self, prompt: str):
+        """Update the system prompt."""
+        self.context.system_prompt = prompt
+        logger.debug("System prompt updated")
+
+    @property
+    def providers(self) -> dict:
+        """Get information about configured providers."""
+        return {
+            "asr": self.asr.name,
+            "llm": self.llm.name,
+            "tts": self.tts.name,
+        }

core/types.py

@@ -0,0 +1,87 @@
+"""Shared types and data structures for the audio engine."""
+
+from dataclasses import dataclass, field
+from typing import Optional
+from enum import Enum
+
+
+class AudioFormat(Enum):
+    """Supported audio formats."""
+
+    PCM_16K = "pcm_16k"  # 16-bit PCM at 16kHz
+    PCM_24K = "pcm_24k"  # 16-bit PCM at 24kHz
+    PCM_44K = "pcm_44k"  # 16-bit PCM at 44.1kHz
+    WAV = "wav"
+    MP3 = "mp3"
+    OGG = "ogg"
+
+
+@dataclass
+class AudioChunk:
+    """A chunk of audio data."""
+
+    data: bytes
+    sample_rate: int = 16000
+    channels: int = 1
+    format: AudioFormat = AudioFormat.PCM_16K
+    timestamp_ms: Optional[int] = None
+    is_final: bool = False
+
+
+@dataclass
+class TranscriptChunk:
+    """A chunk of transcribed text from ASR."""
+
+    text: str
+    is_final: bool = False
+    confidence: Optional[float] = None
+    timestamp_ms: Optional[int] = None
+
+
+@dataclass
+class ResponseChunk:
+    """A chunk of LLM response text."""
+
+    text: str
+    is_final: bool = False
+    timestamp_ms: Optional[int] = None
+
+
+@dataclass
+class ConversationMessage:
+    """A message in the conversation history."""
+
+    role: str  # "user" or "assistant"
+    content: str
+    timestamp_ms: Optional[int] = None
+
+
+@dataclass
+class ConversationContext:
+    """Maintains conversation state and history."""
+
+    messages: list[ConversationMessage] = field(default_factory=list)
+    system_prompt: Optional[str] = None
+    max_history: int = 20
+
+    def add_message(self, role: str, content: str, timestamp_ms: Optional[int] = None):
+        """Add a message to the conversation history."""
+        self.messages.append(
+            ConversationMessage(role=role, content=content, timestamp_ms=timestamp_ms)
+        )
+        # Trim history if needed
+        if len(self.messages) > self.max_history:
+            self.messages = self.messages[-self.max_history :]
+
+    def get_messages_for_llm(self) -> list[dict]:
+        """Get messages formatted for LLM API calls."""
+        result = []
+        if self.system_prompt:
+            result.append({"role": "system", "content": self.system_prompt})
+        for msg in self.messages:
+            result.append({"role": msg.role, "content": msg.content})
+        return result
+
+    def clear(self):
+        """Clear conversation history."""
+        self.messages = []

integrations/__init__.py

@@ -0,0 +1,5 @@
+"""External system integrations."""
+
+from integrations.geneface import GeneFaceIntegration
+
+__all__ = ["GeneFaceIntegration"]

integrations/geneface.py

@@ -0,0 +1,297 @@
+"""GeneFace++ integration for face animation from audio."""
+
+import asyncio
+import logging
+import tempfile
+import os
+from pathlib import Path
+from typing import Optional, AsyncIterator
+from dataclasses import dataclass
+
+from core.types import AudioChunk
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class GeneFaceConfig:
+    """Configuration for GeneFace++ integration."""
+
+    geneface_path: str  # Path to ai-geneface-realtime directory
+    checkpoint_path: Optional[str] = None  # Path to trained model
+    output_resolution: tuple[int, int] = (512, 512)
+    fps: int = 25
+    device: str = "cuda"
+
+
+class GeneFaceIntegration:
+    """
+    Integration with GeneFace++ for generating animated face videos from audio.
+
+    This wraps the GeneFace++ inference system to generate talking face videos
+    from the audio output of the conversation pipeline.
+
+    Example:
+        ```python
+        geneface = GeneFaceIntegration(
+            config=GeneFaceConfig(
+                geneface_path="/path/to/ai-geneface-realtime",
+                checkpoint_path="/path/to/model.ckpt"
+            )
+        )
+
+        # Generate video from audio
+        video_path = await geneface.generate_video(audio_bytes)
+        ```
+    """
+
+    def __init__(self, config: GeneFaceConfig):
+        """
+        Initialize GeneFace++ integration.
+
+        Args:
+            config: GeneFace configuration
+        """
+        self.config = config
+        self._infer = None
+        self._initialized = False
+
+    async def initialize(self):
+        """
+        Initialize the GeneFace++ inference system.
+
+        This loads the models and prepares for inference.
+        """
+        if self._initialized:
+            return
+
+        # Add GeneFace path to Python path
+        import sys
+
+        geneface_path = Path(self.config.geneface_path)
+        if str(geneface_path) not in sys.path:
+            sys.path.insert(0, str(geneface_path))
+
+        try:
+            # Import GeneFace modules
+            from inference.genefacepp_infer import GeneFace2Infer
+
+            # Initialize inference object
+            # Note: This will load models which takes time
+            self._infer = GeneFace2Infer(
+                audio2secc_dir=self.config.checkpoint_path,
+                device=self.config.device,
+            )
+
+            self._initialized = True
+            logger.info("GeneFace++ integration initialized")
+
+        except ImportError as e:
+            logger.error(f"Failed to import GeneFace++: {e}")
+            raise ImportError(
+                f"Could not import GeneFace++. Ensure it's installed at {self.config.geneface_path}"
+            ) from e
+
+    async def generate_video(
+        self,
+        audio: bytes,
+        sample_rate: int = 16000,
+        output_path: Optional[str] = None,
+    ) -> str:
+        """
+        Generate a talking face video from audio.
+
+        Args:
+            audio: Audio bytes (PCM format)
+            sample_rate: Sample rate of the audio
+            output_path: Optional output video path. If not provided,
+                        a temporary file will be created.
+
+        Returns:
+            Path to the generated video file
+        """
+        if not self._initialized:
+            await self.initialize()
+
+        # Save audio to temporary file
+        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
+            temp_audio_path = f.name
+            # Write WAV header and data
+            self._write_wav(f, audio, sample_rate)
+
+        try:
+            # Determine output path
+            if output_path is None:
+                output_path = tempfile.mktemp(suffix=".mp4")
+
+            # Run GeneFace++ inference in executor to not block
+            loop = asyncio.get_event_loop()
+            await loop.run_in_executor(
+                None,
+                self._run_inference,
+                temp_audio_path,
+                output_path,
+            )
+
+            logger.info(f"Generated video at: {output_path}")
+            return output_path
+
+        finally:
+            # Cleanup temp audio file
+            if os.path.exists(temp_audio_path):
+                os.unlink(temp_audio_path)
+
+    def _run_inference(self, audio_path: str, output_path: str):
+        """Run GeneFace++ inference (blocking)."""
+        self._infer.infer_once(
+            inp={
+                "drv_audio": audio_path,
+                "out_name": output_path,
+            }
+        )
+
+    def _write_wav(self, file, audio: bytes, sample_rate: int):
+        """Write audio bytes as a WAV file."""
+        import struct
+
+        # WAV header
+        channels = 1
+        bits_per_sample = 16
+        byte_rate = sample_rate * channels * bits_per_sample // 8
+        block_align = channels * bits_per_sample // 8
+        data_size = len(audio)
+
+        # Write RIFF header
+        file.write(b"RIFF")
+        file.write(struct.pack("<I", 36 + data_size))
+        file.write(b"WAVE")
+
+        # Write fmt chunk
+        file.write(b"fmt ")
+        file.write(struct.pack("<I", 16))  # Chunk size
+        file.write(struct.pack("<H", 1))  # Audio format (PCM)
+        file.write(struct.pack("<H", channels))
+        file.write(struct.pack("<I", sample_rate))
+        file.write(struct.pack("<I", byte_rate))
+        file.write(struct.pack("<H", block_align))
+        file.write(struct.pack("<H", bits_per_sample))
+
+        # Write data chunk
+        file.write(b"data")
+        file.write(struct.pack("<I", data_size))
+        file.write(audio)
+
+    async def generate_video_stream(
+        self,
+        audio_stream: AsyncIterator[AudioChunk],
+        output_path: Optional[str] = None,
+    ) -> str:
+        """
+        Generate video from streaming audio.
+
+        Buffers audio chunks until stream completes, then generates video.
+        For true real-time video streaming, additional work would be needed
+        to chunk the inference.
+
+        Args:
+            audio_stream: Async iterator of audio chunks
+            output_path: Optional output video path
+
+        Returns:
+            Path to the generated video file
+        """
+        # Buffer all audio chunks
+        audio_buffer = bytearray()
+        sample_rate = 16000
+
+        async for chunk in audio_stream:
+            audio_buffer.extend(chunk.data)
+            sample_rate = chunk.sample_rate
+
+        return await self.generate_video(
+            bytes(audio_buffer),
+            sample_rate,
+            output_path,
+        )
+
+
+class GeneFacePipelineWrapper:
+    """
+    Wrapper that adds face animation to an audio pipeline.
+
+    Example:
+        ```python
+        from audio_engine import Pipeline
+        from audio_engine.integrations.geneface import GeneFacePipelineWrapper
+
+        # Create base pipeline
+        pipeline = Pipeline(asr=..., llm=..., tts=...)
+
+        # Wrap with face animation
+        wrapped = GeneFacePipelineWrapper(
+            pipeline=pipeline,
+            geneface_config=GeneFaceConfig(...)
+        )
+
+        # Now returns both audio and video
+        audio, video_path = await wrapped.process_with_video(input_audio)
+        ```
+    """
+
+    def __init__(self, pipeline, geneface_config: GeneFaceConfig):
+        """
+        Initialize the wrapper.
+
+        Args:
+            pipeline: Pipeline instance
+            geneface_config: GeneFace configuration
+        """
+        self.pipeline = pipeline
+        self.geneface = GeneFaceIntegration(geneface_config)
+
+    async def connect(self):
+        """Initialize all components."""
+        await asyncio.gather(
+            self.pipeline.connect(),
+            self.geneface.initialize(),
+        )
+
+    async def disconnect(self):
+        """Clean up all components."""
+        await self.pipeline.disconnect()
+
+    async def __aenter__(self):
+        await self.connect()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self.disconnect()
+
+    async def process_with_video(
+        self,
+        audio: bytes,
+        sample_rate: int = 16000,
+        video_output_path: Optional[str] = None,
+    ) -> tuple[bytes, str]:
+        """
+        Process audio and generate both response audio and face video.
+
+        Args:
+            audio: Input audio bytes
+            sample_rate: Sample rate of input
+            video_output_path: Optional path for output video
+
+        Returns:
+            Tuple of (response_audio_bytes, video_path)
+        """
+        # Get audio response from pipeline
+        response_audio = await self.pipeline.process(audio, sample_rate)
+
+        # Generate face animation video
+        video_path = await self.geneface.generate_video(
+            response_audio,
+            self.pipeline.tts.sample_rate,
+            video_output_path,
+        )
+
+        return response_audio, video_path

llm/__init__.py

@@ -0,0 +1,38 @@
+"""LLM (Large Language Model) providers."""
+
+from core.config import LLMConfig
+
+from .base import BaseLLM
+from .groq import GroqLLM
+
+__all__ = ["BaseLLM", "GroqLLM", "get_llm_from_config"]
+
+
+def get_llm_from_config(config: LLMConfig) -> BaseLLM:
+    """
+    Instantiate LLM provider from config.
+
+    Args:
+        config: LLMConfig object with provider name and settings
+
+    Returns:
+        Initialized BaseLLM provider instance
+
+    Raises:
+        ValueError: If provider name is not recognized
+    """
+    provider_name = config.provider.lower()
+
+    if provider_name == "groq":
+        return GroqLLM(
+            api_key=config.api_key,
+            model=config.model or "llama-3.1-8b-instant",
+            temperature=config.temperature,
+            max_tokens=config.max_tokens,
+            system_prompt=config.system_prompt,
+            **config.extra,
+        )
+    else:
+        raise ValueError(
+            f"Unknown LLM provider: {config.provider}. " f"Supported: groq"
+        )