diff --git a/Makefile b/Makefile index a373f79..4d8b7d9 100644 --- a/Makefile +++ b/Makefile @@ -7,4 +7,4 @@ lint: test: pytest -check: lint test \ No newline at end of file +check: lint test diff --git a/pyproject.toml b/pyproject.toml index 72f1390..e2558d5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,6 +25,9 @@ strict_optional = true warn_unused_ignores = true warn_return_any = true warn_unused_configs = true +exclude = [ + "^tests/", +] [tool.pytest.ini_options] testpaths = ["tests"] diff --git a/src/core/bus.py b/src/core/bus.py new file mode 100644 index 0000000..003abb5 --- /dev/null +++ b/src/core/bus.py @@ -0,0 +1,94 @@ +import asyncio +import logging +from collections import defaultdict +from typing import Any, AsyncGenerator, Coroutine, cast + +from .events import EventData +from .module import Module + +logger = logging.getLogger("ray.serve") + + +class EventGraph: + """ + Asynchronous event routing system for HuRI modules. + + The EventGraph is responsible for: + - Registering module subscribers + - Routing events between modules + - Executing module pipelines asynchronously + - Handling coroutine and async generator outputs + + Modules subscribe to events through their `input_type`. + When an event is published, all subscribed modules are executed + concurrently. + + Supports: + - Coroutine-based modules + - Async generator streaming modules + - Recursive event propagation + + :subscribers: + Dictionary mapping event topics to subscribed modules. + """ + + def __init__(self): + self.subscribers = defaultdict(list) + + def register(self, module: Module): + self.subscribers[module.input_type].append(module) + + async def publish(self, event_topic, data): + subs = self.subscribers[event_topic] + if event_topic not in ("audio_in",): # skip mic-frame spam + logger.info( + "[GRAPH] publish topic=%r subscribers=%s", + event_topic, + [type(m).__name__ for m in subs], + ) + for module in subs: + asyncio.create_task(self._run(module, data)) + + async def _run(self, module: Module, data): + try: + result = module.process(data) + + if hasattr(result, "__aiter__"): + try: + generator = cast(AsyncGenerator[EventData | None, None], result) + async for item in generator: + if item is None: + continue + logger.info( + "[GRAPH] %s -> %r: %s", + type(module).__name__, + module.output_type, + item.summarize(), + ) + await self.publish(module.output_type, item) + except Exception: + logger.exception( + "[GRAPH] async generator failed in %s", type(module).__name__ + ) + + else: + coroutine = cast(Coroutine[Any, Any, EventData | None], result) + try: + value = await coroutine + if value is not None: + logger.info( + "[GRAPH] %s -> %r: %s", + type(module).__name__, + module.output_type, + value.summarize(), + ) + await self.publish(module.output_type, value) + except Exception: + logger.exception( + "[GRAPH] coroutine failed in %s", type(module).__name__ + ) + + except Exception: + logger.exception( + "[GRAPH] process() call failed in %s", type(module).__name__ + ) diff --git a/src/core/client.py b/src/core/client.py index 9df62d0..cfd84e2 100644 --- a/src/core/client.py +++ b/src/core/client.py @@ -4,15 +4,14 @@ import struct from collections import defaultdict from dataclasses import asdict -from typing import Any, Dict, Generic, List, Type, TypeVar +from typing import Any, Dict, Generic, List, Mapping, Type, TypeVar -import numpy as np import websockets from src.core.dataclasses.config import ClientConfig from src.core.events import EventData -T = TypeVar("T", bound=EventData | bytes) +T = TypeVar("T", bound=EventData) class ClientSender(Generic[T]): @@ -45,16 +44,20 @@ async def _send_bytes(self, ws: websockets.ClientConnection, data: bytes): await ws.send(packet) - async def _send_event_data(self, ws: websockets.ClientConnection, data: EventData): - packet = json.dumps({"topic": self.topic, "data": asdict(data)}) + async def _send_event_data( + self, ws: websockets.ClientConnection, data: Mapping[str, Any] + ): + packet = json.dumps({"topic": self.topic, "data": data}) await ws.send(packet) async def send(self, ws: websockets.ClientConnection, data: T): - if isinstance(data, bytes): - await self._send_bytes(ws, data) + wire = data.to_wire() + + if isinstance(wire, bytes): + await self._send_bytes(ws, wire) else: - await self._send_event_data(ws, data) + await self._send_event_data(ws, wire) class ClientHook(Generic[T]): @@ -120,40 +123,12 @@ async def _receive_loop(self, ws: websockets.ClientConnection): topic = msg[2 : 2 + topic_len].decode() data = msg[2 + topic_len :] - if topic == "audio" and len(data) >= 13: - sample_rate, end, pts = struct.unpack(">IBd", data[:13]) - # Samples are native-endian float32 - # (Sender uses ndarray.tobytes()). - samples = np.frombuffer(data[13:], dtype=np.float32) - data = { - "sample_rate": sample_rate, - "end": end, - "pts": pts, - "data": samples, - } - elif topic == "motion" and len(data) >= 16: - pts, fps, n_frames = struct.unpack(">dII", data[:16]) - print(f"<< motion: pts={pts:.3f}s frames={n_frames} @ {fps}fps") - data = { - "poses": np.ndarray(0), - "expressions": np.ndarray(0), - "trans": np.ndarray(0), - "fps": fps, - "pts": pts, - } - else: - print(f"<< {topic}: bytes ({len(data)}B)") else: event = json.loads(msg) topic = event["topic"] data = event["data"] for hook in self.hooks[topic]: - # Hydrate via from_wire (not a bare **data splat) so events - # with nested EventData fields — e.g. RAGQuestion.transcript / - # .emotion — are rebuilt as dataclasses, mirroring the server's - # EventDataFactory. Build a fresh instance per hook so a topic - # with several hooks doesn't re-splat an already-built event. hook_data = ( data if isinstance(data, bytes) diff --git a/src/core/events.py b/src/core/events.py index 862760b..68ee7bd 100644 --- a/src/core/events.py +++ b/src/core/events.py @@ -1,121 +1,50 @@ -import asyncio import logging -from collections import defaultdict -from dataclasses import dataclass -from typing import Any, Mapping - -import numpy as np - -from .module import Module +from dataclasses import asdict, dataclass +from typing import Any, Generic, Mapping, TypeVar logger = logging.getLogger("ray.serve") +WireT = TypeVar("WireT", Mapping[str, Any], bytes) + @dataclass -class EventData: +class EventData(Generic[WireT]): """An event data must be derived from this class, and use @dataclass decorator. Or they can be bytes.""" @classmethod - def from_wire(cls, data: Mapping[str, Any]) -> "EventData": - """Build an event from a JSON payload sent by a client. - - Default: keyword-splat the payload onto the dataclass. Events whose fields - are nested dataclasses (or that accept a simpler external shape than the - in-pipeline one) override this — e.g. RAGQuestion accepts a bare - ``{"text": ...}`` typed question. In-process producers construct the - dataclass directly and never go through this path. - """ - return cls(**data) - - -class EventGraph: - """ - Asynchronous event routing system for HuRI modules. - - The EventGraph is responsible for: - - Registering module subscribers - - Routing events between modules - - Executing module pipelines asynchronously - - Handling coroutine and async generator outputs + def from_wire(cls, data: WireT) -> "EventData[WireT]": + raise NotImplementedError - Modules subscribe to events through their `input_type`. - When an event is published, all subscribed modules are executed - concurrently. + def to_wire(self) -> WireT: + raise NotImplementedError - Supports: - - Coroutine-based modules - - Async generator streaming modules - - Recursive event propagation + def summarize(self) -> str: + cls = type(self).__name__ + return f"{cls}({self!r})" - :subscribers: - Dictionary mapping event topics to subscribed modules. - """ - def __init__(self): - self.subscribers = defaultdict(list) - - def register(self, module: Module): - self.subscribers[module.input_type].append(module) - - async def publish(self, event_topic, data): - subs = self.subscribers[event_topic] - if event_topic not in ("audio_in",): # skip mic-frame spam - logger.info( - "[GRAPH] publish topic=%r subscribers=%s", - event_topic, - [type(m).__name__ for m in subs], - ) - for module in subs: - asyncio.create_task(self._run(module, data)) +@dataclass +class JsonEvent(EventData[Mapping[str, Any]]): + @classmethod + def from_wire(cls, data: Mapping[str, Any]) -> "JsonEvent": + return cls(**data) - async def _run(self, module: Module, data): - try: - result = module.process(data) + def to_wire(self) -> Mapping[str, Any]: + return asdict(self) - if hasattr(result, "__aiter__"): - try: - async for item in result: - if item is None: - continue - logger.info( - "[GRAPH] %s -> %r: %s", - type(module).__name__, - module.output_type, - _summarize(item), - ) - await self.publish(module.output_type, item) - except Exception: - logger.exception( - "[GRAPH] async generator failed in %s", type(module).__name__ - ) - else: - try: - value = await result - if value is not None: - logger.info( - "[GRAPH] %s -> %r: %s", - type(module).__name__, - module.output_type, - _summarize(value), - ) - await self.publish(module.output_type, value) - except Exception: - logger.exception( - "[GRAPH] coroutine failed in %s", type(module).__name__ - ) +@dataclass +class BytesEvent(EventData[bytes]): + data: bytes - except Exception: - logger.exception( - "[GRAPH] process() call failed in %s", type(module).__name__ - ) + @classmethod + def from_wire(cls, data: bytes) -> "BytesEvent": + return cls(data=data) + def to_wire(self) -> bytes: + return self.data -def _summarize(item) -> str: - """Short repr that avoids dumping full numpy arrays into the log.""" - cls = type(item).__name__ - data = getattr(item, "data", None) - if isinstance(data, np.ndarray): - return f"{cls}(shape={data.shape}, dtype={data.dtype})" - return f"{cls}({item!r})" + def summarize(self) -> str: + cls = type(self).__name__ + return f"{cls}(len:{len(self.data)})" diff --git a/src/core/huri.py b/src/core/huri.py index 6b92e12..a38d71b 100644 --- a/src/core/huri.py +++ b/src/core/huri.py @@ -32,7 +32,7 @@ def __init__( self, modules: Dict[str, Type[Module]], handles: Dict[str, handle.DeploymentHandle], - events: Dict[str, Type[EventData | bytes]], + events: Dict[str, Type[EventData]], ) -> None: self.module_factory = ModuleFactory(handles) self.event_factory = EventDataFactory() diff --git a/src/core/module.py b/src/core/module.py index 281a68b..f3144da 100644 --- a/src/core/module.py +++ b/src/core/module.py @@ -2,6 +2,8 @@ from ray.serve import handle +from .events import EventData + class Module: """ @@ -24,7 +26,9 @@ class Module: input_type: str output_type: Optional[str] - def process(self, _) -> Coroutine[Any, Any, Any] | AsyncGenerator[Any, None]: + def process( + self, _ + ) -> Coroutine[Any, Any, EventData | None] | AsyncGenerator[EventData | None, None]: raise NotImplementedError diff --git a/src/core/session.py b/src/core/session.py index 2b64daf..4f09c22 100644 --- a/src/core/session.py +++ b/src/core/session.py @@ -1,6 +1,6 @@ from typing import List -from .events import EventGraph +from .bus import EventGraph from .module import Module diff --git a/src/interfaces/cli_interface.py b/src/interfaces/cli_interface.py index 3cf2282..10949ec 100644 --- a/src/interfaces/cli_interface.py +++ b/src/interfaces/cli_interface.py @@ -11,13 +11,14 @@ from scipy.signal import resample from src.core.client import ClientHook, ClientSender +from src.core.events import BytesEvent from src.core.interface import Interface -from src.modules.rag.events import RAGQuestion, RAGResult +from src.modules.rag.events import RAGQuestion from src.modules.speech_to_text.events import Transcript from src.modules.text_to_speech.events import Audio, Token -class AudioSender(ClientSender[bytes]): +class AudioSender(ClientSender[BytesEvent]): def __init__( self, sample_rate: int = 16000, frame_duration: float = 0.030, **kwargs ): @@ -43,7 +44,7 @@ def callback(indata: np.ndarray, frames, time, status): ): while True: chunk = await queue.get() - await self.send(ws, chunk.tobytes()) + await self.send(ws, BytesEvent(data=chunk.tobytes())) class TextSender(ClientSender[RAGQuestion]): @@ -74,9 +75,9 @@ class AudioHook(ClientHook[Audio]): def __init__( self, + save_audio_dir: str, sample_rate=48000, incoming_sample_rate=16000, - save_audio_dir: Optional[str] = None, **kwargs, ): super().__init__(**kwargs) @@ -162,14 +163,14 @@ async def hook(self, data: Audio): self._collect_audio(data.data, data.sample_rate, bool(data.end)) -class TextHook(ClientHook[RAGResult]): - input_type = RAGResult +class TextHook(ClientHook[RAGQuestion]): + input_type = RAGQuestion def __init__(self, **kwargs): super().__init__(**kwargs) - async def hook(self, data: RAGResult): - print("<<", data.answer) + async def hook(self, data: RAGQuestion): + print("<<", data.transcript, data.emotion) class TokenHook(ClientHook[Token]): diff --git a/src/modules/events.py b/src/modules/events.py index 8cdc066..099174d 100644 --- a/src/modules/events.py +++ b/src/modules/events.py @@ -1,17 +1,17 @@ from typing import Dict, Type -from src.core.events import EventData +from src.core.events import BytesEvent, EventData from src.modules.emotion.events import Emotion from src.modules.gesture.events import Motion from src.modules.rag.events import PartialQuestion, RAGQuestion from src.modules.speech_to_text.events import Transcript, Voice -from src.modules.text_to_speech.events import Token +from src.modules.text_to_speech.events import Audio, Token -def get_events() -> Dict[str, Type[EventData | bytes]]: - events: Dict[str, Type[EventData | bytes]] = { - "audio_in": bytes, # inbound mic frames (raw int16 PCM) - "audio": bytes, +def get_events() -> Dict[str, Type[EventData]]: + events: Dict[str, Type[EventData]] = { + "audio.in": BytesEvent, # inbound mic frames (raw int16 PCM) + "audio.out": Audio, "voice": Voice, "transcript": Transcript, "emotion": Emotion, diff --git a/src/modules/factory.py b/src/modules/factory.py index 315dd6e..1d843ee 100644 --- a/src/modules/factory.py +++ b/src/modules/factory.py @@ -9,9 +9,9 @@ class EventDataFactory: def __init__(self): - self._registry: Dict[str, Type[EventData | bytes]] = {} + self._registry: Dict[str, Type[EventData]] = {} - def register(self, topic: str, event_cls: Type[EventData | bytes] | None) -> None: + def register(self, topic: str, event_cls: Type[EventData] | None) -> None: if topic in self._registry: if event_cls is None or event_cls == self._registry[topic]: return @@ -23,23 +23,16 @@ def register(self, topic: str, event_cls: Type[EventData | bytes] | None) -> Non self._registry[topic] = event_cls - def create(self, topic: str, data: Mapping[str, Any] | bytes) -> EventData | bytes: + def create(self, topic: str, data: Mapping[str, Any]) -> EventData: if topic not in self._registry: raise RuntimeError(f"unknown event topic {topic}") event_cls = self._registry[topic] - if isinstance(data, bytes): - if issubclass(event_cls, bytes): - return data - else: - raise RuntimeError(f"mismatched event data type: \ -{event_cls} is not type bytes but should be.") + if issubclass(event_cls, EventData): + return event_cls.from_wire(data) else: - if issubclass(event_cls, EventData): - return event_cls.from_wire(data) - else: - raise RuntimeError(f"mismatched event data type: \ + raise RuntimeError(f"mismatched event data type: \ {event_cls} is not derived from EventData but should be.") diff --git a/src/modules/gesture/events.py b/src/modules/gesture/events.py index 7793de1..754fde3 100644 --- a/src/modules/gesture/events.py +++ b/src/modules/gesture/events.py @@ -1,3 +1,5 @@ +import logging +import struct from dataclasses import dataclass import numpy as np @@ -6,11 +8,42 @@ _EMAGE_FPS = 30 +logger = logging.getLogger("ray.serve") + @dataclass -class Motion(EventData): +class Motion(EventData[bytes]): poses: np.ndarray # (t, 165) SMPL-X axis-angle, 55 joints × 3 expressions: np.ndarray # (t, 100) facial expression coefficients trans: np.ndarray # (t, 3) global root translation fps: int = _EMAGE_FPS pts: float = 0.0 # presentation timestamp in seconds, paired with Audio.pts + + def to_wire(self) -> bytes: + n_frames = self.poses.shape[0] + logger.info( + "Motion frames=%d fps=%d pts=%.3fs", + n_frames, + self.fps, + self.pts, + ) + header = struct.pack(">dII", self.pts, self.fps, n_frames) + body = ( + self.poses.astype(np.float32).tobytes() + + self.expressions.astype(np.float32).tobytes() + + self.trans.astype(np.float32).tobytes() + ) + return header + body + + @classmethod + def from_wire(cls, data: bytes) -> "Motion": + pts, fps, n_frames = struct.unpack(">dII", data[:16]) + print(f"<< motion: pts={pts:.3f}s frames={n_frames} @ {fps}fps") + + return cls( + poses=np.ndarray(0), + expressions=np.ndarray(0), + trans=np.ndarray(0), + fps=fps, + pts=pts, + ) diff --git a/src/modules/gesture/gesture.py b/src/modules/gesture/gesture.py index a39704a..7c533a2 100644 --- a/src/modules/gesture/gesture.py +++ b/src/modules/gesture/gesture.py @@ -83,24 +83,24 @@ def __init__( print("[Gesture] loading face_vq...") face_vq = EmageVQVAEConv.from_pretrained(hf_repo, subfolder="emage_vq/face").to( - self.device # type: ignore[arg-type] + self.device ) print("[Gesture] loading upper_vq...") upper_vq = EmageVQVAEConv.from_pretrained( hf_repo, subfolder="emage_vq/upper" - ).to(self.device) # type: ignore[arg-type] + ).to(self.device) print("[Gesture] loading lower_vq...") lower_vq = EmageVQVAEConv.from_pretrained( hf_repo, subfolder="emage_vq/lower" - ).to(self.device) # type: ignore[arg-type] + ).to(self.device) print("[Gesture] loading hands_vq...") hands_vq = EmageVQVAEConv.from_pretrained( hf_repo, subfolder="emage_vq/hands" - ).to(self.device) # type: ignore[arg-type] + ).to(self.device) print("[Gesture] loading global_ae...") global_ae = EmageVAEConv.from_pretrained( hf_repo, subfolder="emage_vq/global" - ).to(self.device) # type: ignore[arg-type] + ).to(self.device) self.motion_vq = EmageVQModel( face_model=face_vq, @@ -112,9 +112,7 @@ def __init__( self.motion_vq.eval() print("[Gesture] loading EmageAudioModel...") - self.model = EmageAudioModel.from_pretrained(hf_repo).to( - self.device # type: ignore[arg-type] - ) + self.model = EmageAudioModel.from_pretrained(hf_repo).to(self.device) self.model.eval() self._warmup() @@ -299,7 +297,7 @@ class Gesture(ModuleWithHandle): """ _handle_cls = GestureDeployment - input_type = "audio" + input_type = "audio.out" output_type = "motion" def __init__( @@ -341,9 +339,7 @@ def _end_utterance(self) -> None: self._buf_start = 0 self._emitted = 0 - async def process( # type: ignore[override] - self, audio: Audio - ) -> AsyncGenerator[Motion, None]: + async def process(self, audio: Audio) -> AsyncGenerator[Motion, None]: # Each chunk arrives as its own process() task on the shared per-session # instance, so serialise under a lock to keep the buffer ordered. async with self._lock: diff --git a/src/modules/rag/events.py b/src/modules/rag/events.py index 0873376..955ff40 100644 --- a/src/modules/rag/events.py +++ b/src/modules/rag/events.py @@ -1,13 +1,13 @@ from dataclasses import dataclass, field from typing import Any, Mapping, Optional -from src.core.events import EventData +from src.core.events import JsonEvent from src.modules.emotion.events import Emotion from src.modules.speech_to_text.events import Transcript @dataclass -class RAGResult(EventData): +class RAGResult(JsonEvent): """What RAGHandle returns.""" answer: str @@ -15,7 +15,7 @@ class RAGResult(EventData): @dataclass -class PartialQuestion(EventData): +class PartialQuestion(JsonEvent): """Partial question used to aggregate a sentence to an emotion.""" transcript: Optional[Transcript] @@ -23,7 +23,7 @@ class PartialQuestion(EventData): @dataclass -class RAGQuestion(EventData): +class RAGQuestion(JsonEvent): """Fully aggregated question to send to the RAG.""" transcript: Transcript diff --git a/src/modules/rag/rag.py b/src/modules/rag/rag.py index 23fddee..782ca8b 100644 --- a/src/modules/rag/rag.py +++ b/src/modules/rag/rag.py @@ -426,9 +426,7 @@ def base_strength(payload: dict) -> float: vec = await self._embed(merged) vector_size = len(vec) now = datetime.now().isoformat() - imp = min( - max((p.payload or {}).get("importance", 3) for p in weak) + 1, 10 - ) + imp = min(max((p.payload or {}).get("importance", 3) for p in weak) + 1, 10) self._qdrant.upsert( collection_name=self._cfg.memory_collection, points=[ @@ -580,7 +578,7 @@ def _build_prompt( # slot for a formatting rule. open-mistral-nemo skips the persona-level # no-Ah/Oh rule often enough that we restate it right at the tail. user_prompt += ( - '\n\nStart your answer straight on the substance — do not open with ' + "\n\nStart your answer straight on the substance — do not open with " '"Ah" or "Oh".' ) @@ -780,9 +778,7 @@ def __init__( self._max_history_turns = max_history_turns self.history: list[dict] = [] - async def process( # type: ignore[override] - self, data: RAGQuestion - ) -> AsyncGenerator[Token, None]: + async def process(self, data: RAGQuestion) -> AsyncGenerator[Token, None]: """ Called when a "question" event arrives through the event bus. Packages _user_id + question, sends to the stateless RAGHandle. diff --git a/src/modules/speech_to_text/events.py b/src/modules/speech_to_text/events.py index fc04674..e53fb9d 100644 --- a/src/modules/speech_to_text/events.py +++ b/src/modules/speech_to_text/events.py @@ -3,15 +3,24 @@ import numpy as np -from src.core.events import EventData +from src.core.events import JsonEvent @dataclass -class Transcript(EventData): +class Transcript(JsonEvent): text: str end: bool @dataclass -class Voice(EventData): +class Voice(JsonEvent): data: Optional[np.ndarray] + + def summarize(self) -> str: + """Short repr that avoids dumping full numpy arrays into the log.""" + + if self.data: + cls = type(self).__name__ + return f"{cls}(shape={self.data.shape}, dtype={self.data.dtype})" + else: + return super().summarize() diff --git a/src/modules/speech_to_text/microphone_vad.py b/src/modules/speech_to_text/microphone_vad.py index eb52f74..2142198 100644 --- a/src/modules/speech_to_text/microphone_vad.py +++ b/src/modules/speech_to_text/microphone_vad.py @@ -3,6 +3,7 @@ import numpy as np import webrtcvad +from src.core.events import BytesEvent from src.core.module import Module from .events import Voice @@ -13,7 +14,7 @@ class MIC(Module): Detect voice and silence using WebRTC VAD. - input: audio_in, + input: audio.in, output: voice :vad_agressiveness: from 0 (low) to 3 (high, can distord audio). @@ -26,7 +27,7 @@ class MIC(Module): # Inbound microphone frames travel on their own topic so the TTS-output # "audio" topic (consumed by Gesture and the client Sender) never collides # with mic input — otherwise raw mic bytes get echoed back to the client. - input_type = "audio_in" + input_type = "audio.in" output_type = "voice" def __init__( @@ -49,11 +50,11 @@ def __init__( self.vad = webrtcvad.Vad(vad_agressiveness) - async def process(self, data: bytes) -> Optional[Voice]: - if self.vad.is_speech(data, self.sample_rate) is True: + async def process(self, data: BytesEvent) -> Optional[Voice]: + if self.vad.is_speech(data.data, self.sample_rate) is True: self.silence_frames_count = 0 - audio_array = np.frombuffer(data, dtype=np.int16) + audio_array = np.frombuffer(data.data, dtype=np.int16) audio_array_float = audio_array.astype(np.float32) / 32768.0 return Voice(audio_array_float) diff --git a/src/modules/text_to_speech/events.py b/src/modules/text_to_speech/events.py index dceb269..434c14f 100644 --- a/src/modules/text_to_speech/events.py +++ b/src/modules/text_to_speech/events.py @@ -1,19 +1,47 @@ +import logging +import struct from dataclasses import dataclass import numpy as np -from src.core.events import EventData +from src.core.events import EventData, JsonEvent + +logger = logging.getLogger("ray.serve") @dataclass -class Token(EventData): +class Token(JsonEvent): text: str end: bool @dataclass -class Audio(EventData): +class Audio(EventData[bytes]): data: np.ndarray sample_rate: int end: bool = False pts: float = 0.0 # presentation timestamp in seconds from utterance start + + def to_wire(self) -> bytes: + logger.info( + "Audio samples=%d sr=%d end=%s pts=%.3fs", + self.data.shape[0], + self.sample_rate, + self.end, + self.pts, + ) + header = struct.pack(">IBd", self.sample_rate, int(self.end), self.pts) + return header + self.data.tobytes() + + @classmethod + def from_wire(cls, data: bytes) -> "Audio": + sample_rate, end, pts = struct.unpack(">IBd", data[:13]) + samples = np.frombuffer(data[13:], dtype=np.float32) + + return cls(sample_rate=sample_rate, end=end, pts=pts, data=samples) + + def summarize(self) -> str: + """Short repr that avoids dumping full numpy arrays into the log.""" + + cls = type(self).__name__ + return f"{cls}(shape={self.data.shape}, dtype={self.data.dtype})" diff --git a/src/modules/text_to_speech/text_to_speech.py b/src/modules/text_to_speech/text_to_speech.py index 63e1b46..c080792 100644 --- a/src/modules/text_to_speech/text_to_speech.py +++ b/src/modules/text_to_speech/text_to_speech.py @@ -232,7 +232,7 @@ class TTS(ModuleWithHandle): _handle_cls = TTSDeployment input_type = "token" - output_type = "audio" + output_type = "audio.out" def __init__(self, _handle: handle.DeploymentHandle): super().__init__(_handle) @@ -248,9 +248,7 @@ def __init__(self, _handle: handle.DeploymentHandle): # and silently drop trailing words). self._push_lock = asyncio.Lock() - async def process( # type: ignore[override] - self, token: Token - ) -> AsyncGenerator[Audio, None]: + async def process(self, token: Token) -> AsyncGenerator[Audio, None]: # Acquire BEFORE any await so lock-acquisition order matches token order. # Setup + push happen under the lock; only the first token of an # utterance goes on to drain/yield audio (outside the lock, so pushes of diff --git a/src/modules/utils/sender.py b/src/modules/utils/sender.py index 434f457..255678a 100644 --- a/src/modules/utils/sender.py +++ b/src/modules/utils/sender.py @@ -1,14 +1,10 @@ import logging import struct -from dataclasses import asdict -import numpy as np from fastapi import WebSocket from src.core.events import EventData from src.core.module import Module -from src.modules.gesture.events import Motion -from src.modules.text_to_speech.events import Audio logger = logging.getLogger("ray.serve") @@ -41,40 +37,21 @@ def __init__(self, ws: WebSocket, type: str): self.ws: WebSocket = ws self.input_type = type - async def process(self, data: EventData | bytes): + async def process(self, data: EventData): logger.info("[Sender:%s] received %s", self.input_type, type(data).__name__) - if isinstance(data, bytes): - await self.ws.send_bytes(self._prefix(data)) - elif isinstance(data, Audio): - logger.info( - "[Sender:%s] Audio samples=%d sr=%d end=%s pts=%.3fs", - self.input_type, - data.data.shape[0], - data.sample_rate, - data.end, - data.pts, - ) - header = struct.pack(">IBd", data.sample_rate, int(data.end), data.pts) - await self.ws.send_bytes(self._prefix(header + data.data.tobytes())) - elif isinstance(data, Motion): - n_frames = data.poses.shape[0] - logger.info( - "[Sender:%s] Motion frames=%d fps=%d pts=%.3fs", - self.input_type, - n_frames, - data.fps, - data.pts, - ) - header = struct.pack(">dII", data.pts, data.fps, n_frames) - body = ( - data.poses.astype(np.float32).tobytes() - + data.expressions.astype(np.float32).tobytes() - + data.trans.astype(np.float32).tobytes() - ) - await self.ws.send_bytes(self._prefix(header + body)) + + wire = data.to_wire() + if isinstance(wire, bytes): + await self.ws.send_bytes(self._prefix(wire)) else: - await self.ws.send_json({"topic": self.input_type, "data": asdict(data)}) + await self.ws.send_json( + { + "topic": self.input_type, + "data": wire, + } + ) def _prefix(self, payload: bytes) -> bytes: + """Encode topic and topic len and adds it as a prefix to the payload""" topic_bytes = self.input_type.encode() return struct.pack(">H", len(topic_bytes)) + topic_bytes + payload diff --git a/tests/core/test_events.py b/tests/core/test_events.py index c4485b1..9c2430b 100644 --- a/tests/core/test_events.py +++ b/tests/core/test_events.py @@ -1,11 +1,11 @@ import asyncio -import logging from dataclasses import dataclass import numpy as np import pytest -from src.core.events import EventData, EventGraph, _summarize +from src.core.bus import EventGraph +from src.core.events import EventData # --------------------------------------------------------------------------- # Helpers @@ -206,25 +206,3 @@ async def process(self, data): await graph.publish("a", DummyEvent(1)) await asyncio.wait_for(done.wait(), timeout=1) - - -# --------------------------------------------------------------------------- -# summarize() -# --------------------------------------------------------------------------- - - -def test_summarize_numpy(): - event = ArrayEvent(np.zeros((3, 4), dtype=np.float32)) - - text = _summarize(event) - - assert "shape=(3, 4)" in text - assert "float32" in text - - -def test_summarize_regular(): - event = DummyEvent(123) - - text = _summarize(event) - - assert "DummyEvent" in text