""" YorubaApp STT API - Speech-to-Text para Yoruba usando Facebook MMS-1b-all Deploy: Hugging Face Spaces (Docker) """ from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel from transformers import Wav2Vec2ForCTC, AutoProcessor import torch import numpy as np import base64 import logging import asyncio import io import tempfile import os from typing import Optional import soundfile as sf import librosa # Configurar logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) app = FastAPI( title="YorubaApp STT API", description="Speech-to-Text para Yoruba usando Facebook MMS-1b-all", version="1.0.0" ) # CORS para permitir conexões do app app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # Variáveis globais para o modelo processor: Optional[AutoProcessor] = None model: Optional[Wav2Vec2ForCTC] = None device = "cuda" if torch.cuda.is_available() else "cpu" @app.on_event("startup") async def load_model(): """Carrega o modelo MMS na inicialização""" global processor, model logger.info("Carregando modelo MMS-1b-all...") MODEL_ID = "facebook/mms-1b-all" try: processor = AutoProcessor.from_pretrained(MODEL_ID) model = Wav2Vec2ForCTC.from_pretrained(MODEL_ID) # Configurar para Yoruba processor.tokenizer.set_target_lang("yor") model.load_adapter("yor") model.to(device) model.eval() logger.info(f"Modelo carregado com sucesso! Device: {device}") except Exception as e: logger.error(f"Erro ao carregar modelo: {e}") raise def process_audio(audio_data: bytes) -> str: """Processa áudio e retorna transcrição""" global processor, model if processor is None or model is None: raise RuntimeError("Modelo não carregado") try: # Tenta detectar e converter o formato do áudio audio_np = convert_audio_to_pcm(audio_data) if audio_np is None or len(audio_np) == 0: logger.warning("Áudio vazio após conversão") return "" # Verifica se há áudio suficiente if len(audio_np) < 1600: # Menos de 0.1s logger.warning(f"Áudio muito curto: {len(audio_np)} samples") return "" logger.info(f"Processando {len(audio_np)} samples de áudio") # Processa com MMS inputs = processor( audio_np, sampling_rate=16000, return_tensors="pt" ).to(device) with torch.no_grad(): outputs = model(**inputs).logits ids = torch.argmax(outputs, dim=-1)[0] transcription = processor.decode(ids) return transcription.strip() except Exception as e: logger.error(f"Erro no processamento de áudio: {e}") return "" def convert_audio_to_pcm(audio_data: bytes) -> Optional[np.ndarray]: """ Converte áudio de qualquer formato (WebM, MP3, M4A, WAV) para PCM 16kHz mono float32. Retorna numpy array normalizado [-1.0, 1.0] ou None em caso de erro. """ try: # Primeiro, tenta detectar se é PCM raw (como antes) # PCM 16-bit geralmente tem tamanho par e valores válidos if len(audio_data) >= 3200 and len(audio_data) % 2 == 0: # Tenta como PCM 16-bit try: pcm_data = np.frombuffer(audio_data, dtype=np.int16).astype(np.float32) / 32768.0 # Verifica se parece áudio válido (não apenas ruído) if np.abs(pcm_data).max() < 0.01: logger.info("PCM parece silêncio, tentando como arquivo codificado") else: # Verifica se tem variação razoável std = np.std(pcm_data) if std > 0.001 and std < 0.5: logger.info("Detectado como PCM raw 16-bit") return pcm_data except: pass # Tenta carregar como arquivo de áudio codificado (WebM, MP3, M4A, etc.) # Usa arquivo temporário porque algumas bibliotecas não aceitam bytes diretamente with tempfile.NamedTemporaryFile(suffix='.audio', delete=False) as tmp: tmp.write(audio_data) tmp_path = tmp.name try: # Usa librosa para carregar e resamplear automaticamente para 16kHz audio_np, sr = librosa.load(tmp_path, sr=16000, mono=True) logger.info(f"Áudio carregado via librosa: {len(audio_np)} samples, sr={sr}") return audio_np except Exception as e1: logger.warning(f"librosa falhou: {e1}") # Fallback: tenta com soundfile try: audio_np, sr = sf.read(io.BytesIO(audio_data)) # Converte para mono se estéreo if len(audio_np.shape) > 1: audio_np = np.mean(audio_np, axis=1) # Resamplea para 16kHz se necessário if sr != 16000: audio_np = librosa.resample(audio_np, orig_sr=sr, target_sr=16000) logger.info(f"Áudio carregado via soundfile: {len(audio_np)} samples") return audio_np.astype(np.float32) except Exception as e2: logger.warning(f"soundfile falhou: {e2}") # Último fallback: assume PCM raw logger.info("Fallback para PCM raw") return np.frombuffer(audio_data, dtype=np.int16).astype(np.float32) / 32768.0 finally: # Limpa arquivo temporário try: os.unlink(tmp_path) except: pass except Exception as e: logger.error(f"Erro na conversão de áudio: {e}") return None # ============== WebSocket Endpoint (Streaming) ============== @app.websocket("/ws/transcribe") async def websocket_transcribe(websocket: WebSocket): """ WebSocket endpoint para transcrição em streaming. Protocolo: - Cliente envia: {"type": "audio_chunk", "audio": ""} - Cliente envia: {"type": "end"} para finalizar - Servidor responde: {"type": "partial", "text": "..."} durante streaming - Servidor responde: {"type": "final", "text": "..."} ao finalizar """ await websocket.accept() audio_buffer = bytearray() chunk_count = 0 logger.info("WebSocket conectado") try: while True: data = await websocket.receive_json() msg_type = data.get("type") if msg_type == "audio_chunk": # Decodifica chunk base64 e adiciona ao buffer try: chunk = base64.b64decode(data.get("audio", "")) audio_buffer.extend(chunk) chunk_count += 1 logger.debug(f"Chunk {chunk_count}: +{len(chunk)} bytes, total: {len(audio_buffer)}") # Para streaming, processar a cada ~2 segundos de áudio (estimado) # Mas só para partial results - o final será processado no "end" if len(audio_buffer) >= 64000 and chunk_count % 5 == 0: transcription = process_audio(bytes(audio_buffer)) if transcription: await websocket.send_json({ "type": "partial", "text": transcription }) logger.info(f"Parcial: {transcription}") except Exception as e: logger.error(f"Erro ao processar chunk: {e}") elif msg_type == "end": # Processa áudio final completo logger.info(f"Finalizando - {chunk_count} chunks, {len(audio_buffer)} bytes") if audio_buffer: final_text = process_audio(bytes(audio_buffer)) logger.info(f"Transcrição final: '{final_text}' (len={len(final_text)})") await websocket.send_json({ "type": "final", "text": final_text }) else: logger.warning("Buffer vazio no end") await websocket.send_json({ "type": "final", "text": "" }) # Limpa buffer para próxima sessão audio_buffer = bytearray() chunk_count = 0 elif msg_type == "ping": # Keep-alive await websocket.send_json({"type": "pong"}) except WebSocketDisconnect: logger.info("WebSocket desconectado") except Exception as e: logger.error(f"Erro no WebSocket: {e}") # ============== REST Endpoint (Fallback) ============== class STTRequest(BaseModel): audio_base64: str class STTResponse(BaseModel): text: str success: bool error: Optional[str] = None @app.post("/stt", response_model=STTResponse) async def speech_to_text(request: STTRequest): """ Endpoint REST para transcrição (fallback se WebSocket não funcionar). Body: - audio_base64: Áudio em base64 (PCM 16-bit, 16kHz, mono) Returns: - text: Transcrição em Yoruba - success: Se a operação foi bem sucedida """ try: audio_data = base64.b64decode(request.audio_base64) transcription = process_audio(audio_data) return STTResponse( text=transcription, success=True ) except Exception as e: logger.error(f"Erro no endpoint REST: {e}") return STTResponse( text="", success=False, error=str(e) ) # ============== Health Check ============== @app.get("/health") async def health_check(): """Health check endpoint""" return { "status": "healthy", "model_loaded": model is not None, "device": device } @app.get("/") async def root(): """Root endpoint com informações da API""" return { "name": "YorubaApp STT API", "version": "1.0.0", "model": "facebook/mms-1b-all", "language": "Yoruba (yor)", "endpoints": { "websocket": "/ws/transcribe", "rest": "/stt", "health": "/health" } }