scripts/render_examples.py

"""Pre-render the 'Try an Example' showcase bundles → examples/.

Drives the LIVE Space end to end for each persona in src.examples.EXAMPLES and
writes the cached bundle the app serves with zero quota:
  examples/<id>.lyrics.txt   examples/<id>.trace.json   examples/<id>.song.mp3
and, IF a headshot is provided, examples/<id>.video.mp4 (demo-video footage).

Drop examples/<id>.photo.(jpg|png) before running to get the full talking-head
video for that persona; otherwise it renders song-only (no photo needed).

Run (after the Space is deployed with this code, on the new render):
  python scripts/render_examples.py            # all personas
  python scripts/render_examples.py --only datasci_unhinged
Costs a few ZeroGPU minutes total (one lyrics call + one render per persona).
"""

import argparse
import glob
import json
import os
import shutil
import sys

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)

from src.examples import EX_DIR, EXAMPLES  # noqa: E402

DEFAULT_SPACE = "build-small-hackathon/open-to-work-musical"
VOICE = "🎲 Surprise me"


def _photo_for(eid):
    for ext in ("jpg", "jpeg", "png", "webp"):
        hits = glob.glob(os.path.join(EX_DIR, f"{eid}.photo.{ext}"))
        if hits:
            return hits[0]
    return None


def render_one(client, ex):
    from gradio_client import handle_file

    eid = ex["id"]
    print(f"\n=== {eid} · level {ex['level']} · {ex['genre']} ===", flush=True)
    # new on_write_lyrics signature: (resume, jd, genre, level, tiny_mode, voice_type)
    # returns (lyrics, trace_html)  ← trace is rendered HTML now, not a dict
    lyrics, trace_html = client.predict(
        ex["resume"], ex["jd"], ex["genre"], ex["level"], False, VOICE,
        api_name="/on_write_lyrics",
    )
    if not lyrics or not lyrics.strip():
        print("  !! no lyrics produced — skipping", flush=True)
        return False
    os.makedirs(EX_DIR, exist_ok=True)
    with open(os.path.join(EX_DIR, f"{eid}.lyrics.txt"), "w", encoding="utf-8") as f:
        f.write(lyrics)
    with open(os.path.join(EX_DIR, f"{eid}.trace.html"), "w", encoding="utf-8") as f:
        f.write(trace_html or "")
    print("  lyrics+trace saved", flush=True)

    photo = _photo_for(eid)
    print(f"  rendering {'song+video' if photo else 'song-only'}…", flush=True)
    # new on_render signature: (lyrics, photo, voice_sample, voice_consent, level, genre, voice_type)
    try:
        song, video, _dbg = client.predict(
            lyrics, handle_file(photo) if photo else None, None, False,
            ex["level"], ex["genre"], VOICE,
            api_name="/on_render",
        )
    except Exception as e:
        msg = str(e).splitlines()[-1][:200]
        print(f"  !! render skipped (lyrics+trace kept): {msg}", flush=True)
        return False
    if song and os.path.isfile(song):
        shutil.copy(song, os.path.join(EX_DIR, f"{eid}.song.mp3"))
        print("  song saved.", flush=True)
    else:
        print("  !! no song produced", flush=True)
        return False
    if video and os.path.isfile(video):
        shutil.copy(video, os.path.join(EX_DIR, f"{eid}.video.mp4"))
        print("  video saved.", flush=True)
    return True


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--space", default=DEFAULT_SPACE)
    ap.add_argument("--only", help="render a single example id")
    args = ap.parse_args()

    from gradio_client import Client
    from huggingface_hub import get_token

    client = Client(args.space, token=get_token(), verbose=False)
    todo = [e for e in EXAMPLES if not args.only or e["id"] == args.only]
    ok = sum(render_one(client, e) for e in todo)
    print(f"\nDone: {ok}/{len(todo)} bundles rendered into {EX_DIR}", flush=True)
    sys.exit(0 if ok == len(todo) else 1)


if __name__ == "__main__":
    main()