Realtime-Venus-Omni/memory_adapter/model_adapter.py

# Copyright 2026 The Realtime-Venus Team. All rights reserved.

"""Memory configuration and routing for Realtime-Venus-Omni Chat and Duplex.

Convert conversation inputs, manage memory sessions, and coordinate text
and speech generation through the model interfaces.
"""

from __future__ import annotations

from copy import deepcopy
from dataclasses import replace
import logging
import math
from pathlib import Path
import sys
from typing import Any, Mapping, Sequence

import numpy as np

from .omni_memory.config import MemoryConfig, load_memory_config
from .omni_memory.duplex.proactive_adapter import DuplexMemoryAdapter
from .omni_memory.duplex.runtime import MemoryDuplexRuntime
from .omni_memory.schema import MediaChunk, Question
from .omni_memory.session import MediaMemorySession


_CONTROLLER_ATTRIBUTE = "_realtime_venus_omni_memory_controller"
_DEFAULT_CHAT_QUERY = "Describe the video."
_MEMORY_MEDIA_SLOT_TYPE = "_realtime_venus_omni_memory_media"
_LOGGER = logging.getLogger(__name__)


def configure_memory(
    model: Any,
    *,
    memory_minutes: float | None = None,
) -> "RealtimeVenusOmniMemoryController":
    """Enable Memory with an optional Duplex archive duration in minutes."""
    overrides = None
    if memory_minutes is not None:
        if isinstance(memory_minutes, bool) or not isinstance(
            memory_minutes, (int, float)
        ):
            raise TypeError(
                "memory_minutes must be a positive finite number of minutes"
            )
        try:
            seconds = memory_minutes * 60.0
        except OverflowError as error:
            raise ValueError("memory_minutes must be positive and finite") from error
        if not math.isfinite(seconds) or seconds <= 0:
            raise ValueError("memory_minutes must be positive and finite")
        overrides = {"duplex": {"archive": {"retention_seconds": seconds}}}
    controller = RealtimeVenusOmniMemoryController(
        model=model,
        config=load_memory_config(overrides),
    )
    setattr(model, _CONTROLLER_ATTRIBUTE, controller)
    return controller


