eval.py

"""
eval.py — before/after LoRA comparison
=======================================
Produces a side-by-side comparison video for a held-out test clip:

  LEFT   : original frames (raw camera)
  MIDDLE : naive fal edit  (frame-by-frame, no temporal conditioning → flickery)
  RIGHT  : lora output     (temporally consistent — run after training)

Usage
-----
  # Step 1 — extract frames + run naive fal edits (the "before" baseline)
  python3 eval.py baseline --video data/test/92s-RpM7Tks_6.mp4

  # Step 2 — run LoRA inference on the test frames (after training)
  python3 eval.py lora --video data/test/92s-RpM7Tks_6.mp4 --lora path/to/lora.safetensors

  # Step 3 — render the comparison video (requires both steps above)
  python3 eval.py compare --video data/test/92s-RpM7Tks_6.mp4

  # All-in-one (baseline + compare, no LoRA yet)
  python3 eval.py baseline --video data/test/92s-RpM7Tks_6.mp4 --compare
"""

from __future__ import annotations

import argparse
import logging
import os
import random
import sys
from pathlib import Path
from typing import Any

import ffmpeg
import fal_client
from PIL import Image

# ---------------------------------------------------------------------------
# Import shared helpers from pipeline
# ---------------------------------------------------------------------------
sys.path.insert(0, str(Path(__file__).parent))
from pipeline import DEFAULT_PROMPT, VIDEO_PROMPTS, FAL_MODEL, _upload_and_edit, _download

FAL_LORA_MODEL = "fal-ai/flux-2/klein/4b/base/edit/lora"

# ---------------------------------------------------------------------------
# Defaults
# ---------------------------------------------------------------------------
DEFAULT_TEST_DIR  = Path("data/test")
DEFAULT_FPS       = 8
OUTPUT_VIDEO_FPS  = 8   # fps for the rendered comparison video
DEFAULT_LORA_URL_FILE = Path("data/lora_url.txt")
DEFAULT_K_MIN    = 5   # minimum lookback frames
DEFAULT_K_MAX    = 20  # maximum lookback frames

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s  %(levelname)-8s  %(message)s",
    datefmt="%H:%M:%S",
)
log = logging.getLogger(__name__)

# ===========================================================================
# Helpers
# ===========================================================================
def _stem(video_path: Path) -> str:
    return video_path.stem

def _test_dirs(video_path: Path, test_dir: Path) -> tuple[Path, Path, Path, Path]:
    """Return (frames_dir, naive_dir, lora_dir, output_dir) for a test video."""
    stem = _stem(video_path)
    base = test_dir / stem
    return (
        base / "frames",
        base / "naive",
        base / "lora",
        base / "output",
    )

def _extract_frames(video_path: Path, frames_dir: Path, fps: int) -> list[Path]:
    frames_dir.mkdir(parents=True, exist_ok=True)
    existing = sorted(frames_dir.glob("*.jpg"))
    if existing:
        log.info("SKIP extract — %d frames already in %s", len(existing), frames_dir)
        return existing

    log.info("Extracting frames from %s at %d fps", video_path.name, fps)
    pattern = str(frames_dir / "%06d.jpg")
    ffmpeg.input(str(video_path)).filter("fps", fps=fps).output(
        pattern, qscale=2
    ).overwrite_output().run(quiet=True)

    frames = sorted(frames_dir.glob("*.jpg"))
    log.info("Extracted %d frames", len(frames))
    return frames

def _prompt_for(stem: str) -> str:
    return VIDEO_PROMPTS.get(stem, DEFAULT_PROMPT)

# ===========================================================================
# Stage 1 — Naive fal baseline
# ===========================================================================
def run_baseline(video_path: Path, test_dir: Path, fps: int, dry_run: bool) -> None:
    """Extract frames and edit each one naively (no temporal conditioning)."""
    frames_dir, naive_dir, _, _ = _test_dirs(video_path, test_dir)
    naive_dir.mkdir(parents=True, exist_ok=True)

    frames = _extract_frames(video_path, frames_dir, fps)
    prompt = _prompt_for(_stem(video_path))
    log.info("Naive edit prompt: %s", prompt)

    for frame_path in frames:
        dest = naive_dir / frame_path.name
        if dest.exists():
            log.debug("SKIP naive %s", frame_path.name)
            continue

        log.info("EDIT (naive) %s", frame_path.name)
        if dry_run:
            dest.write_bytes(frame_path.read_bytes())
            continue

        url = _upload_and_edit(frame_path, prompt)
        _download(url, dest)

    log.info("Naive edits done: %d frames in %s", len(frames), naive_dir)

