Create app.py

app.py

@@ -0,0 +1,138 @@
+from fastapi import FastAPI, File, UploadFile, HTTPException, Form
+from fastapi.responses import JSONResponse, PlainTextResponse
+import whisper
+import tempfile
+import os
+from pathlib import Path
+from typing import Optional
+import subprocess
+import json
+
+app = FastAPI(title="Subtitle Generator API")
+
+# 加载Whisper模型（首次启动时会下载模型）
+model_size = os.getenv("WHISPER_MODEL", "base")  # tiny, base, small, medium, large
+model = whisper.load_model(model_size)
+
+def format_timestamp(seconds: float) -> str:
+    """将秒数转换为SRT时间戳格式"""
+    hours = int(seconds // 3600)
+    minutes = int((seconds % 3600) // 60)
+    secs = int(seconds % 60)
+    millisecs = int((seconds - int(seconds)) * 1000)
+    return f"{hours:02d}:{minutes:02d}:{secs:02d},{millisecs:03d}"
+
+def generate_srt(segments) -> str:
+    """生成SRT格式字幕"""
+    srt_content = []
+    for i, segment in enumerate(segments, 1):
+        start_time = format_timestamp(segment['start'])
+        end_time = format_timestamp(segment['end'])
+        text = segment['text'].strip()
+        srt_content.append(f"{i}\n{start_time} --> {end_time}\n{text}\n")
+    return "\n".join(srt_content)
+
+def generate_vtt(segments) -> str:
+    """生成WebVTT格式字幕"""
+    vtt_content = ["WEBVTT\n"]
+    for i, segment in enumerate(segments):
+        start_time = format_timestamp(segment['start']).replace(',', '.')
+        end_time = format_timestamp(segment['end']).replace(',', '.')
+        text = segment['text'].strip()
+        vtt_content.append(f"{start_time} --> {end_time}\n{text}\n")
+    return "\n".join(vtt_content)
+
+@app.get("/")
+async def root():
+    return {
+        "message": "Subtitle Generator API",
+        "endpoints": {
+            "/generate_subtitle": "POST - Upload audio/video and get subtitles",
+            "/health": "GET - Health check"
+        },
+        "supported_formats": ["srt", "vtt", "txt", "json"],
+        "model": model_size
+    }
+
+@app.get("/health")
+async def health_check():
+    return {"status": "healthy", "model": model_size}
+
+@app.post("/generate_subtitle")
+async def generate_subtitle(
+    file: UploadFile = File(...),
+    output_format: str = Form("srt"),
+    language: Optional[str] = Form(None)
+):
+    """
+    生成字幕文件
+    
+    - file: 音频或视频文件
+    - output_format: 输出格式 (srt, vtt, txt, json)
+    - language: 语言代码 (如 en, zh, ja)，不指定则自动检测
+    """
+    
+    # 验证格式
+    valid_formats = ["srt", "vtt", "txt", "json"]
+    if output_format not in valid_formats:
+        raise HTTPException(
+            status_code=400, 
+            detail=f"Invalid format. Choose from: {valid_formats}"
+        )
+    
+    # 保存上传文件到临时目录
+    with tempfile.NamedTemporaryFile(delete=False, suffix=Path(file.filename).suffix) as tmp_file:
+        content = await file.read()
+        tmp_file.write(content)
+        tmp_path = tmp_file.name
+    
+    try:
+        # 转码选项
+        transcribe_options = {}
+        if language:
+            transcribe_options["language"] = language
+        
+        # 使用Whisper转录音频
+        result = model.transcribe(tmp_path, **transcribe_options)
+        
+        # 根据请求格式生成字幕
+        if output_format == "srt":
+            subtitle_text = generate_srt(result["segments"])
+            media_type = "text/plain"
+            filename = "subtitle.srt"
+        elif output_format == "vtt":
+            subtitle_text = generate_vtt(result["segments"])
+            media_type = "text/plain"
+            filename = "subtitle.vtt"
+        elif output_format == "txt":
+            subtitle_text = result["text"]
+            media_type = "text/plain"
+            filename = "subtitle.txt"
+        elif output_format == "json":
+            subtitle_text = json.dumps({
+                "text": result["text"],
+                "segments": result["segments"],
+                "language": result["language"]
+            }, ensure_ascii=False, indent=2)
+            media_type = "application/json"
+            filename = "subtitle.json"
+        
+        return PlainTextResponse(
+            content=subtitle_text,
+            media_type=media_type,
+            headers={
+                "Content-Disposition": f"attachment; filename={filename}"
+            }
+        )
+    
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Transcription failed: {str(e)}")
+    
+    finally:
+        # 清理临时文件
+        if os.path.exists(tmp_path):
+            os.unlink(tmp_path)
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)