class RealtimeVenusOmniMemoryController:
    """Manage memory sessions for Chat and Duplex on a loaded model."""

    def __init__(
        self,
        *,
        model: Any,
        config: MemoryConfig,
    ) -> None:
        self.model = model
        self.config = config
        self.default_chat_query = _DEFAULT_CHAT_QUERY

    def chat(self, official_call: Mapping[str, Any]) -> Any:
        """Convert conversation media into one-second chunks and run Memory Chat."""
        call = _normalize_official_chat_call(official_call)
        _validate_memory_chat_options(call)
        media_chunks, query, prompt_context = _media_chunks_and_query(
            call.get("msgs"),
            default_query=self.default_chat_query,
            image=call.get("image"),
            call=call,
            normalize_content=_official_content_normalizer(self.model),
        )
        max_new_tokens = _positive_integer(
            call.get("max_new_tokens", self.config.chat.max_new_tokens),
            "max_new_tokens",
        )
        resolved_config = replace(
            self.config,
            chat=replace(self.config.chat, max_new_tokens=max_new_tokens),
        )
        processor = getattr(self.model, "processor", None)
        requested_processor = call.get("processor")
        requested_tokenizer = call.get("tokenizer")
        if (
            processor is None
            or requested_processor is not None
            or requested_tokenizer is not None
        ):
            prepare_processor = getattr(self.model, "prepare_processor", None)
            if not callable(prepare_processor):
                raise RuntimeError(
                    "Memory Chat requires model.processor or prepare_processor()"
                )
            prepare_processor(
                processor=requested_processor,
                tokenizer=requested_tokenizer,
            )
            processor = getattr(self.model, "processor", None)
        if processor is None:
            raise RuntimeError(
                "Realtime-Venus-Omni processor initialization returned no processor"
            )
        session = MediaMemorySession(
            model=self.model,
            processor=processor,
            use_memory=True,
            config=resolved_config,
            generate_audio=False,
            prompt_context=prompt_context,
        )
        rendered_prompt: str | None = None
        try:
            answer = session.run(
                media_chunks=media_chunks,
                questions=(Question(query_id="chat", text=query),),
                generation_options=_official_generation_options(call),
            )[0]
            rendered_prompt = session.last_rendered_prompt
        finally:
            session.close()
        if (
            bool(call.get("generate_audio", False))
            and bool(prompt_context["use_tts_template"])
            and call.get("output_audio_path")
        ):
            self._write_official_chat_audio(call, answer.text)
        if bool(call.get("return_prompt", False)):
            if not isinstance(rendered_prompt, str):
                raise RuntimeError(
                    "Memory Chat did not retain the prompt used for generation"
                )
            return answer.text, rendered_prompt
        return answer.text

    def _write_official_chat_audio(
        self,
        call: Mapping[str, Any],
        answer: str,
    ) -> None:
        """Generate speech from the answer text and write it to the requested path."""
        source_msgs = call.get("msgs")
        if not isinstance(source_msgs, Sequence):
            raise TypeError("msgs must be a sequence")
        tts_msgs = [
            deepcopy(message)
            for message in source_msgs
            if isinstance(message, Mapping) and message.get("role") == "system"
        ]
        tts_msgs.append({"role": "assistant", "content": f"<|tts_bos|>{answer}<|tts_eos|>"})
        if not hasattr(self.model, _CONTROLLER_ATTRIBUTE):
            raise RuntimeError("Memory controller state disappeared during Chat TTS")
        delattr(self.model, _CONTROLLER_ATTRIBUTE)
        try:
            self.model.chat(
                msgs=tts_msgs,
                max_new_tokens=1,
                do_sample=False,
                max_inp_length=call.get("max_inp_length", 8192),
                max_slice_nums=call.get("max_slice_nums"),
                use_image_id=call.get("use_image_id"),
                enable_thinking=False,
                use_tts_template=True,
                generate_audio=True,
                output_audio_path=call.get("output_audio_path"),
                output_tts_inputs_embeds_path=call.get("output_tts_inputs_embeds_path"),
                omni_mode=bool(call.get("omni_mode", False)),
                teacher_forcing=True,
                tts_proj_layer=call.get("tts_proj_layer", -1),
                tts_sampling_params=call.get("tts_sampling_params"),
                merge_audio_from_same_content=call.get(
                    "merge_audio_from_same_content", True
                ),
                tokenizer=call.get("tokenizer"),
                processor=call.get("processor"),
            )
        finally:
            setattr(self.model, _CONTROLLER_ATTRIBUTE, self)

    def duplex_factory_kwargs(self, kwargs: Mapping[str, Any]) -> dict[str, Any]:
        """Configure the decoder window mode required by Duplex Memory."""
        resolved = dict(kwargs)
        expected = "off"
        requested = resolved.get("sliding_window_mode", expected)
        if requested != expected:
            raise ValueError(
                f"Memory Duplex requires sliding_window_mode={expected!r}, "
                f"got {requested!r}"
            )
        resolved["sliding_window_mode"] = expected
        return resolved

    def wrap_duplex(self, official_duplex: Any) -> Any:
        """Attach long-term memory to an existing Duplex instance."""
        adapter = DuplexMemoryAdapter(official_duplex)
        runtime = MemoryDuplexRuntime(
            model=self.model,
            config=self.config.to_duplex_engine_config(),
            generate_audio=bool(getattr(official_duplex, "generate_audio", True)),
            adapter=adapter,
            auto_prepare=False,
        )
        return HuggingFaceMemoryDuplex(runtime=runtime, official=official_duplex)


