Update app.py

app.py

@@ -1,6 +1,7 @@
 import torch
 import gradio as gr
 import librosa
+import spaces
 
 from espnet2.bin.s2t_inference_language import Speech2Language
 from espnet2.bin.s2t_inference import Speech2Text as ARSpeech2Text
@@ -75,49 +76,11 @@ Please consider citing the following papers if you find our work helpful.
 # if not torch.cuda.is_available():
 #     raise RuntimeError("Please use GPU for better inference speed.")
 
-# device = "cuda"
+device = "cuda"
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# device = "cuda" if torch.cuda.is_available() else "cpu"
 
 
-try:
-    s2l = Speech2Language.from_pretrained(
-        model_tag=f"espnet/owsm_v4_medium_1B",
-        device=device,
-        nbest=1,
-    )
-except Exception as e:
-    print("File downloaded")
-
-# 2. Remove unrequired file
-import yaml
-from pathlib import Path
-import espnet_model_zoo
-
-d = "models--espnet--owsm_v4_medium_1B/snapshots/471418ddaf0b03c9ab1fd75f1f5d26fc3aea3aa9/exp/s2t_train_conv2d8_size1024_e18_d18_mel128_raw_bpe50000/config.yaml"
-p = Path(espnet_model_zoo.__file__)
-config_path = p.parent / d
-
-def remove_key(obj, key="gradient_checkpoint_layers"):
-    if isinstance(obj, dict):
-        if key in obj:
-            del obj[key]
-        for k, v in list(obj.items()):
-            remove_key(v, key)
-    elif isinstance(obj, list):
-        for item in obj:
-            remove_key(item, key)
-
-with open(config_path, "r") as f:
-    config = yaml.safe_load(f)
-
-remove_key(config)
-
-with open(config_path, "w") as f:
-    yaml.safe_dump(config, f, sort_keys=False, allow_unicode=True)
-
-print("Done! All 'gradient_checkpoint_layers' keys removed.")
-
 s2l = Speech2Language.from_pretrained(
     model_tag=f"espnet/owsm_v4_medium_1B",
     device=device,
@@ -138,12 +101,12 @@ s2t_ar = ARSpeech2Text.from_pretrained(
 
 # CTC looks okay.
 s2t_ctc = CTCSpeech2Text.from_pretrained(
-            model_tag=f"espnet/owsm_ctc_v4_1B",
-            device=device,
-            lang_sym="<eng>",
-            task_sym="<asr>",
-            predict_time=False,
-        )
+    model_tag=f"espnet/owsm_ctc_v4_1B",
+    device=device,
+    lang_sym="<eng>",
+    task_sym="<asr>",
+    predict_time=False,
+)
 
 
 iso_codes = ['abk', 'afr', 'amh', 'ara', 'asm', 'ast', 'aze', 'bak', 'bas', 'bel', 'ben', 'bos', 'bre', 'bul', 'cat', 'ceb', 'ces', 'chv', 'ckb', 'cmn', 'cnh', 'cym', 'dan', 'deu', 'dgd', 'div', 'ell', 'eng', 'epo', 'est', 'eus', 'fas', 'fil', 'fin', 'fra', 'frr', 'ful', 'gle', 'glg', 'grn', 'guj', 'hat', 'hau', 'heb', 'hin', 'hrv', 'hsb', 'hun', 'hye', 'ibo', 'ina', 'ind', 'isl', 'ita', 'jav', 'jpn', 'kab', 'kam', 'kan', 'kat', 'kaz', 'kea', 'khm', 'kin', 'kir', 'kmr', 'kor', 'lao', 'lav', 'lga', 'lin', 'lit', 'ltz', 'lug', 'luo', 'mal', 'mar', 'mas', 'mdf', 'mhr', 'mkd', 'mlt', 'mon', 'mri', 'mrj', 'mya', 'myv', 'nan', 'nep', 'nld', 'nno', 'nob', 'npi', 'nso', 'nya', 'oci', 'ori', 'orm', 'ory', 'pan', 'pol', 'por', 'pus', 'quy', 'roh', 'ron', 'rus', 'sah', 'sat', 'sin', 'skr', 'slk', 'slv', 'sna', 'snd', 'som', 'sot', 'spa', 'srd', 'srp', 'sun', 'swa', 'swe', 'swh', 'tam', 'tat', 'tel', 'tgk', 'tgl', 'tha', 'tig', 'tir', 'tok', 'tpi', 'tsn', 'tuk', 'tur', 'twi', 'uig', 'ukr', 'umb', 'urd', 'uzb', 'vie', 'vot', 'wol', 'xho', 'yor', 'yue', 'zho', 'zul']
@@ -187,6 +150,7 @@ def format_timestamp(
     )
 
 
+@spaces.GPU
 def predict(audio_path, src_lang: str, task: str, model_name: str, beam_size, long_form: bool, text_prev: str,):
     task_sym = f'<{task2code[task]}>'