# ===========================================================================
# Stage 2 — LoRA inference
# ===========================================================================
def _infer_lora_frame(
    reference_path: Path,
    original_path: Path,
    prompt: str,
    lora_url: str,
    lora_scale: float = 1.0,
) -> str:
    """
    Run the edit model with the trained LoRA.

    Inputs match the training format:
      reference_path = edited[t-k]  (past edit, ROOT_start2 during training)
      original_path  = original[t]  (raw frame,  ROOT_start  during training)

    The LoRA learned:  original[t] + ref edited[t-k]  →  edited[t]
    """
    ref_url  = fal_client.upload_file(str(reference_path))
    orig_url = fal_client.upload_file(str(original_path))

    result: Any = fal_client.run(
        FAL_LORA_MODEL,
        arguments={
            "image_urls": [orig_url, ref_url],   # start, start2
            "prompt":     prompt,
            "loras":      [{"path": lora_url, "scale": lora_scale}],
        },
    )

    images = result.get("images") or result.get("image")
    if isinstance(images, list) and images:
        url = images[0].get("url") or images[0].get("image_url")
    elif isinstance(images, dict):
        url = images.get("url") or images.get("image_url")
    else:
        raise ValueError(f"Unexpected fal response shape: {result}")
    if not url:
        raise ValueError(f"No URL in fal response: {result}")
    return url

def run_lora(
    video_path: Path,
    test_dir: Path,
    lora_url: str | None,
    fps: int,
    k_min: int,
    k_max: int,
    dry_run: bool,
) -> None:
    """
    Run LoRA inference on the test frames, using the autoregressive strategy from training:

        edited[t-k]  +  original[t]  →  edited[t]

    For each frame t:
      - reference = our own previous LoRA output at t-k  (autoregressive)
        where k is randomly selected from [k_min, k_max] range
      - original  = raw camera frame at t
    Both are passed to the fine-tuned edit model, which produces edited[t].
    """
    frames_dir, naive_dir, lora_dir, _ = _test_dirs(video_path, test_dir)
    lora_dir.mkdir(parents=True, exist_ok=True)

    frames = sorted(frames_dir.glob("*.jpg"))
    if not frames:
        log.error("No frames found — run 'baseline' stage first")
        sys.exit(1)

    # Resolve LoRA URL -------------------------------------------------
    if lora_url is None:
        if DEFAULT_LORA_URL_FILE.exists():
            lora_url = DEFAULT_LORA_URL_FILE.read_text().strip()
            log.info("Loaded LoRA URL from %s", DEFAULT_LORA_URL_FILE)
        else:
            log.error(
                "No --lora-url given and %s not found. Run train.py first.",
                DEFAULT_LORA_URL_FILE,
            )
            sys.exit(1)

    prompt = f"CONSISTENTEDIT {_prompt_for(_stem(video_path))}"
    log.info("LoRA URL:  %s", lora_url)
    log.info("Prompt:    %s", prompt)
    log.info("k_range=[%d,%d]  frames=%d  output=%s", k_min, k_max, len(frames), lora_dir)

    for t, frame_path in enumerate(frames):
        dest = lora_dir / frame_path.name
        if dest.exists():
            log.debug("SKIP lora %s", frame_path.name)
            continue

        # Reference = previous LoRA output at t-k (autoregressive)
        # k is randomly selected from [k_min, k_max] range
        if t >= k_max:
            # Randomly select k from the range
            k = random.randint(k_min, k_max)
            ref_path = lora_dir / frames[t - k].name
        elif t >= k_min:
            # Within valid range but below maximum lookback distance
            k = random.randint(k_min, min(t, k_max))
            if (lora_dir / frames[t - k].name).exists():
                ref_path = lora_dir / frames[t - k].name
            else:
                # Bootstrap: use naive edit
                bootstrap = naive_dir / frame_path.name
                ref_path = bootstrap if bootstrap.exists() else frame_path
        else:
            # Bootstrap: use naive edit for first k_min frames
            bootstrap = naive_dir / frame_path.name
            ref_path = bootstrap if bootstrap.exists() else frame_path

        log.info("EDIT (lora) %s  ref=%s", frame_path.name, ref_path.name)

        if dry_run:
            dest.write_bytes(frame_path.read_bytes())
            continue

        url = _infer_lora_frame(ref_path, frame_path, prompt, lora_url)
        _download(url, dest)

    log.info("LoRA edits done: %d frames in %s", len(frames), lora_dir)

