Update inference.py

inference.py

@@ -2,6 +2,7 @@ import os
 import lightning as L
 import torch
 import time
+import spaces
 from snac import SNAC
 from litgpt import Tokenizer
 from litgpt.utils import (
@@ -380,6 +381,7 @@ class OmniInference:
             download_model(ckpt_dir)
         self.fabric, self.model, self.text_tokenizer, self.snacmodel, self.whispermodel = load_model(ckpt_dir, device)
 
+    @spaces.GPU
     def warm_up(self, sample='./data/samples/output1.wav'):
         for _ in self.run_AT_batch_stream(sample):
             pass