controller/controller.py

import gc

import torch
import torch.nn.functional as F

from einops import repeat, rearrange
from vidtome import merge
from utils.flow_utils import flow_warp, coords_grid

# AdaIn


def calc_mean_std(feat, eps=1e-5):
    # eps is a small value added to the variance to avoid divide-by-zero.
    size = feat.size()
    assert (len(size) == 4)
    N, C = size[:2]
    feat_var = feat.view(N, C, -1).var(dim=2) + eps
    feat_std = feat_var.sqrt().view(N, C, 1, 1)
    feat_mean = feat.view(N, C, -1).mean(dim=2).view(N, C, 1, 1)
    return feat_mean, feat_std


class AttentionControl():

    def __init__(self,
                 warp_period=(0.0, 0.0),
                 merge_period=(0.0, 0.0),
                 merge_ratio=(0.3, 0.3),
                 ToMe_period=(0.0, 1.0),
                 mask_period=(0.0, 0.0),
                 cross_period=(0.0, 0.0),
                 ada_period=(0.0, 0.0),
                 inner_strength=1.0,
                 loose_cfatnn=False,
                 flow_merge=True, 
                ):
        
        self.cur_frame_idx = 0

        self.step_store = self.get_empty_store()
        self.cur_step = 0
        self.total_step = 0
        self.cur_index = 0
        self.init_store = False
        self.restore = False
        self.update = False
        self.flow = None
        self.mask = None
        self.cldm = None
        self.decoded_imgs = []
        self.restorex0 = True
        self.updatex0 = False
        self.inner_strength = inner_strength
        self.cross_period = cross_period
        self.mask_period = mask_period
        self.ada_period = ada_period
        self.warp_period = warp_period
        self.ToMe_period = ToMe_period
        self.merge_period = merge_period
        self.merge_ratio = merge_ratio
        self.keyframe_idx = 0
        self.flow_merge = flow_merge
        self.distances = {}
        self.flow_correspondence = {}
        self.non_pad_ratio = (1.0, 1.0)
        self.up_resolution = 1280 if loose_cfatnn else 1281

    @staticmethod
    def get_empty_store():
        return {
            'first': [],
            'previous': [],
            'x0_previous': [],
            'first_ada': [],
            'pre_x0': [],
            "pre_keyframe_lq": None,
            "flows": None,
            "occ_masks": None,
            "flow_confids": None,
            "merge": None,
            "unmerge": None,
            "corres_scores": None,
            "flows2": None,
            "flow_confids2": None,
        }

    def forward(self, context, is_cross: bool, place_in_unet: str):
        cross_period = (self.total_step * self.cross_period[0],
                        self.total_step * self.cross_period[1])
        if not is_cross and place_in_unet == 'up' and context.shape[
                2] < self.up_resolution:
            if self.init_store:
                self.step_store['first'].append(context.detach())
                self.step_store['previous'].append(context.detach())
            if self.update:
                tmp = context.clone().detach()
            if self.restore and self.cur_step >= cross_period[0] and \
                    self.cur_step <= cross_period[1]:
                # context = torch.cat(
                #     (self.step_store['first'][self.cur_index],
                #      self.step_store['previous'][self.cur_index]),
                #     dim=1).clone()
                context = self.step_store['previous'][self.cur_index].clone()
            if self.update:
                self.step_store['previous'][self.cur_index] = tmp
            self.cur_index += 1
        # print(is_cross, place_in_unet, context.shape[2])
        # import ipdb; ipdb.set_trace()
        return context

    def update_x0(self, x0, cur_frame=0):
        # if self.init_store:
        #     self.step_store['x0_previous'].append(x0.detach())
        #     style_mean, style_std = calc_mean_std(x0.detach())
        #     self.step_store['first_ada'].append(style_mean.detach())
        #     self.step_store['first_ada'].append(style_std.detach())
        # if self.updatex0:
        #     tmp = x0.clone().detach()
        if self.restorex0:
            # if self.cur_step >= self.total_step * self.ada_period[
            #         0] and self.cur_step <= self.total_step * self.ada_period[
            #             1]:
            #     x0 = F.instance_norm(x0) * self.step_store['first_ada'][
            #         2 * self.cur_step +
            #         1] + self.step_store['first_ada'][2 * self.cur_step]
            if self.cur_step >= self.total_step * self.warp_period[
                    0] and self.cur_step < int(self.total_step * self.warp_period[1]):
                
                # mid_x = repeat(x[mid][None], 'b c h w -> (repeat b) c h w', repeat=x.shape[0])
                mid = x0.shape[0] // 2
                if len(self.step_store["pre_x0"]) == int(self.total_step * self.warp_period[1]):
                    print(f"[INFO] keyframe latent warping @ step {self.cur_step}...")
                    x0[mid] = (1 - self.step_store["occ_masks"][mid]) * x0[mid] + \
                        flow_warp(self.step_store["pre_x0"][self.cur_step][None], self.step_store["flows"][mid], mode='nearest')[0] * self.step_store["occ_masks"][mid] 
                    
                print(f"[INFO] local latent warping @ step {self.cur_step}...")
                for i in range(x0.shape[0]):
                    if i == mid:
                        continue
                    x0[i] = (1 - self.step_store["occ_masks"][i]) * x0[i] + \
                        flow_warp(x0[mid][None], self.step_store["flows"][i], mode='nearest')[0] * self.step_store["occ_masks"][i] 
                # x = rearrange(x, 'b c h w -> b (h w) c', h=64) 
                # self.step_store['x0_previous'][self.cur_step] = tmp
                # print(f"[INFO] storeing {self.cur_frame_idx} th frame x0 for step {self.cur_step}...")
                if len(self.step_store["pre_x0"]) < int(self.total_step * self.warp_period[1]):
                    self.step_store['pre_x0'].append(x0[mid])
                else:
                    self.step_store['pre_x0'][self.cur_step] = x0[mid]

        return x0

    def merge_x0(self, x0, merge_ratio):
        # print(f"[INFO] {self.total_step * self.merge_period[0]} {self.cur_step} {int(self.total_step * self.merge_period[1])} ...")
        if self.cur_step >= self.total_step * self.merge_period[0] and \
            self.cur_step < int(self.total_step * self.merge_period[1]):
            print(f"[INFO] latent merging @ step {self.cur_step}...")

            B, C, H, W = x0.shape
            non_pad_ratio_h, non_pad_ratio_w = self.non_pad_ratio
            padding_size_w = W - int(W * non_pad_ratio_w)
            padding_size_h = H - int(H * non_pad_ratio_h)
            non_pad_w = W - padding_size_w
            non_pad_h = H - padding_size_h
            padding_mask = torch.zeros((H, W), device=x0.device, dtype=torch.bool)
            if padding_size_w:
                padding_mask[:, -padding_size_w:] = 1
            if padding_size_h:
                padding_mask[-padding_size_h:, :] = 1
            padding_mask = rearrange(padding_mask, 'h w -> (h w)')
            
            idx_buffer = torch.arange(H*W, device=x0.device, dtype=torch.int64)
            non_pad_idx = idx_buffer[None, ~padding_mask, None]
            del idx_buffer, padding_mask
            x0 = rearrange(x0, 'b c h w -> b (h w) c', h=H)
            x_non_pad = torch.gather(x0, dim=1, index=non_pad_idx.expand(B, -1, C))
            # import ipdb; ipdb.set_trace()
            # merge.visualize_correspondence(x_non_pad[0][None], x_non_pad[B//2][None], ratio=0.3, H=H, out="latent_correspondence.png")

            # m, u, ret_dict = merge.bipartite_soft_matching_randframe(
            #                     x_non_pad, B, merge_ratio, 0, target_stride=B)
            import copy
            flows = copy.deepcopy(self.step_store["flows"])
            for i in range(B):
                if flows[i] is not None:
                    flows[i] = flows[i][:, :, :non_pad_h, :non_pad_w]
            # merge.visualize_flow_correspondence(x_non_pad[1][None], x_non_pad[B // 2][None], flow=flows[1], flow_confid=self.step_store["flow_confids"][1], \
            #                                  ratio=0.8, H=H, out=f"flow_correspondence_08.png")
            # import ipdb; ipdb.set_trace()
            x_non_pad = rearrange(x_non_pad, 'b a c -> 1 (b a) c')
            m, u, ret_dict = merge.bipartite_soft_matching_randframe(
                    x_non_pad, B, merge_ratio, 0, target_stride=B, 
                    H=H,
                    flow=flows, 
                    flow_confid=self.step_store["flow_confids"],
                    )
            x_non_pad = u(m(x_non_pad))
            # x_non_pad = self.step_store["unmerge"](self.step_store["merge"](x_non_pad))
            x_non_pad = rearrange(x_non_pad, '1 (b a) c -> b a c', b=B)
            # print(torch.mean(x0[0]).item(), torch.mean(x0[1]).item(), torch.mean(x0[2]).item(), torch.mean(x0[3]).item(), torch.mean(x0[4]).item())
            # print(torch.std(x0[0]).item(), torch.std(x0[1]).item(), torch.std(x0[2]).item(), torch.std(x0[3]).item(), torch.std(x0[4]).item())
            # import ipdb; ipdb.set_trace()
            x0.scatter_(dim=1, index=non_pad_idx.expand(B, -1, C), src=x_non_pad)
            x0 = rearrange(x0, 'b (h w) c -> b c h w ', h=H)
            # import ipdb; ipdb.set_trace()
        
        return x0
    
    def merge_x0_scores(self, x0, merge_ratio, merge_mode="replace"):
        # print(f"[INFO] {self.total_step * self.merge_period[0]} {self.cur_step} {int(self.total_step * self.merge_period[1])} ...")
        # import ipdb; ipdb.set_trace()
        if self.cur_step >= self.total_step * self.merge_period[0] and \
            self.cur_step < int(self.total_step * self.merge_period[1]):
            print(f"[INFO] latent merging @ step {self.cur_step}...")

            B, C, H, W = x0.shape
            non_pad_ratio_h, non_pad_ratio_w = self.non_pad_ratio
            padding_size_w = W - int(W * non_pad_ratio_w)
            padding_size_h = H - int(H * non_pad_ratio_h)
            padding_mask = torch.zeros((H, W), device=x0.device, dtype=torch.bool)
            if padding_size_w:
                padding_mask[:, -padding_size_w:] = 1
            if padding_size_h:
                padding_mask[-padding_size_h:, :] = 1
            padding_mask = rearrange(padding_mask, 'h w -> (h w)')
            
            idx_buffer = torch.arange(H*W, device=x0.device, dtype=torch.int64)
            non_pad_idx = idx_buffer[None, ~padding_mask, None]
            x0 = rearrange(x0, 'b c h w -> b (h w) c', h=H)
            x_non_pad = torch.gather(x0, dim=1, index=non_pad_idx.expand(B, -1, C))
            x_non_pad_A, x_non_pad_N = x_non_pad.shape[1], x_non_pad.shape[1] * B
            mid = B // 2
            
            x_non_pad_ = x_non_pad.clone()
            x_non_pad = rearrange(x_non_pad, 'b a c -> 1 (b a) c')
            # import ipdb; ipdb.set_trace()

            idx_buffer = torch.arange(x_non_pad_N, device=x0.device, dtype=torch.int64)
            randf = torch.tensor(B // 2, dtype=torch.int).to(x0.device)
            # print(f"[INFO] {randf.item()} th frame as target")
            dst_select = ((torch.div(idx_buffer, x_non_pad_A, rounding_mode='floor')) % B == randf).to(torch.bool)
            # a_idx: src index. b_idx: dst index
            a_idx = idx_buffer[None, ~dst_select, None] 
            b_idx = idx_buffer[None, dst_select, None]
            del idx_buffer, padding_mask
            num_dst = b_idx.shape[1]
            # b, _, _ = x_non_pad.shape
            b = 1
            src = torch.gather(x_non_pad, dim=1, index=a_idx.expand(b, x_non_pad_N - num_dst, C))
            tar = torch.gather(x_non_pad, dim=1, index=b_idx.expand(b, num_dst, C))
            # tar = x_non_pad[mid][None]
            # src = torch.cat((x_non_pad[:mid], x_non_pad[mid+1:]), dim=0)
            # src = rearrange(src, 'b n c -> 1 (b n) c')
            # print(f"[INFO] {x_non_pad.shape} {src.shape} {tar.shape} ...")
            # print(f"[INFO] maximum score {torch.max(self.step_store['corres_scores'])} ...")
            flow_src_idx = self.flow_correspondence[H][0]
            flow_tar_idx = self.flow_correspondence[H][1]
            flow_confid = self.step_store["flow_confids"][:mid] + self.step_store["flow_confids"][mid+1:]
            flow_confid = torch.cat(flow_confid, dim=0) 
            flow_confid = rearrange(flow_confid, 'b h w -> 1 (b h w)')
            scores = F.normalize(self.step_store["corres_scores"], p=2, dim=-1)

            flow_confid -= (torch.max(flow_confid) - torch.max(scores)) 

            # merge.visualize_correspondence_score(x_non_pad_[0][None], x_non_pad_[mid][None], 
            #                                    score=scores[:,:x_non_pad_A],
            #                                    ratio=0.2, H=H-padding_size_h, out="latent_correspondence.png")
            # import ipdb; ipdb.set_trace()
            scores[:, flow_src_idx[0, :, 0], flow_tar_idx[0, :, 0]] += (flow_confid[:, flow_src_idx[0, :, 0]] * 0.3)
            # merge.visualize_correspondence_score(x_non_pad_[0][None], x_non_pad_[mid][None], 
            #                                    score=scores[:,:x_non_pad_A],
            #                                    ratio=0.2, H=H-padding_size_h, out="latent_correspondence_flow.png")

            # import ipdb; ipdb.set_trace()
            r = min(src.shape[1], int(src.shape[1] * merge_ratio))
            node_max, node_idx = scores.max(dim=-1)
            edge_idx = node_max.argsort(dim=-1, descending=True)[..., None]
            unm_idx = edge_idx[..., r:, :]  # Unmerged Tokens
            src_idx = edge_idx[..., :r, :]  # Merged Tokens 
            tar_idx = torch.gather(node_idx[..., None], dim=-2, index=src_idx)
            unm = torch.gather(src, dim=-2, index=unm_idx.expand(-1, -1, C))
            if merge_mode != "replace":
                src = torch.gather(src, dim=-2, index=src_idx.expand(-1, -1, C))
                # In other mode such as mean, combine matched src and dst tokens.
                tar = tar.scatter_reduce(-2, tar_idx.expand(-1, -1, C),
                                        src, reduce=merge_mode, include_self=True)
            # In replace mode, just cat unmerged tokens and tar tokens. Ignore src tokens.
            # token = torch.cat([unm, tar], dim=1)

            # unm_len = unm_idx.shape[1]
            # unm, tar = token[..., :unm_len, :], token[..., unm_len:, :]
            src = torch.gather(tar, dim=-2, index=tar_idx.expand(-1, -1, C))
            # Combine back to the original shape
            # x_non_pad = torch.zeros(b, x_non_pad_N, C, device=x0.device, dtype=x0.dtype)
            # Scatter dst tokens
            x_non_pad.scatter_(dim=-2, index=b_idx.expand(b, -1, C), src=tar)
            # Scatter unmerged tokens
            x_non_pad.scatter_(dim=-2, index=torch.gather(a_idx.expand(b, -1, 1),
                        dim=1, index=unm_idx).expand(-1, -1, C), src=unm)
            # Scatter src tokens
            x_non_pad.scatter_(dim=-2, index=torch.gather(a_idx.expand(b, -1, 1),
                        dim=1, index=src_idx).expand(-1, -1, C), src=src)

            x_non_pad = rearrange(x_non_pad, '1 (b a) c -> b a c', a=x_non_pad_A)
            x0.scatter_(dim=1, index=non_pad_idx.expand(B, -1, C), src=x_non_pad)
            x0 = rearrange(x0, 'b (h w) c -> b c h w ', h=H)
        
        return x0

    def set_distance(self, B, H, W, radius, device):
        y, x = torch.meshgrid(torch.arange(H), torch.arange(W))
        coords = torch.stack((y, x), dim=-1).float().to(device)
        coords = rearrange(coords, 'h w c -> (h w) c')

        # Calculate the Euclidean distance between all pixels
        distances = torch.cdist(coords, coords)
        # radius = W // 30
        radius = 1 if radius == 0 else radius
        # print(f"[INFO]  W: {W} Radius: {radius} ")
        distances //= radius
        distances = torch.exp(-distances)
        # distances += torch.diag_embed(torch.ones(A)).to(metric.device)
        distances = repeat(distances, 'h a -> 1 (b h) a', b=B)
        self.distances[H] = distances
    
    def set_flow_correspondence(self, B, H, W, key_idx, flow_confid, flow):

        if len(flow) != B - 1:
                flow_confid = flow_confid[:key_idx] + flow_confid[key_idx+1:]
                flow = flow[:key_idx] + flow[key_idx+1:]

        flow_confid = torch.cat(flow_confid, dim=0) 
        flow = torch.cat(flow, dim=0) 
        flow_confid = rearrange(flow_confid, 'b h w -> 1 (b h w)')
        
        edge_idx = flow_confid.argsort(dim=-1, descending=True)[..., None]

        src_idx = edge_idx[..., :, :]  # Merged Tokens 

        A = H * W
        src_idx_tensor = src_idx[0, : ,0]
        f = src_idx_tensor // A
        id = src_idx_tensor % A
        x = id % W
        y = id // W

        # Stack the results into a 2D tensor
        src_fxy = torch.stack((f, x, y), dim=1)
        # import ipdb; ipdb.set_trace()
        grid = coords_grid(B-1, H, W).to(flow.device) + flow  # [F-1, 2, H, W]

        x = grid[src_fxy[:, 0], 0, src_fxy[:, 2], src_fxy[:, 1]].clamp(0, W-1).long()
        y = grid[src_fxy[:, 0], 1, src_fxy[:, 2], src_fxy[:, 1]].clamp(0, H-1).long()
        tar_xy = torch.stack((x, y), dim=1)
        tar_idx = y * W + x
        tar_idx = rearrange(tar_idx, ' d -> 1 d 1')

        self.flow_correspondence[H] = (src_idx, tar_idx)

    def set_merge(self, merge, unmerge):
        self.step_store["merge"] = merge
        self.step_store["unmerge"] = unmerge

    def set_warp(self, flows, masks, flow_confids=None):
        self.step_store["flows"] = flows
        self.step_store["occ_masks"] = masks
        if flow_confids is not None:
            self.step_store["flow_confids"] = flow_confids

    def set_warp2(self, flows, flow_confids):
        self.step_store["flows2"] = flows
        self.step_store["flow_confids2"] = flow_confids

    def set_pre_keyframe_lq(self, pre_keyframe_lq):
        self.step_store["pre_keyframe_lq"] = pre_keyframe_lq

    def __call__(self, context, is_cross: bool, place_in_unet: str):
        context = self.forward(context, is_cross, place_in_unet)
        return context

    def set_cur_frame_idx(self, frame_idx):
        self.cur_frame_idx = frame_idx

    def set_step(self, step):
        self.cur_step = step

    def set_total_step(self, total_step):
        self.total_step = total_step
        self.cur_index = 0

    def clear_store(self):
        del self.step_store
        torch.cuda.empty_cache()
        gc.collect()
        self.step_store = self.get_empty_store()

    def set_task(self, task, restore_step=1.0):
        self.init_store = False
        self.restore = False
        self.update = False
        self.cur_index = 0
        self.restore_step = restore_step
        self.updatex0 = False
        self.restorex0 = False
        if 'initfirst' in task:
            self.init_store = True
            self.clear_store()
        if 'updatestyle' in task:
            self.update = True
        if 'keepstyle' in task:
            self.restore = True
        if 'updatex0' in task: 
            self.updatex0 = True
        if 'keepx0' in task:
            self.restorex0 = True