Update float_lipsync.py

float_lipsync.py

@@ -229,6 +229,7 @@ class FloatLipsync:
             img = self._crop_face(img)
 
         self.preprocessed_ref_image = self.transform(image=img)['image'].unsqueeze(0).to(self.device)
+        self._original_ref_image = self.preprocessed_ref_image.clone()
         logger.info(f"[FLOAT] Reference image preprocessed: {self.preprocessed_ref_image.shape}")
 
     def _crop_face(self, img: np.ndarray) -> np.ndarray:
@@ -273,7 +274,7 @@ class FloatLipsync:
             return False
 
     @torch.no_grad()
-    def generate(self, audio_path: str, output_path: str = None, emo: str = 'S2E') -> str:
+    def generate(self, audio_path: str, output_path: str = None, emo: str = 'S2E', chain_frames: bool = False) -> str:
         """
         Generate a lipsync video from an audio file.
 
@@ -281,6 +282,7 @@ class FloatLipsync:
             audio_path: Path to the input WAV audio file
             output_path: Where to save the output video (auto-generated if None)
             emo: Emotion mode ('S2E' for speech-to-emotion auto-detect)
+            chain_frames: If True, use the last frame as reference for next call
 
         Returns:
             Path to the generated video file
@@ -298,6 +300,11 @@ class FloatLipsync:
 
         # 1. Load and preprocess audio
         speech_array, sr = librosa.load(audio_path, sr=self.opt.sampling_rate)
+
+        # Pad with 0.5s silence so FLOAT eases the face back to neutral at the end
+        pad_samples = int(0.5 * sr)
+        speech_array = np.concatenate([speech_array, np.zeros(pad_samples, dtype=speech_array.dtype)])
+
         processed_audio = self.wav2vec_preprocessor(
             speech_array, sampling_rate=sr, return_tensors='pt'
         ).input_values[0].unsqueeze(0).to(self.device)
@@ -324,13 +331,25 @@ class FloatLipsync:
         )['d_hat']
         logger.info(f"[FLOAT] Model inference: {time.time() - gen_start:.2f}s")
 
-        # 3. Save video with audio
+        # 3. Chain frames: use last frame as reference for next chunk
+        if chain_frames and d_hat.shape[0] > 0:
+            # d_hat is (T, C, H, W) in [-1, 1] range — same as preprocessed_ref_image
+            self.preprocessed_ref_image = d_hat[-1:].clone().to(self.device)
+            logger.info(f"[FLOAT] Chained last frame as next reference")
+
+        # 4. Save video with audio
         self._save_video(d_hat, output_path, audio_path)
 
         logger.info(f"[FLOAT] ✓ Total generation: {time.time() - inference_start:.2f}s")
         logger.info(f"[FLOAT] ✓ Output: {output_path}")
         return output_path
 
+    def reset_reference(self):
+        """Reset the reference image back to the original ref.png."""
+        if hasattr(self, '_original_ref_image') and self._original_ref_image is not None:
+            self.preprocessed_ref_image = self._original_ref_image.clone()
+            logger.info("[FLOAT] Reference reset to original")
+
     def _save_video(self, vid_tensor: torch.Tensor, video_path: str, audio_path: str):
         """Save video frames tensor to mp4 with audio using ffmpeg."""
         # Prepare frames: (T, C, H, W) -> (T, H, W, C) uint8