streaming.py

"""Full duplex streaming mode for MiniCPM-o 4.5 MLX.

Captures screen video + system audio, processes through the model in real-time,
and outputs text analysis with optional TTS playback.

Architecture:
    [Screen 1fps] + [Audio 16kHz] -> ChunkSynchronizer -> DuplexGenerator -> TTSPlayback
"""

import queue
import threading
import time
from typing import Optional

import mlx.core as mx
import numpy as np


class ScreenCapture:
    """Capture screen region at 1fps using mss.

    Produces (H, W, C) float32 frames resized to 448x448.
    """

    def __init__(
        self,
        out_queue: queue.Queue,
        region: Optional[tuple] = None,
        fps: float = 1.0,
        target_size: int = 448,
    ):
        self.out_queue = out_queue
        self.region = region  # (x, y, w, h) or None for primary monitor
        self.fps = fps
        self.target_size = target_size
        self._stop = threading.Event()
        self._thread: Optional[threading.Thread] = None

    def start(self):
        self._stop.clear()
        self._thread = threading.Thread(target=self._run, daemon=True)
        self._thread.start()

    def stop(self):
        self._stop.set()
        if self._thread:
            self._thread.join(timeout=2)

    def _run(self):
        import mss
        from PIL import Image

        with mss.mss() as sct:
            if self.region:
                x, y, w, h = self.region
                monitor = {"left": x, "top": y, "width": w, "height": h}
            else:
                monitor = sct.monitors[1]  # Primary monitor

            while not self._stop.is_set():
                t0 = time.time()
                screenshot = sct.grab(monitor)
                # Convert to PIL Image, resize, convert to float32
                img = Image.frombytes("RGB", screenshot.size, screenshot.rgb)
                img = img.resize(
                    (self.target_size, self.target_size), Image.BILINEAR
                )
                frame = np.array(img, dtype=np.float32) / 255.0  # (H, W, 3)

                try:
                    self.out_queue.put_nowait(
                        {"type": "video", "frame": frame, "time": time.time()}
                    )
                except queue.Full:
                    pass  # Drop frame if queue full

                elapsed = time.time() - t0
                sleep_time = max(0, (1.0 / self.fps) - elapsed)
                if sleep_time > 0:
                    self._stop.wait(sleep_time)


class AudioCapture:
    """Capture system audio at 16kHz using sounddevice.

    Uses BlackHole virtual audio device for system audio loopback on macOS.
    Produces 1-second mono float32 audio chunks.
    """

    def __init__(
        self,
        out_queue: queue.Queue,
        device: Optional[str] = None,
        sample_rate: int = 16000,
        chunk_seconds: float = 1.0,
    ):
        self.out_queue = out_queue
        self.device = device  # Device name or index
        self.sample_rate = sample_rate
        self.chunk_seconds = chunk_seconds
        self.chunk_samples = int(sample_rate * chunk_seconds)
        self._stop = threading.Event()
        self._thread: Optional[threading.Thread] = None

    def start(self):
        self._stop.clear()
        self._thread = threading.Thread(target=self._run, daemon=True)
        self._thread.start()

    def stop(self):
        self._stop.set()
        if self._thread:
            self._thread.join(timeout=2)

    def _find_device(self):
        """Find audio device by name."""
        import sounddevice as sd

        if self.device is None:
            return None  # Use default

        if isinstance(self.device, int):
            return self.device

        devices = sd.query_devices()
        for i, d in enumerate(devices):
            if self.device.lower() in d["name"].lower() and d["max_input_channels"] > 0:
                return i

        print(f"Warning: Audio device '{self.device}' not found, using default.")
        return None

    def _run(self):
        import sounddevice as sd

        device_id = self._find_device()
        buffer = np.array([], dtype=np.float32)

        def callback(indata, frames, time_info, status):
            nonlocal buffer
            if status:
                pass  # Ignore overflow/underflow
            mono = indata.mean(axis=1) if indata.ndim > 1 else indata.flatten()
            buffer = np.concatenate([buffer, mono])

        try:
            with sd.InputStream(
                device=device_id,
                channels=1,
                samplerate=self.sample_rate,
                blocksize=1024,
                callback=callback,
            ):
                while not self._stop.is_set():
                    if len(buffer) >= self.chunk_samples:
                        chunk = buffer[: self.chunk_samples].copy()
                        buffer = buffer[self.chunk_samples :]
                        try:
                            self.out_queue.put_nowait(
                                {
                                    "type": "audio",
                                    "data": chunk,
                                    "time": time.time(),
                                }
                            )
                        except queue.Full:
                            pass
                    else:
                        self._stop.wait(0.05)
        except Exception as e:
            print(f"Audio capture error: {e}")


