#!/usr/bin/env python3
"""Generate sequential keyframes with the VBVR-Pro Neo-Unify checkpoint."""

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
import torch
from PIL import Image
from transformers import AutoModel, AutoTokenizer


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--model",
        required=True,
        help="Local model directory or Hugging Face repository ID.",
    )
    parser.add_argument("--input", type=Path, required=True, help="Initial RGB image.")
    parser.add_argument("--prompt", required=True, help="Text instruction.")
    parser.add_argument("--output-dir", type=Path, default=Path("outputs"))
    parser.add_argument("--num-images", type=int, default=1)
    parser.add_argument("--width", type=int, default=512)
    parser.add_argument("--height", type=int, default=512)
    parser.add_argument("--num-steps", type=int, default=50)
    parser.add_argument("--cfg-scale", type=float, default=1.0)
    parser.add_argument("--img-cfg-scale", type=float, default=1.0)
    parser.add_argument("--timestep-shift", type=float, default=1.0)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--device", default="cuda:0")
    args = parser.parse_args()

    if args.num_images <= 0:
        parser.error("--num-images must be positive")
    if args.num_steps <= 0:
        parser.error("--num-steps must be positive")
    if args.width <= 0 or args.height <= 0:
        parser.error("--width and --height must be positive")
    if args.width % 32 or args.height % 32:
        parser.error("--width and --height must be multiples of 32")
    if args.cfg_scale < 0 or args.img_cfg_scale < 0:
        parser.error("CFG scales must be non-negative")
    if not args.device.startswith("cuda"):
        parser.error("the included inference path requires a CUDA device")
    return args


def tensor_to_image(frame: torch.Tensor) -> Image.Image:
    """Convert one model output in [-1, 1] to an RGB PIL image."""
    if frame.ndim != 4 or frame.shape[0] != 1 or frame.shape[1] != 3:
        raise ValueError(f"unexpected generated tensor shape: {tuple(frame.shape)}")
    image = (frame.detach().float() * 0.5 + 0.5).clamp(0, 1)
    array = (
        image[0].permute(1, 2, 0).cpu().numpy() * 255.0
    ).round().astype(np.uint8)
    return Image.fromarray(array).convert("RGB")


def main() -> int:
    args = parse_args()
    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is required for this model")
    if not args.input.is_file():
        raise FileNotFoundError(f"input image does not exist: {args.input}")

    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)

    tokenizer = AutoTokenizer.from_pretrained(
        args.model,
        trust_remote_code=True,
    )
    model = AutoModel.from_pretrained(
        args.model,
        torch_dtype=torch.bfloat16,
        trust_remote_code=True,
    ).to(args.device)
    model.eval()

    with Image.open(args.input) as image_file:
        first_frame = image_file.convert("RGB")
    prompt = args.prompt.replace("<image>", "").strip()
    if not prompt:
        raise ValueError("prompt is empty after removing <image> placeholders")

    with torch.inference_mode():
        frames = model.interleave_gen_image_only(
            tokenizer,
            prompt,
            gt_text="<image>" * args.num_images,
            images=[first_frame],
            image_size=(args.width, args.height),
            max_images=args.num_images,
            num_steps=args.num_steps,
            cfg_scale=args.cfg_scale,
            img_cfg_scale=args.img_cfg_scale,
            timestep_shift=args.timestep_shift,
        )

    if len(frames) != args.num_images:
        raise RuntimeError(
            f"model returned {len(frames)} frames; expected {args.num_images}"
        )
    args.output_dir.mkdir(parents=True, exist_ok=True)
    for index, frame in enumerate(frames, start=1):
        output = args.output_dir / f"frame_{index}.png"
        tensor_to_image(frame).save(output)
        print(output)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())