Spaces:
Sleeping
Sleeping
File size: 1,536 Bytes
3af511d 9792cba 3af511d 3b048f8 9792cba |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 |
import os
import tempfile
import numpy as np
import torch
from fastapi import FastAPI, File, UploadFile
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
os.environ["TRANSFORMERS_CACHE"] = "/app/cache"
app = FastAPI(
title = "Whisper API",
redirect_slashes=False
)
# Device configuration
device = "cuda:0" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
# Load Whisper model
model_id = "openai/whisper-large-v3-turbo"
model = AutoModelForSpeechSeq2Seq.from_pretrained(
model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
).to(device)
processor = AutoProcessor.from_pretrained(model_id)
pipe = pipeline(
"automatic-speech-recognition",
model=model,
tokenizer=processor.tokenizer,
feature_extractor=processor.feature_extractor,
torch_dtype=torch_dtype,
device=device
)
@app.get("/")
async def root():
return {"message": "Welcome to Whisper API!"}
@app.post("/transcribe/")
async def transcribe_audio(file: UploadFile = File(...)):
try:
# Save the uploaded file temporarily
with tempfile.NamedTemporaryFile(delete=True, suffix=".wav") as temp_audio:
temp_audio.write(await file.read())
temp_audio.flush()
# Transcribe the audio
result = pipe(temp_audio.name, return_timestamps="word")
return {"transcription": result["chunks"]}
except Exception as e:
return {"error": str(e)}
|