commentary.py

import json
import re

import spaces
import torch
from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration

MODEL_ID = "Qwen/Qwen2.5-VL-7B-Instruct"
DEVICE   = "cuda" if torch.cuda.is_available() else "cpu"

STYLES = {
    "Football (Premier League hype)":
        "an over-the-top English football commentator, dramatic and poetic",
    "MasterChef judge":
        "a brutally dramatic cooking-show judge who treats every moment like a "
        "high-stakes elimination round — stern, intense, occasionally impressed",
    "Nature documentary":
        "a hushed, awestruck nature narrator observing these humans like rare "
        "wildlife in their natural habitat",
    "Boxing announcer":
        "a booming boxing announcer treating every move as championship-defining",
    "Diva Hour":
        "a fabulous, shady red-carpet diva commentator narrating like everyone "
        "is a celebrity arriving at a star-studded gala, full of glamour and sass",
}

# Vibe → frontend key mapping (used by app.py)
VIBE_TO_STYLE = {
    "football":  "Football (Premier League hype)",
    "masterchef":"MasterChef judge",
    "wildlife":  "Nature documentary",
    "boxing":    "Boxing announcer",
    "diva":      "Diva Hour",
}

SYSTEM_PROMPT = """You are {persona}, calling live play-by-play over a short clip.

You'll receive {n_frames} keyframes in time order, each tagged "t=<seconds>".
Some people have their NAME burned in a box above their head — use those exact
names. Refer to anyone unnamed by what you see ("the one in the red shirt").
People we already know in this clip: {roster}.

NARRATIVE ARC — mandatory:
  Frame 1-2  : establish the scene and mood (who, where, what's at stake).
  Middle     : build tension — notice details, call the action, raise stakes.
  Final frame: pay off the arc with a punchy closing line.
Each line must feel like it follows the one before. No generic filler.

Write EXACTLY one line per keyframe — {n_frames} lines, in the same order.

Voice rules:
- React to what is ACTUALLY in each frame — never recycle a phrase.
- Vary rhythm: short punchy lines for action, longer for atmosphere.
- Drop names only when it lands — not in every line.
- Stay fully in character as {persona} from first line to last.
- CAPS and exclamation marks only for the BIG beats — overuse kills the punch.
- Keep every line under {max_words} words — it must be spoken before the next beat.

Output rules (strict):
- Output ONLY a JSON array. No markdown, no text outside it.
- Schema: [{{"time": <float seconds>, "text": "<one line>"}}]
- Use the EXACT "t=" value of each keyframe as its "time"."""

# Max words per line — tuned per vibe so audio fits in the gap
MAX_WORDS = {
    "football":   14,
    "masterchef": 13,
    "wildlife":   15,
    "boxing":     12,
    "diva":       14,
}
_DEFAULT_MAX_WORDS = 14

# Minimum gap (seconds) between events — events closer than this get merged
MIN_GAP_SEC = 2.5


print(f"[info] Loading Qwen2.5-VL-7B (first boot: download ~16GB + load) -> {DEVICE}")
_model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    MODEL_ID, torch_dtype=torch.bfloat16)
_model.to(DEVICE)
_processor = AutoProcessor.from_pretrained(MODEL_ID)
print("[info] Qwen ready.")


def _merge_close_events(event_frames, min_gap=MIN_GAP_SEC):
    """Merge keyframes that are too close together.

    When two events are < min_gap seconds apart the commentary lines would
    collide in audio. We keep the one with higher motion score (second of
    the pair, which usually shows the peak) and drop the earlier one.
    Returns a filtered list of (timestamp, image) tuples.
    """
    if len(event_frames) <= 1:
        return event_frames
    merged = [event_frames[0]]
    for ts, img in event_frames[1:]:
        prev_ts, _ = merged[-1]
        if ts - prev_ts < min_gap:
            # keep the later frame (usually the action peak)
            merged[-1] = (ts, img)
        else:
            merged.append((ts, img))
    return merged


@spaces.GPU(duration=120)
def generate_commentary(event_frames, roster, style_key, vibe=None):
    """event_frames: [(timestamp_sec, PIL.Image), ...] (annotated key events)
       roster:       list of names the user assigned
       style_key:    key into STYLES dict
       vibe:         optional vibe slug for per-vibe word limit

    Returns [{"time": float, "text": str}, ...] sorted by time.
    """
    from qwen_vl_utils import process_vision_info

    # Merge events that are too close — prevents audio collision at gen time
    event_frames = _merge_close_events(event_frames)

    n          = len(event_frames)
    max_words  = MAX_WORDS.get(vibe or "", _DEFAULT_MAX_WORDS)
    content    = []

    for i, (ts, img) in enumerate(event_frames):
        content.append({"type": "text",
                        "text": f"Keyframe {i+1}/{n} — t={ts:.2f}s:"})
        content.append({"type": "image", "image": img})
    content.append({"type": "text",
                    "text": f"Now write the {n}-line commentary JSON."})

    persona    = STYLES.get(style_key, list(STYLES.values())[0])
    roster_str = ", ".join(roster) if roster else "nobody named yet"

    messages = [
        {"role": "system",
         "content": SYSTEM_PROMPT.format(
             persona=persona, roster=roster_str,
             n_frames=n, max_words=max_words)},
        {"role": "user", "content": content},
    ]

    text          = _processor.apply_chat_template(messages, tokenize=False,
                                                   add_generation_prompt=True)
    image_inputs, _ = process_vision_info(messages)
    inputs        = _processor(text=[text], images=image_inputs,
                               return_tensors="pt").to(_model.device)

    out = _model.generate(**inputs, max_new_tokens=600, do_sample=True,
                          temperature=0.8, top_p=0.95)
    new = out[:, inputs.input_ids.shape[1]:]
    raw = _processor.batch_decode(new, skip_special_tokens=True)[0]

    valid_times = [round(ts, 2) for ts, _ in event_frames]
    return _parse(raw, valid_times, max_words)


def _trim_to_words(text: str, max_words: int) -> str:
    words = text.split()
    if len(words) <= max_words:
        return text
    # trim at sentence boundary if possible
    trimmed = " ".join(words[:max_words])
    for punct in (".", "!", "?", "—", ","):
        idx = trimmed.rfind(punct)
        if idx > len(trimmed) // 2:
            return trimmed[:idx + 1]
    return trimmed + "…"


def _parse(raw: str, valid_times: list[float], max_words: int = 14) -> list[dict]:
    raw = re.sub(r"```(?:json)?", "", raw).strip()
    m   = re.search(r"\[.*\]", raw, re.DOTALL)
    try:
        items = json.loads(m.group(0) if m else raw)
    except (json.JSONDecodeError, AttributeError):
        items = []

    script = []
    for i, it in enumerate(items):
        try:
            txt = str(it["text"]).strip()
            t   = float(it.get("time", valid_times[min(i, len(valid_times)-1)]))
        except (KeyError, TypeError, ValueError):
            continue
        t   = min(valid_times, key=lambda v: abs(v - t)) if valid_times else t
        txt = _trim_to_words(txt, max_words)
        if txt:
            script.append({"time": t, "text": txt})

    if not script and valid_times:
        script = [{"time": valid_times[0],
                   "text": "WHAT A MOMENT, LADIES AND GENTLEMEN!"}]
    script.sort(key=lambda x: x["time"])
    return script