# ===========================================================================
# Stage 3 — Render comparison video
# ===========================================================================
def render_comparison(
    video_path: Path,
    test_dir: Path,
    fps: int,
    include_lora: bool = False,
) -> Path:
    """
    Stack columns side by side and encode to MP4:
      always:        original | naive
      if lora ready: original | naive | lora

    Uses PIL to composite frames (avoids ffmpeg JPEG-LS decode issues)
    then pipes raw frames into ffmpeg for H.264 encoding.
    """
    import subprocess

    frames_dir, naive_dir, lora_dir, output_dir = _test_dirs(video_path, test_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    frames = sorted(frames_dir.glob("*.jpg"))
    naives = sorted(naive_dir.glob("*.jpg"))
    loras  = sorted(lora_dir.glob("*.jpg")) if include_lora else []

    if not frames:
        log.error("No original frames — run baseline first"); sys.exit(1)
    if not naives:
        log.error("No naive edits — run baseline first");     sys.exit(1)
    if include_lora and not loras:
        log.error("No lora frames — run lora stage first");   sys.exit(1)

    n = min(len(frames), len(naives))
    if include_lora:
        n = min(n, len(loras))

    cols = 3 if include_lora else 2
    log.info("Rendering %d frames (%d columns) via PIL…", n, cols)

    out_name = "comparison_with_lora.mp4" if include_lora else "comparison_baseline.mp4"
    out_path = output_dir / out_name

    # Open first frame to determine canvas size
    h = 720
    sample = Image.open(frames[0]).convert("RGB")
    w0 = int(sample.width * h / sample.height)  # orig column width

    sample_fal = Image.open(naives[0]).convert("RGB")
    w1 = int(sample_fal.width * h / sample_fal.height)  # fal column width

    total_w = w0 + w1 * (cols - 1)

    # Pipe raw RGB frames into ffmpeg
    ffmpeg_cmd = [
        "ffmpeg", "-y",
        "-f", "rawvideo",
        "-vcodec", "rawvideo",
        "-s", f"{total_w}x{h}",
        "-pix_fmt", "rgb24",
        "-r", str(fps),
        "-i", "pipe:0",
        "-vcodec", "libx264",
        "-crf", "18",
        "-pix_fmt", "yuv420p",
        str(out_path),
    ]
    proc = subprocess.Popen(ffmpeg_cmd, stdin=subprocess.PIPE,
                            stderr=subprocess.DEVNULL)

    for i in range(n):
        canvas = Image.new("RGB", (total_w, h))
        orig_img = Image.open(frames[i]).convert("RGB").resize((w0, h), Image.LANCZOS)
        canvas.paste(orig_img, (0, 0))

        naive_img = Image.open(naives[i]).convert("RGB").resize((w1, h), Image.LANCZOS)
        canvas.paste(naive_img, (w0, 0))

        if include_lora:
            lora_img = Image.open(loras[i]).convert("RGB").resize((w1, h), Image.LANCZOS)
            canvas.paste(lora_img, (w0 + w1, 0))

        proc.stdin.write(canvas.tobytes())

    proc.stdin.close()
    proc.wait()

    log.info("Saved comparison video → %s", out_path)
    return out_path

# ===========================================================================
# CLI
# ===========================================================================
def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Evaluate LoRA before/after on a held-out test video",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    parser.add_argument(
        "stage",
        choices=["baseline", "lora", "compare"],
        help="Which stage to run",
    )
    parser.add_argument(
        "--video",
        type=Path,
        required=True,
        help="Path to the test video file",
    )
    parser.add_argument(
        "--test-dir",
        type=Path,
        default=DEFAULT_TEST_DIR,
        metavar="DIR",
        help=f"Root dir for test outputs (default: {DEFAULT_TEST_DIR})",
    )
    parser.add_argument(
        "--fps",
        type=int,
        default=DEFAULT_FPS,
        help=f"FPS for frame extraction (default: {DEFAULT_FPS})",
    )
    parser.add_argument(
        "--lora-url",
        type=str,
        default=None,
        metavar="URL",
        help=(
            "fal.ai LoRA URL from train.py. "
            f"Defaults to contents of {DEFAULT_LORA_URL_FILE} if present."
        ),
    )
    parser.add_argument(
        "--k-min",
        type=int,
        default=DEFAULT_K_MIN,
        help=f"Minimum temporal lookback distance in frames (default: {DEFAULT_K_MIN})",
    )
    parser.add_argument(
        "--k-max",
        type=int,
        default=DEFAULT_K_MAX,
        help=f"Maximum temporal lookback distance in frames (default: {DEFAULT_K_MAX})",
    )
    parser.add_argument(
        "--compare",
        action="store_true",
        help="Auto-render comparison video after the stage completes",
    )
    parser.add_argument(
        "--with-lora",
        action="store_true",
        help="Include the lora column in comparison video",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Skip fal API calls (copy originals as placeholders)",
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="Enable DEBUG logging",
    )
    return parser.parse_args()

def main() -> None:
    args = parse_args()
    if args.verbose:
        logging.getLogger().setLevel(logging.DEBUG)

    if not args.video.exists():
        log.error("Video not found: %s", args.video)
        sys.exit(1)

    if args.stage == "baseline":
        if not os.environ.get("FAL_KEY") and not args.dry_run:
            log.error("FAL_KEY not set — export it or use --dry-run")
            sys.exit(1)
        run_baseline(args.video, args.test_dir, args.fps, args.dry_run)
        if args.compare:
            render_comparison(args.video, args.test_dir, args.fps, include_lora=False)

    elif args.stage == "lora":
        if not os.environ.get("FAL_KEY") and not args.dry_run:
            log.error("FAL_KEY not set — export it or use --dry-run")
            sys.exit(1)
        run_lora(
            args.video, args.test_dir, args.lora_url,
            args.fps, args.k_min, args.k_max, args.dry_run,
        )
        if args.compare:
            render_comparison(args.video, args.test_dir, args.fps, include_lora=True)

    elif args.stage == "compare":
        render_comparison(
            args.video, args.test_dir, args.fps, include_lora=args.with_lora
        )

if __name__ == "__main__":
    main()