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