class ChunkSynchronizer:
    """Synchronize video frames and audio into 1-second chunks.

    Pairs the latest video frame with each 1-second audio chunk.
    Runs mel processing on the audio.
    """

    def __init__(
        self,
        raw_queue: queue.Queue,
        sync_queue: queue.Queue,
        mel_processor,
    ):
        self.raw_queue = raw_queue
        self.sync_queue = sync_queue
        self.mel_processor = mel_processor
        self._stop = threading.Event()
        self._thread: Optional[threading.Thread] = None
        self._latest_frame: Optional[np.ndarray] = None

    def start(self):
        self._stop.clear()
        self._thread = threading.Thread(target=self._run, daemon=True)
        self._thread.start()

    def stop(self):
        self._stop.set()
        if self._thread:
            self._thread.join(timeout=2)

    def _run(self):
        while not self._stop.is_set():
            try:
                item = self.raw_queue.get(timeout=0.1)
            except queue.Empty:
                continue

            if item["type"] == "video":
                self._latest_frame = item["frame"]
            elif item["type"] == "audio":
                self.mel_processor.add_audio(item["data"])
                mel_chunk = self.mel_processor.get_mel_chunk()
                if mel_chunk is not None:
                    try:
                        self.sync_queue.put_nowait(
                            {
                                "video_frame": self._latest_frame,
                                "mel_chunk": mel_chunk,
                                "time": item["time"],
                            }
                        )
                    except queue.Full:
                        pass  # Drop if consumer is slow


