Blame for float_streamer.py

"""
FLOAT Streaming Engine - Continuous real-time lipsync generation.

Generates frames continuously in a background thread:
- Idle mode: silent audio → natural idle motion (breathing, blinking, micro-movements)
- Speech mode: TTS audio features injected → lip-synced speech motion
"""
import contextlib
import os
import sys
import math
import time
import threading
import queue
import logging
import random
import numpy as np

import torch
import torch.nn.functional as F
import cv2
import librosa

import timeline

logger = logging.getLogger(__name__)

# Max idle frames buffered ahead while a reply is between TTS chunks (0.2s at 25fps)
MID_REPLY_IDLE_LEAD_FRAMES = 5

FLOAT_REPO_PATH = "/app/float_repo"
if FLOAT_REPO_PATH not in sys.path:
    sys.path.insert(0, FLOAT_REPO_PATH)


class FloatStreamer:
    def __init__(self):
        self.model = None
        self.device = None
        self.opt = None

        self.s_r = None          
        self.r_s = None          
        self.s_r_feats = None    

        self.prev_x = None       
        self.prev_wa = None      

        self.wav2vec_preprocessor = None

        self._frame_buffer = queue.Queue(maxsize=75) 
        self._speech_queue = queue.Queue()
        # Out-of-band control events (e.g. audio preload hints) sent to the
        # client immediately, bypassing the frame buffer's playback queue.
        self._event_queue = queue.Queue(maxsize=50)

        self._running = False
        self._thread = None
        self._lock = threading.Lock()
        self._is_speaking = False

        self.fps = 25.0
        self.frames_per_chunk = 50      
        self.num_prev_frames = 10
        self.sampling_rate = 16000

        self.ready = False
        self.current_dampening = 0.35
        
        # --- Procedural Idle State Machine Variables ---
        self._idle_frame_counter = 0
        self._idle_is_active_state = False
        self._idle_state_frames_left = 125 # Start with 5 seconds of calm
        
        # Target values we want to reach
        self._idle_target_mid = 0.32
        self._idle_target_amp = 0.06
        
        # Current values we are smoothly animating
        self._idle_current_mid = 0.32
        self._idle_current_amp = 0.06

    def initialize(self, model, device, opt, ref_image_tensor, wav2vec_preprocessor):
        self.model = model
        self.device = device
        self.opt = opt
        self.wav2vec_preprocessor = wav2vec_preprocessor

        self.fps = opt.fps
        self.frames_per_chunk = int(opt.wav2vec_sec * opt.fps)
        self.num_prev_frames = opt.num_prev_frames
        self.sampling_rate = opt.sampling_rate

        self._encode_reference(ref_image_tensor)

        with torch.no_grad():
            silence_audio = torch.zeros(1, int(opt.wav2vec_sec * self.sampling_rate)).to(device)
            T_silence = self.frames_per_chunk
            self._silence_wa = self.model.audio_encoder.inference(silence_audio, seq_len=T_silence)
            self._silence_we = self.model.emotion_encoder.predict_emotion(silence_audio).unsqueeze(1)
            logger.info(f"[STREAMER] Silence features encoded: wa={self._silence_wa.shape}")

        self.prev_x = torch.zeros(1, self.num_prev_frames, opt.dim_w).to(device)
        self.prev_wa = torch.zeros(1, self.num_prev_frames, opt.dim_w).to(device)

        # Dedicated RNG rather than the device's default generator, which Breeze
        # TTS's CUDA graphs also use — sharing it with another model's graph
        # capture/replay is what raised "Offset increment outside graph capture".
        self._rng = torch.Generator(device=device)
        self._rng.manual_seed(0)

        # Render on a high-priority CUDA stream. Breeze TTS runs on the default
        # (normal-priority) stream in the same process, and while it generates
        # the tail chunks of a reply its big-model kernels were starving FLOAT's
        # many small ones — the first clip rendered ~2.6-3x slower than real time
        # and played slowed/out of sync. Priority lets the GPU schedule FLOAT's
        # kernels first whenever both are waiting.
        self._stream = None
        if torch.cuda.is_available() and str(device).startswith("cuda"):
            _, high = torch.cuda.Stream.priority_range()
            self._stream = torch.cuda.Stream(device=device, priority=high)

        self.ready = True
        logger.info("[STREAMER] Initialized and ready")

    def _encode_reference(self, ref_tensor):
        with torch.no_grad():
            s = ref_tensor.to(self.device)
            self.s_r, r_s_lambda, self.s_r_feats = self.model.encode_image_into_latent(s)
            self.r_s = self.model.motion_autoencoder.dec.direction(r_s_lambda)

    def update_reference(self, ref_tensor):
        with self._lock:
            self._encode_reference(ref_tensor)
            self.prev_x = torch.zeros(1, self.num_prev_frames, self.opt.dim_w).to(self.device)
            self.prev_wa = torch.zeros(1, self.num_prev_frames, self.opt.dim_w).to(self.device)
        logger.info("[STREAMER] Reference updated")

    def start(self):
        if self._running: return
        self._running = True
        self._thread = threading.Thread(target=self._generation_loop, daemon=True)
        self._thread.start()
        logger.info("[STREAMER] Generation loop started")

    def stop(self):
        self._running = False
        if self._thread: self._thread.join(timeout=5)

    def inject_speech(self, audio_path: str, audio_url: str = None, clear_buffer: bool = True, is_last: bool = True,
                      tag: str = None, turn: int = None):
        """tag ("c2/4") and turn label this TTS chunk's timeline entries all the
        way through render, WS delivery and client playback."""
        t0 = time.time()
        
        if clear_buffer:
            flushed = self._frame_buffer.qsize()
            self.drain_buffer()
            timeline.mark(f"{tag} flushed {flushed} buffered idle frames so speech plays immediately", turn)

        speech_array, sr = librosa.load(audio_path, sr=self.sampling_rate)

        if is_last:
            pad_samples = int(0.5 * self.sampling_rate)
            speech_array = np.concatenate([speech_array, np.zeros(pad_samples, dtype=speech_array.dtype)])
            
        audio_tensor = torch.FloatTensor(speech_array).unsqueeze(0).to(self.device)

        with torch.no_grad():
            T = math.ceil(audio_tensor.shape[-1] * self.fps / self.sampling_rate)
            wa_full = self.model.audio_encoder.inference(audio_tensor, seq_len=T)
            we = self.model.emotion_encoder.predict_emotion(audio_tensor).unsqueeze(1)

        num_chunks = math.ceil(T / self.frames_per_chunk)

        for i in range(num_chunks):
            start_frame = i * self.frames_per_chunk
            end_frame = min(start_frame + self.frames_per_chunk, T)
            wa_chunk = wa_full[:, start_frame:end_frame]

            if wa_chunk.shape[1] < self.frames_per_chunk:
                wa_chunk = F.pad(wa_chunk, (0, 0, 0, self.frames_per_chunk - wa_chunk.shape[1]), mode='replicate')

            chunk_data = {
                "wa": wa_chunk, "we": we, "chunk_index": i,
                "total_chunks": num_chunks, "actual_frames": end_frame - start_frame,
                "tag": tag, "turn": turn,
            }

            if i == 0:
                chunk_data["speech_start"] = {
                    "type": "speech_start", "audio_url": audio_url,
                    "duration": len(speech_array) / self.sampling_rate, "num_chunks": num_chunks,
                    "tag": tag, "turn": turn,
                }
            self._speech_queue.put(chunk_data)

        if is_last:
            self._speech_queue.put({"type": "speech_end", "turn": turn})

        timeline.mark(
            f"{tag} audio features encoded in {time.time()-t0:.2f}s → {num_chunks} render chunk(s) "
            f"({T} frames) queued{' + speech_end' if is_last else ''}", turn)

    def push_event(self, event: dict):
        try: self._event_queue.put_nowait(event)
        except queue.Full: pass  # no client draining it — drop

    def get_event(self):
        try: return self._event_queue.get_nowait()
        except queue.Empty: return None

    def get_frame(self, timeout=0.1):
        try: return self._frame_buffer.get(timeout=timeout)
        except queue.Empty: return None

    @contextlib.contextmanager
    def _render_stream(self):
        """Run one chunk's GPU work on the high-priority render stream. Inputs
        (audio features, reference latents) are produced on the default stream by
        other threads, so wait for it first; each chunk ends with per-frame .cpu()
        copies, which sync this stream before the inputs can be freed."""
        if self._stream is None:
            yield
            return
        self._stream.wait_stream(torch.cuda.default_stream(self.device))
        with torch.cuda.stream(self._stream):
            yield

    def _generation_loop(self):
        logger.info("[STREAMER] Generation loop running")
        self._idle_cfg_scale = 1.0  

        while self._running:
            try:
                speech_data = None
                try: speech_data = self._speech_queue.get_nowait()
                except queue.Empty: pass

                if speech_data and speech_data.get("type") == "speech_end":
                    self._is_speaking = False
                    self._idle_cfg_scale = self.opt.a_cfg_scale
                    timeline.mark(
                        f"speech_end queued behind {self._frame_buffer.qsize()} buffered frames — back to idle",
                        speech_data.get("turn"))
                    try:
                        self._frame_buffer.put({"type": "speech_end", "turn": speech_data.get("turn")}, timeout=2.0)
                    except queue.Full:
                        logger.warning("[STREAMER] Frame buffer full, dropping speech_end event")
                    continue

                with self._render_stream():
                    if speech_data and "wa" in speech_data:
                        self._is_speaking = True
                        speech_start_event = speech_data.get("speech_start")

                        self._generate_chunk(
                            wa=speech_data["wa"], we=speech_data["we"],
                            actual_frames=speech_data.get("actual_frames", self.frames_per_chunk),
                            is_speech=True, nfe=5,
                            speech_start_event=speech_start_event,
                            tag=f"{speech_data.get('tag')} render {speech_data['chunk_index']+1}/{speech_data['total_chunks']}",
                            turn=speech_data.get("turn"),
                        )

                    else:
                        self._generate_idle_chunk()

            except Exception as e:
                logger.error(f"[STREAMER] Generation error: {e}", exc_info=True)
                time.sleep(0.5)

    def _generate_idle_chunk(self):
        cfg = self._idle_cfg_scale
        if cfg > 1.2: self._idle_cfg_scale = max(1.2, cfg - 0.2)

        dampening = torch.zeros(1, self.frames_per_chunk, 1, device=self.device)
        
        for t in range(self.frames_per_chunk):
            # 1. State Machine Timer
            self._idle_state_frames_left -= 1
            if self._idle_state_frames_left <= 0:
                # Toggle state
                self._idle_is_active_state = not self._idle_is_active_state
                
                if self._idle_is_active_state:
                    # Switch to Dynamic/Active (High swaying)
                    self._idle_target_mid = random.uniform(0.48, 0.52)
                    self._idle_target_amp = random.uniform(0.12, 0.16)
                    self._idle_state_frames_left = random.randint(125, 250) # Hold for 5-10s
                else:
                    # Switch to Static/Calm (Subtle breathing)
                    self._idle_target_mid = random.uniform(0.28, 0.34)
                    self._idle_target_amp = random.uniform(0.04, 0.08)
                    self._idle_state_frames_left = random.randint(125, 375) # Hold for 5-15s

            # 2. Smooth Interpolation (Lerp)
            # This smoothly glides the current values toward the target over ~2 seconds
            lerp_speed = 0.02
            self._idle_current_mid += (self._idle_target_mid - self._idle_current_mid) * lerp_speed
            self._idle_current_amp += (self._idle_target_amp - self._idle_current_amp) * lerp_speed

            # 3. Apply to Sine Wave
            global_t = self._idle_frame_counter + t
            dampening[0, t, 0] = self._idle_current_mid + self._idle_current_amp * math.sin(global_t * 0.05)
            
        self._idle_frame_counter += self.frames_per_chunk

        self._generate_chunk(
            wa=self._silence_wa, we=self._silence_we, actual_frames=self.frames_per_chunk,
            a_cfg_scale=cfg, e_cfg_scale=1.0, is_speech=False, nfe=3,
            dynamic_dampening=dampening
        )

    @torch.no_grad()
    def _generate_chunk(self, wa, we, actual_frames=None, a_cfg_scale=None, e_cfg_scale=None, is_speech=False, nfe=None, speech_start_event=None, dynamic_dampening=None, tag=None, turn=None):
        if actual_frames is None: actual_frames = self.frames_per_chunk
        if a_cfg_scale is None: a_cfg_scale = self.opt.a_cfg_scale
        if e_cfg_scale is None: e_cfg_scale = self.opt.e_cfg_scale
        if nfe is None: nfe = self.opt.nfe

        t0 = time.time()

        with self._lock:
            r_s = self.r_s
            s_r = self.s_r
            s_r_feats = self.s_r_feats
            prev_x = self.prev_x
            prev_wa = self.prev_wa

        x0 = torch.randn(1, self.frames_per_chunk, self.opt.dim_w, device=self.device, generator=self._rng)
        time_steps = torch.linspace(0, 1, nfe, device=self.device)

        def sample_chunk(tt, zt):
            out = self.model.fmt.forward_with_cfv(
                t=tt.unsqueeze(0), x=zt, wa=wa, wr=r_s, we=we,
                prev_x=prev_x, prev_wa=prev_wa,
                a_cfg_scale=a_cfg_scale, r_cfg_scale=self.opt.r_cfg_scale, e_cfg_scale=e_cfg_scale,
            )
            return out[:, self.num_prev_frames:]

        from torchdiffeq import odeint
        trajectory = odeint(sample_chunk, x0, time_steps, atol=self.opt.ode_atol, rtol=self.opt.ode_rtol, method=self.opt.torchdiffeq_ode_method)
        sample = trajectory[-1]  

        t_ode = time.time() - t0

        if not is_speech:
            sp = sample.permute(0, 2, 1) 
            sp = F.avg_pool1d(F.pad(sp, (1, 1), mode='replicate'), kernel_size=3, stride=1)
            sample = sp.permute(0, 2, 1)

        if is_speech:
            target_dampening = 1.0
            if isinstance(self.current_dampening, torch.Tensor):
                start_val = self.current_dampening[0, -1, 0].item()
                damp_curve = torch.linspace(start_val, 1.0, sample.shape[1], device=self.device).view(1, -1, 1)
                sample_display = sample * damp_curve
            elif abs(self.current_dampening - target_dampening) > 0.01:
                damp_curve = torch.linspace(self.current_dampening, target_dampening, sample.shape[1], device=self.device).view(1, -1, 1)
                sample_display = sample * damp_curve
            else:
                sample_display = sample * target_dampening
            end_dampening = 1.0
        else:
            if dynamic_dampening is not None:
                if isinstance(self.current_dampening, float):
                    fade = torch.linspace(1.0, 0.0, sample.shape[1], device=self.device).view(1, -1, 1)
                    damp_curve = self.current_dampening * fade + dynamic_dampening * (1.0 - fade)
                    sample_display = sample * damp_curve
                else:
                    damp_curve = dynamic_dampening
                    sample_display = sample * dynamic_dampening
                end_dampening = damp_curve
            else:
                sample_display = sample * 0.35
                end_dampening = 0.35

        t_dec = time.time()
        frames_pushed = 0
        t_wait = 0.0  # time blocked on a full frame buffer — not render cost

        if speech_start_event:
            # Frames already buffered ahead of speech_start play before this
            # chunk's audio is even requested — that's added reply latency.
            ahead = self._frame_buffer.qsize()
            timeline.mark(
                f"{speech_start_event.get('tag')} speech_start queued behind {ahead} buffered frames "
                f"(~{ahead / self.fps:.2f}s of video before it)", turn)
            try:
                self._frame_buffer.put(speech_start_event, timeout=2.0)
            except queue.Full:
                logger.warning("[STREAMER] Frame buffer full, dropping speech_start event")

        n_shown = 0  # frames actually rendered (a cut-short idle chunk shows only some)
        for t in range(min(actual_frames, sample.shape[1])):
            
            if not is_speech:
                # Mid-reply (between TTS chunks — _is_speaking stays True until
                # speech_end), keep only a tiny idle lead: the next chunk's frames
                # queue BEHIND whatever idle is buffered, so pre-rendering seconds
                # of idle here (the buffer holds 3s) became a pause of up to 3s
                # between chunks whenever TTS for the next chunk finished after
                # the current chunk's video had already rendered.
                while (self._is_speaking and self._running and self._speech_queue.empty()
                       and self._frame_buffer.qsize() >= MID_REPLY_IDLE_LEAD_FRAMES):
                    time.sleep(0.01)
                if not self._speech_queue.empty():
                    break

            s_r_d_t = s_r + sample_display[:, t]
            img_t, _ = self.model.motion_autoencoder.dec(s_r_d_t, alpha=None, feats=s_r_feats)

            frame = img_t.squeeze().permute(1, 2, 0).detach().clamp(-1, 1)
            frame = ((frame + 1) / 2 * 255).to(torch.uint8).cpu().numpy()
            _, jpeg_data = cv2.imencode('.jpg', cv2.cvtColor(frame, cv2.COLOR_RGB2BGR), [cv2.IMWRITE_JPEG_QUALITY, 85])
            frame_bytes = jpeg_data.tobytes()
            n_shown = t + 1

            if is_speech:
                try:
                    t_put = time.time()
                    self._frame_buffer.put(frame_bytes, timeout=2.0)
                    t_wait += time.time() - t_put
                    frames_pushed += 1
                except queue.Full:
                    pass
            else:
                try:
                    self._frame_buffer.put_nowait(frame_bytes)
                    frames_pushed += 1
                except queue.Full:
                    pass  

        t_dec_done = time.time()

        self._commit_continuity(sample, wa, n_shown, end_dampening)

        if is_speech:
            # Rendering slower than the video it produces means the frame buffer
            # drains and video falls behind the audio.
            video_s = frames_pushed / self.fps
            # Exclude time spent waiting for buffer space: a full buffer means
            # rendering is AHEAD of playback, not behind it.
            took = t_dec_done - t0 - t_wait
            slow = f" ⚠ SLOWER THAN REAL TIME ({took/video_s:.2f}x)" if video_s and took > video_s else ""
            timeline.mark(
                f"{tag} rendered {frames_pushed} frames ({video_s:.2f}s video) in {took:.2f}s "
                f"[ode {t_ode:.2f}s, decode {t_dec_done-t_dec-t_wait:.2f}s, waited {t_wait:.2f}s for buffer space]{slow}", turn)
        else:
            # Idle chunks render back-to-back forever — keep them out of INFO.
            logger.debug(
                f"[STREAMER] Idle chunk rendered | frames={frames_pushed} | "
                f"ode={t_ode:.2f}s | encode={t_dec_done-t_dec:.2f}s | total={t_dec_done-t0:.2f}s"
            )

    def _commit_continuity(self, sample, wa, n_shown, end_dampening):
        """Hand the next chunk the motion state at the last frame that was
        actually SHOWN — its prev_x/prev_wa conditioning and the dampening level
        it blends from. Using the end of the full computed chunk instead made a
        visible jump whenever a chunk was cut short: e.g. the idle filler between
        TTS chunks of a reply, interrupted after ~5 frames when the next chunk
        arrives (the next chunk then continued from idle frame 50's pose and
        ~0.3 motion strength), or a clip's last partial chunk (padding frames)."""
        if n_shown <= 0:
            return  # nothing shown — the previous state is still what's on screen
        n = self.num_prev_frames
        with self._lock:
            if n_shown >= n:
                self.prev_x = sample[:, n_shown - n:n_shown].clone()
                self.prev_wa = wa[:, n_shown - n:n_shown].clone()
            else:
                self.prev_x = torch.cat([self.prev_x, sample[:, :n_shown]], dim=1)[:, -n:].clone()
                self.prev_wa = torch.cat([self.prev_wa, wa[:, :n_shown]], dim=1)[:, -n:].clone()
        if isinstance(end_dampening, torch.Tensor):
            self.current_dampening = end_dampening[:, :n_shown].clone()
        else:
            self.current_dampening = end_dampening

    def drain_buffer(self):
        while not self._frame_buffer.empty():
            try: self._frame_buffer.get_nowait()
            except queue.Empty: break
        while not self._speech_queue.empty():
            try: self._speech_queue.get_nowait()
            except queue.Empty: break

    def interrupt(self):
        """Barge-in: drop any queued/buffered speech and snap back to idle
        immediately, instead of waiting for a speech_end marker to work through
        the queue naturally. A chunk already mid-render can't be aborted, but
        nothing queued behind it will play."""
        self.drain_buffer()
        self._is_speaking = False
        self._idle_cfg_scale = self.opt.a_cfg_scale
        timeline.mark("streamer interrupted — queued speech dropped, back to idle")

    def signal_speech_end(self):
        """Tell the generation loop speech is over with no audio to inject — used
        when a reply's final TTS chunk fails to synthesize, so speech_end still
        reaches the client instead of leaving it waiting for one that never comes."""
        self._speech_queue.put({"type": "speech_end"})

_streamer = None

def get_streamer() -> FloatStreamer:
    global _streamer
    if _streamer is None:
        _streamer = FloatStreamer()
    return _streamer