class HuggingFaceMemoryDuplex:
    """Route Duplex calls through the memory runtime."""

    def __init__(self, *, runtime: MemoryDuplexRuntime, official: Any) -> None:
        self._memory_runtime = runtime
        self._official_duplex = official

    @property
    def memory_runtime(self) -> MemoryDuplexRuntime:
        """Return the underlying memory runtime."""
        return self._memory_runtime

    @property
    def official_duplex(self) -> Any:
        """Return the underlying Duplex instance."""
        return self._official_duplex

    def prepare(
        self,
        prefix_system_prompt: str | None = None,
        ref_audio: Any | None = None,
        prompt_wav_path: str | Path | None = None,
        context_previous_marker: str = "\n\nprevious: ",
        **kwargs: Any,
    ) -> Any:
        """Prepare a session and return the rendered system prompt."""
        return self._memory_runtime.prepare(
            prefix_system_prompt=prefix_system_prompt,
            ref_audio=ref_audio,
            prompt_wav_path=prompt_wav_path,
            context_previous_marker=context_previous_marker,
            **kwargs,
        )

    def streaming_prefill(
        self,
        audio_waveform: Any | None = None,
        frame_list: list[Any] | None = None,
        text_list: list[Any] | None = None,
        max_slice_nums: Any = 1,
        batch_vision_feed: bool = False,
    ) -> dict[str, Any]:
        """Prefill media and route non-empty text_list inputs to Memory."""
        return self._memory_runtime.streaming_prefill(
            audio_waveform=audio_waveform,
            frame_list=frame_list,
            text_list=text_list,
            max_slice_nums=max_slice_nums,
            batch_vision_feed=batch_vision_feed,
        )

    def streaming_generate(
        self,
        prompt_wav_path: str | Path | None = None,
        max_new_speak_tokens_per_chunk: int = 6,
        decode_mode: str = "sampling",
        temperature: float = 0.7,
        top_k: int = 100,
        top_p: float = 0.8,
        listen_prob_scale: float = 1.0,
        listen_top_k: int | None = None,
        text_repetition_penalty: float = 1.05,
        text_repetition_window_size: int = 512,
    ) -> dict[str, Any]:
        """Generate a chunk, including pending answers after the end of media."""
        return self._memory_runtime.streaming_generate(
            prompt_wav_path=(None if prompt_wav_path is None else str(prompt_wav_path)),
            max_new_speak_tokens_per_chunk=max_new_speak_tokens_per_chunk,
            decode_mode=decode_mode,
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
            listen_prob_scale=listen_prob_scale,
            listen_top_k=listen_top_k,
            text_repetition_penalty=text_repetition_penalty,
            text_repetition_window_size=text_repetition_window_size,
        )

    def as_simplex(
        self,
        reset_session: bool = True,
        reset_token2wav_cache: bool = False,
    ) -> Any:
        """Close both runtimes while preserving the original cleanup exception."""
        try:
            self._memory_runtime.close()
        except BaseException:
            try:
                self._official_duplex.as_simplex(
                    reset_session=reset_session,
                    reset_token2wav_cache=reset_token2wav_cache,
                )
            except BaseException:
                _LOGGER.exception(
                    "Official Duplex cleanup also failed after Memory close failed; "
                    "re-raising the original Memory error"
                )
            raise
        return self._official_duplex.as_simplex(
            reset_session=reset_session,
            reset_token2wav_cache=reset_token2wav_cache,
        )

    def __getattr__(self, name: str) -> Any:
        """Delegate attributes not defined here to the underlying Duplex instance."""
        official = object.__getattribute__(self, "_official_duplex")
        return getattr(official, name)


def _normalize_official_chat_call(official_call: Mapping[str, Any]) -> dict[str, Any]:
    """Expand nested keyword arguments and reject duplicate fields."""
    call = dict(official_call)
    call.pop("self", None)
    call.pop("memory_controller", None)
    extras = call.pop("kwargs", {})
    if not isinstance(extras, Mapping):
        raise TypeError("official chat kwargs must be a mapping")
    duplicates = sorted(set(call).intersection(extras))
    if duplicates:
        raise ValueError("duplicate Chat arguments: " + ", ".join(duplicates))
    call.update(extras)
    return call


def _validate_memory_chat_options(call: Mapping[str, Any]) -> None:
    """Reject options incompatible with embedding-based Memory Chat prefill."""
    if call.get("vision_hidden_states") is not None:
        raise ValueError("Memory Chat does not accept precomputed vision_hidden_states")
    for name in ("stream", "stream_input", "teacher_forcing"):
        if bool(call.get(name, False)):
            raise NotImplementedError(f"Memory Chat does not support {name}=True")
    max_slice_nums = call.get("max_slice_nums")
    if isinstance(max_slice_nums, bool) or max_slice_nums not in (None, 1):
        raise NotImplementedError(
            "Memory Chat currently supports max_slice_nums=None or 1; HD slicing "
            "requires a variable visual-token store and must not be silently "
            "treated as one slice"
        )
    _positive_integer(call.get("max_inp_length", 8192), "max_inp_length")


