update

inference.py

@@ -138,6 +138,7 @@ def get_input_ids_whisper_ATBatch(mel, leng, whispermodel, device):
     return torch.stack([audio_feature, audio_feature]), stacked_inputids
 
 
+@spaces.GPU
 def load_audio(path):
     audio = whisper.load_audio(path)
     duration_ms = (len(audio) / 16000) * 1000
@@ -357,7 +358,7 @@ def load_model(ckpt_dir, device):
     config.post_adapter = False
 
     with fabric.init_module(empty_init=False):
-        model = GPT(config)
+        model = GPT(config, device=device)
 
     # model = fabric.setup(model)
     state_dict = lazy_load(ckpt_dir + "/lit_model.pth")
@@ -401,8 +402,8 @@ 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)
+        # 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)