Update main.py

main.py

@@ -16,6 +16,13 @@ import logging
 import ffmpeg
 import tempfile
 
+# Force HF cache to /tmp
+os.environ["HF_HOME"] = "/tmp/huggingface"
+os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface/transformers"
+os.environ["HF_HUB_CACHE"] = "/tmp/huggingface/hub"
+
+os.makedirs(os.environ["HF_HOME"], exist_ok=True)
+
 # Silence all transformers and huggingface logging
 logging.getLogger("transformers").setLevel(logging.ERROR)
 logging.getLogger("urllib3").setLevel(logging.ERROR)
@@ -108,8 +115,10 @@ def correct_grammar(text, grammar_pipeline):
     return '. '.join([r['generated_text'] for r in results])
 
 # ========== Initialize Models ==========
-processor, whisper_model = load_whisper_model("small")
-grammar_pipeline = load_grammar_model()
+# processor, whisper_model = load_whisper_model("small")
+# grammar_pipeline = load_grammar_model()
+processor = WhisperProcessor.from_pretrained(model_name)
+model = WhisperForConditionalGeneration.from_pretrained(model_name)
 
 # ========== Warm-Up Models ==========
 def warm_up_models():