Update inference.py

inference.py

@@ -147,8 +147,8 @@ def load_audio(path):
 
 def A1_A2_batch(fabric, audio_feature, input_ids, leng, model, text_tokenizer, step,
                 snacmodel, out_dir=None):
-    with fabric.init_tensor():
-        model.set_kv_cache(batch_size=2)
+
+    model.set_kv_cache(batch_size=2)
     tokenlist = generate_TA_BATCH(
         model,
         audio_feature,
@@ -191,8 +191,8 @@ def A1_A2_batch(fabric, audio_feature, input_ids, leng, model, text_tokenizer, s
 
 
 def A1_T2(fabric, audio_feature, input_ids, leng, model, text_tokenizer, step):
-    with fabric.init_tensor():
-        model.set_kv_cache(batch_size=1)
+
+    model.set_kv_cache(batch_size=1)
     tokenlist = generate_AT(
         model,
         audio_feature,
@@ -214,8 +214,8 @@ def A1_T2(fabric, audio_feature, input_ids, leng, model, text_tokenizer, step):
 
 def A1_A2(fabric, audio_feature, input_ids, leng, model, text_tokenizer, step,
           snacmodel, out_dir=None):
-    with fabric.init_tensor():
-        model.set_kv_cache(batch_size=1)
+
+    model.set_kv_cache(batch_size=1)
     tokenlist = generate_AA(
         model,
         audio_feature,
@@ -256,8 +256,8 @@ def A1_A2(fabric, audio_feature, input_ids, leng, model, text_tokenizer, step,
 
 
 def A1_T1(fabric, audio_feature, input_ids, leng, model, text_tokenizer, step):
-    with fabric.init_tensor():
-        model.set_kv_cache(batch_size=1)
+
+    model.set_kv_cache(batch_size=1)
     tokenlist = generate_ASR(
         model,
         audio_feature,
@@ -280,8 +280,8 @@ def A1_T1(fabric, audio_feature, input_ids, leng, model, text_tokenizer, step):
 
 def T1_A2(fabric, input_ids, model, text_tokenizer, step,
           snacmodel, out_dir=None):
-    with fabric.init_tensor():
-        model.set_kv_cache(batch_size=1)
+
+    model.set_kv_cache(batch_size=1)
     tokenlist = generate_TA(
         model,
         None,
@@ -325,8 +325,8 @@ def T1_A2(fabric, input_ids, model, text_tokenizer, step,
 
 def T1_T2(fabric, input_ids, model, text_tokenizer, step):
 
-    with fabric.init_tensor():
-        model.set_kv_cache(batch_size=1)
+
+    model.set_kv_cache(batch_size=1)
     tokenlist = generate_TT(
         model,
         None,
@@ -386,6 +386,7 @@ class OmniInference:
             pass
 
     @torch.inference_mode()
+    @spaces.GPU
     def run_AT_batch_stream(self, 
                             audio_path, 
                             stream_stride=4,
@@ -400,8 +401,7 @@ class OmniInference:
         assert os.path.exists(audio_path), f"audio file {audio_path} not found"
         model = self.model
 
-        with self.fabric.init_tensor():
-            model.set_kv_cache(batch_size=2)
+        model.set_kv_cache(batch_size=2)
 
         mel, leng = load_audio(audio_path)
         audio_feature, input_ids = get_input_ids_whisper_ATBatch(mel, leng, self.whispermodel, self.device)