sdafd commited on
Commit
2077444
·
verified ·
1 Parent(s): 9f28c12

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +117 -0
app.py ADDED
@@ -0,0 +1,117 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # app.py
2
+ import time
3
+ import whisper
4
+ from fastapi import FastAPI, UploadFile, File, HTTPException
5
+ from fastapi.responses import FileResponse
6
+ from typing import Optional
7
+ import os
8
+ import psutil
9
+
10
+ app = FastAPI()
11
+ start_time = time.time()
12
+
13
+ # Load model during startup
14
+ @app.on_event("startup")
15
+ def load_model():
16
+ try:
17
+ app.state.model = whisper.load_model("large")
18
+ print("Model loaded successfully")
19
+ except Exception as e:
20
+ print(f"Error loading model: {e}")
21
+ raise
22
+
23
+ def format_time(seconds: float) -> str:
24
+ """Convert seconds to SRT time format"""
25
+ milliseconds = int((seconds - int(seconds)) * 1000)
26
+ hours = int(seconds // 3600)
27
+ minutes = int((seconds % 3600) // 60)
28
+ seconds = int(seconds % 60)
29
+ return f"{hours:02d}:{minutes:02d}:{seconds:02d},{milliseconds:03d}"
30
+
31
+ def generate_srt(transcript: dict) -> str:
32
+ """Generate SRT content from Whisper transcript"""
33
+ srt_content = []
34
+ index = 1
35
+ for segment in transcript['segments']:
36
+ for word in segment.get('words', []):
37
+ start = word['start']
38
+ end = word['end']
39
+ start_time = format_time(start)
40
+ end_time = format_time(end)
41
+ srt_content.append(
42
+ f"{index}\n"
43
+ f"{start_time} --> {end_time}\n"
44
+ f"{word['word'].strip()}\n\n"
45
+ )
46
+ index += 1
47
+ return "".join(srt_content)
48
+
49
+ @app.post("/transcribe")
50
+ async def transcribe_audio(
51
+ file: UploadFile = File(..., description="Audio/video file to transcribe"),
52
+ task_token: Optional[str] = None
53
+ ):
54
+ """Endpoint for submitting transcription tasks"""
55
+ try:
56
+ # Save uploaded file temporarily
57
+ temp_file = f"temp_{file.filename}"
58
+ with open(temp_file, "wb") as buffer:
59
+ content = await file.read()
60
+ buffer.write(content)
61
+
62
+ # Transcribe audio
63
+ result = app.state.model.transcribe(
64
+ temp_file,
65
+ word_timestamps=True
66
+ )
67
+
68
+ # Generate SRT file
69
+ srt_content = generate_srt(result)
70
+ srt_file = f"{temp_file}.srt"
71
+ with open(srt_file, "w") as f:
72
+ f.write(srt_content)
73
+
74
+ # Clean up temporary files
75
+ os.remove(temp_file)
76
+
77
+ return FileResponse(
78
+ srt_file,
79
+ media_type='application/x-subrip',
80
+ filename=f"{file.filename}.srt"
81
+ )
82
+
83
+ except Exception as e:
84
+ raise HTTPException(status_code=500, detail=str(e))
85
+ finally:
86
+ if os.path.exists(temp_file):
87
+ os.remove(temp_file)
88
+ if os.path.exists(srt_file):
89
+ os.remove(srt_file)
90
+
91
+ @app.get("/status")
92
+ async def get_status():
93
+ """Get server health status"""
94
+ process = psutil.Process(os.getpid())
95
+ return {
96
+ "status": "OK",
97
+ "uptime": round(time.time() - start_time, 2),
98
+ "memory_usage": f"{process.memory_info().rss / 1024 / 1024:.2f} MB",
99
+ "model_loaded": hasattr(app.state, "model"),
100
+ "active_requests": len(process.connections())
101
+ }
102
+
103
+ @app.get("/model_status")
104
+ async def get_model_status():
105
+ """Get model information"""
106
+ if not hasattr(app.state, "model"):
107
+ return {"model_status": "Not loaded"}
108
+
109
+ return {
110
+ "model_name": "Whisper large",
111
+ "device": app.state.model.device,
112
+ "parameters": f"{sum(p.numel() for p in app.state.model.parameters()):,}"
113
+ }
114
+
115
+ if __name__ == "__main__":
116
+ import uvicorn
117
+ uvicorn.run(app, host="0.0.0.0", port=8000)