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