app.py

import os
import io
import gigaam
import gradio as gr
from mistralai import Mistral
from pydub import AudioSegment
import markdown2
from xhtml2pdf import pisa
import torch
import json
import numpy as np
import tempfile


def load_tts_model():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = torch.package.PackageImporter("v4_ru.pt").load_pickle("tts_models", "model")
    model.to(device)
    return model, device

tts_model, tts_device = load_tts_model()

# === Функция генерации озвучки ===
def synthesize_ssml_parts(parts, speaker="baya", sample_rate=48000):
    audio_segments = []
    for part in parts:
        if isinstance(part, dict):
            text_ssml = part.get("part", "")
        else:
            text_ssml = part


        # генерируем аудио
        audio_tensor = tts_model.apply_tts(
            text=text_ssml,
            speaker=speaker,   # можно поменять: aidar, baya, kseniya, xenia, eugene
            sample_rate=48000,
            put_accent=True,
            put_yo=True,
        )

        # конвертируем в numpy float32
        audio_np = audio_tensor.cpu().numpy()

        # создаём AudioSegment
        segment = AudioSegment(
            (audio_np * 32767).astype(np.int16).tobytes(),
            frame_rate=48000,
            sample_width=2,
            channels=1
        )
        audio_segments.append(segment)

    # объединяем все сегменты
    combined = sum(audio_segments)

    # сохраняем в BytesIO
    buffer = io.BytesIO()
    combined.export(buffer, format="wav")
    buffer.seek(0)
    return buffer


# === Определение длительности аудио ===
def get_audio_duration(file_path: str) -> float:
    audio = AudioSegment.from_file(file_path)
    return audio.duration_seconds


# === Транскрибация с прогресс-баром ===
def transcribe_audio(audio_file: str, progress_bar) -> str:
    os.environ["HF_TOKEN"] = os.getenv("HF_TOKEN")
    model = gigaam.load_model("v2_rnnt")

    total_duration = get_audio_duration(audio_file)
    recognition_result = model.transcribe_longform(audio_file)

    all_text = []
    last_progress = 0

    for utterance in recognition_result:
        transcription = utterance["transcription"]
        start, end = utterance["boundaries"]

        all_text.append(f"[{gigaam.format_time(start)} - {gigaam.format_time(end)}]: {transcription}")

        # обновляем прогресс
        current_progress = int((end / total_duration) * 100 * 0.9)
        if current_progress > last_progress:
            progress_bar.progress(current_progress, text="⏳ Транскрибируем аудио...")
            last_progress = current_progress

    return "\n".join(all_text)


# === Генерация PDF из Markdown ===
def create_pdf_abstract(markdown_text: str) -> bytes:
    html = markdown2.markdown(markdown_text)

    buffer = io.BytesIO()
    pisa.CreatePDF(io.StringIO(html), dest=buffer)
    buffer.seek(0)
    return buffer.read()


# === Суммаризация ===
def summarize_text(text: str, style: str, length: str) -> str:
    MISTRAL_API_KEY = os.getenv("MISTRAL_API_KEY")
    client = Mistral(api_key=MISTRAL_API_KEY)

    prompt = f"""
    Ты — умный помощник.
    Сделай {length} {style} конспект по этому тексту (на русском языке):

    {text}
    """

    response = client.chat.complete(
        model="mistral-large-latest",
        messages=[
            {"role": "system", "content": "Ты создаёшь структурированные конспекты в формате Markdown."},
            {"role": "user", "content": prompt},
        ],
    )
    return response.choices[0].message.content


def convert_summarize_text_with_ssml(text: str, style: str, length: str) -> list:
    MISTRAL_API_KEY = os.getenv("MISTRAL_API_KEY")
    client = Mistral(api_key=MISTRAL_API_KEY)

    prompt = f"""
    Ты — умный помощник.
    Разбей текст на части, где каждая часть не больше 1000 символов.
    Каждую часть оберни в SSML тег <speak>.
    - Оберни текст в тег <speak> ... </speak>.
    - **Знаки препинания не должны произноситься словами.** Вместо этого:
    - запятая → вставь `<break time="220ms"/>` в месте запятой;
    - точка / конец предложения → вставь `<break time="450ms"/>` после предложения;
    - двоеточие / точка с запятой → `<break time="350ms"/>`;
    - длинная пауза / переход к новому абзацу → `<break time="700ms"/>`.
    - **Вопросительные предложения**: в конце вопроса НЕ вставляй слово «вопросительный знак». Вместо этого оберни заключительную часть вопроса в `<prosody pitch="+12%" rate="95%">...</prosody>` чтобы задать подъём интонации, и затем `<break time="450ms"/>`.
    - **Восклицательные предложения**: выдели ключевую фразу с помощью `<emphasis level="strong">...</emphasis>` и / или `<prosody pitch="+15%" rate="105%">...</prosody>`, затем `<break time="450ms"/>`.
    - **Кавычки / прямые речи**: при открытии цитаты добавь небольшую паузу `<break time="200ms"/>`, затем внутри цитаты можно использовать `<emphasis level="moderate">` или слегка поднять `pitch` для выразительности, после цитаты — пауза `<break time="300ms"/>`. Не произноси слово «кавычки».
    - **Числа**: записывай цифры буквенно; если невозможно — используй `<say-as interpret-as="cardinal">...</say-as>` для чисел (но приоритет — слова).
    - **Не вставляй** никаких дополнительных SSML-тегов, которые могут быть не поддержаны (например, vendor-specific `<amazon:...>`). Используй только: `<speak>`, `<break>`, `<prosody>`, `<emphasis>`, `<say-as>`.
    Цифры запиши буквенно.
    Пеши без сокращений слов(не г. а год/года)
    Верни результат в JSON формате: список объектов с полем 'part'.
    Пример:
    [
    {{"part": "<speak>Текст части 1...</speak>"}},
    {{"part": "<speak>Текст части 2...</speak>"}}
    ]
    Только JSON, никаких объяснений.
    {text}

    RETURN ONLY JSON
    """

    response = client.chat.complete(
        model="pixtral-12b-2409",
        messages=[
            {"role": "system", "content": "Ты создаёшь структурированные конспекты для TTS с SSML."},
            {"role": "user", "content": f"{prompt}"},
        ],
        response_format={"type": "json_object"}
    )

    print(response.choices[0].message.content)
    # Получаем JSON
    json_text = response.choices[0].message.content
    return json.loads(json_text)

