src/amuseme/transcriber.py

"""
transcriber.py — Audio → word-level timestamps + frame metadata generation using Whisper and MiniCPM5-1B
"""
import os
import re
import math
from dataclasses import dataclass

from faster_whisper import WhisperModel

from amuseme.logger import get_logger
from .frame_metadata import FrameAnim, Frame, SongFrames

logger = get_logger("transcriber")


@dataclass
class Word:
    text: str
    start: float
    end: float


@dataclass
class FrameMeta:
    words: list[Word]
    frame_animation: FrameAnim


# Gap in seconds between words that triggers a new line fallback
SILENCE_GAP = 0.5
# Max words per line before forcing a break fallback
MAX_WORDS_PER_LINE = 7
# Lines shorter than this get merged into a neighbor so the on-screen text
# changes less often and reads as a complete phrase rather than a fragment.
MIN_WORDS_PER_LINE = 4


# ─── Observability helpers ──────────────────────────────────────────────────
# Human-readable renderings of each step's input/output, written to the shared
# logger (logs/model_io.log + console) so the whole pipeline is inspectable.

def format_transcript(words: list["Word"]) -> str:
    """Whisper output as an indexed, timestamped, readable transcript."""
    if not words:
        return "(no words)"
    sentence = " ".join(w.text for w in words)
    rows = [f"  [{i:>3}] {w.start:6.2f}-{w.end:6.2f}s  {w.text}" for i, w in enumerate(words)]
    return f"transcript: {sentence}\n" + "\n".join(rows)


def format_frames(metas: list["FrameMeta"]) -> str:
    """FrameMeta list as readable lines with their frame-level effects."""
    if not metas:
        return "(no frames)"
    out = []
    for fi, m in enumerate(metas):
        text = " ".join(w.text for w in m.words)
        head = f"  Frame {fi:>2}: \"{text}\""
        if m.frame_animation != FrameAnim.none:
            head += f"  [frame_animation={m.frame_animation.value}]"
        out.append(head)
    return "\n".join(out)


_model = None

def _load_model(model_size: str = "large-v3"):
    global _model
    if _model is None:
        if model_size == "turbo":
            model_size = "large-v3-turbo"

        device = "cpu" if os.environ.get("FORCE_CPU") == "1" else "cuda"
        logger.info(f"Loading Whisper {model_size} on {device}...")
        compute_type = "float16" if device == "cuda" else "int8"
        try:
            _model = WhisperModel(model_size, device=device, compute_type=compute_type)
        except Exception as e:
            if device == "cuda":
                # CUDA/cuDNN library mismatch (e.g. on ZeroGPU). Retry on CPU.
                logger.warning(f"CUDA load failed ({e}). Falling back to CPU int8.")
                try:
                    _model = WhisperModel(model_size, device="cpu", compute_type="int8")
                except Exception as e2:
                    logger.error(f"CPU fallback also failed: {e2}")
                    raise
            else:
                logger.warning(f"Failed to load with {compute_type}: {e}. Retrying with float32.")
                _model = WhisperModel(model_size, device=device, compute_type="float32")
    return _model


_demucs_model = None

def _separate_vocals(audio_path: str) -> str:
    """Isolate vocals from a song using Demucs. Returns path to vocals wav."""
    import tempfile
    import torch
    import torchaudio
    from demucs.pretrained import get_model
    from demucs.apply import apply_model

    global _demucs_model
    if _demucs_model is None:
        logger.info("Loading Demucs htdemucs model...")
        _demucs_model = get_model('htdemucs')
        device = "cpu" if os.environ.get("FORCE_CPU") == "1" else ("cuda" if torch.cuda.is_available() else "cpu")
        _demucs_model.to(device)

    device = next(_demucs_model.parameters()).device
    wav, sr = torchaudio.load(audio_path)

    if sr != _demucs_model.samplerate:
        wav = torchaudio.functional.resample(wav, sr, _demucs_model.samplerate)
        sr = _demucs_model.samplerate

    wav = wav.to(device)

    logger.info(f"Separating vocals from {audio_path}...")
    with torch.no_grad():
        sources = apply_model(_demucs_model, wav.unsqueeze(0), split=True, overlap=0.25)[0]

    vocals_idx = _demucs_model.sources.index('vocals')
    vocals = sources[vocals_idx].cpu()

    vocals_path = os.path.join(tempfile.gettempdir(), "amuseme_vocals.wav")
    torchaudio.save(vocals_path, vocals, sr)
    logger.info(f"Vocal separation complete. Saved to {vocals_path}")
    return vocals_path


