from typing import Callable, Dict, List, Optional, Union import gc import numpy as np import torch import torch.nn.functional as F from diffusers.pipelines.stable_video_diffusion.pipeline_stable_video_diffusion import ( _resize_with_antialiasing, StableVideoDiffusionPipeline, retrieve_timesteps, ) from diffusers.utils import logging from kornia.utils import create_meshgrid from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution logger = logging.get_logger(__name__) # pylint: disable=invalid-name @torch.no_grad() def normalize_point_map(point_map, valid_mask): # T,H,W,3 T,H,W norm_factor = (point_map[..., 2] * valid_mask.float()).mean() / (valid_mask.float().mean() + 1e-8) norm_factor = norm_factor.clip(min=1e-3) return point_map / norm_factor def point_map_xy2intrinsic_map(point_map_xy): # *,h,w,2 height, width = point_map_xy.shape[-3], point_map_xy.shape[-2] assert height % 2 == 0 assert width % 2 == 0 mesh_grid = create_meshgrid( height=height, width=width, normalized_coordinates=True, device=point_map_xy.device, dtype=point_map_xy.dtype )[0] # h,w,2 assert mesh_grid.abs().min() > 1e-4 # *,h,w,2 mesh_grid = mesh_grid.expand_as(point_map_xy) nc = point_map_xy.mean(dim=-2).mean(dim=-2) # *, 2 nc_map = nc[..., None, None, :].expand_as(point_map_xy) nf = ((point_map_xy - nc_map) / mesh_grid).mean(dim=-2).mean(dim=-2) nf_map = nf[..., None, None, :].expand_as(point_map_xy) # print((mesh_grid * nf_map + nc_map - point_map_xy).abs().max()) return torch.cat([nc_map, nf_map], dim=-1) def robust_min_max(tensor, quantile=0.99): T, H, W = tensor.shape min_vals = [] max_vals = [] for i in range(T): min_vals.append(torch.quantile(tensor[i], q=1-quantile, interpolation='nearest').item()) max_vals.append(torch.quantile(tensor[i], q=quantile, interpolation='nearest').item()) return min(min_vals), max(max_vals) class GeometryCrafterDiffPipeline(StableVideoDiffusionPipeline): @torch.inference_mode() def encode_video( self, video: torch.Tensor, chunk_size: int = 14, ) -> torch.Tensor: """ :param video: [b, c, h, w] in range [-1, 1], the b may contain multiple videos or frames :param chunk_size: the chunk size to encode video :return: image_embeddings in shape of [b, 1024] """ video_224 = _resize_with_antialiasing(video.float(), (224, 224)) video_224 = (video_224 + 1.0) / 2.0 # [-1, 1] -> [0, 1] embeddings = [] for i in range(0, video_224.shape[0], chunk_size): emb = self.feature_extractor( images=video_224[i : i + chunk_size], do_normalize=True, do_center_crop=False, do_resize=False, do_rescale=False, return_tensors="pt", ).pixel_values.to(video.device, dtype=video.dtype) embeddings.append(self.image_encoder(emb).image_embeds) # [b, 1024] embeddings = torch.cat(embeddings, dim=0) # [t, 1024] return embeddings @torch.inference_mode() def encode_vae_video( self, video: torch.Tensor, chunk_size: int = 14, ): """ :param video: [b, c, h, w] in range [-1, 1], the b may contain multiple videos or frames :param chunk_size: the chunk size to encode video :return: vae latents in shape of [b, c, h, w] """ video_latents = [] for i in range(0, video.shape[0], chunk_size): video_latents.append( self.vae.encode(video[i : i + chunk_size]).latent_dist.mode() ) video_latents = torch.cat(video_latents, dim=0) return video_latents @torch.inference_mode() def produce_priors(self, prior_model, frame, chunk_size=8): T, _, H, W = frame.shape # frame = (frame + 1) / 2 pred_point_maps = [] pred_masks = [] for i in range(0, len(frame), chunk_size): pred_p, pred_m = prior_model.forward_image(frame[i:i+chunk_size]) pred_point_maps.append(pred_p) pred_masks.append(pred_m) pred_point_maps = torch.cat(pred_point_maps, dim=0) pred_masks = torch.cat(pred_masks, dim=0) pred_masks = pred_masks.float() * 2 - 1 # T,H,W,3 T,H,W pred_point_maps = normalize_point_map(pred_point_maps, pred_masks > 0) pred_disps = 1.0 / pred_point_maps[..., 2].clamp_min(1e-3) pred_disps = pred_disps * (pred_masks > 0) min_disparity, max_disparity = robust_min_max(pred_disps) pred_disps = ((pred_disps - min_disparity) / (max_disparity - min_disparity+1e-4)).clamp(0, 1) pred_disps = pred_disps * 2 - 1 pred_point_maps[..., :2] = pred_point_maps[..., :2] / (pred_point_maps[..., 2:3] + 1e-7) pred_point_maps[..., 2] = torch.log(pred_point_maps[..., 2] + 1e-7) * (pred_masks > 0) # [x/z, y/z, log(z)] pred_intr_maps = point_map_xy2intrinsic_map(pred_point_maps[..., :2]).permute(0,3,1,2) # T,H,W,2 pred_point_maps = pred_point_maps.permute(0,3,1,2) return pred_disps, pred_masks, pred_point_maps, pred_intr_maps @torch.inference_mode() def encode_point_map(self, point_map_vae, disparity, valid_mask, point_map, intrinsic_map, chunk_size=8): T, _, H, W = point_map.shape latents = [] psedo_image = disparity[:, None].repeat(1,3,1,1) intrinsic_map = torch.norm(intrinsic_map[:, 2:4], p=2, dim=1, keepdim=False) for i in range(0, T, chunk_size): latent_dist = self.vae.encode(psedo_image[i : i + chunk_size].to(self.vae.dtype)).latent_dist latent_dist = point_map_vae.encode( torch.cat([ intrinsic_map[i:i+chunk_size, None], point_map[i:i+chunk_size, 2:3], disparity[i:i+chunk_size, None], valid_mask[i:i+chunk_size, None]], dim=1), latent_dist ) if isinstance(latent_dist, DiagonalGaussianDistribution): latent = latent_dist.mode() else: latent = latent_dist assert isinstance(latent, torch.Tensor) latents.append(latent) latents = torch.cat(latents, dim=0) latents = latents * self.vae.config.scaling_factor return latents @torch.no_grad() def decode_point_map(self, point_map_vae, latents, chunk_size=8, force_projection=True, force_fixed_focal=True, use_extract_interp=False, need_resize=False, height=None, width=None): T = latents.shape[0] rec_intrinsic_maps = [] rec_depth_maps = [] rec_valid_masks = [] for i in range(0, T, chunk_size): lat = latents[i:i+chunk_size] rec_imap, rec_dmap, rec_vmask = point_map_vae.decode( lat, num_frames=lat.shape[0], ) rec_intrinsic_maps.append(rec_imap) rec_depth_maps.append(rec_dmap) rec_valid_masks.append(rec_vmask) rec_intrinsic_maps = torch.cat(rec_intrinsic_maps, dim=0) rec_depth_maps = torch.cat(rec_depth_maps, dim=0) rec_valid_masks = torch.cat(rec_valid_masks, dim=0) if need_resize: rec_depth_maps = F.interpolate(rec_depth_maps, (height, width), mode='nearest-exact') if use_extract_interp else F.interpolate(rec_depth_maps, (height, width), mode='bilinear', align_corners=False) rec_valid_masks = F.interpolate(rec_valid_masks, (height, width), mode='nearest-exact') if use_extract_interp else F.interpolate(rec_valid_masks, (height, width), mode='bilinear', align_corners=False) rec_intrinsic_maps = F.interpolate(rec_intrinsic_maps, (height, width), mode='bilinear', align_corners=False) H, W = rec_intrinsic_maps.shape[-2], rec_intrinsic_maps.shape[-1] mesh_grid = create_meshgrid( H, W, normalized_coordinates=True ).to(rec_intrinsic_maps.device, rec_intrinsic_maps.dtype, non_blocking=True) # 1,h,w,2 rec_intrinsic_maps = torch.cat([rec_intrinsic_maps * W / np.sqrt(W**2+H**2), rec_intrinsic_maps * H / np.sqrt(W**2+H**2)], dim=1) # t,2,h,w mesh_grid = mesh_grid.permute(0,3,1,2) rec_valid_masks = rec_valid_masks.squeeze(1) > 0 if force_projection: if force_fixed_focal: nfx = (rec_intrinsic_maps[:, 0, :, :] * rec_valid_masks.float()).mean() / (rec_valid_masks.float().mean() + 1e-4) nfy = (rec_intrinsic_maps[:, 1, :, :] * rec_valid_masks.float()).mean() / (rec_valid_masks.float().mean() + 1e-4) rec_intrinsic_maps = torch.tensor([nfx, nfy], device=rec_intrinsic_maps.device)[None, :, None, None].repeat(T, 1, 1, 1) else: nfx = (rec_intrinsic_maps[:, 0, :, :] * rec_valid_masks.float()).mean(dim=[-1, -2]) / (rec_valid_masks.float().mean(dim=[-1, -2]) + 1e-4) nfy = (rec_intrinsic_maps[:, 1, :, :] * rec_valid_masks.float()).mean(dim=[-1, -2]) / (rec_valid_masks.float().mean(dim=[-1, -2]) + 1e-4) rec_intrinsic_maps = torch.stack([nfx, nfy], dim=-1)[:, :, None, None] # t,2,1,1 rec_point_maps = torch.cat([rec_intrinsic_maps * mesh_grid, rec_depth_maps], dim=1).permute(0,2,3,1) xy, z = rec_point_maps.split([2, 1], dim=-1) z = torch.clamp_max(z, 10) # for numerical stability z = torch.exp(z) rec_point_maps = torch.cat([xy * z, z], dim=-1) return rec_point_maps, rec_valid_masks @torch.no_grad() def __call__( self, video: Union[np.ndarray, torch.Tensor], point_map_vae, prior_model, height: int = 320, width: int = 640, num_inference_steps: int = 5, guidance_scale: float = 1.0, window_size: Optional[int] = 14, noise_aug_strength: float = 0.02, decode_chunk_size: Optional[int] = None, generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, latents: Optional[torch.FloatTensor] = None, callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, callback_on_step_end_tensor_inputs: List[str] = ["latents"], overlap: int = 4, force_projection: bool = True, force_fixed_focal: bool = True, use_extract_interp: bool = False, track_time: bool = False, ): # video: in shape [t, h, w, c] if np.ndarray or [t, c, h, w] if torch.Tensor, in range [0, 1] # 0. Default height and width to unet if isinstance(video, np.ndarray): video = torch.from_numpy(video.transpose(0, 3, 1, 2)) else: assert isinstance(video, torch.Tensor) height = height or video.shape[-2] width = width or video.shape[-1] original_height = video.shape[-2] original_width = video.shape[-1] num_frames = video.shape[0] decode_chunk_size = decode_chunk_size if decode_chunk_size is not None else 8 if num_frames <= window_size: window_size = num_frames overlap = 0 stride = window_size - overlap # 1. Check inputs. Raise error if not correct assert height % 64 == 0 and width % 64 == 0 if original_height != height or original_width != width: need_resize = True else: need_resize = False # 2. Define call parameters batch_size = 1 device = self._execution_device # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` # corresponds to doing no classifier free guidance. self._guidance_scale = guidance_scale if track_time: start_event = torch.cuda.Event(enable_timing=True) prior_event = torch.cuda.Event(enable_timing=True) encode_event = torch.cuda.Event(enable_timing=True) denoise_event = torch.cuda.Event(enable_timing=True) decode_event = torch.cuda.Event(enable_timing=True) start_event.record() # 3. Encode input video pred_disparity, pred_valid_mask, pred_point_map, pred_intrinsic_map = self.produce_priors( prior_model, video.to(device=device, dtype=torch.float32), chunk_size=decode_chunk_size ) # T,H,W T,H,W T,3,H,W T,2,H,W if need_resize: pred_disparity = F.interpolate(pred_disparity.unsqueeze(1), (height, width), mode='bilinear', align_corners=False).squeeze(1) pred_valid_mask = F.interpolate(pred_valid_mask.unsqueeze(1), (height, width), mode='bilinear', align_corners=False).squeeze(1) pred_point_map = F.interpolate(pred_point_map, (height, width), mode='bilinear', align_corners=False) pred_intrinsic_map = F.interpolate(pred_intrinsic_map, (height, width), mode='bilinear', align_corners=False) if track_time: prior_event.record() torch.cuda.synchronize() elapsed_time_ms = start_event.elapsed_time(prior_event) print(f"Elapsed time for computing per-frame prior: {elapsed_time_ms} ms") else: gc.collect() torch.cuda.empty_cache() # 3. Encode input video if need_resize: video = F.interpolate(video, (height, width), mode="bicubic", align_corners=False, antialias=True).clamp(0, 1) video = video.to(device=device, dtype=self.dtype) video = video * 2.0 - 1.0 # [0,1] -> [-1,1], in [t, c, h, w] video_embeddings = self.encode_video(video, chunk_size=decode_chunk_size).unsqueeze(0) prior_latents = self.encode_point_map( point_map_vae, pred_disparity, pred_valid_mask, pred_point_map, pred_intrinsic_map, chunk_size=decode_chunk_size ).unsqueeze(0).to(video_embeddings.dtype) # 1,T,C,H,W # 4. Encode input image using VAE # pdb.set_trace() needs_upcasting = ( self.vae.dtype == torch.float16 and self.vae.config.force_upcast ) if needs_upcasting: self.vae.to(dtype=torch.float32) video_latents = self.encode_vae_video( video.to(self.vae.dtype), chunk_size=decode_chunk_size, ).unsqueeze(0).to(video_embeddings.dtype) # [1, t, c, h, w] torch.cuda.empty_cache() if track_time: encode_event.record() torch.cuda.synchronize() elapsed_time_ms = prior_event.elapsed_time(encode_event) print(f"Elapsed time for encode prior and frames: {elapsed_time_ms} ms") else: gc.collect() torch.cuda.empty_cache() # cast back to fp16 if needed if needs_upcasting: self.vae.to(dtype=torch.float16) # 5. Get Added Time IDs added_time_ids = self._get_add_time_ids( 7, 127, noise_aug_strength, video_embeddings.dtype, batch_size, 1, False, ) # [1 or 2, 3] added_time_ids = added_time_ids.to(device) # 6. Prepare timesteps timesteps, num_inference_steps = retrieve_timesteps( self.scheduler, num_inference_steps, device, None, None ) num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order self._num_timesteps = len(timesteps) # 7. Prepare latent variables # num_channels_latents = self.unet.config.in_channels - prior_latents.shape[1] num_channels_latents = 8 latents_init = self.prepare_latents( batch_size, window_size, num_channels_latents, height, width, video_embeddings.dtype, device, generator, latents, ) # [1, t, c, h, w] latents_all = None idx_start = 0 if overlap > 0: weights = torch.linspace(0, 1, overlap, device=device) weights = weights.view(1, overlap, 1, 1, 1) else: weights = None while idx_start < num_frames - overlap: idx_end = min(idx_start + window_size, num_frames) self.scheduler.set_timesteps(num_inference_steps, device=device) # 9. Denoising loop # latents_init = latents_init.flip(1) latents = latents_init[:, : idx_end - idx_start].clone() latents_init = torch.cat( [latents_init[:, -overlap:], latents_init[:, :stride]], dim=1 ) video_latents_current = video_latents[:, idx_start:idx_end] prior_latents_current = prior_latents[:, idx_start:idx_end] video_embeddings_current = video_embeddings[:, idx_start:idx_end] with self.progress_bar(total=num_inference_steps) as progress_bar: for i, t in enumerate(timesteps): if latents_all is not None and i == 0: latents[:, :overlap] = ( latents_all[:, -overlap:] + latents[:, :overlap] / self.scheduler.init_noise_sigma * self.scheduler.sigmas[i] ) latent_model_input = latents latent_model_input = self.scheduler.scale_model_input( latent_model_input, t ) # [1 or 2, t, c, h, w] latent_model_input = torch.cat( [latent_model_input, video_latents_current, prior_latents_current], dim=2 ) noise_pred = self.unet( latent_model_input, t, encoder_hidden_states=video_embeddings_current, added_time_ids=added_time_ids, return_dict=False, )[0] # pdb.set_trace() # perform guidance if self.do_classifier_free_guidance: latent_model_input = latents latent_model_input = self.scheduler.scale_model_input( latent_model_input, t ) latent_model_input = torch.cat( [latent_model_input, torch.zeros_like(latent_model_input), torch.zeros_like(latent_model_input)], dim=2, ) noise_pred_uncond = self.unet( latent_model_input, t, encoder_hidden_states=torch.zeros_like( video_embeddings_current ), added_time_ids=added_time_ids, return_dict=False, )[0] noise_pred = noise_pred_uncond + self.guidance_scale * ( noise_pred - noise_pred_uncond ) latents = self.scheduler.step(noise_pred, t, latents).prev_sample if callback_on_step_end is not None: callback_kwargs = {} for k in callback_on_step_end_tensor_inputs: callback_kwargs[k] = locals()[k] callback_outputs = callback_on_step_end( self, i, t, callback_kwargs ) latents = callback_outputs.pop("latents", latents) if i == len(timesteps) - 1 or ( (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0 ): progress_bar.update() if latents_all is None: latents_all = latents.clone() else: if overlap > 0: latents_all[:, -overlap:] = latents[ :, :overlap ] * weights + latents_all[:, -overlap:] * (1 - weights) latents_all = torch.cat([latents_all, latents[:, overlap:]], dim=1) idx_start += stride latents_all = 1 / self.vae.config.scaling_factor * latents_all.squeeze(0).to(torch.float32) if track_time: denoise_event.record() torch.cuda.synchronize() elapsed_time_ms = encode_event.elapsed_time(denoise_event) print(f"Elapsed time for denoise latent: {elapsed_time_ms} ms") else: gc.collect() torch.cuda.empty_cache() point_map, valid_mask = self.decode_point_map( point_map_vae, latents_all, chunk_size=decode_chunk_size, force_projection=force_projection, force_fixed_focal=force_fixed_focal, use_extract_interp=use_extract_interp, need_resize=need_resize, height=original_height, width=original_width) if track_time: decode_event.record() torch.cuda.synchronize() elapsed_time_ms = denoise_event.elapsed_time(decode_event) print(f"Elapsed time for decode latent: {elapsed_time_ms} ms") else: gc.collect() torch.cuda.empty_cache() self.maybe_free_model_hooks() # t,h,w,3 t,h,w return point_map, valid_mask