File size: 4,256 Bytes
39e9e02 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 | #!/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())
|