def _media_chunks_and_query(
    msgs: Any,
    *,
    default_query: str,
    image: Any | None = None,
    call: Mapping[str, Any] | None = None,
    normalize_content: Any | None = None,
) -> tuple[tuple[MediaChunk, ...], str, dict[str, Any]]:
    """Extract media chunks while preserving the conversation structure."""
    if not isinstance(msgs, Sequence) or isinstance(msgs, (str, bytes)) or not msgs:
        raise ValueError("Memory Chat requires a non-empty msgs sequence")
    if isinstance(msgs[0], Sequence) and not isinstance(msgs[0], Mapping):
        raise NotImplementedError("Memory Chat does not support batched msgs")
    try:
        from PIL import Image
    except ImportError as error:
        raise RuntimeError("Memory Chat requires Pillow") from error

    source_messages = []
    for message in msgs:
        if not isinstance(message, Mapping):
            raise TypeError("each Chat message must be a mapping")
        copied = dict(message)
        content = copied.get("content")
        if isinstance(content, Sequence) and not isinstance(content, (str, bytes)):
            copied["content"] = list(content)
        source_messages.append(copied)
    if image is not None and source_messages:
        first = source_messages[0]
        if not isinstance(first, Mapping):
            raise TypeError("each Chat message must be a mapping")
        if isinstance(first.get("content"), str):
            first["content"] = [image, first["content"]]

    prompt_messages: list[dict[str, Any]] = []
    chunks: list[MediaChunk] = []
    latest_user_text = ""
    inserted_memory_slot = False
    memory_message_index: int | None = None
    saw_audio = False
    for message_index, message in enumerate(source_messages):
        if not isinstance(message, Mapping):
            raise TypeError("each Chat message must be a mapping")
        role = message.get("role")
        if role not in {"system", "user", "assistant"}:
            raise ValueError(f"unsupported Chat role: {role!r}")
        if message_index == 0 and role not in {"system", "user"}:
            raise ValueError("the first Chat message must be system or user")
        content = message.get("content")
        content_was_text = isinstance(content, str)
        if callable(normalize_content):
            items = list(normalize_content(content))
        elif content_was_text:
            items = [content]
        elif isinstance(content, Sequence) and not isinstance(content, (str, bytes)):
            items = list(_normalize_content_items(content))
        else:
            raise TypeError("message content must be text or a content sequence")

        if role == "user":
            text_parts = [
                item.strip() for item in items if isinstance(item, str) and item.strip()
            ]
            if text_parts:
                latest_user_text = "\n".join(text_parts)

        preserved: list[Any] = []
        cursor = 0
        video_region_started = False
        video_region_closed = False
        while cursor < len(items):
            item = items[cursor]
            if isinstance(item, np.ndarray):
                saw_audio = True
            is_video_unit = (
                role == "user"
                and isinstance(item, Image.Image)
                and cursor + 1 < len(items)
                and isinstance(items[cursor + 1], np.ndarray)
            )
            if not is_video_unit:
                if video_region_started:
                    video_region_closed = True
                preserved.append(item)
                cursor += 1
                continue

            if video_region_closed:
                raise NotImplementedError(
                    "Memory Chat requires video frame/audio units to form one "
                    "contiguous region inside their user message"
                )
            video_region_started = True

            main_frame = item
            audio = items[cursor + 1]
            if memory_message_index is None:
                memory_message_index = message_index
            elif memory_message_index != message_index:
                raise NotImplementedError(
                    "Memory Chat requires all video frame/audio units to remain "
                    "inside one user message so their original conversation turn "
                    "is unambiguous"
                )
            saw_audio = True
            cursor += 2
            frames = [main_frame]
            if cursor < len(items) and isinstance(items[cursor], Image.Image):
                next_is_stacked = cursor + 1 == len(items) or not isinstance(
                    items[cursor + 1], np.ndarray
                )
                if next_is_stacked:
                    frames.append(items[cursor])
                    cursor += 1
            index = len(chunks)
            frame_times = tuple(
                float(index) + offset / len(frames) for offset in range(len(frames))
            )
            chunks.append(
                MediaChunk(
                    sequence_number=index,
                    start_seconds=float(index),
                    end_seconds=float(index + 1),
                    frames=tuple(frames),
                    frame_timestamps_seconds=frame_times,
                    audio_waveform=audio,
                    is_final=False,
                )
            )
            if not inserted_memory_slot:
                preserved.append({"type": _MEMORY_MEDIA_SLOT_TYPE})
                inserted_memory_slot = True

        prompt_messages.append(
            {
                "role": role,
                "content": (
                    preserved[0]
                    if content_was_text
                    and len(preserved) == 1
                    and isinstance(preserved[0], str)
                    else preserved
                ),
            }
        )

    if not chunks:
        raise ValueError(
            "Memory Chat found no [PIL frame, numpy audio] pairs in user msgs"
        )
    chunks[-1] = replace(chunks[-1], is_final=True)
    query = latest_user_text.strip() or default_query
    options = {} if call is None else dict(call)
    prompt_context = {
        "msgs": tuple(prompt_messages),
        "memory_media_slot_type": _MEMORY_MEDIA_SLOT_TYPE,
        "omni_mode": bool(options.get("omni_mode", False)),
        "use_tts_template": bool(options.get("use_tts_template", False)) or saw_audio,
        "enable_thinking": bool(options.get("enable_thinking", False)),
        "max_slice_nums": 1,
        "use_image_id": options.get("use_image_id"),
        "max_input_length": options.get("max_inp_length", 8192),
        "merge_audio_from_same_content": bool(
            options.get("merge_audio_from_same_content", True)
        ),
    }
    return tuple(chunks), query, prompt_context


