From d8224326763895bb96fda66eab4f8e766955db74 Mon Sep 17 00:00:00 2001
From: supermario_leo <leo.stack@outlook.com>
Date: Wed, 29 Apr 2026 09:25:57 +0800
Subject: [PATCH] fix(demo): free GPU tensors between batches to prevent VRAM
 accumulation

In transcribe_batch(), the large `inputs` and `output_ids` tensors were
retained in scope until the next GC cycle, preventing the CUDA allocator
from reclaiming their memory between batches. On a 24 GB GPU this causes
OOM when processing 3+ files of 25-min audio even though each individual
batch fits (issue #368).

Changes:
- Explicitly delete `inputs` and `output_ids` after the decode loop in
  transcribe_batch() and flush the allocator with empty_cache().
- Add _release_device_cache() helper (empty_cache + gc.collect) called
  after each batch in transcribe_with_batching() as well.
- Add close() method to release the model and processor at end-of-use.
- Wrap the transcribe_with_batching() call in main() with try/finally so
  close() runs even when an exception is raised mid-batch.
---
 demo/vibevoice_asr_inference_from_file.py | 53 +++++++++++++++++------
 1 file changed, 40 insertions(+), 13 deletions(-)

diff --git a/demo/vibevoice_asr_inference_from_file.py b/demo/vibevoice_asr_inference_from_file.py
index bf4f75df..da1b8abe 100644
--- a/demo/vibevoice_asr_inference_from_file.py
+++ b/demo/vibevoice_asr_inference_from_file.py
@@ -6,6 +6,7 @@
 between batch processing and single-sample processing.
 """
 
+import gc
 import os
 import sys
 import torch
@@ -202,9 +203,31 @@ def transcribe_batch(
         
         print(f"  Total generation time: {generation_time:.2f}s")
         print(f"  Average time per sample: {generation_time/batch_size:.2f}s")
-        
+
+        # Release the large input/output tensors back to the CUDA allocator so that
+        # VRAM does not accumulate across consecutive batches when processing many files
+        # (see https://github.com/microsoft/VibeVoice/issues/368).
+        del inputs, output_ids
+        self._release_device_cache()
+
         return results
-    
+
+    def _release_device_cache(self) -> None:
+        """Return any cached device memory to the allocator and run a GC cycle."""
+        if self.device == "cuda" and torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        gc.collect()
+
+    def close(self) -> None:
+        """Release the model and processor from device memory.
+
+        Call this when the instance is no longer needed — especially after processing
+        many files on GPU — to return all VRAM to the allocator before loading another
+        model in the same process.
+        """
+        del self.model, self.processor
+        self._release_device_cache()
+
     def transcribe_with_batching(
         self,
         audio_inputs: List,
@@ -246,7 +269,8 @@ def transcribe_with_batching(
                 num_beams=num_beams,
             )
             all_results.extend(batch_results)
-        
+            self._release_device_cache()
+
         return all_results
 
 
@@ -558,16 +582,19 @@ def main():
     print(f"Processing {len(all_audio_inputs)} audio(s)")
     print("="*80)
     
-    all_results = asr.transcribe_with_batching(
-        all_audio_inputs,
-        batch_size=args.batch_size,
-        max_new_tokens=args.max_new_tokens,
-        temperature=args.temperature,
-        top_p=args.top_p,
-        do_sample=do_sample,
-        num_beams=args.num_beams,
-    )
-    
+    try:
+        all_results = asr.transcribe_with_batching(
+            all_audio_inputs,
+            batch_size=args.batch_size,
+            max_new_tokens=args.max_new_tokens,
+            temperature=args.temperature,
+            top_p=args.top_p,
+            do_sample=do_sample,
+            num_beams=args.num_beams,
+        )
+    finally:
+        asr.close()
+
     # Print results
     print("\n" + "="*80)
     print("Results")