# Небольшая "заглушка" прогресс-бара, чтобы можно было вызвать transcribe_audio (он ожидает progress_bar)
class DummyProgress:
    def progress(self, *args, **kwargs):
        return None

progress_dummy = DummyProgress()

# Обёртки (НЕ меняют логику твоих функций)
def transcribe_wrapper(audio_filepath):
    if audio_filepath is None or audio_filepath == "":
        return ""
    # Gradio даёт путь к временному файлу — передаём напрямую в твою функцию
    try:
        return transcribe_audio(audio_filepath, progress_dummy)
    except Exception as e:
        return f"Transcription error: {e}"

def summarize_wrapper(audio_filepath, style, compression):
    # получаем транскрипт (если пользователь подаёт текст в дальнейшем — можно расширить)
    transcript = transcribe_wrapper(audio_filepath)
    if transcript.startswith("Transcription error"):
        return transcript
    try:
        summary = summarize_text(transcript, style, compression)
        return summary
    except Exception as e:
        return f"Summarization error: {e}"

def pdf_wrapper(markdown_text):
    try:
        pdf_bytes = create_pdf_abstract(markdown_text)
        # Gradio принимает bytes для File/Download
        return ("abstract.pdf", pdf_bytes)
    except Exception as e:
        return None

def ssml_and_tts_wrapper(markdown_text, style, compression, speaker):
    try:
        # Получаем части с SSML (твоя функция)
        parts = convert_summarize_text_with_ssml(markdown_text, style, compression)
        # Генерируем аудио (твоя функция возвращает BytesIO)
        audio_buffer = synthesize_ssml_parts(parts, speaker)
        audio_buffer.seek(0)
        # Для gr.Audio можно возвращать bytes либо путь к файлу.
        return audio_buffer.read()
    except Exception as e:
        # В случае ошибки возвращаем строку с сообщением (Gradio покажет)
        return f"TTS error: {e}"

# Запускаем интерфейс Gradio
with gr.Blocks() as demo:
    gr.Markdown("# 🎙️ Аудио-конспекты (Gradio)")

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("## 1) Загрузка аудио и создание конспекта")
            upload = gr.Audio(label="Загрузите аудио (mp3/wav)", type="filepath")
            style = gr.Dropdown(["структурированный", "в виде списка", "подробный", "короткий"], value="структурированный", label="Стиль конспекта")
            compression = gr.Slider(0, 100, 50, label="Уровень сжатия (0 — подробно, 100 — кратко)")
            btn_summarize = gr.Button("✨ Сделать конспект")
            transcript_out = gr.Textbox(label="Транскрипт (результат распознавания)", lines=6)
            summary_md = gr.Textbox(label="Конспект (Markdown)", lines=15)

            # Кнопка создаёт транскрипт и конспект
            def on_summarize(audio_fp, stl, cmp):
                tr = transcribe_wrapper(audio_fp)
                if tr.startswith("Transcription error"):
                    return tr, ""
                summary = summarize_text(tr, stl, cmp)
                return tr, summary

            btn_summarize.click(
                fn=on_summarize,
                inputs=[upload, style, compression],
                outputs=[transcript_out, summary_md],
            )

        with gr.Column(scale=1):
            gr.Markdown("## 2) Озвучка (SSML → TTS)")
            speaker = gr.Dropdown(["aidar", "baya", "kseniya", "xenia", "eugene"], value="baya", label="Выберите голос")
            btn_tts = gr.Button("🔊 Сгенерировать озвучку")
            audio_out = gr.Audio(label="Озвучка (WAV)", type="numpy")
            tts_file = gr.File(label="Скачать WAV")

            def on_tts(markdown_text, stl, cmp, sp):
                """
                Возвращает:
                - путь (str) для gr.Audio (type="filepath")
                - путь (str) для gr.File
                В случае ошибки возвращает (None, None).
                """
                try:
                    # получаем SSML-части (твоя функция)
                    parts = convert_summarize_text_with_ssml(markdown_text, stl, cmp)

                    # генерируем BytesIO с wav (твоя функция)
                    buf = synthesize_ssml_parts(parts, sp)
                    buf.seek(0)
                    data = buf.read()
                    if not data:
                        print("on_tts: audio buffer is empty")
                        return None, None

                    # записываем во временный файл и возвращаем путь
                    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
                    try:
                        tmp.write(data)
                        tmp.flush()
                    finally:
                        tmp.close()

                    # Опционально: можно вернуть имя файла без пути для отображения,
                    # но Gradio требует полный путь чтобы прочитать файл, поэтому возвращаем tmp.name
                    return tmp.name, tmp.name

                except Exception as e:
                    # полезный лог в консоль для отладки
                    print("TTS generation error:", repr(e))
                    return None, None

            # Подключаем обработчик к кнопке
            btn_tts.click(
                fn=on_tts,
                inputs=[summary_md, style, compression, speaker],
                outputs=[audio_out, tts_file],
            )


demo.launch()