USER
import asyncio
from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.responses import FileResponse
import torch
import io
import os
import uuid
import aiofiles
import edge_tts
from pydub import AudioSegment
from pydub.utils import which
AudioSegment.converter = which("ffmpeg") or "ffmpeg.exe"
from transformers import (
WhisperProcessor,
WhisperForConditionalGeneration,
MT5Tokenizer,
AutoModelForSeq2SeqLM
)
from inference_wav2lip import run_wav2lip
from concurrent.futures import ThreadPoolExecutor
# Tạo thư mục lưu tạm nếu chưa có
UPLOAD_DIR = "uploads"
OUTPUT_DIR = "outputs"
os.makedirs(UPLOAD_DIR, exist_ok=True)
os.makedirs(OUTPUT_DIR, exist_ok=True)
app = FastAPI()
# Load Whisper ASR model + processor (lưu ý: có thể cần CUDA nếu máy bạn có GPU)
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()
# Load chatbot Vietnamese mT5
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()
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()
return await loop.run_in_executor(executor, _speech2text_sync, audio_bytes)
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()
return await loop.run_in_executor(executor, _chat_llm_sync, text)
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)
return buffer
@app.post("/chat")
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:
# Lưu file mặt
async with aiofiles.open(face_path, 'wb') as f_face:
await f_face.write(await face.read())
# Lưu file audio upload
async with aiofiles.open(audio_path, 'wb') as f_audio:
audio_bytes = await audio.read()
await f_audio.write(audio_bytes)
# Chạy ASR chuyển giọng nói sang text
transcription = await speech2text(audio_bytes)
if not transcription.strip():
raise HTTPException(status_code=400, detail="No speech detected")
# Chatbot trả lời
reply_text = await chat_llm(transcription)
# TTS tạo audio mp3
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())
# Convert mp3 sang wav để Wav2Lip xử lý
tts_wav_path = os.path.join(UPLOAD_DIR, f"{unique_id}_tts.wav")
audio = AudioSegment.from_file(tts_mp3_path, format="mp3")
audio.export(tts_wav_path, format="wav")
# Chạy Wav2Lip lip-sync chạy trong executor tránh block event loop
loop = asyncio.get_event_loop()
await loop.run_in_executor(executor, run_wav2lip, face_path, tts_wav_path, output_video_path)
# Trả về file video mp4
return FileResponse(output_video_path, media_type="video/mp4")
except Exception as e:
print("Error:", e)
raise HTTPException(status_code=500, detail=str(e))
finally:
# Xoá file tạm nếu tồn tại
for fpath in [face_path, ]:
if os.path.exists(fpath):
os.remove(fpath)
for fpath in [tts_mp3_path, tts_wav_path]:
if fpath and os.path.exists(fpath):
os.remove(fpath)
# (Nếu bạn cần phục vụ static files hoặc giao diện web, mount static files ở đây)
from fastapi.staticfiles import StaticFiles
app.mount("/", StaticFiles(directory="static", html=True), name="static")
# Chạy server
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=8000, log_level="info")
print("Lưu file mặt...")
print("Lưu file audio...")
print(f"Nhận file audio, size={len(audio_bytes)} bytes")
print(f"Phần chuyển speech2text...")
print(f"Transcript: {transcription}")
print(f"Phần chat bot trả lời...")
print(f"Reply: {reply_text}")
print(f"Phần text2speech...")
print(f"Phần lưu file TTS...")
print(f"Chạy Wav2Lip...")
print(f"Trả video...")ASSISTANT
Dưới đây là phiên bản tối ưu và sạch hơn của code bạn, với các điểm cải tiến:
- **Tách log rõ trong API** (dùng logging thay vì print, tránh gọi print ngoài hàm `__main__`)
- **Xóa toàn bộ file tạm (face, audio, tts mp3, wav, output video) trong `finally`** tránh sót file
- **Dùng async/await với ThreadPoolExecutor đúng cách**
- **Thêm kiểm tra và xử lý lỗi chi tiết, tránh lỗi unbound biến**
- **Cải thiện import, setup ffmpeg cho pydub**
---
```python
import asyncio
import logging
import os
import io
import uuid
from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles
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
# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# Setup ffmpeg path for pydub
AudioSegment.converter = which("ffmpeg") or "ffmpeg.exe"
# Directories
UPLOAD_DIR = "uploads"
OUTPUT_DIR = "outputs"
os.makedirs(UPLOAD_DIR, exist_ok=True)
os.makedirs(OUTPUT_DIR, exist_ok=True)
# FastAPI app
app = FastAPI()
# Mount static files if needed
app.mount("/", StaticFiles(directory="static", html=True), name="static")
# Load models once
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 for blocking calls
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")
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():
logger.warning("No speech detected")
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 = AudioSegment.from_file(tts_mp3_path, format="mp3")
audio.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"Return video file: {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=str(e))
finally:
# Xóa file tạm, kiểm tra tồn tại trước khi xóa
logger.info("Cleaning up temporary files...")
for path in [face_path, audio_path, tts_mp3_path, tts_wav_path]:
if path and os.path.exists(path):
os.remove(path)
if os.path.exists(output_video_path):
os.remove(output_video_path)
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=8000, log_level="info")
```
---
### Những điểm tối ưu chính
- Dùng Logging chuẩn (`logging`), dễ kiểm soát, không in lộn xộn ra console
- Đảm bảo các file tạm luôn được xóa sạch
- Chạy các hàm blocking trong executor tránh block event loop
- Không in log debug ngoài hàm main mà đặt trong các function
- Dùng `asyncio.get_event_loop()` một lần trong hàm async gọi executor
- Xử lý lỗi chi tiết, log stacktrace rõ ràng giúp debug dễ dàng
Nếu bạn cần thêm tối ưu về mặt hiệu năng hoặc cache model, bạn có thể mở rộng thêm nhé!