Update app.py

app.py

@@ -10,20 +10,15 @@ MAX_AUDIO_MINS = 30  # maximum audio input in minutes
 
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
 torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
-use_flash_attention_2 = is_flash_attn_2_available()
+attn_implementation = "flash_attention_2" if is_flash_attn_2_available() else "sdpa"
 
 model = AutoModelForSpeechSeq2Seq.from_pretrained(
-    "openai/whisper-large-v2", torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True, use_flash_attention_2=use_flash_attention_2
+    "openai/whisper-large-v2", torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True, attn_implementation=attn_implementation
 )
 distilled_model = AutoModelForSpeechSeq2Seq.from_pretrained(
-    "distil-whisper/distil-large-v2", torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True, use_flash_attention_2=use_flash_attention_2
+    "distil-whisper/distil-large-v2", torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True, attn_implementation=attn_implementation
 )
 
-if not use_flash_attention_2:
-    # use flash attention from pytorch sdpa
-    model = model.to_bettertransformer()
-    distilled_model = distilled_model.to_bettertransformer()
-
 processor = AutoProcessor.from_pretrained("openai/whisper-large-v2")
 
 model.to(device)