流式传输

services/a2f_api/api.py

@@ -1,34 +1,73 @@
-from flask import Flask, request, jsonify
-from flask_cors import CORS
+import json
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import StreamingResponse
+from pydantic import BaseModel
 from text_to_blendshapes_service import TextToBlendShapesService
 
-app = Flask(__name__)
-CORS(app)
+app = FastAPI()
 
-@app.route('/health', methods=['GET'])
-def health():
-    return jsonify({'status': 'ok'})
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
 
-@app.route('/text-to-blendshapes', methods=['POST'])
-def text_to_blendshapes():
+class TextRequest(BaseModel):
+    text: str
+    language: str = 'zh-CN'
+    segment: bool = False
+    split_punctuations: str = None
+    max_sentence_length: int = None
+    first_sentence_split_size: int = None
+
+@app.get('/health')
+async def health():
+    return {'status': 'ok'}
+
+@app.post('/text-to-blendshapes')
+async def text_to_blendshapes(request: TextRequest):
     try:
-        data = request.get_json()
-        if not data or 'text' not in data:
-            return jsonify({'success': False, 'error': 'Missing text'}), 400
-
-        text = data['text']
-        language = data.get('language', 'zh-CN')
-
-        service = TextToBlendShapesService(lang=language)
-        result = service.text_to_blend_shapes(text)
-
-        return jsonify(result)
-
+        service = TextToBlendShapesService(lang=request.language)
+        result = service.text_to_blend_shapes(
+            request.text,
+            segment=request.segment,
+            split_punctuations=request.split_punctuations,
+            max_sentence_length=request.max_sentence_length
+        )
+        return result
     except Exception as e:
         import traceback
         traceback.print_exc()
-        return jsonify({'success': False, 'error': str(e)}), 500
+        return {'success': False, 'error': str(e)}
+
+@app.post('/text-to-blendshapes/stream')
+async def text_to_blendshapes_stream(request: TextRequest):
+    async def generate():
+        service = TextToBlendShapesService(lang=request.language)
+        try:
+            for message in service.iter_text_to_blend_shapes_stream(
+                request.text,
+                split_punctuations=request.split_punctuations,
+                max_sentence_length=request.max_sentence_length,
+                first_sentence_split_size=request.first_sentence_split_size
+            ):
+                yield json.dumps(message) + "\n"
+        except Exception as e:
+            yield json.dumps({'type': 'error', 'message': str(e)}) + "\n"
+
+    return StreamingResponse(
+        generate(),
+        media_type='application/x-ndjson',
+        headers={
+            'Cache-Control': 'no-cache',
+            'X-Accel-Buffering': 'no'
+        }
+    )
 
 if __name__ == '__main__':
+    import uvicorn
     print("Text to BlendShapes API: http://localhost:5001")
-    app.run(host='0.0.0.0', port=5001, debug=True)
+    uvicorn.run(app, host='0.0.0.0', port=5001)