Kirito_Sub / main.py
abhiswork's picture
Update main.py
8047cb5 verified
Raw
History Blame Contribute Delete
6.92 kB
from fastapi import FastAPI, UploadFile, File, Depends, HTTPException, status
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from fastapi.responses import HTMLResponse
from faster_whisper import WhisperModel
import tempfile
import os
app = FastAPI()
security = HTTPBearer()
API_KEY = os.environ.get("API_KEY")
# 1. Authentication Dependency
def verify_api_key(credentials: HTTPAuthorizationCredentials = Depends(security)):
if credentials.credentials != API_KEY:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="Invalid API Key"
)
return credentials.credentials
# 2. Load Model globally (Runs on CPU with int8 quantization for speed)
print("Loading model...")
model = WhisperModel("base", device="cpu", compute_type="int8")
print("Model loaded!")
# 3. Transcription Endpoint
@app.post("/v1/transcribe")
async def transcribe(
file: UploadFile = File(...),
token: str = Depends(verify_api_key)
):
# Save uploaded file to a temporary location
with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
tmp.write(await file.read())
tmp_path = tmp.name
try:
# Transcribe with word_timestamps=True to get word-level timing
segments, info = model.transcribe(tmp_path, beam_size=5, word_timestamps=True)
# Convert the segments generator into a list of dictionaries
segments_data = []
full_text = []
for segment in segments:
full_text.append(segment.text)
seg_dict = {
"text": segment.text,
"start": segment.start,
"end": segment.end,
"words": []
}
# If word timestamps are available, include them
if segment.words:
for word in segment.words:
seg_dict["words"].append({
"word": word.word,
"start": word.start,
"end": word.end
})
segments_data.append(seg_dict)
return {
"text": " ".join(full_text).strip(),
"language": info.language,
"segments": segments_data
}
finally:
# Clean up the temp file
os.remove(tmp_path)
# 4. Health Check Endpoint (For Automated Ping)
@app.get("/health")
def health_check():
return {"status": "awake"}
# 5. Simple UI for Testing
HTML_TEMPLATE = """
<!DOCTYPE html>
<html>
<head>
<title>Kirito_Sub Transcription Tester</title>
<style>
body { font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif; padding: 40px; max-width: 600px; margin: auto; background-color: #0b0f19; color: #e1e7ef; }
.container { background: #1e293b; padding: 25px; border-radius: 12px; box-shadow: 0 4px 6px rgba(0, 0, 0, 0.3); }
h2 { margin-top: 0; color: #38bdf8; }
label { display: block; margin-top: 15px; font-weight: bold; font-size: 14px; color: #94a3b8; }
input[type="password"], input[type="file"] {
width: 100%; padding: 10px; margin-top: 5px;
background: #0f172a; border: 1px solid #334155; color: white; border-radius: 6px; box-sizing: border-box;
}
button {
margin-top: 20px; padding: 12px; width: 100%;
background: #0ea5e9; color: white; border: none; border-radius: 6px;
font-size: 16px; font-weight: bold; cursor: pointer; transition: background 0.2s;
}
button:hover { background: #0284c7; }
#result { margin-top: 20px; padding: 15px; background: #0f172a; border-left: 4px solid #10b981; border-radius: 4px; display: none; }
#loading { display: none; margin-top: 15px; color: #fbbf24; font-weight: bold; text-align: center; }
</style>
</head>
<body>
<div class="container">
<h2>Transcription Tester</h2>
<p style="font-size: 14px; color: #94a3b8;">Test your Whisper model directly in the browser.</p>
<label>Secret API Key:</label>
<input type="password" id="apiKey" placeholder="Enter the key from your Space secrets">
<label>Audio File:</label>
<input type="file" id="audioFile" accept="audio/*,video/*">
<button onclick="testTranscription()">Transcribe Audio</button>
<div id="loading">Processing audio on CPU... This may take a minute! ⏳</div>
<div id="result">
<strong style="color: #10b981;">Result:</strong>
<p id="transcriptionText" style="margin-top: 10px; line-height: 1.5;"></p>
<small id="langInfo" style="color: #64748b;"></small>
</div>
</div>
<script>
async function testTranscription() {
const apiKey = document.getElementById('apiKey').value;
const fileInput = document.getElementById('audioFile');
const loadingDiv = document.getElementById('loading');
const resultDiv = document.getElementById('result');
const textOutput = document.getElementById('transcriptionText');
if (!apiKey || fileInput.files.length === 0) {
alert("Please enter your API key and select an audio file.");
return;
}
// Reset UI
resultDiv.style.display = "none";
loadingDiv.style.display = "block";
textOutput.innerText = "";
// Prepare form data
const formData = new FormData();
formData.append("file", fileInput.files[0]);
try {
const response = await fetch('/v1/transcribe', {
method: 'POST',
headers: { 'Authorization': 'Bearer ' + apiKey },
body: formData
});
if (!response.ok) {
const errText = await response.text();
throw new Error(`HTTP ${response.status}: ${errText}`);
}
const data = await response.json();
// Display results
loadingDiv.style.display = "none";
resultDiv.style.display = "block";
textOutput.innerText = data.text;
document.getElementById('langInfo').innerText = "Detected Language: " + data.language;
} catch (err) {
loadingDiv.style.display = "none";
resultDiv.style.display = "block";
resultDiv.style.borderLeftColor = "#ef4444"; // Red for error
textOutput.innerText = "Error: " + err.message;
}
}
</script>
</body>
</html>
"""
@app.get("/", response_class=HTMLResponse)
def serve_home():
return HTML_TEMPLATE