class DuplexGenerator:
    """Main processing loop for full duplex streaming.

    Dequeues synchronized chunks, runs model inference, generates text responses,
    and optionally queues TTS audio for playback.
    """

    def __init__(
        self,
        model,
        processor,
        sync_queue: queue.Queue,
        tts_queue: Optional[queue.Queue] = None,
        temperature: float = 0.0,
        max_tokens_per_chunk: int = 50,
        enable_tts: bool = False,
    ):
        self.model = model
        self.processor = processor
        self.sync_queue = sync_queue
        self.tts_queue = tts_queue
        self.temperature = temperature
        self.max_tokens = max_tokens_per_chunk
        self.enable_tts = enable_tts
        self._stop = threading.Event()
        self._thread: Optional[threading.Thread] = None
        self.ctx = None
        self.chunk_count = 0
        self.on_text = None  # callback(text: str)
        self.on_status = None  # callback(status: dict)

    def start(self):
        self._stop.clear()
        self._thread = threading.Thread(target=self._run, daemon=True)
        self._thread.start()

    def stop(self):
        self._stop.set()
        if self._thread:
            self._thread.join(timeout=5)

    def _build_chunk_prompt(self, has_video: bool, has_audio: bool):
        """Build prompt tokens for one streaming chunk.

        Returns:
            dict with input_ids, image_bound, audio_bound
        """
        tokenizer = self.processor.tokenizer

        parts = []
        parts.append("<|im_start|>user\n")

        image_bound = []
        audio_bound = []

        # Video placeholder
        if has_video:
            # 64 query tokens for resampled image
            n_img_tokens = self.model.config.query_num  # 64
            img_placeholder = "<image>" + "<unk>" * n_img_tokens + "</image>"
            parts.append(img_placeholder)

        # Audio placeholder
        if has_audio:
            # Approximate audio tokens: ~10 after pooling for 1 second
            n_audio_tokens = 10
            audio_placeholder = (
                "<|audio_start|>" + "<unk>" * n_audio_tokens + "<|audio_end|>"
            )
            parts.append(audio_placeholder)

        parts.append("\nDescribe what you see and hear.<|im_end|>\n")
        parts.append("<|im_start|>assistant\n")

        text = "".join(parts)
        tokenized = tokenizer(text, return_tensors="np")
        input_ids = mx.array(tokenized["input_ids"])

        # Find image_bound and audio_bound positions
        ids_list = tokenized["input_ids"][0].tolist()
        unk_id = tokenizer.convert_tokens_to_ids("<unk>")

        if has_video:
            img_start_id = tokenizer.convert_tokens_to_ids("<image>")
            img_end_id = tokenizer.convert_tokens_to_ids("</image>")
            in_img = False
            start_idx = None
            for i, tok in enumerate(ids_list):
                if tok == img_start_id:
                    in_img = True
                    start_idx = i + 1
                elif tok == img_end_id and in_img:
                    image_bound.append((start_idx, i))
                    in_img = False

        if has_audio:
            audio_start_id = tokenizer.convert_tokens_to_ids("<|audio_start|>")
            audio_end_id = tokenizer.convert_tokens_to_ids("<|audio_end|>")
            in_audio = False
            start_idx = None
            for i, tok in enumerate(ids_list):
                if tok == audio_start_id:
                    in_audio = True
                    start_idx = i + 1
                elif tok == audio_end_id and in_audio:
                    audio_bound.append((start_idx, i))
                    in_audio = False

        return {
            "input_ids": input_ids,
            "image_bound": image_bound if image_bound else None,
            "audio_bound": audio_bound if audio_bound else None,
        }

    def _prepare_video_frame(self, frame: np.ndarray):
        """Prepare a video frame for model input.

        Args:
            frame: (H, W, 3) float32 frame

        Returns:
            (pixel_values, tgt_sizes, patch_attention_mask)
        """
        # Frame is already (448, 448, 3) float32
        # Add batch dimension: (1, H, W, 3)
        pv = mx.array(frame[np.newaxis, ...])

        # Compute patch sizes
        h_patches = frame.shape[0] // 14  # 32
        w_patches = frame.shape[1] // 14  # 32
        tgt_sizes = mx.array([[h_patches, w_patches]], dtype=mx.int32)

        total_patches = h_patches * w_patches
        patch_attention_mask = mx.ones((1, total_patches), dtype=mx.bool_)

        return pv, tgt_sizes, patch_attention_mask

    def _run(self):
        # Initialize streaming context
        self.ctx = self.model.init_streaming()
        self.chunk_count = 0

        while not self._stop.is_set():
            try:
                chunk = self.sync_queue.get(timeout=0.5)
            except queue.Empty:
                continue

            t0 = time.time()
            self.chunk_count += 1

            video_frame = chunk.get("video_frame")
            mel_chunk = chunk.get("mel_chunk")

            has_video = video_frame is not None
            has_audio = mel_chunk is not None

            if not has_video and not has_audio:
                continue

            # Build prompt for this chunk
            prompt = self._build_chunk_prompt(has_video, has_audio)

            # Prepare video
            pixel_values = None
            tgt_sizes = None
            patch_attention_mask = None
            if has_video:
                pixel_values, tgt_sizes, patch_attention_mask = (
                    self._prepare_video_frame(video_frame)
                )

            # Process chunk through model
            logits = self.model.process_streaming_chunk(
                ctx=self.ctx,
                video_frame=pixel_values,
                audio_chunk=mel_chunk,
                prompt_tokens=prompt["input_ids"],
                image_bound=prompt["image_bound"],
                audio_bound=prompt["audio_bound"],
                tgt_sizes=tgt_sizes,
                patch_attention_mask=patch_attention_mask,
            )

            # Generate text response
            tokens = self.model.streaming_generate(
                ctx=self.ctx,
                logits=logits,
                tokenizer=self.processor.tokenizer,
                max_tokens=self.max_tokens,
                temperature=self.temperature,
            )

            elapsed = time.time() - t0

            if tokens:
                text = self.processor.tokenizer.decode(
                    tokens, skip_special_tokens=True
                )
                if self.on_text and text.strip():
                    self.on_text(text.strip())

                # TTS if enabled
                if self.enable_tts and self.tts_queue and tokens:
                    self.tts_queue.put_nowait(
                        {"tokens": tokens, "text": text}
                    )

            if self.on_status:
                self.on_status(
                    {
                        "chunk": self.chunk_count,
                        "mode": self.ctx.mode,
                        "cache_tokens": self.ctx.total_tokens,
                        "latency_ms": int(elapsed * 1000),
                        "mem_gb": mx.get_peak_memory() / 1e9,
                    }
                )


