Update app.py

app.py

@@ -6,7 +6,7 @@ import torch
 import gradio as gr
 import pytube as pt
 import spaces
-from transformers import pipeline
+from transformers import WhisperForConditionalGeneration, WhisperProcessor, pipeline
 from huggingface_hub import model_info
 try:
     import flash_attn
@@ -25,9 +25,14 @@ print(f"Using device: {device}")
 
 @spaces.GPU(duration=60)
 def pipe(file, return_timestamps=False):
+    model = WhisperForConditionalGeneration.from_pretrained(MODEL_NAME, torch_dtype=torch.float16, low_cpu_mem_usage=True)
+    model.to(device)
+    processor = WhisperProcessor.from_pretrained(MODEL_NAME)
+    model.forward = torch.compile(model.forward, mode="reduce-overhead", fullgraph=True)
+    model.generation_config.cache_implementation = "static"
     asr = pipeline(
         task="automatic-speech-recognition",
-        model=MODEL_NAME,
+        model=model,
         chunk_length_s=30,
         device=device,
         token=auth_token,