whisper ar2

app.py

@@ -5,23 +5,22 @@ import torchaudio
 from transformers import AutoProcessor, pipeline
 import io
 from pydub import AudioSegment
-from optimum.onnxruntime import ORTModelForSpeechSeq2Seq
+# from optimum.onnxruntime import ORTModelForSpeechSeq2Seq
 import numpy as np
 import uvicorn
 import time
 app = FastAPI()
+from transformers import WhisperForConditionalGeneration, WhisperProcessor
 
 # Device configuration
 device = "cuda" if torch.cuda.is_available() else "cpu"
 print(device)
-torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
-
 # Load the model and processor
-model_id = "WajeehAzeemX/whisper-small-ar2_onnx"
-model = ORTModelForSpeechSeq2Seq.from_pretrained(
-    model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
+model_id = "whitefox123/whisper-small-ar2"
+model = WhisperForConditionalGeneration.from_pretrained(
+    model_id
 )
-processor = AutoProcessor.from_pretrained(model_id)
+processor = WhisperProcessor.from_pretrained(model_id)
 
 
 pipe = pipeline(
@@ -29,7 +28,6 @@ pipe = pipeline(
     model=model,
     tokenizer=processor.tokenizer,
     feature_extractor=processor.feature_extractor,
-    torch_dtype=torch_dtype,
 )
 
 @app.post("/transcribe/")