def transcribe(audio_path: str, lyrics_override: str = "", model_size: str = "large-v3",
               use_demucs: bool = True, condition_on_previous_text: bool = True, use_vad: bool = True,
               theme: str = "Dark", visual_prompt: str = "") -> list[FrameMeta]:
    """
    Transcribe audio and return timestamped display frames with metadata.
    """
    if use_demucs:
        vocals_path = _separate_vocals(audio_path)
    else:
        logger.info(f"Skipping Demucs separation, using original audio.")
        vocals_path = audio_path

    model = _load_model(model_size)
    logger.info(
        "=== WHISPER INPUT ===\n"
        f"  audio={audio_path}\n"
        f"  model={model_size}  demucs={use_demucs}  vad={use_vad}  "
        f"condition_on_previous_text={condition_on_previous_text}\n"
        f"  lyrics_override={'yes' if lyrics_override.strip() else 'no'}"
    )

    segments_gen, info = model.transcribe(
        vocals_path, 
        word_timestamps=True,
        condition_on_previous_text=condition_on_previous_text,
        compression_ratio_threshold=1.50,
        temperature=(0.0, 0.2),
        no_speech_threshold=0.8,
        vad_filter=use_vad,
        vad_parameters=dict(threshold=0.05, min_silence_duration_ms=2000, speech_pad_ms=2000, min_speech_duration_ms=50) if use_vad else None,
        initial_prompt=lyrics_override if lyrics_override.strip() else None
    )
    
    segments = list(segments_gen)
    
    raw_words: list[Word] = []
    for segment in segments:
        for w in getattr(segment, "words", []):
            clean = re.sub(r"[^\w\s''-]", "", w.word).strip()
            if clean:
                raw_words.append(Word(text=clean, start=w.start, end=w.end))

    if not raw_words:
        logger.warning(f"No words extracted from {audio_path}")
        return []

    logger.info(
        f"=== WHISPER OUTPUT === {len(raw_words)} words "
        f"(detected language: {getattr(info, 'language', '?')})\n"
        f"{format_transcript(raw_words)}"
    )

    return _generate_frame_metadata(raw_words, theme, visual_prompt)


_minicpm_model = None
_minicpm_generator = None
_minicpm_tokenizer = None

def _load_minicpm():
    global _minicpm_model, _minicpm_generator, _minicpm_tokenizer
    if _minicpm_model is None:
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer
        from outlines import from_transformers, Generator

        device = "cpu" if os.environ.get("FORCE_CPU") == "1" else ("cuda" if torch.cuda.is_available() else "cpu")
        logger.info(f"Loading MiniCPM5-1B on {device}...")
        try:
            # MiniCPM5-1B is a standard LlamaForCausalLM — no trust_remote_code needed.
            hf_model = AutoModelForCausalLM.from_pretrained(
                "openbmb/MiniCPM5-1B",
                torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32,
            ).to(device)
            _minicpm_tokenizer = AutoTokenizer.from_pretrained("openbmb/MiniCPM5-1B")

            _minicpm_model = from_transformers(hf_model, _minicpm_tokenizer)
            _minicpm_generator = Generator(_minicpm_model, SongFrames)
        except Exception as e:
            logger.error(f"Error loading MiniCPM5-1B via outlines: {e}")
            _minicpm_model = "failed"
    return _minicpm_model, _minicpm_generator, _minicpm_tokenizer


# Words per LLM chunk. Kept small so the model only has to produce a couple of
# {count, frame_animation} objects per call — a much easier task than
# partitioning a long chunk into several lines.
CHUNK_SIZE = 10


