Update float_streamer.py

float_streamer.py

@@ -52,6 +52,9 @@ class FloatStreamer:
 
         self.ready = False
         self.current_dampening = 0.35
+        
+        # Global counter to keep sine waves continuous across chunks
+        self._idle_frame_counter = 0
 
     def initialize(self, model, device, opt, ref_image_tensor, wav2vec_preprocessor):
         self.model = model
@@ -112,8 +115,6 @@ class FloatStreamer:
 
         speech_array, sr = librosa.load(audio_path, sr=self.sampling_rate)
 
-        # BUG FIX: ONLY pad with silence if this is the very last sentence of the entire response.
-        # Otherwise, the avatar's mouth will awkwardly clamp shut between every sentence.
         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)])
@@ -147,7 +148,6 @@ class FloatStreamer:
                 }
             self._speech_queue.put(chunk_data)
 
-        # BUG FIX: ONLY tell the streamer to transition back to the Idle phase if this is the final sentence.
         if is_last:
             self._speech_queue.put({"type": "speech_end"})
             
@@ -195,13 +195,25 @@ class FloatStreamer:
         cfg = self._idle_cfg_scale
         if cfg > 1.0: self._idle_cfg_scale = max(1.0, cfg - 0.35)
 
+        # Ported from batch_idle.py: Recreate the lively "breathing" using a continuous sine wave!
+        dampening = torch.zeros(1, self.frames_per_chunk, 1, device=self.device)
+        damp_mid = 0.50
+        damp_amp = 0.15
+        
+        for t in range(self.frames_per_chunk):
+            global_t = self._idle_frame_counter + t
+            dampening[0, t, 0] = damp_mid + damp_amp * math.sin(global_t * 0.06)
+            
+        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, 
+            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):
+    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
@@ -237,15 +249,36 @@ class FloatStreamer:
             self.prev_x = sample[:, -self.num_prev_frames:].clone()
             self.prev_wa = wa[:, -self.num_prev_frames:].clone()
 
-        target_dampening = 1.0 if is_speech else 0.35
-
-        if abs(self.current_dampening - target_dampening) > 0.01:
-            damp_curve = torch.linspace(self.current_dampening, target_dampening, sample.shape[1], device=self.device)
-            damp_curve = damp_curve.view(1, -1, 1) 
-            sample_display = sample * damp_curve
-            self.current_dampening = target_dampening
+        # ---- FLUID DAMPENING LOGIC ----
+        if is_speech:
+            target_dampening = 1.0
+            if isinstance(self.current_dampening, torch.Tensor):
+                # Transitioning from organic idle wave -> full speech
+                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:
+                # Transitioning from flat float -> float
+                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:
-            sample_display = sample * target_dampening
+            # Applying organic breathing curve from batch_idle.py
+            if dynamic_dampening is not None:
+                if isinstance(self.current_dampening, float):
+                    # Transitioning from speech -> down to the organic idle wave
+                    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:
+                    # Continuous organic idle wave
+                    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