Update app.py

app.py

@@ -3,6 +3,8 @@ import time
 
 from importlib.metadata import version
 
+import spaces
+
 import torch
 import torchaudio
 import torchaudio.transforms as T
@@ -11,6 +13,16 @@ import gradio as gr
 
 from transformers import AutoModelForCTC, Wav2Vec2BertProcessor
 
+use_cuda = torch.cuda.is_available()
+
+if use_cuda:
+    print('CUDA is available, setting correct inference_device variable.')
+    device = 'cuda'
+    torch_dtype = torch.float16
+else:
+    device = 'cpu'
+    torch_dtype = torch.float32
+
 # Config
 model_name = "Yehor/w2v-bert-2.0-uk-v2.1"
 
@@ -20,10 +32,6 @@ max_duration = 60
 concurrency_limit = 5
 use_torch_compile = False
 
-# Torch
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
-
 # Load the model
 asr_model = AutoModelForCTC.from_pretrained(model_name, torch_dtype=torch_dtype, device_map=device)
 processor = Wav2Vec2BertProcessor.from_pretrained(model_name)
@@ -117,6 +125,7 @@ tech_libraries = f"""
 """.strip()
 
 
+@spaces.GPU
 def inference(audio_path, progress=gr.Progress()):
     if not audio_path:
         raise gr.Error("Please upload an audio file.")