""" Example usage: `python genie/evaluate.py --checkpoint_dir 1x-technologies/GENIE_35M` """ import argparse import time import os import sys from collections import defaultdict from pathlib import Path import accelerate import wandb import lpips import torch import transformers from accelerate import DataLoaderConfiguration from einops import rearrange from torch.utils.data import DataLoader from tqdm import tqdm from transformers import default_data_collator import numpy as np sys.path.append(os.getcwd()) import re from data import RawTokenDataset from visualize import decode_latents_wrapper from genie.st_mask_git import STMaskGIT from skimage import metrics as image_metrics from cont_data import RawFeatureDataset from raw_image_data import RawImageDataset from genie.st_mar import STMAR from datasets import utils from common.fid_score import calculate_fid from common.calculate_fvd import calculate_fvd from common.eval_utils import decode_tokens, decode_features, compute_lpips, AvgMetric, compute_loss wandb.login(key='4c1540ebf8cb9964703ac212a937c00848a79b67') # Hardcoded values for the v1.1 dataset WINDOW_SIZE = 12 STRIDE = 15 # Data is 30 Hz so with stride 15, video is 2 Hz SVD_SCALE = 0.18215 def parse_args(): parser = argparse.ArgumentParser(description="Evaluate GENIE-style models.") parser.add_argument( "--val_data_dir", type=str, default="data/1x_humanoid_magvit_traj10_val", help="A directory with video data, should have a `metadata.json` and `video.bin`." ) parser.add_argument( "--checkpoint_dir", type=str, help="Path to a HuggingFace-style checkpoint." ) parser.add_argument( "--batch_size", type=int, default=4, help="Batch size, current script only supports a single GPU." ) parser.add_argument( "--maskgit_steps", type=int, default=4, help="Number of MaskGIT sampling steps." ) parser.add_argument( "--temperature", type=float, default=0, help="Sampling temperature. If `temperature` <= 1e-8, will do greedy sampling." ) parser.add_argument( "--save_outputs_dir", type=str, help="Debug option. If specified, will save model predictions and ground truths to this directory. " "Specifically, will save `{pred_frames,pred_logits,gtruth_frames,gtruth_tokens}.pt`" ) parser.add_argument( "--max_examples", type=int, default=200, help="If specified, will stop evaluation early after `max_examples` examples." ) parser.add_argument( "--autoregressive_time", action="store_true", help="If True, autoregressive generation in time dimension." ) parser.add_argument( "--add_action_input", action="store_true", help="If True, uses action in the video output." ) parser.add_argument( "--perturbation_type", type=str, default="gaussian", help="Type of perturbation to apply to the action input. Options: gaussian " ) parser.add_argument( "--perturbation_scale", type=float, default=0.1, help="Perturbation applied to each action dimension." ) parser.add_argument( "--project_prefix", type=str, default="", help="Project suffix." ) parser.add_argument( "--use_feature", action="store_true", help="visualize the features rather than tokens" ) parser.add_argument( "--use_raw_image", action="store_true", help="use raw images as inputs", default=True ) return parser.parse_args() def get_model_step(checkpoint_dir): if os.path.exists(f"{checkpoint_dir}/scheduler.bin"): sch = torch.load(f"{checkpoint_dir}/scheduler.bin") return sch['_step_count'] return 0 class GenieEvaluator: def __init__(self, args, decode_latents, device="cuda"): super().__init__() if not os.path.exists(args.checkpoint_dir + "/config.json"): # search and find the latest modified checkpoint folder dirs = [os.path.join(args.checkpoint_dir, f.name) for f in os.scandir(args.checkpoint_dir) if f.is_dir()] dirs.sort(key=os.path.getctime) if len(dirs) > 3 and os.path.join(args.checkpoint_dir, "epoch_1") in dirs: dirs.remove(os.path.join(args.checkpoint_dir, "epoch_1")) if len(dirs) == 0: exit(f"No checkpoint found in {args.checkpoint_dir}") paths = dirs[:-3] # only keep the last 3 for path in paths: print(f"evaluation: remove rm -rf {path}") os.system(f"rm -rf {path}") args.checkpoint_dir = dirs[-1] print("Loading model from:", args.checkpoint_dir) self.model = STMAR.from_pretrained(args.checkpoint_dir) self.model_step = get_model_step(args.checkpoint_dir) self.model = self.model.to(device=device) self.model.eval() self.decode_latents = decode_latents self.device = device self.args = args def predict_zframe_logits(self, input_ids: torch.Tensor, action_ids: torch.Tensor = None, domains = None, skip_normalization: bool = False) -> tuple[torch.LongTensor, torch.FloatTensor]: """ Conditioned on each prefix: [frame_0], [frame_0, frame_1], ..., [frame_0, frame_1, ... frame_{T-1}], predict the tokens in the following frame: [pred_frame_1, pred_frame_2, ..., pred_frame_T]. Image logits are denoised in parallel across spatial dimension and teacher-forced across the time dimension. To compute logits, we save both the samples and logits as we do MaskGIT generation. Total number of forward passes is (T-1) * maskgit steps. Args: input_ids: Tensor of size (B, T*H*W) corresponding to flattened, tokenized images. Returns: (samples_THW, factored_logits) samples_THW: size (B, T, H, W) corresponding to the token ids of the predicted frames. May differ from the argmax of `factored_logits` if not greedy sampling. factored_logits: size (B, 512, 2, T-1, H, W) corresponding to the predicted logits. Note that we are factorizing the 2**18 vocabulary into two separate vocabularies of size 512 each. """ inputs_THW = rearrange(input_ids, "b (t h w) ... -> b t h w ...", t=WINDOW_SIZE, h=self.args.latent_h, w=self.args.latent_w).to(self.device) all_samples = [] all_logits = [] samples_HW = inputs_THW.clone() for timestep in range(1, WINDOW_SIZE): print(f"Generating frame {timestep}") inputs_masked = inputs_THW.clone() if self.args.autoregressive_time: if timestep > self.model.config.num_prompt_frames: inputs_masked[:, timestep-1] = samples_HW.clone() inputs_masked[:, timestep:] = self.model.mask_token # MaskGIT sampling samples_HW, factored_logits, _ = self.model.maskgit_generate( inputs_masked, out_t=timestep, maskgit_steps=self.args.maskgit_steps, temperature=self.args.temperature, action_ids=action_ids, domain=domains, skip_normalization=skip_normalization ) all_samples.append(samples_HW) all_logits.append(factored_logits) samples_THW = torch.stack(all_samples, dim=1) return samples_THW, torch.stack(all_logits, dim=3) def predict_next_frames(self, samples_THW) -> torch.Tensor: """ All model submissions should have this defined. Like predict_next_frames, this is teacher-forced along spatial dimension, autoregressive along time dimension. Conditioned on each prefix: [frame_0], [frame_0, frame_1], ..., [frame_0, frame_1, ..., frame_{T-1}], predict the following frame: [pred_frame_1, pred_frame_2, ..., pred_frame_T]. For this model, the frames are generated by using the argmax of `predict_zframe_logits` and decoding the quantized latent space tokens back to the original image space. Args: samples_THW: LongTensor of size (B, T, H, W) corresponding to sampled images in the quantized latent space. Returns: LongTensor of size (B, T-1, 3, 256, 256) corresponding to the predicted frames. """ return decode_features(samples_THW.cpu() / SVD_SCALE, self.decode_latents) @torch.no_grad() def main(): transformers.set_seed(42) args = parse_args() # allow different batch sizes in final batch accelerator = accelerate.Accelerator(dataloader_config=DataLoaderConfiguration(even_batches=False)) # if "robomimic" in args.val_data_dir: # dataset = "robomimic" # save the results to wandb. hardcoded the input dataset to have magvit and will change later dataset = re.search(r"data/(.*?)_magvit", args.val_data_dir).group(1) # rtrim the last / and get the last part of the path args.checkpoint_dir = args.checkpoint_dir.rstrip('/') name = args.checkpoint_dir.split('/')[-1] decode_latents = decode_latents_wrapper(device=accelerator.device, encoder_name_or_path="stabilityai/stable-video-diffusion-img2vid", encoder_type="temporalvae") evaluator = GenieEvaluator(args, decode_latents) action_d = len(evaluator.model.action_preprocessor[dataset].mean) action_d_horizon = evaluator.model.config.d_actions[evaluator.model.config.action_domains.index(dataset)] stride = action_d_horizon // action_d print("model stride:", stride) if accelerator.is_main_process: wandb.teardown() wandb.init(project='video_val', resume="allow", id=f"{args.project_prefix}{name}", name=f"{args.project_prefix}{name}", settings=wandb.Settings(start_method="thread")) with_action_input = True if args.use_raw_image: args.val_data_dir = args.val_data_dir.replace("magvit", "image") val_dataset = RawImageDataset(args.val_data_dir, window_size=WINDOW_SIZE, compute_stride_from_freq_table=False, stride=stride, filter_overlaps=True, use_actions=with_action_input) else: # args.val_data_dir = args.val_data_dir.replace("magvit", "vae") args.val_data_dir = args.val_data_dir.replace("magvit_traj1000000", "noquant_temporalvae_shard0_of_1") val_dataset = RawFeatureDataset(args.val_data_dir, window_size=WINDOW_SIZE, compute_stride_from_freq_table=False, stride=stride, filter_overlaps=True, use_actions=with_action_input) dataset_metadata = val_dataset.metadata assert hasattr(evaluator, "model"), "Expected Evaluator to have attribute `model`." evaluator.model = accelerator.prepare_model(evaluator.model, evaluation_mode=True) # No DDP with_action_input = evaluator.model.config.use_actions # hack to reset lpips_alex = lpips.LPIPS(net="alex") # Calculate LPIPS w/ AlexNet, which is the fastest model out of their options random_samples = None if args.max_examples is not None: val_dataset.valid_start_inds = val_dataset.valid_start_inds[:args.max_examples] dataloader = DataLoader(val_dataset, collate_fn=default_data_collator, batch_size=args.batch_size) metrics = defaultdict(AvgMetric) batch_idx = 0 latent_side_len = 32 # hardcoded args.latent_h = args.latent_w = latent_side_len dataloader = accelerator.prepare(dataloader) gt_full_sequence = [] generated_full_sequence = [] for batch in tqdm(dataloader): batch_idx += 1 if args.use_raw_image: # token the batches on the fly images = batch["images"].detach().cpu().numpy().astype(np.uint8) outputs = [] for context in images: output = [] for image_t in context: output_t = utils.get_vae_image_embeddings( image_t, encoder_type="temporalvae", encoder_name_or_path="stabilityai/stable-video-diffusion-img2vid", ) output.append(output_t) outputs.append(output) batch["input_ids"] = torch.FloatTensor(outputs).to(evaluator.device) batch["input_ids"] = rearrange(batch["input_ids"], "b t c h w -> b (t h w) c") * SVD_SCALE batch["labels"] = batch["input_ids"].clone() batch_size = batch["input_ids"].size(0) reshaped_input_ids = rearrange(batch["input_ids"], "b (t h w) ... -> b t h w ...", t=WINDOW_SIZE, h=latent_side_len, w=latent_side_len) start_time = time.time() if not with_action_input: samples, _ = evaluator.predict_zframe_logits(batch["input_ids"].to(evaluator.device), domains=[val_dataset.name]) else: samples, _ = evaluator.predict_zframe_logits(batch["input_ids"].to(evaluator.device), batch["action_ids"].to(evaluator.device), [val_dataset.name]) frames_per_batch = (WINDOW_SIZE - 1) * batch["input_ids"].size(0) metrics["gen_time"].update((time.time() - start_time) / frames_per_batch, batch_size) start_time = time.time() pred_frames = evaluator.predict_next_frames(samples) metrics["dec_time"].update((time.time() - start_time) / frames_per_batch, batch_size) decoded_gtruth = decode_features(reshaped_input_ids / SVD_SCALE, decode_latents) decoded_gtruth_clone = batch['images'].permute(0, 1, 4, 2, 3)[:len(decoded_gtruth)] if args.use_raw_image: # key: use raw image as the groundtruth decoded_gtruth = batch['images'].permute(0, 1, 4, 2, 3)[:len(decoded_gtruth)].long().cpu().detach() metrics["pred_lpips"].update_list(compute_lpips(decoded_gtruth[:, 1:], pred_frames, lpips_alex)) gt_frames_numpy = decoded_gtruth[:, 1:].detach().cpu().numpy() pred_frames_numpy = pred_frames.detach().cpu().numpy() # save the image to wandb # if accelerator.is_main_process: # for i in range(gt_frames_numpy.shape[0] // 4): # wandb.log({ # f"{dataset}/gt_{i}": [wandb.Image(np.transpose(gt_frames_numpy[i][j], (1,2,0))) for j in range(gt_frames_numpy.shape[1] // 2)], # f"{dataset}/pred_{i}": [wandb.Image(np.transpose(pred_frames_numpy[i][j], (1,2,0))) for j in range(pred_frames_numpy.shape[1] // 2)] # }) psnr = [image_metrics.peak_signal_noise_ratio( gt_frames_numpy[i][-1] / 255., pred_frames_numpy[i][-1] / 255., data_range=1.0) for i in range(gt_frames_numpy.shape[0])] ssim = [np.mean([image_metrics.structural_similarity( gt_frames_numpy[i][j] / 255., pred_frames_numpy[i][j] / 255., data_range=1.0, channel_axis=0) \ for i in range(gt_frames_numpy.shape[0])]) for j in range(gt_frames_numpy.shape[1])] metrics["ssim"].update_list(ssim) metrics["psnr"].update_list(psnr) gt_full_sequence.append(decoded_gtruth_clone[:, 1:]) generated_full_sequence.append(pred_frames) # metrics["fvd"].update_list(calculate_fvd(.float().to(accelerator.device) / 255., # pred_frames.float().to(accelerator.device) / 255, device=accelerator.device)) # try: # metrics["fvd"].update_list(calculate_fvd(decoded_gtruth[:, 1:], pred_frames)) # except Exception as e: # print(f"Error calculating FVD: {e}") # As in Genie. we also compute psnr_delta = PSNR(x_t, x_t_hat) - PSNR(x_t, x_t_hatprime) where x_t_hatprime samples random actions # this difference in PSNR measures the controllability # actions need to be just uniform random actions if with_action_input: # for computing delta psnr N_TRIALS = 5 psnr_delta_mean = np.zeros(gt_frames_numpy.shape[0]) for _ in range(N_TRIALS): # action_mean, action_std = val_dataset.action_stat # action_std = torch.tensor(action_std).to(evaluator.device) # action_mean = torch.tensor(action_mean).to(evaluator.device) action_mean = evaluator.model.action_preprocessor[dataset].mean.repeat(stride) action_std = evaluator.model.action_preprocessor[dataset].std.repeat(stride) random_action_ids = torch.randn_like(batch["action_ids"]) * action_std + action_mean random_samples, _ = evaluator.predict_zframe_logits(batch["input_ids"].to(evaluator.device), random_action_ids.to(evaluator.device), [val_dataset.name], skip_normalization=False) random_pred_frames = evaluator.predict_next_frames(random_samples) random_pred_frames_numpy = random_pred_frames.detach().cpu().numpy() # random subtracts groundtruth psnr_delta = [psnr[i] - image_metrics.peak_signal_noise_ratio( gt_frames_numpy[i][-1] / 255., random_pred_frames_numpy[i][-1] / 255., data_range=1.0) for i in range(gt_frames_numpy.shape[0])] psnr_delta_mean += np.array(psnr_delta) / N_TRIALS metrics[f"psnr_delta"].update_list(psnr_delta_mean) print(f"=== dataset {dataset} model: {name}") print({key: f"{val.mean():.4f}" for key, val in metrics.items()}) if batch_idx > args.max_examples: break generated_full_sequence = torch.cat(generated_full_sequence, dim=0) / 255. gt_full_sequence = torch.cat(gt_full_sequence, dim=0) / 255. gt_full_sequence.detach().cpu().numpy().tofile(args.checkpoint_dir + "/gt_video.bin") generated_full_sequence.detach().cpu().numpy().tofile(args.checkpoint_dir + "/generated_video.bin") # save the generated and groundtruth sequences # import IPython; IPython.embed() metrics["fid"].update_list([calculate_fid(gt_full_sequence, generated_full_sequence, device=accelerator.device)]) metrics["fvd"].update_list([calculate_fvd(gt_full_sequence, generated_full_sequence, device=accelerator.device)]) for key, val in metrics.items(): agg_total, agg_count = accelerator.reduce( torch.tensor([val.total, val.count], device=accelerator.device) ) accelerator.print(f"{key}: {agg_total / agg_count:.4f}") if accelerator.is_main_process: prefix = "teacher_force" if not args.autoregressive_time else "autoregressive" for key, val in metrics.items(): try: wandb.log({f"{dataset}/{prefix}_{key}": val.mean()}) wandb.log({f"{prefix}_{key}": val.mean()}) except Exception as e: print(e) wandb.log({f"{dataset}/num_examples": len(val_dataset)}) wandb.log({f"{dataset}/perturbation_scale": args.perturbation_scale}) wandb.log({f"model_step": evaluator.model_step}) # model training steps dataset_metadata = { f"{dataset}/dataset_name": f"{dataset}", f"{dataset}/num_examples": len(val_dataset), f"{dataset}/num_features": len(val_dataset[0]) if val_dataset else 0, f"{dataset}/sample_data": val_dataset[0] if len(val_dataset) > 0 else "N/A", f"{dataset}/model_step": evaluator.model_step } for k, v in dataset_metadata.items(): wandb.run.summary[k] = v wandb.finish() if __name__ == "__main__": main()