geometrycrafter/diff_ppl.py

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