| import os |
| import sys |
| sys.path.append(os.getcwd()) |
|
|
| import argparse |
| import warnings |
| warnings.filterwarnings("ignore") |
|
|
| import time |
| from contextlib import nullcontext |
| from omegaconf import OmegaConf |
| from pathlib import Path |
| from tqdm import tqdm |
|
|
| import numpy as np |
| import torch |
| import decord |
| from einops import rearrange |
| from lightning.pytorch import seed_everything |
| from torch import autocast |
| from torchvision import transforms |
| from torchvision.io import write_video |
|
|
| from vidtok.modules.util import print0 |
| from scripts.inference_evaluate import load_model_from_config |
|
|
|
|
| class SingleVideoDataset(torch.utils.data.Dataset): |
| def __init__( |
| self, |
| video_path, |
| input_height=128, |
| input_width=128, |
| sample_fps=8, |
| chunk_size=16, |
| is_causal=True, |
| read_long_video=False |
| ): |
| decord.bridge.set_bridge("torch") |
| self.video_path = video_path |
| self.transform = transforms.Compose( |
| [ |
| transforms.Resize(input_height, antialias=True), |
| transforms.CenterCrop((input_height, input_width)), |
| transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)), |
| ] |
| ) |
|
|
| self.video_reader = decord.VideoReader(video_path, num_threads=0) |
| total_frames = len(self.video_reader) |
| fps = self.video_reader.get_avg_fps() |
|
|
| interval = round(fps / sample_fps) |
| frame_ids = list(range(0, total_frames, interval)) |
| self.frame_ids_batch = [] |
| if read_long_video: |
| video_length = len(frame_ids) |
| if is_causal and video_length > chunk_size: |
| self.frame_ids_batch.append(frame_ids[:chunk_size * ((video_length - 1) // chunk_size) + 1]) |
| elif not is_causal and video_length >= chunk_size: |
| self.frame_ids_batch.append(frame_ids[:chunk_size * (video_length // chunk_size)]) |
| else: |
| num_frames_per_batch = chunk_size + 1 if is_causal else chunk_size |
| for x in range(0, len(frame_ids), num_frames_per_batch): |
| if len(frame_ids[x : x + num_frames_per_batch]) == num_frames_per_batch: |
| self.frame_ids_batch.append(frame_ids[x : x + num_frames_per_batch]) |
|
|
| def __len__(self): |
| return len(self.frame_ids_batch) |
|
|
| def __getitem__(self, idx): |
| frame_ids = self.frame_ids_batch[idx] |
| frames = self.video_reader.get_batch(frame_ids).permute(0, 3, 1, 2).float() / 255.0 |
| frames = self.transform(frames).permute(1, 0, 2, 3) |
| return frames |
|
|
|
|
| def tensor_to_uint8(tensor): |
| tensor = torch.clamp(tensor, -1.0, 1.0) |
| tensor = (tensor + 1.0) / 2.0 |
| tensor = (tensor.cpu().numpy() * 255).astype(np.uint8) |
| return tensor |
|
|
|
|
| def main(): |
| def str2bool(v): |
| if isinstance(v, bool): |
| return v |
| if v.lower() in ("yes", "true", "t", "y", "1"): |
| return True |
| elif v.lower() in ("no", "false", "f", "n", "0"): |
| return False |
| else: |
| raise argparse.ArgumentTypeError("Boolean value expected.") |
|
|
| parser = argparse.ArgumentParser() |
|
|
| parser.add_argument( |
| "--seed", |
| type=int, |
| default=42, |
| help="the seed (for reproducible sampling)", |
| ) |
| parser.add_argument( |
| "--precision", type=str, help="evaluate at this precision", choices=["full", "autocast"], default="full" |
| ) |
| parser.add_argument( |
| "--config", |
| type=str, |
| default="configs/vidtok_kl_causal_488_4chn.yaml", |
| help="path to config which constructs model", |
| ) |
| parser.add_argument( |
| "--ckpt", |
| type=str, |
| default="checkpoints/vidtok_kl_causal_488_4chn.ckpt", |
| help="path to checkpoint of model", |
| ) |
| parser.add_argument( |
| "--output_video_dir", |
| type=str, |
| default="tmp", |
| help="path to save the outputs", |
| ) |
| parser.add_argument( |
| "--input_video_path", |
| type=str, |
| default="assets/example.mp4", |
| help="path to the input video", |
| ) |
| parser.add_argument( |
| "--input_height", |
| type=int, |
| default=256, |
| help="height of the input video", |
| ) |
| parser.add_argument( |
| "--input_width", |
| type=int, |
| default=256, |
| help="width of the input video", |
| ) |
| parser.add_argument( |
| "--sample_fps", |
| type=int, |
| default=30, |
| help="sample fps", |
| ) |
| parser.add_argument( |
| "--chunk_size", |
| type=int, |
| default=16, |
| help="the size of a chunk - we split a long video into several chunks", |
| ) |
| parser.add_argument( |
| "--read_long_video", |
| action='store_true' |
| ) |
| parser.add_argument( |
| "--pad_gen_frames", |
| action="store_true", |
| help="Used only in causal mode. If True, pad frames generated in the last batch, else replicate the first frame instead", |
| ) |
| parser.add_argument( |
| "--concate_input", |
| type=str2bool, |
| const=True, |
| default=True, |
| nargs="?", |
| help="", |
| ) |
|
|
| args = parser.parse_args() |
| seed_everything(args.seed) |
|
|
| print0(f"[bold red]\[scripts.inference_reconstruct][/bold red] Evaluating model {args.ckpt}") |
| device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") |
| precision_scope = autocast if args.precision == "autocast" else nullcontext |
| config = OmegaConf.load(args.config) |
|
|
| os.makedirs(args.output_video_dir, exist_ok=True) |
|
|
| model = load_model_from_config(args.config, args.ckpt) |
| model.to(device).eval() |
| assert args.chunk_size % model.encoder.time_downsample_factor == 0 |
|
|
| if args.read_long_video: |
| assert hasattr(model, 'use_tiling'), "Tiling inference is needed to conduct long video reconstruction." |
| print(f"Using tiling inference to save memory usage...") |
| model.use_tiling = True |
| model.t_chunk_enc = args.chunk_size |
| model.t_chunk_dec = model.t_chunk_enc // model.encoder.time_downsample_factor |
| model.use_overlap = True |
| |
| dataset = SingleVideoDataset( |
| video_path=args.input_video_path, |
| input_height=args.input_height, |
| input_width=args.input_width, |
| sample_fps=args.sample_fps, |
| chunk_size=args.chunk_size, |
| is_causal=model.is_causal, |
| read_long_video=args.read_long_video |
| ) |
| dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False) |
|
|
| inputs = [] |
| outputs = [] |
| with torch.no_grad(), precision_scope("cuda"): |
| tic = time.time() |
| for i, input in tqdm(enumerate(dataloader)): |
| input = input.to(device) |
| |
| if model.is_causal and not args.read_long_video and args.pad_gen_frames: |
| if i == 0: |
| _, xrec, _ = model(input) |
| else: |
| _, xrec, _ = model(torch.cat([last_gen_frames, input], dim=2)) |
| xrec = xrec[:, :, -input.shape[2]:].clamp(-1, 1) |
| last_gen_frames = xrec[:, :, (1 - model.encoder.time_downsample_factor):, :, :] |
| else: |
| _, xrec, _ = model(input) |
| |
| input = rearrange(input, "b c t h w -> (b t) c h w") |
| inputs.append(input) |
| xrec = rearrange(xrec.clamp(-1, 1), "b c t h w -> (b t) c h w") |
| outputs.append(xrec) |
|
|
| toc = time.time() |
|
|
| |
| inputs = tensor_to_uint8(torch.cat(inputs, dim=0)) |
| inputs = rearrange(inputs, "t c h w -> t h w c") |
| outputs = tensor_to_uint8(torch.cat(outputs, dim=0)) |
| outputs = rearrange(outputs, "t c h w -> t h w c") |
| min_len = min(inputs.shape[0], outputs.shape[0]) |
| final = np.concatenate([inputs[:min_len], outputs[:min_len]], axis=2) if args.concate_input else outputs[:min_len] |
|
|
| output_video_path = os.path.join(args.output_video_dir, f"{Path(args.input_video_path).stem}_reconstructed.mp4") |
| write_video(output_video_path, final, args.sample_fps) |
|
|
| print0(f"[bold red]Results saved in: {output_video_path}[/bold red]") |
| print0(f"[bold red]\[scripts.inference_reconstruct][/bold red] Time taken: {toc - tic:.2f}s") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
