"""VBVR-Pro-Wan2.2-I2V-A14B inference example.

Usage:
    python example.py --model_path ./VBVR-Pro-Wan2.2-I2V-A14B \
        --image input.png --prompt "Your video instruction"
"""

import argparse

import torch
from diffusers import AutoencoderKLWan, WanImageToVideoPipeline
from diffusers.utils import export_to_video, load_image


parser = argparse.ArgumentParser()
parser.add_argument("--model_path", type=str, default="VBVR-Pro-Wan2.2-I2V-A14B")
parser.add_argument("--image", type=str, required=True, help="Path or URL to input image")
parser.add_argument("--prompt", type=str, required=True, help="Video instruction")
parser.add_argument(
    "--negative_prompt",
    type=str,
    default="Bright tones, overexposed, static, blurred details, subtitles, low quality",
)
parser.add_argument("--output", type=str, default="output.mp4")
parser.add_argument("--width", type=int, default=832)
parser.add_argument("--height", type=int, default=480)
parser.add_argument("--num_frames", type=int, default=81)
parser.add_argument("--steps", type=int, default=50)
parser.add_argument("--guidance_scale", type=float, default=5.0)
parser.add_argument("--fps", type=int, default=15)
parser.add_argument("--seed", type=int, default=42)
args = parser.parse_args()

print(f"Loading model from: {args.model_path}")
vae = AutoencoderKLWan.from_pretrained(
    args.model_path, subfolder="vae", torch_dtype=torch.float32
)
pipe = WanImageToVideoPipeline.from_pretrained(
    args.model_path, vae=vae, torch_dtype=torch.bfloat16
)
pipe.enable_model_cpu_offload()

image = load_image(args.image).convert("RGB").resize((args.width, args.height))
print(f"Input image: {args.image} ({image.size[0]}x{image.size[1]})")

frames = pipe(
    image=image,
    prompt=args.prompt,
    negative_prompt=args.negative_prompt,
    height=args.height,
    width=args.width,
    num_frames=args.num_frames,
    num_inference_steps=args.steps,
    guidance_scale=args.guidance_scale,
    generator=torch.manual_seed(args.seed),
).frames[0]

export_to_video(frames, args.output, fps=args.fps)
print(f"Saved to: {args.output}")