| |
| """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()) |
|
|