def _build_prompt(chunk_words: list[Word], theme: str, visual_prompt: str) -> str:
    """Build the user message describing one chunk of words."""
    text_parts = []
    for i, w in enumerate(chunk_words):
        if i > 0:
            silence = w.start - chunk_words[i - 1].end
            if silence > 0.5:
                text_parts.append(f"[PAUSE {silence:.1f}s]")
        word_str = f"'{w.text}'"
        duration = w.end - w.start
        if duration > 0.6:
            # Flag long-held/sustained notes — these can stand alone as a line.
            word_str += f"({duration:.1f}s)"
        text_parts.append(word_str)
    context_str = " ".join(text_parts)

    return f"""LYRICS (in order):
{context_str}

THEME: {theme}
AI BACKGROUND ART STYLE (a separate image model paints a new background every two lines, using this style): {visual_prompt if visual_prompt else 'None'}

Split these {len(chunk_words)} words into short on-screen lines and return JSON for the SongFrames schema.

Each frame has:
- "count": how many of the NEXT words (in order) belong to this line.
- "frame_animation": one effect for the whole line, matching its mood.

RULES:
1. The counts must add up to exactly {len(chunk_words)} (every word used once, in order).
2. Each line should read as a complete phrase — aim for 4 to 7 words per line, never more than 7. Avoid 1-2 word lines unless they are the very last words or a word marked with a long held duration like (1.5s).
3. Use [PAUSE] markers as strong hints for where a line should end.
4. Pick "frame_animation" to match the line's mood and complement the AI background art style above (e.g. "zoom_in" for emphasis, "flash" for a dramatic hit, "fade_to_black" for a quiet ending, "pan_left"/"pan_right" for gentle movement) — most lines should be "none"."""


SYSTEM_PROMPT = (
    "You are a lyric video director. You split song lyrics into short on-screen "
    "lines and pick one mood-matching animation effect per line. Your output "
    "drives the timing of a kinetic-typography video that also has AI-generated "
    "background art behind the text, so choose line breaks and animations that "
    "complement that art without crowding it. "
    "You never change, add, remove or reorder words."
)


def _generate_frame_metadata(words: list[Word], theme: str, visual_prompt: str) -> list[FrameMeta]:
    model, generator, tokenizer = _load_minicpm()
    if model == "failed" or model is None:
        logger.warning("LLM unavailable, falling back to rule-based annotator.")
        return _smooth_short_lines(_fallback_annotator(words))

    all_metas: list[FrameMeta] = []
    total_words = len(words)
    start_idx = 0

    while start_idx < total_words:
        end_idx = min(start_idx + CHUNK_SIZE, total_words)
        chunk_words = words[start_idx:end_idx]

        user_prompt = _build_prompt(chunk_words, theme, visual_prompt)
        prompt = tokenizer.apply_chat_template(
            [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": user_prompt},
            ],
            tokenize=False,
            add_generation_prompt=True,
            enable_thinking=False,  # annotation task, no reasoning needed
        )

        logger.info(
            f"=== LLM INPUT [chunk {start_idx}-{end_idx - 1}] === (MiniCPM5-1B)\n{user_prompt}"
        )

        try:
            # Outlines guarantees schema-valid JSON; with overrides removed the
            # output is just a couple of {count, frame_animation} objects.
            result_str = generator(prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_p=0.95)
            logger.info(f"=== LLM RAW OUTPUT [chunk {start_idx}-{end_idx - 1}] ===\n{result_str}")
            result = SongFrames.model_validate_json(result_str)
            chunk_metas = _reconcile_chunk(words, start_idx, end_idx, result.frames)
            logger.info(
                f"=== LLM PARSED FRAMES [chunk {start_idx}-{end_idx - 1}] ===\n"
                f"{format_frames(chunk_metas)}"
            )
        except Exception as e:
            logger.error(
                f"Generation failed for chunk {start_idx}-{end_idx - 1}: {e}. "
                "Using rule-based fallback annotator for this chunk."
            )
            chunk_metas = _fallback_annotator(words[start_idx:end_idx])
            logger.info(
                f"=== FALLBACK FRAMES [chunk {start_idx}-{end_idx - 1}] ===\n"
                f"{format_frames(chunk_metas)}"
            )

        all_metas.extend(chunk_metas)
        start_idx = end_idx

    all_metas = _smooth_short_lines(all_metas)

    logger.info(
        f"=== FINAL FRAME METADATA === {len(all_metas)} frames total\n"
        f"{format_frames(all_metas)}"
    )
    return all_metas


