ASSISTANT
Dưới đây là đoạn code FastAPI đầy đủ chuẩn, đảm bảo endpoint `/chat` chỉ nhận POST, bạn có thể test dễ dàng bằng Postman hoặc curl với method POST.
Mình cũng bổ sung logging cơ bản và cấu trúc file tạm được xóa sạch tránh rác, đồng thời chú thích rõ từng phần.
---
```python
import asyncio
import logging
import os
import io
import uuid
from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.responses import FileResponse
import torch
import aiofiles
import edge_tts
from pydub import AudioSegment
from pydub.utils import which
from transformers import (
WhisperProcessor,
WhisperForConditionalGeneration,
MT5Tokenizer,
AutoModelForSeq2SeqLM
)
from inference_wav2lip import run_wav2lip
from concurrent.futures import ThreadPoolExecutor
# Thiết lập logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# Thiết lập path ffmpeg cho pydub nếu cần
AudioSegment.converter = which("ffmpeg") or "ffmpeg.exe"
# Thư mục lưu trữ tệp
UPLOAD_DIR = "uploads"
OUTPUT_DIR = "outputs"
os.makedirs(UPLOAD_DIR, exist_ok=True)
os.makedirs(OUTPUT_DIR, exist_ok=True)
# Khởi tạo FastAPI app
app = FastAPI()
# Load các model một lần lúc khởi động server
logger.info("Loading ASR model...")
model_asr_name = "hoangdeeptry/whisper-vietnamese-3"
processor_asr = WhisperProcessor.from_pretrained(model_asr_name)
model_asr = WhisperForConditionalGeneration.from_pretrained(model_asr_name)
model_asr.eval()
logger.info("Loading chatbot model...")
model_chat_name = "google/mt5-base"
tokenizer_chat = MT5Tokenizer.from_pretrained(model_chat_name, use_fast=False)
model_chat = AutoModelForSeq2SeqLM.from_pretrained(model_chat_name)
model_chat.eval()
# ThreadPool để chạy các hàm blocking không block event loop
executor = ThreadPoolExecutor(max_workers=2)
async def speech2text(audio_bytes: bytes) -> str:
def _speech2text_sync(audio_bytes_sync):
import soundfile as sf
import torchaudio
audio_input, sample_rate = sf.read(io.BytesIO(audio_bytes_sync))
if sample_rate != 16000:
waveform, sr = torchaudio.load(io.BytesIO(audio_bytes_sync))
resampler = torchaudio.transforms.Resample(sr, 16000)
audio_input = resampler(waveform).squeeze().numpy()
inputs = processor_asr(audio_input, sampling_rate=16000, return_tensors="pt")
with torch.no_grad():
predicted_ids = model_asr.generate(inputs.input_features)
transcription = processor_asr.batch_decode(predicted_ids, skip_special_tokens=True)[0]
return transcription
loop = asyncio.get_event_loop()
transcription = await loop.run_in_executor(executor, _speech2text_sync, audio_bytes)
logger.info(f"ASR transcription: {transcription}")
return transcription
async def chat_llm(text: str) -> str:
def _chat_llm_sync(text_sync):
inputs = tokenizer_chat.encode(text_sync, return_tensors="pt", max_length=512, truncation=True)
outputs = model_chat.generate(inputs, max_length=128, num_beams=5)
reply = tokenizer_chat.decode(outputs[0], skip_special_tokens=True)
return reply
loop = asyncio.get_event_loop()
reply = await loop.run_in_executor(executor, _chat_llm_sync, text)
logger.info(f"Chatbot reply: {reply}")
return reply
async def text2speech(text: str) -> io.BytesIO:
communicate = edge_tts.Communicate(text, voice="vi-VN-Nam-AiNeural")
buffer = io.BytesIO()
async for chunk in communicate.stream():
buffer.write(chunk)
buffer.seek(0)
logger.info(f"TTS audio generated, size={buffer.getbuffer().nbytes} bytes")
return buffer
@app.post("/chat") # Chỉ nhận POST
async def chat_api(face: UploadFile = File(...), audio: UploadFile = File(...)):
unique_id = str(uuid.uuid4())
face_path = os.path.join(UPLOAD_DIR, f"{unique_id}_{face.filename}")
audio_path = os.path.join(UPLOAD_DIR, f"{unique_id}_{audio.filename}")
output_video_path = os.path.join(OUTPUT_DIR, f"{unique_id}_out.mp4")
tts_mp3_path = None
tts_wav_path = None
try:
logger.info("Saving face file...")
async with aiofiles.open(face_path, "wb") as f_face:
await f_face.write(await face.read())
logger.info("Saving audio file...")
async with aiofiles.open(audio_path, "wb") as f_audio:
audio_bytes = await audio.read()
await f_audio.write(audio_bytes)
logger.info("Running ASR...")
transcription = await speech2text(audio_bytes)
if not transcription.strip():
raise HTTPException(status_code=400, detail="No speech detected")
logger.info("Running chatbot LLM...")
reply_text = await chat_llm(transcription)
logger.info("Running TTS...")
tts_audio = await text2speech(reply_text)
tts_mp3_path = os.path.join(UPLOAD_DIR, f"{unique_id}_tts.mp3")
with open(tts_mp3_path, "wb") as f_tts:
f_tts.write(tts_audio.getbuffer())
logger.info("Converting TTS mp3 to wav...")
tts_wav_path = os.path.join(UPLOAD_DIR, f"{unique_id}_tts.wav")
audio_seg = AudioSegment.from_file(tts_mp3_path, format="mp3")
audio_seg.export(tts_wav_path, format="wav")
logger.info("Running Wav2Lip lip-sync...")
loop = asyncio.get_event_loop()
await loop.run_in_executor(executor, run_wav2lip, face_path, tts_wav_path, output_video_path)
logger.info(f"Returning output video at {output_video_path} ...")
return FileResponse(output_video_path, media_type="video/mp4")
except Exception as e:
logger.error(f"Error in /chat endpoint: {e}", exc_info=True)
raise HTTPException(status_code=500, detail=f"Server error: {str(e)}")
finally:
# Xóa file tạm
for p in [face_path, audio_path, tts_mp3_path, tts_wav_path, output_video_path]:
if p and os.path.exists(p):
os.remove(p)
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=8000, log_level="info")
```
---
## Hướng dẫn test
- Gửi **POST** request đến `http://localhost:8000/chat`.
- Body dạng `form-data` gồm 2 trường:
- `face`: file ảnh (jpg, png,...)
- `audio`: file audio (wav, mp3,...)
---
Nếu bạn gọi API bằng trình duyệt (nhập URL):
- Kết quả sẽ là lỗi **Method Not Allowed**, vì trình duyệt gửi GET, trong khi server chỉ nhận POST với `/chat`.
---
Nếu còn thắc mắc hay lỗi gì, cứ báo mình nhé!