app.py

import nemo.collections.asr as nemo_asr
from urllib.request import urlretrieve
import os
import json
from flask import Flask, request, jsonify
import logging
import torchaudio
import torch
import time
from torchaudio.backend.sox_io_backend import info

# Настройка логирования
logging.basicConfig(filename="/tmp/transcribe.log", level=logging.INFO, format='%(asctime)s %(levelname)s: %(message)s')

# Инициализация Flask
app = Flask(__name__)

# Загружаем модель
logging.info("Начало загрузки модели Parakeet...")
print("Начало загрузки модели Parakeet...")
try:
    model = nemo_asr.models.ASRModel.from_pretrained("nvidia/parakeet-tdt-0.6b-v2")
    logging.info("Модель успешно загружена")
    print("Модель успешно загружена")
except Exception as e:
    logging.error(f"Ошибка загрузки модели: {str(e)}")
    print(f"Ошибка загрузки модели: {str(e)}")
    raise e

# Функция конвертации стерео в моно и ресэмплинга
def convert_to_mono(audio_path):
    try:
        waveform, sample_rate = torchaudio.load(audio_path)
        # Конвертация в моно
        if waveform.shape[0] > 1:
            waveform = torch.mean(waveform, dim=0, keepdim=True)
        # Ресэмплинг до 16000 Гц
        if sample_rate != 16000:
            waveform = torchaudio.transforms.Resample(sample_rate, 16000)(waveform)
        mono_path = "/tmp/audio_mono.wav"
        torchaudio.save(mono_path, waveform, 16000)
        # Проверка длительности
        audio_info = info(mono_path)
        duration = audio_info.num_frames / audio_info.sample_rate
        logging.info(f"Длительность аудио: {duration:.2f} секунд")
        if duration > 300:  # 5 минут
            logging.warning("Аудио длиннее 5 минут, транскрипция может занять много времени")
        return mono_path
    except Exception as e:
        logging.error(f"Ошибка конвертации аудио: {str(e)}")
        raise e

# Flask-эндпоинт для транскрипции
@app.route('/transcribe', methods=['POST'])
def transcribe():
    try:
        data = request.get_json()
        if not data or 'audio_url' not in data:
            logging.error("Отсутствует audio_url в JSON")
            return jsonify({"error": "Missing audio_url in JSON"}), 400

        audio_url = data['audio_url']
        audio_path = "/tmp/audio.mp3"

        # Скачиваем аудиофайл
        logging.info(f"Скачивание аудиофайла: {audio_url}")
        print(f"Скачивание аудиофайла: {audio_url}")
        urlretrieve(audio_url, audio_path)

        if not os.path.exists(audio_path):
            logging.error("Не удалось скачать аудиофайл")
            return jsonify({"error": "Failed to download audio file"}), 500

        # Конвертируем в моно
        logging.info("Конвертация в моно...")
        print("Конвертация в моно...")
        mono_audio_path = convert_to_mono(audio_path)

        # Транскрибируем
        logging.info("Начало транскрипции...")
        print("Начало транскрипции...")
        start_time = time.time()
        transcription = model.transcribe([mono_audio_path], timestamps=True)
        end_time = time.time()
        logging.info(f"Транскрипция завершена за {end_time - start_time:.2f} секунд")

        result = {
            "text": transcription[0]["text"],
            "segments": transcription[0]["segments"]
        }

        # Сохраняем для отладки
        with open("/tmp/transcription.json", "w") as f:
            json.dump(result, f, indent=2, ensure_ascii=False)

        logging.info("Транскрипция успешно завершена")
        return jsonify(result)

    except Exception as e:
        logging.error(f"Ошибка транскрипции: {str(e)}")
        return jsonify({"error": str(e)}), 500

# Запуск Flask
if __name__ == "__main__":
    logging.info("Запуск Flask-сервера...")
    print("Запуск Flask-сервера...")
    app.run(host='0.0.0.0', port=7860, debug=False)