class TTSPlayback:
    """Dequeue TTS tokens, convert to audio, and play back.

    Uses Token2wav vocoder for audio synthesis and sounddevice for playback.
    """

    def __init__(self, tts_queue: queue.Queue, sample_rate: int = 24000):
        self.tts_queue = tts_queue
        self.sample_rate = sample_rate
        self._stop = threading.Event()
        self._thread: Optional[threading.Thread] = None
        self._vocoder = None

    def start(self):
        self._stop.clear()
        self._thread = threading.Thread(target=self._run, daemon=True)
        self._thread.start()

    def stop(self):
        self._stop.set()
        if self._thread:
            self._thread.join(timeout=2)

    def _run(self):
        import sounddevice as sd

        # Try loading vocoder
        try:
            from stepaudio2 import Token2wav
            self._vocoder = Token2wav()
        except ImportError:
            print("TTSPlayback: Token2wav not available, TTS disabled.")
            return

        while not self._stop.is_set():
            try:
                item = self.tts_queue.get(timeout=0.5)
            except queue.Empty:
                continue

            tokens = item.get("tokens", [])
            if not tokens:
                continue

            try:
                import io
                import soundfile as sf

                wav_bytes = self._vocoder(tokens, None)
                waveform, sr = sf.read(io.BytesIO(wav_bytes))
                sd.play(waveform, sr, blocking=False)
            except Exception as e:
                print(f"TTS playback error: {e}")


def run_live_mode(model, processor, args):
    """Run full duplex streaming mode.

    Args:
        model: loaded MiniCPM-o model
        processor: tokenizer/processor
        args: argparse namespace with capture_region, audio_device, tts options
    """
    from mlx_vlm.models.minicpmo.audio import StreamingMelProcessor

    print("Starting live streaming mode...")
    print("Press Ctrl+C to stop.\n")

    # Create queues
    raw_queue = queue.Queue(maxsize=30)
    sync_queue = queue.Queue(maxsize=10)
    tts_queue = queue.Queue(maxsize=10) if args.tts else None

    # Create mel processor
    mel_processor = StreamingMelProcessor(sample_rate=16000)

    # Parse capture region
    region = None
    if hasattr(args, "capture_region") and args.capture_region:
        parts = args.capture_region.split(",")
        if len(parts) == 4:
            region = tuple(int(p) for p in parts)

    # Create threads
    screen = ScreenCapture(raw_queue, region=region, fps=1.0)
    audio_dev = getattr(args, "audio_device", "BlackHole")
    audio = AudioCapture(raw_queue, device=audio_dev, sample_rate=16000)
    sync = ChunkSynchronizer(raw_queue, sync_queue, mel_processor)

    generator = DuplexGenerator(
        model,
        processor,
        sync_queue,
        tts_queue=tts_queue,
        temperature=getattr(args, "temp", 0.0),
        max_tokens_per_chunk=getattr(args, "max_tokens", 50),
        enable_tts=getattr(args, "tts", False),
    )

    tts_playback = None
    if tts_queue:
        tts_playback = TTSPlayback(tts_queue)

    # Set up callbacks
    def on_text(text):
        print(f"[{generator.chunk_count}] {text}")

    def on_status(status):
        print(
            f"  >> chunk={status['chunk']} mode={status['mode']} "
            f"cache={status['cache_tokens']}tok "
            f"latency={status['latency_ms']}ms "
            f"mem={status['mem_gb']:.1f}GB",
            flush=True,
        )

    generator.on_text = on_text
    generator.on_status = on_status

    # Start all threads
    screen.start()
    audio.start()
    sync.start()
    generator.start()
    if tts_playback:
        tts_playback.start()

    print("Live mode active. Capturing screen + audio...\n")

    try:
        while True:
            time.sleep(0.5)
    except KeyboardInterrupt:
        print("\nStopping live mode...")
    finally:
        screen.stop()
        audio.stop()
        sync.stop()
        generator.stop()
        if tts_playback:
            tts_playback.stop()
        print("Live mode stopped.")