Spaces:
Paused
Paused
JasonZhai-9936
Claude Opus 5.5
Hand motion continuity to the next chunk from the last SHOWN frame, not the end of the computed chunk
0e35357 | """ | |
| 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 | |
| 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 | |
| ) | |
| 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 |