def _official_content_normalizer(model: Any) -> Any | None:
    """Get the content normalizer from the loaded model module."""
    module = sys.modules.get(type(model).__module__)
    normalizer = None if module is None else getattr(module, "normalize_content", None)
    return normalizer if callable(normalizer) else None


def _normalize_content_items(content: Sequence[Any]) -> tuple[Any, ...]:
    """Normalize supported content formats into model input items."""
    normalized: list[Any] = []
    for item in content:
        if not isinstance(item, Mapping):
            normalized.append(item)
            continue
        item_type = item.get("type")
        if item_type in {"text", "input_text"}:
            normalized.append(item.get("text", ""))
        elif item_type in {"image", "image_url", "input_image"}:
            normalized.append(item.get("image", item.get("image_url")))
        elif item_type in {"audio", "input_audio"}:
            normalized.append(item.get("audio", item.get("input_audio")))
        else:
            raise ValueError(f"unsupported structured content type: {item_type!r}")
    return tuple(normalized)


def _official_generation_options(call: Mapping[str, Any]) -> dict[str, Any]:
    """Extract text generation options, excluding media and TTS controls."""
    controls = {
        "image",
        "msgs",
        "vision_hidden_states",
        "max_new_tokens",
        "max_inp_length",
        "max_slice_nums",
        "use_image_id",
        "enable_thinking",
        "use_tts_template",
        "generate_audio",
        "output_audio_path",
        "output_tts_inputs_embeds_path",
        "omni_mode",
        "teacher_forcing",
        "return_prompt",
        "tts_proj_layer",
        "tts_sampling_params",
        "merge_audio_from_same_content",
        "stream",
        "stream_input",
        "tokenizer",
        "processor",
    }
    return {key: value for key, value in call.items() if key not in controls}


def _positive_integer(value: Any, name: str) -> int:
    """Validate a positive integer parameter, excluding boolean values."""
    if isinstance(value, bool) or not isinstance(value, int):
        raise TypeError(f"{name} must be an integer")
    if value <= 0:
        raise ValueError(f"{name} must be positive")
    return value


__all__ = [
    "HuggingFaceMemoryDuplex",
    "RealtimeVenusOmniMemoryController",
    "configure_memory",
]