Slice the kvcache when getting the model state for a given voice (#32)

Related to #14

pocket_tts/models/tts_model.py

@@ -261,6 +261,66 @@ class TTSModel(nn.Module):
         conditioning = F.linear(latents, self.flow_lm.speaker_proj_weight)
         return conditioning
 
+    def _slice_kv_cache(self, model_state: dict, num_frames: int) -> None:
+        """Slice KV cache to only keep the first num_frames elements.
+
+        This optimizes memory usage when caching voice states by discarding
+        unused cache capacity beyond the actual audio prompt length.
+
+        Args:
+            model_state: The model state dict containing KV caches for all modules
+            num_frames: Number of frames to keep in the KV cache
+        """
+        original_size = 0
+        sliced_size = 0
+        for module_name, module_state in model_state.items():
+            if "cache" in module_state:
+                # KV cache has shape [2, batch_size, sequence_length, num_heads, dim_per_head]
+                cache = module_state["cache"]
+                original_size += cache.numel() * cache.element_size()
+                # Slice to keep only the first num_frames positions
+                module_state["cache"] = cache[:, :, :num_frames, :, :].clone()
+                sliced_size += module_state["cache"].numel() * module_state["cache"].element_size()
+
+        memory_saved_mb = (original_size - sliced_size) / (1024 * 1024)
+        logger.info(
+            f"Sliced KV cache from {original_size / (1024 * 1024):.1f} MB to {sliced_size / (1024 * 1024):.1f} MB "
+            f"(saved {memory_saved_mb:.1f} MB)"
+        )
+
+    def _expand_kv_cache(self, model_state: dict, sequence_length: int) -> None:
+        """Expand KV cache back to full sequence_length for generation.
+
+        When a model state is retrieved from cache with sliced KV caches,
+        this method expands them back to the full size needed for generation.
+
+        Args:
+            model_state: The model state dict containing potentially sliced KV caches
+            sequence_length: Target sequence length to expand caches to
+        """
+        for module_name, module_state in model_state.items():
+            if "cache" in module_state:
+                cache = module_state["cache"]
+                # KV cache has shape [2, batch_size, current_length, num_heads, dim_per_head]
+                current_length = cache.shape[2]
+                if current_length < sequence_length:
+                    # Create expanded cache filled with NaN for unused positions
+                    expanded_cache = torch.full(
+                        (
+                            cache.shape[0],
+                            cache.shape[1],
+                            sequence_length,
+                            cache.shape[3],
+                            cache.shape[4],
+                        ),
+                        float("NaN"),
+                        device=cache.device,
+                        dtype=cache.dtype,
+                    )
+                    # Copy existing data to the beginning
+                    expanded_cache[:, :, :current_length, :, :] = cache
+                    module_state["cache"] = expanded_cache
+
     @torch.no_grad
     def _decode_audio_worker(self, latents_queue: queue.Queue, result_queue: queue.Queue):
         """Worker thread function for decoding audio latents from queue with immediate streaming."""
@@ -417,6 +477,9 @@ class TTSModel(nn.Module):
         if copy_state:
             model_state = copy.deepcopy(model_state)
 
+        # Expand sliced KV caches back to full size for generation
+        self._expand_kv_cache(model_state, sequence_length=1000)
+
         # Set up multithreaded generation and decoding
         latents_queue = queue.Queue()
         result_queue = queue.Queue()
@@ -632,6 +695,10 @@ class TTSModel(nn.Module):
         with display_execution_time("Prompting audio"):
             self._run_flow_lm_and_increment_step(model_state=model_state, audio_conditioning=prompt)
 
+        # Optimize memory by slicing KV cache to only keep frames from the audio prompt
+        num_audio_frames = prompt.shape[1]
+        self._slice_kv_cache(model_state, num_audio_frames)
+
         return model_state