"""
wan_flf2v_infer.py — run Wan2.1 FLF2V over the 4 consecutive keyframe pairs
(kf0-kf1, kf1-kf2, kf2-kf3, kf3-kf4) produced by pipe14.py, on a single
16GB RTX A4000 using sequential CPU offload.

UNVERIFIED beyond the diffusers signature check — this has not been run yet.
Checkpoint ID below is a best guess (see comment) — confirm the exact repo
string on huggingface.co before running; if wrong, from_pretrained will 404.

Usage:
    python wan_flf2v_infer.py \
        --frames-dir /path/to/frames/eat2 \
        --prompt "a hand reaching to pick up a spoon and bring it to the mouth" \
        --out-dir ./wan_out/eat2 \
        --num-frames 25 --steps 20   # start SMALL, see note below
"""

import argparse
import os

import torch
from diffusers import WanImageToVideoPipeline, AutoencoderKLWan
from diffusers.utils import export_to_video
from PIL import Image


# BEST GUESS — confirm on huggingface.co/Wan-AI before trusting this.
# The plain I2V checkpoint (Wan2.1-I2V-14B-*) is NOT the same thing as this —
# last_image conditioning needs a checkpoint actually trained for FLF2V.
MODEL_ID = "Wan-AI/Wan2.1-FLF2V-14B-720P-Diffusers"


def load_pipeline(model_id=MODEL_ID):
    # Wan's VAE is commonly recommended in fp32 for numerical stability —
    # [Guessing, from general Wan model-card guidance, not verified against
    # this exact diffusers release]. If this OOMs or errors, try dropping
    # torch_dtype here to torch.bfloat16 and see if quality holds up.
    vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)
    pipe = WanImageToVideoPipeline.from_pretrained(model_id, vae=vae, torch_dtype=torch.bfloat16)

    # Sequential offload, not just enable_model_cpu_offload — this is the
    # aggressive layer-by-layer streaming mode needed to fit a 14B model in
    # 16GB VRAM at all. It will be slow. That's expected, not a bug.
    pipe.enable_sequential_cpu_offload()

    return pipe


def run_segment(pipe, first_frame_path, last_frame_path, prompt, out_path,
                 height=480, width=832, num_frames=25, steps=20, seed=None):
    """
    height/width: MUST match (or be close to) what the checkpoint was
    trained at (480P checkpoint -> 480x832-ish, 720P checkpoint -> larger).
    Using the wrong resolution for the checkpoint variant degrades quality
    silently rather than erroring — [Guessing this is true here the same
    way it is for most diffusion checkpoints; not confirmed for Wan specifically].

    num_frames=25, steps=20 here are DELIBERATELY reduced from the Wan
    defaults (81 frames / 50 steps) for a first test — verify this actually
    runs and fits in VRAM/time before scaling up. Going straight to 81/50
    on a first attempt on offloaded 14B risks a very long first debug cycle.
    """
    image = Image.open(first_frame_path).convert("RGB").resize((width, height))
    last_image = Image.open(last_frame_path).convert("RGB").resize((width, height))

    generator = torch.Generator(device="cpu")
    if seed is not None:
        generator = generator.manual_seed(seed)

    output = pipe(
        image=image,
        last_image=last_image,
        prompt=prompt,
        height=height,
        width=width,
        num_frames=num_frames,
        num_inference_steps=steps,
        generator=generator,
    )
    frames = output.frames[0]
    export_to_video(frames, out_path, fps=16)  # fps is a guess for a short test clip — not Wan's canonical value
    print(f"wrote {out_path}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--frames-dir", required=True, help="dir containing kf0.png ... kf4.png")
    ap.add_argument("--prompt", required=True)
    ap.add_argument("--out-dir", required=True)
    ap.add_argument("--num-frames", type=int, default=25)
    ap.add_argument("--steps", type=int, default=20)
    ap.add_argument("--height", type=int, default=480)
    ap.add_argument("--width", type=int, default=832)
    ap.add_argument("--seed", type=int, default=None)
    ap.add_argument("--segment", type=int, default=None,
                     help="Run only ONE segment (0-3) instead of all 4 — "
                          "use this for the single-segment style-drift test "
                          "before attempting the full 4-segment stitch.")
    args = ap.parse_args()

    os.makedirs(args.out_dir, exist_ok=True)
    pipe = load_pipeline()

    kf_paths = [os.path.join(args.frames_dir, f"kf{i}.png") for i in range(5)]
    for p in kf_paths:
        if not os.path.exists(p):
            raise FileNotFoundError(f"missing keyframe: {p}")

    segments = list(range(4)) if args.segment is None else [args.segment]
    for i in segments:
        out_path = os.path.join(args.out_dir, f"segment_{i}_{i+1}.mp4")
        print(f"\n--- segment {i}: kf{i} -> kf{i+1} ---")
        run_segment(
            pipe, kf_paths[i], kf_paths[i + 1], args.prompt, out_path,
            height=args.height, width=args.width,
            num_frames=args.num_frames, steps=args.steps, seed=args.seed,
        )


if __name__ == "__main__":
    main()