def _reconcile_chunk(words: list[Word], start_idx: int, end_idx: int, frames: list[Frame]) -> list[FrameMeta]:
    """
    Turn count-based LLM frames into FrameMeta covering exactly [start_idx, end_idx).

    Indices are assigned sequentially from the counts, so contiguity/ordering/coverage
    are guaranteed regardless of what the model emitted. Counts are clamped to
    MAX_WORDS_PER_LINE so a single line can never balloon to the whole chunk; any
    words left over (from undershoot or clamping) are packed into additional
    MAX_WORDS_PER_LINE-sized lines with frame_animation=none.
    """
    metas: list[FrameMeta] = []
    cursor = start_idx
    for fr in frames:
        if cursor >= end_idx:
            break
        count = max(1, min(fr.count, MAX_WORDS_PER_LINE))
        stop = min(cursor + count, end_idx)
        metas.append(FrameMeta(words=words[cursor:stop], frame_animation=fr.frame_animation))
        cursor = stop

    # Any words the model failed to account for (undershoot, or overshoot
    # clamped above): pack into additional clean, short lines.
    while cursor < end_idx:
        stop = min(cursor + MAX_WORDS_PER_LINE, end_idx)
        metas.append(FrameMeta(words=words[cursor:stop], frame_animation=FrameAnim.none))
        cursor = stop

    return metas


# A short line whose words span at least this long (e.g. one held/sustained
# note) reads fine on its own — don't merge it away just for being short.
SUSTAINED_LINE_DURATION = 1.2


def _line_duration(meta: FrameMeta) -> float:
    if not meta.words:
        return 0.0
    return meta.words[-1].end - meta.words[0].start


def _smooth_short_lines(metas: list[FrameMeta]) -> list[FrameMeta]:
    """
    Merge lines shorter than MIN_WORDS_PER_LINE into a neighboring line,
    unless the short line is actually a sustained/held word (long duration),
    in which case it's left standing on its own.

    Short lines (especially single words) cause the background/text to change
    on almost every word, which feels frantic and breaks the sentence up.
    Words stay in chronological order, so a merged FrameMeta's on-screen
    window (via get_frame_times) still spans correctly.
    """
    if len(metas) <= 1:
        return metas

    result: list[FrameMeta] = []
    i = 0
    while i < len(metas):
        cur = metas[i]
        if len(cur.words) < MIN_WORDS_PER_LINE and _line_duration(cur) < SUSTAINED_LINE_DURATION:
            nxt = metas[i + 1] if i + 1 < len(metas) else None
            # Prefer merging forward (keeps reading order natural).
            if nxt is not None and len(cur.words) + len(nxt.words) <= MAX_WORDS_PER_LINE:
                merged_anim = cur.frame_animation if cur.frame_animation != FrameAnim.none else nxt.frame_animation
                result.append(FrameMeta(words=cur.words + nxt.words, frame_animation=merged_anim))
                i += 2
                continue
            # Too big to merge fully — borrow just enough words from the
            # front of the next line so `cur` reaches MIN_WORDS_PER_LINE.
            if nxt is not None:
                moved = min(MIN_WORDS_PER_LINE - len(cur.words), len(nxt.words) - 1)
                if moved > 0:
                    cur = FrameMeta(words=cur.words + nxt.words[:moved], frame_animation=cur.frame_animation)
                    nxt.words = nxt.words[moved:]
                    result.append(cur)
                    i += 1
                    continue
            # Otherwise fold backward into the line we just emitted.
            if result and len(result[-1].words) + len(cur.words) <= MAX_WORDS_PER_LINE:
                result[-1].words.extend(cur.words)
                i += 1
                continue
            # Too big to fold fully backward — borrow from the end of the
            # previous line instead.
            if result:
                prev = result[-1]
                moved = min(MIN_WORDS_PER_LINE - len(cur.words), len(prev.words) - 1)
                if moved > 0:
                    cur.words = prev.words[-moved:] + cur.words
                    prev.words = prev.words[:-moved]
        result.append(cur)
        i += 1
    return result


def _fallback_annotator(words: list[Word]) -> list[FrameMeta]:
    """
    Rule-based annotator used when the LLM is unavailable or a chunk fails.
    Groups words into lines by silence gaps / max line length, with no
    frame-level animation.
    """
    if not words:
        return []

    metas: list[FrameMeta] = []
    current: list[Word] = [words[0]]

    for i in range(1, len(words)):
        silence = words[i].start - words[i - 1].end
        if len(current) >= MAX_WORDS_PER_LINE or silence > SILENCE_GAP:
            metas.append(FrameMeta(words=current, frame_animation=FrameAnim.none))
            current = []
        current.append(words[i])

    if current:
        metas.append(FrameMeta(words=current, frame_animation=FrameAnim.none))

    return metas