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 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

logger = logging.getLogger(__name__)

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()

        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)

        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):
        t0 = time.time()
        
        if clear_buffer:
            self.drain_buffer()
            logger.info("[STREAMER] Buffer flushed for immediate speech playback")

        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,
            }

            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,
                }
            self._speech_queue.put(chunk_data)

        if is_last:
            self._speech_queue.put({"type": "speech_end"})
            
        logger.info(f"[STREAMER] Speech injected: {T} frames, {num_chunks} chunks, is_last={is_last}")

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

    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  
                    logger.info("[STREAMER] Speech ended, transitioning to idle")
                    continue

                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
                    )

                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):
        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)
        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)

        with self._lock:
            self.prev_x = sample[:, -self.num_prev_frames:].clone()
            self.prev_wa = wa[:, -self.num_prev_frames:].clone()

        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
            self.current_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:
                    sample_display = sample * dynamic_dampening
                self.current_dampening = dynamic_dampening
            else:
                sample_display = sample * 0.35
                self.current_dampening = 0.35

        t_dec = time.time()
        frames_pushed = 0

        if speech_start_event:
            try:
                self._frame_buffer.put(speech_start_event, timeout=2.0)
            except queue.Full:
                logger.warning("[STREAMER] Frame buffer full, dropping speech_start event")

        for t in range(min(actual_frames, sample.shape[1])):
            
            if not is_speech and not self._speech_queue.empty():
                if self._frame_buffer.qsize() > 25:
                    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()

            if is_speech:
                try:
                    self._frame_buffer.put(frame_bytes, timeout=2.0)
                    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()

    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

_streamer = None

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