Update app.py

app.py

@@ -7,6 +7,7 @@ import torch
 from transformers import WhisperForConditionalGeneration, WhisperProcessor
 from indic_transliteration import sanscript
 from indic_transliteration.sanscript import transliterate
+import spaces
 
 # ---------------- CONFIG ---------------- #
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
@@ -63,17 +64,18 @@ SENTENCE_BANK = {
     ]
 }
 
-# ---------------- LOAD MODELS ---------------- #
-print("Loading Whisper models...")
+# Global variables for models (will be loaded lazily)
 whisper_models = {}
 whisper_processors = {}
 
-for lang, model_id in MODEL_CONFIGS.items():
-    print(f"Loading {lang} model: {model_id}")
-    whisper_models[lang] = WhisperForConditionalGeneration.from_pretrained(model_id).to(DEVICE)
-    whisper_processors[lang] = WhisperProcessor.from_pretrained(model_id)
-
-print("All models loaded successfully!")
+def load_model(language_choice):
+    """Load model for specific language if not already loaded"""
+    if language_choice not in whisper_models:
+        model_id = MODEL_CONFIGS[language_choice]
+        print(f"Loading {language_choice} model: {model_id}")
+        whisper_models[language_choice] = WhisperForConditionalGeneration.from_pretrained(model_id).to(DEVICE)
+        whisper_processors[language_choice] = WhisperProcessor.from_pretrained(model_id)
+        print(f"{language_choice} model loaded successfully!")
 
 # ---------------- HELPERS ---------------- #
 def get_random_sentence(language_choice):
@@ -91,7 +93,11 @@ def transliterate_to_hk(text, lang_choice):
     }
     return transliterate(text, mapping[lang_choice], sanscript.HK) if mapping[lang_choice] else text
 
+@spaces.GPU
 def transcribe_once(audio_path, language_choice, initial_prompt, beam_size, temperature, condition_on_previous_text):
+    # Load model if not already loaded
+    load_model(language_choice)
+    
     # Get the appropriate model and processor for the language
     model = whisper_models[language_choice]
     processor = whisper_processors[language_choice]
@@ -151,6 +157,7 @@ def char_level_highlight(ref, hyp):
     return "".join(out)
 
 # ---------------- MAIN ---------------- #
+@spaces.GPU
 def compare_pronunciation(audio, language_choice, intended_sentence,
                           pass1_beam, pass1_temp, pass1_condition):
     if audio is None or not intended_sentence.strip():