algorithms/worldmem/pose_prediction.py

from omegaconf import DictConfig
import torch
from lightning.pytorch.utilities.types import STEP_OUTPUT
from algorithms.common.metrics import (
    FrechetInceptionDistance,
    LearnedPerceptualImagePatchSimilarity,
    FrechetVideoDistance,
)
from .df_base import DiffusionForcingBase
from utils.logging_utils import log_video, get_validation_metrics_for_videos
from .models.vae import VAE_models
from .models.dit import DiT_models
from einops import rearrange
from torch import autocast
import numpy as np
from tqdm import tqdm
import torch.nn.functional as F
from .models.pose_prediction import PosePredictionNet
import torchvision.transforms.functional as TF
import random
from torchvision.transforms import InterpolationMode
from PIL import Image
import math
from packaging import version as pver
import torch.distributed as dist
import matplotlib.pyplot as plt

import torch
import math
import wandb

import torch.nn as nn
from algorithms.common.base_pytorch_algo import BasePytorchAlgo

class PosePrediction(BasePytorchAlgo):

    def __init__(self, cfg: DictConfig):

        super().__init__(cfg)

    def _build_model(self):
        self.pose_prediction_model = PosePredictionNet()
        vae = VAE_models["vit-l-20-shallow-encoder"]()
        self.vae = vae.eval()

    def training_step(self, batch, batch_idx) -> STEP_OUTPUT:
        xs, conditions, pose_conditions= batch
        pose_conditions[:,:,3:] = pose_conditions[:,:,3:] // 15
        xs = self.encode(xs)

        b,f,c,h,w = xs.shape
        xs = xs[:,:-1].reshape(-1, c, h, w)
        conditions = conditions[:,1:].reshape(-1, 25)
        offset_gt = pose_conditions[:,1:] - pose_conditions[:,:-1]
        pose_conditions = pose_conditions[:,:-1].reshape(-1, 5)
        offset_gt = offset_gt.reshape(-1, 5)
        offset_gt[:, 3][offset_gt[:, 3]==23] = -1
        offset_gt[:, 3][offset_gt[:, 3]==-23] = 1
        offset_gt[:, 4][offset_gt[:, 4]==23] = -1
        offset_gt[:, 4][offset_gt[:, 4]==-23] = 1

        offset_pred = self.pose_prediction_model(xs, conditions, pose_conditions)
        criterion = nn.MSELoss()
        loss = criterion(offset_pred, offset_gt)
        if batch_idx % 200 == 0:
            self.log("training/loss", loss.cpu())
        output_dict = {
            "loss": loss}
        return output_dict

    def encode(self, x):
        # vae encoding
        B = x.shape[1]
        T = x.shape[0]
        H, W = x.shape[-2:]
        scaling_factor = 0.07843137255

        x = rearrange(x, "t b c h w -> (t b) c h w")
        with torch.no_grad():
            with autocast("cuda", dtype=torch.half):
                x = self.vae.encode(x * 2 - 1).mean * scaling_factor
        x = rearrange(x, "(t b) (h w) c -> t b c h w", t=T, h=H // self.vae.patch_size, w=W // self.vae.patch_size)
        # x = x[:, :n_prompt_frames]
        return x

    def decode(self, x):
        total_frames = x.shape[0]
        scaling_factor = 0.07843137255
        x = rearrange(x, "t b c h w -> (t b) (h w) c")
        with torch.no_grad():
            with autocast("cuda", dtype=torch.half):
                x = (self.vae.decode(x / scaling_factor) + 1) / 2

        x = rearrange(x, "(t b) c h w-> t b c h w", t=total_frames)
        return x

    def validation_step(self, batch, batch_idx, namespace="validation") -> STEP_OUTPUT:
        xs, conditions, pose_conditions= batch
        pose_conditions[:,:,3:] = pose_conditions[:,:,3:] // 15
        xs = self.encode(xs)

        b,f,c,h,w = xs.shape
        xs = xs[:,:-1].reshape(-1, c, h, w)
        conditions = conditions[:,1:].reshape(-1, 25)
        offset_gt = pose_conditions[:,1:] - pose_conditions[:,:-1]
        pose_conditions = pose_conditions[:,:-1].reshape(-1, 5)
        offset_gt = offset_gt.reshape(-1, 5)
        offset_gt[:, 3][offset_gt[:, 3]==23] = -1
        offset_gt[:, 3][offset_gt[:, 3]==-23] = 1
        offset_gt[:, 4][offset_gt[:, 4]==23] = -1
        offset_gt[:, 4][offset_gt[:, 4]==-23] = 1

        offset_pred = self.pose_prediction_model(xs, conditions, pose_conditions)

        criterion = nn.MSELoss()
        loss = criterion(offset_pred, offset_gt)

        if batch_idx % 200 == 0:
            self.log("validation/loss", loss.cpu())
        output_dict = {
            "loss": loss}
        return

    @torch.no_grad()
    def interactive(self, batch, context_frames, device):
        with torch.cuda.amp.autocast():
            condition_similar_length = self.condition_similar_length
            # xs_raw, conditions, pose_conditions, c2w_mat, masks, frame_idx = self._preprocess_batch(batch)

            first_frame, new_conditions, new_pose_conditions, new_c2w_mat, new_frame_idx = batch

            if self.frames is None:
                first_frame_encode = self.encode(first_frame[None, None].to(device))
                self.frames = first_frame_encode.to(device)
                self.actions = new_conditions[None, None].to(device)
                self.poses = new_pose_conditions[None, None].to(device)
                self.memory_c2w = new_c2w_mat[None, None].to(device)
                self.frame_idx = torch.tensor([[new_frame_idx]]).to(device)
                return first_frame
            else:
                self.actions = torch.cat([self.actions, new_conditions[None, None].to(device)])
                self.poses = torch.cat([self.poses, new_pose_conditions[None, None].to(device)])
                self.memory_c2w = torch.cat([self.memory_c2w, new_c2w_mat[None, None].to(device)])
                self.frame_idx = torch.cat([self.frame_idx, torch.tensor([[new_frame_idx]]).to(device)])

            conditions = self.actions.clone()
            pose_conditions = self.poses.clone()
            c2w_mat = self.memory_c2w .clone()
            frame_idx = self.frame_idx.clone()


            curr_frame = 0
            horizon = 1
            batch_size = 1
            n_frames = curr_frame + horizon
            # context
            n_context_frames = context_frames // self.frame_stack
            xs_pred = self.frames[:n_context_frames].clone()
            curr_frame += n_context_frames

            pbar = tqdm(total=n_frames, initial=curr_frame, desc="Sampling")

            # generation on frame
            scheduling_matrix = self._generate_scheduling_matrix(horizon)
            chunk = torch.randn((horizon, batch_size, *xs_pred.shape[2:])).to(xs_pred.device)
            chunk = torch.clamp(chunk, -self.clip_noise, self.clip_noise)

            xs_pred = torch.cat([xs_pred, chunk], 0)

            # sliding window: only input the last n_tokens frames
            start_frame = max(0, curr_frame + horizon - self.n_tokens)

            pbar.set_postfix(
                {
                    "start": start_frame,
                    "end": curr_frame + horizon,
                }
            )

            if condition_similar_length:

                if curr_frame < condition_similar_length:
                    random_idx = [i for i in range(curr_frame)] + [0] * (condition_similar_length-curr_frame)
                    random_idx = np.repeat(np.array(random_idx)[:,None], xs_pred.shape[1], -1)
                else:
                    num_samples = 10000
                    radius = 30
                    samples = torch.rand((num_samples, 1), device=pose_conditions.device)
                    angles = 2 * np.pi * torch.rand((num_samples,), device=pose_conditions.device)
                    # points = radius * torch.sqrt(samples) * torch.stack((torch.cos(angles), torch.sin(angles)), dim=1)
                    
                    points = generate_points_in_sphere(num_samples, radius).to(pose_conditions.device)
                    points = points[:, None].repeat(1, pose_conditions.shape[1], 1)
                    points += pose_conditions[curr_frame, :, :3][None]
                    fov_half_h = torch.tensor(105/2, device=pose_conditions.device)
                    fov_half_v = torch.tensor(75/2, device=pose_conditions.device)
                    # in_fov1 = is_inside_fov(points, pose_conditions[curr_frame, :, [0, 2]], pose_conditions[curr_frame, :, -1], fov_half)

                    in_fov1 = is_inside_fov_3d_hv(points, pose_conditions[curr_frame, :, :3], 
                        pose_conditions[curr_frame, :, -2], pose_conditions[curr_frame, :, -1],
                        fov_half_h, fov_half_v)

                    in_fov_list = []
                    for pc in pose_conditions[:curr_frame]:
                        in_fov_list.append(is_inside_fov_3d_hv(points, pc[:, :3], pc[:, -2], pc[:, -1],
                                                        fov_half_h, fov_half_v))
                    
                    in_fov_list = torch.stack(in_fov_list)
                    # v3
                    random_idx = []

                    for csl in range(self.condition_similar_length // 2):
                        overlap_ratio = ((in_fov1[None].bool() & in_fov_list).sum(1))/in_fov1.sum()
                        # mask = distance > (in_fov1.bool().sum(0) / 4)
                        #_, r_idx = torch.topk(overlap_ratio / tensor_max_with_number((frame_idx[curr_frame] - frame_idx[:curr_frame]), 10), k=1, dim=0)
                        
                        # if csl > self.condition_similar_length:
                        # 	_, r_idx = torch.topk(overlap_ratio, k=1, dim=0)
                        # else:
                        # 	_, r_idx = torch.topk(overlap_ratio / tensor_max_with_number((frame_idx[curr_frame] - frame_idx[:curr_frame]), 10), k=1, dim=0)

                        _, r_idx = torch.topk(overlap_ratio, k=1, dim=0)
                        # _, r_idx = torch.topk(overlap_ratio / tensor_max_with_number((frame_idx[curr_frame] - frame_idx[:curr_frame]), 10), k=1, dim=0)

                        # if curr_frame >=93:
                        #     import pdb;pdb.set_trace()

                        # start_time = time.time()
                        cos_sim = F.cosine_similarity(xs_pred.to(r_idx.device)[r_idx[:, range(in_fov1.shape[1])], 
                            range(in_fov1.shape[1])], xs_pred.to(r_idx.device)[:curr_frame], dim=2)
                        cos_sim = cos_sim.mean((-2,-1))

                        mask_sim = cos_sim>0.9
                        in_fov_list = in_fov_list & ~mask_sim[:,None].to(in_fov_list.device)

                        random_idx.append(r_idx)

                    for bi in range(conditions.shape[1]):
                        if len(torch.nonzero(conditions[:,bi,24] == 1))==0:
                            pass
                        else:
                            last_idx = torch.nonzero(conditions[:,bi,24] == 1)[-1]
                            in_fov_list[:last_idx,:,bi] = False

                    for csl in range(self.condition_similar_length // 2):
                        overlap_ratio = ((in_fov1[None].bool() & in_fov_list).sum(1))/in_fov1.sum()
                        # mask = distance > (in_fov1.bool().sum(0) / 4)
                        #_, r_idx = torch.topk(overlap_ratio / tensor_max_with_number((frame_idx[curr_frame] - frame_idx[:curr_frame]), 10), k=1, dim=0)
                        
                        # if csl > self.condition_similar_length:
                        # 	_, r_idx = torch.topk(overlap_ratio, k=1, dim=0)
                        # else:
                        # 	_, r_idx = torch.topk(overlap_ratio / tensor_max_with_number((frame_idx[curr_frame] - frame_idx[:curr_frame]), 10), k=1, dim=0)

                        _, r_idx = torch.topk(overlap_ratio, k=1, dim=0)
                        # _, r_idx = torch.topk(overlap_ratio / tensor_max_with_number((frame_idx[curr_frame] - frame_idx[:curr_frame]), 10), k=1, dim=0)

                        # if curr_frame >=93:
                        #     import pdb;pdb.set_trace()

                        # start_time = time.time()
                        cos_sim = F.cosine_similarity(xs_pred.to(r_idx.device)[r_idx[:, range(in_fov1.shape[1])], 
                            range(in_fov1.shape[1])], xs_pred.to(r_idx.device)[:curr_frame], dim=2)
                        cos_sim = cos_sim.mean((-2,-1))

                        mask_sim = cos_sim>0.9
                        in_fov_list = in_fov_list & ~mask_sim[:,None].to(in_fov_list.device)

                        random_idx.append(r_idx)
                    
                    random_idx = torch.cat(random_idx).cpu()
                    condition_similar_length = len(random_idx)
                
                xs_pred = torch.cat([xs_pred, xs_pred[random_idx[:,range(xs_pred.shape[1])], range(xs_pred.shape[1])].clone()], 0)

            if condition_similar_length:
                # import pdb;pdb.set_trace()
                padding = torch.zeros((condition_similar_length,) + conditions.shape[1:], device=conditions.device, dtype=conditions.dtype)
                input_condition = torch.cat([conditions[start_frame : curr_frame + horizon], padding], dim=0)
                if self.pose_cond_dim:
                    # if not self.use_plucker:
                    input_pose_condition = torch.cat([pose_conditions[start_frame : curr_frame + horizon], pose_conditions[random_idx[:,range(xs_pred.shape[1])], range(xs_pred.shape[1])]], dim=0).clone()

                if self.use_plucker:
                    if self.all_zero_frame:
                        frame_idx_list = []
                        input_pose_condition = []
                        for i in range(start_frame, curr_frame + horizon):
                            input_pose_condition.append(convert_to_plucker(torch.cat([c2w_mat[i:i+1],c2w_mat[random_idx[:,range(xs_pred.shape[1])], range(xs_pred.shape[1])]]).clone(), 0, focal_length=self.focal_length, is_old_setting=self.old_setting).to(xs_pred.dtype))
                            frame_idx_list.append(torch.cat([frame_idx[i:i+1]-frame_idx[i:i+1], frame_idx[random_idx[:,range(xs_pred.shape[1])], range(xs_pred.shape[1])]-frame_idx[i:i+1]]))
                        input_pose_condition = torch.cat(input_pose_condition)
                        frame_idx_list = torch.cat(frame_idx_list)

                        # print(frame_idx_list[:,0])
                    else:
                        # print(curr_frame-start_frame)
                        # input_pose_condition = torch.cat([c2w_mat[start_frame : curr_frame + horizon], c2w_mat[random_idx[:,range(xs_pred.shape[1])], range(xs_pred.shape[1])]], dim=0).clone()
                        # import pdb;pdb.set_trace()
                        if self.last_frame_refer:
                            input_pose_condition = torch.cat([c2w_mat[start_frame : curr_frame + horizon], c2w_mat[-1:]], dim=0).clone()
                        else:
                            input_pose_condition = torch.cat([c2w_mat[start_frame : curr_frame + horizon], c2w_mat[random_idx[:,range(xs_pred.shape[1])], range(xs_pred.shape[1])]], dim=0).clone()
                        
                        if self.zero_curr:
                            # print("="*50)
                            input_pose_condition = convert_to_plucker(input_pose_condition, curr_frame-start_frame, focal_length=self.focal_length, is_old_setting=self.old_setting)
                            # input_pose_condition[:curr_frame-start_frame] = input_pose_condition[curr_frame-start_frame:curr_frame-start_frame+1]
                        # input_pose_condition = convert_to_plucker(input_pose_condition, -self.condition_similar_length-1, focal_length=self.focal_length)
                        else:
                            input_pose_condition = convert_to_plucker(input_pose_condition, -condition_similar_length, focal_length=self.focal_length, is_old_setting=self.old_setting)
                        frame_idx_list = None
                else:
                    input_pose_condition = torch.cat([pose_conditions[start_frame : curr_frame + horizon], pose_conditions[random_idx[:,range(xs_pred.shape[1])], range(xs_pred.shape[1])]], dim=0).clone()
                    frame_idx_list = None
            else:
                input_condition = conditions[start_frame : curr_frame + horizon]
                input_pose_condition = None
                frame_idx_list = None
                
            for m in range(scheduling_matrix.shape[0] - 1):
                from_noise_levels = np.concatenate((np.zeros((curr_frame,), dtype=np.int64), scheduling_matrix[m]))[
                    :, None
                ].repeat(batch_size, axis=1)
                to_noise_levels = np.concatenate(
                    (
                        np.zeros((curr_frame,), dtype=np.int64),
                        scheduling_matrix[m + 1],
                    )
                )[
                    :, None
                ].repeat(batch_size, axis=1)

                if condition_similar_length:
                    from_noise_levels = np.concatenate([from_noise_levels, np.zeros((condition_similar_length,from_noise_levels.shape[-1]), dtype=np.int32)], axis=0)
                    to_noise_levels = np.concatenate([to_noise_levels, np.zeros((condition_similar_length,from_noise_levels.shape[-1]), dtype=np.int32)], axis=0)

                from_noise_levels = torch.from_numpy(from_noise_levels).to(self.device)
                to_noise_levels = torch.from_numpy(to_noise_levels).to(self.device)


                if input_pose_condition is not None:
                    input_pose_condition = input_pose_condition.to(xs_pred.dtype)
                
                xs_pred[start_frame:] = self.diffusion_model.sample_step(
                    xs_pred[start_frame:],
                    input_condition,
                    input_pose_condition,
                    from_noise_levels[start_frame:],
                    to_noise_levels[start_frame:],
                    current_frame=curr_frame,
                    mode="validation",
                    reference_length=condition_similar_length,
                    frame_idx=frame_idx_list
                )

                # if curr_frame > 14:
                #     import pdb;pdb.set_trace()

                # if xs_pred_back is not None:
                #     xs_pred = torch.cat([xs_pred[:6], xs_pred_back[6:12], xs_pred[6:]], dim=0)
        
            # import pdb;pdb.set_trace()
            if condition_similar_length: # and curr_frame+1!=n_frames:
                xs_pred = xs_pred[:-condition_similar_length]

            curr_frame += horizon
            pbar.update(horizon)

            self.frames = torch.cat([self.frames, xs_pred[n_context_frames:]])

            xs_pred = self.decode(xs_pred[n_context_frames:])

            return xs_pred[-1,0].cpu()