Spaces:
Runtime error
Runtime error
| """ | |
| References: | |
| - Diffusion Forcing: https://github.com/buoyancy99/diffusion-forcing | |
| """ | |
| import torch | |
| from dit import DiT_models | |
| from vae import VAE_models | |
| from torchvision.io import read_video, write_video | |
| from utils import one_hot_actions, sigmoid_beta_schedule | |
| from tqdm import tqdm | |
| from einops import rearrange | |
| from torch import autocast | |
| import os | |
| #assert torch.cuda.is_available() | |
| #device = "cuda:0" | |
| def run_mod(): | |
| device = "cpu" | |
| # load DiT checkpoint | |
| ckpt = torch.load("oasis500m.pt",map_location=torch.device('cpu')) | |
| model = DiT_models["DiT-S/2"]() | |
| model.load_state_dict(ckpt, strict=False) | |
| model = model.to(device).eval() | |
| # load VAE checkpoint | |
| vae_ckpt = torch.load("vit-l-20.pt",map_location=torch.device('cpu')) | |
| vae = VAE_models["vit-l-20-shallow-encoder"]() | |
| vae.load_state_dict(vae_ckpt) | |
| vae = vae.to(device).eval() | |
| # sampling params | |
| B = 1 | |
| total_frames = 32 | |
| max_noise_level = 1000 | |
| ddim_noise_steps = 100 | |
| noise_range = torch.linspace(-1, max_noise_level - 1, ddim_noise_steps + 1) | |
| noise_abs_max = 20 | |
| ctx_max_noise_idx = ddim_noise_steps // 10 * 3 | |
| # get input video | |
| print(os.getcwd()) | |
| video_id = "snippy-chartreuse-mastiff-f79998db196d-20220401-224517.chunk_001" | |
| mp4_path = f"{os.getcwd()}/open_oasis_master/sample_data/{video_id}.mp4" | |
| actions_path = f"{os.getcwd()}/open_oasis_master/sample_data/{video_id}.actions.pt" | |
| video = read_video(mp4_path, pts_unit="sec")[0].float() / 255 | |
| actions = one_hot_actions(torch.load(actions_path,map_location=torch.device('cpu'))) | |
| offset = 100 | |
| video = video[offset:offset+total_frames].unsqueeze(0) | |
| actions = actions[offset:offset+total_frames].unsqueeze(0) | |
| # sampling inputs | |
| n_prompt_frames = 1 | |
| x = video[:, :n_prompt_frames] | |
| x = x.to(device) | |
| actions = actions.to(device) | |
| # vae encoding | |
| scaling_factor = 0.07843137255 | |
| x = rearrange(x, "b t h w c -> (b t) c h w") | |
| H, W = x.shape[-2:] | |
| with torch.no_grad(): | |
| x = vae.encode(x * 2 - 1).mean * scaling_factor | |
| x = rearrange(x, "(b t) (h w) c -> b t c h w", t=n_prompt_frames, h=H//vae.patch_size, w=W//vae.patch_size) | |
| # get alphas | |
| betas = sigmoid_beta_schedule(max_noise_level).to(device) | |
| alphas = 1.0 - betas | |
| alphas_cumprod = torch.cumprod(alphas, dim=0) | |
| alphas_cumprod = rearrange(alphas_cumprod, "T -> T 1 1 1") | |
| # sampling loop | |
| for i in tqdm(range(n_prompt_frames, total_frames)): | |
| chunk = torch.randn((B, 1, *x.shape[-3:]), device=device) | |
| chunk = torch.clamp(chunk, -noise_abs_max, +noise_abs_max) | |
| x = torch.cat([x, chunk], dim=1) | |
| start_frame = max(0, i + 1 - model.max_frames) | |
| for noise_idx in reversed(range(1, ddim_noise_steps + 1)): | |
| # set up noise values | |
| ctx_noise_idx = min(noise_idx, ctx_max_noise_idx) | |
| t_ctx = torch.full((B, i), noise_range[ctx_noise_idx], dtype=torch.long, device=device) | |
| t = torch.full((B, 1), noise_range[noise_idx], dtype=torch.long, device=device) | |
| t_next = torch.full((B, 1), noise_range[noise_idx - 1], dtype=torch.long, device=device) | |
| t_next = torch.where(t_next < 0, t, t_next) | |
| t = torch.cat([t_ctx, t], dim=1) | |
| t_next = torch.cat([t_ctx, t_next], dim=1) | |
| # sliding window | |
| x_curr = x.clone() | |
| x_curr = x_curr[:, start_frame:] | |
| t = t[:, start_frame:] | |
| t_next = t_next[:, start_frame:] | |
| # add some noise to the context | |
| ctx_noise = torch.randn_like(x_curr[:, :-1]) | |
| ctx_noise = torch.clamp(ctx_noise, -noise_abs_max, +noise_abs_max) | |
| x_curr[:, :-1] = alphas_cumprod[t[:, :-1]].sqrt() * x_curr[:, :-1] + (1 - alphas_cumprod[t[:, :-1]]).sqrt() * ctx_noise | |
| # get model predictions | |
| with torch.no_grad(): | |
| with autocast("cpu", dtype=torch.half): | |
| v = model(x_curr, t, actions[:, start_frame : i + 1]) | |
| x_start = alphas_cumprod[t].sqrt() * x_curr - (1 - alphas_cumprod[t]).sqrt() * v | |
| x_noise = ((1 / alphas_cumprod[t]).sqrt() * x_curr - x_start) \ | |
| / (1 / alphas_cumprod[t] - 1).sqrt() | |
| # get frame prediction | |
| x_pred = alphas_cumprod[t_next].sqrt() * x_start + x_noise * (1 - alphas_cumprod[t_next]).sqrt() | |
| x[:, -1:] = x_pred[:, -1:] | |
| # vae decoding | |
| x = rearrange(x, "b t c h w -> (b t) (h w) c") | |
| with torch.no_grad(): | |
| x = (vae.decode(x / scaling_factor) + 1) / 2 | |
| x = rearrange(x, "(b t) c h w -> b t h w c", t=total_frames) | |
| # save video | |
| x = torch.clamp(x, 0, 1) | |
| x = (x * 255).byte() | |
| write_video("video.mp4", x[0], fps=20) | |
| print("generation saved to video.mp4.") | |
| return "video.mp4" | |