Upload 3 files

inference.py

@@ -2,7 +2,6 @@ import os
 import lightning as L
 import torch
 import time
-import spaces
 from snac import SNAC
 from litgpt import Tokenizer
 from litgpt.utils import (
@@ -147,8 +146,8 @@ def load_audio(path):
 
 def A1_A2_batch(fabric, audio_feature, input_ids, leng, model, text_tokenizer, step,
                 snacmodel, out_dir=None):
-
-    model.set_kv_cache(batch_size=2)
+    with fabric.init_tensor():
+        model.set_kv_cache(batch_size=2)
     tokenlist = generate_TA_BATCH(
         model,
         audio_feature,
@@ -191,8 +190,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):
-
-    model.set_kv_cache(batch_size=1)
+    with fabric.init_tensor():
+        model.set_kv_cache(batch_size=1)
     tokenlist = generate_AT(
         model,
         audio_feature,
@@ -214,8 +213,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):
-
-    model.set_kv_cache(batch_size=1)
+    with fabric.init_tensor():
+        model.set_kv_cache(batch_size=1)
     tokenlist = generate_AA(
         model,
         audio_feature,
@@ -256,8 +255,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):
-
-    model.set_kv_cache(batch_size=1)
+    with fabric.init_tensor():
+        model.set_kv_cache(batch_size=1)
     tokenlist = generate_ASR(
         model,
         audio_feature,
@@ -280,8 +279,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):
-
-    model.set_kv_cache(batch_size=1)
+    with fabric.init_tensor():
+        model.set_kv_cache(batch_size=1)
     tokenlist = generate_TA(
         model,
         None,
@@ -325,8 +324,8 @@ def T1_A2(fabric, input_ids, model, text_tokenizer, step,
 
 def T1_T2(fabric, input_ids, model, text_tokenizer, step):
 
-
-    model.set_kv_cache(batch_size=1)
+    with fabric.init_tensor():
+        model.set_kv_cache(batch_size=1)
     tokenlist = generate_TT(
         model,
         None,
@@ -356,13 +355,12 @@ def load_model(ckpt_dir, device):
     config.post_adapter = False
 
     with fabric.init_module(empty_init=False):
-        model = GPT(config, device=device)
+        model = GPT(config)
 
-    # model = fabric.setup(model)
+    model = fabric.setup(model)
     state_dict = lazy_load(ckpt_dir + "/lit_model.pth")
     model.load_state_dict(state_dict, strict=True)
-    model = model.to(device)
-    model.eval()
+    model.to(device).eval()
 
     return fabric, model, text_tokenizer, snacmodel, whispermodel
 
@@ -385,8 +383,7 @@ class OmniInference:
         for _ in self.run_AT_batch_stream(sample):
             pass
 
-    # @torch.inference_mode()
-    @spaces.GPU
+    @torch.inference_mode()
     def run_AT_batch_stream(self, 
                             audio_path, 
                             stream_stride=4,
@@ -401,7 +398,8 @@ class OmniInference:
         assert os.path.exists(audio_path), f"audio file {audio_path} not found"
         model = self.model
 
-        model.set_kv_cache(batch_size=2)
+        with self.fabric.init_tensor():
+            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)
@@ -419,7 +417,7 @@ class OmniInference:
         list_output = [[] for i in range(8)]
         tokens_A, token_T = next_token_batch(
             model,
-            audio_feature.to(torch.float32).to(device),
+            audio_feature.to(torch.float32).to(model.device),
             input_ids,
             [T - 3, T - 3],
             ["A1T2", "A1T2"],