Spaces:
Running
Running
| import numpy as np | |
| import torch | |
| import cv2 | |
| from torchvision.transforms import Resize, InterpolationMode, ToTensor, Compose, CenterCrop | |
| from einops import rearrange | |
| import glob | |
| from diffusers.utils import USE_PEFT_BACKEND | |
| from diffusers.utils import load_image | |
| from natsort import natsorted | |
| def read_mask(mask_dir): | |
| transform = Compose([ | |
| Resize((512, 512), interpolation=InterpolationMode.BILINEAR, antialias=True), | |
| # CenterCrop((512, 512)), | |
| ToTensor()]) | |
| mask_paths = glob.glob(mask_dir + '/*.png') | |
| mask_paths = natsorted(mask_paths) | |
| mask_list = [] | |
| for mask_path in mask_paths: | |
| mask = load_image(mask_path) | |
| mask_torch = transform(mask).bool().unsqueeze(0) # torch.Size([1, 3, 512, 512]) -1~1 | |
| mask_list.append(mask_torch) | |
| return mask_list | |
| def read_rgb(rgb_dir): | |
| transform = Compose([ | |
| Resize((512, 512), interpolation=InterpolationMode.BILINEAR, antialias=True), | |
| # CenterCrop((512, 512)), | |
| ToTensor()]) | |
| rgb_paths = sorted(glob.glob(rgb_dir + '/*.jpg')) | |
| rgb_list = [] | |
| rgb_frame = [] | |
| for rgb_path in rgb_paths: | |
| rgb = load_image(rgb_path); | |
| width, height = rgb.size | |
| file_name = rgb_path.split('/')[-1] | |
| frame_number = int(file_name.split('_')[1].split('.')[0].lstrip('0') or '0') | |
| rgb_frame.append(frame_number) | |
| rgb_torch = transform(rgb).unsqueeze(0) # torch.Size([1, 3, 512, 512]) | |
| rgb_list.append(rgb_torch) | |
| return rgb_list, (width, height), rgb_frame | |
| def read_depth2disparity(depth_dir): | |
| depth_paths = sorted(glob.glob(depth_dir + '/*.npy')) | |
| disparity_list = [] | |
| for depth_path in depth_paths: | |
| depth = np.load(depth_path) | |
| depth = cv2.resize(depth, (512, 512)).reshape((512, 512, 1)) # [512,512,1] | |
| # depth = CenterCrop((512, 512))(torch.from_numpy(depth))[..., None].numpy() # [512,512,1] | |
| disparity = 1 / (depth + 1e-5) | |
| disparity_map = disparity / np.max(disparity) # 0.00233~1 | |
| # disparity_map = disparity_map.astype(np.uint8)[:,:,0] | |
| disparity_map = np.concatenate([disparity_map, disparity_map, disparity_map], axis=2) | |
| disparity_list.append(torch.from_numpy(disparity_map[None]).permute(0, 3, 1, 2).float()) # [1,512,512,3] | |
| return disparity_list | |
| def compute_attn(attn, query, key, value, video_length, ref_frame_index, attention_mask): | |
| key_ref_cross = rearrange(key, "(b f) d c -> b f d c", f=video_length) | |
| key_ref_cross = key_ref_cross[:, ref_frame_index] | |
| key_ref_cross = rearrange(key_ref_cross, "b f d c -> (b f) d c") | |
| value_ref_cross = rearrange(value, "(b f) d c -> b f d c", f=video_length) | |
| value_ref_cross = value_ref_cross[:, ref_frame_index] | |
| value_ref_cross = rearrange(value_ref_cross, "b f d c -> (b f) d c") | |
| key_ref_cross = attn.head_to_batch_dim(key_ref_cross) | |
| value_ref_cross = attn.head_to_batch_dim(value_ref_cross) | |
| attention_probs = attn.get_attention_scores(query, key_ref_cross, attention_mask) | |
| hidden_states_ref_cross = torch.bmm(attention_probs, value_ref_cross) | |
| return hidden_states_ref_cross | |
| class CrossViewAttnProcessor: | |
| def __init__(self, self_attn_coeff, unet_chunk_size=2): | |
| self.unet_chunk_size = unet_chunk_size | |
| self.self_attn_coeff = self_attn_coeff | |
| def __call__( | |
| self, | |
| attn, | |
| hidden_states, | |
| encoder_hidden_states=None, | |
| attention_mask=None, | |
| temb=None, | |
| scale=1.0, ): | |
| residual = hidden_states | |
| args = () if USE_PEFT_BACKEND else (scale,) | |
| if attn.spatial_norm is not None: | |
| hidden_states = attn.spatial_norm(hidden_states, temb) | |
| input_ndim = hidden_states.ndim | |
| if input_ndim == 4: | |
| batch_size, channel, height, width = hidden_states.shape | |
| hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) | |
| batch_size, sequence_length, _ = ( | |
| hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape | |
| ) | |
| attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) | |
| if attn.group_norm is not None: | |
| hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) | |
| query = attn.to_q(hidden_states, *args) | |
| is_cross_attention = encoder_hidden_states is not None | |
| if encoder_hidden_states is None: | |
| encoder_hidden_states = hidden_states | |
| elif attn.norm_cross: | |
| encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) | |
| key = attn.to_k(encoder_hidden_states, *args) | |
| value = attn.to_v(encoder_hidden_states, *args) | |
| query = attn.head_to_batch_dim(query) | |
| # Sparse Attention | |
| if not is_cross_attention: | |
| ################## Perform self attention | |
| key_self = attn.head_to_batch_dim(key) | |
| value_self = attn.head_to_batch_dim(value) | |
| attention_probs = attn.get_attention_scores(query, key_self, attention_mask) | |
| hidden_states_self = torch.bmm(attention_probs, value_self) | |
| ####################################### | |
| video_length = key.size()[0] // self.unet_chunk_size | |
| ref0_frame_index = [0] * video_length | |
| ref1_frame_index = [1] * video_length | |
| ref2_frame_index = [2] * video_length | |
| ref3_frame_index = [3] * video_length | |
| hidden_states_ref0 = compute_attn(attn, query, key, value, video_length, ref0_frame_index, attention_mask) | |
| hidden_states_ref1 = compute_attn(attn, query, key, value, video_length, ref1_frame_index, attention_mask) | |
| hidden_states_ref2 = compute_attn(attn, query, key, value, video_length, ref2_frame_index, attention_mask) | |
| key = rearrange(key, "(b f) d c -> b f d c", f=video_length) | |
| key = key[:, ref3_frame_index] | |
| key = rearrange(key, "b f d c -> (b f) d c") | |
| value = rearrange(value, "(b f) d c -> b f d c", f=video_length) | |
| value = value[:, ref3_frame_index] | |
| value = rearrange(value, "b f d c -> (b f) d c") | |
| key = attn.head_to_batch_dim(key) | |
| value = attn.head_to_batch_dim(value) | |
| attention_probs = attn.get_attention_scores(query, key, attention_mask) | |
| hidden_states_ref3 = torch.bmm(attention_probs, value) | |
| hidden_states = self.self_attn_coeff * hidden_states_self + (1 - self.self_attn_coeff) * torch.mean( | |
| torch.stack([hidden_states_ref0, hidden_states_ref1, hidden_states_ref2, hidden_states_ref3]), | |
| dim=0) if not is_cross_attention else hidden_states_ref3 | |
| hidden_states = attn.batch_to_head_dim(hidden_states) | |
| # linear proj | |
| hidden_states = attn.to_out[0](hidden_states, *args) | |
| # dropout | |
| hidden_states = attn.to_out[1](hidden_states) | |
| if input_ndim == 4: | |
| hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width) | |
| if attn.residual_connection: | |
| hidden_states = hidden_states + residual | |
| hidden_states = hidden_states / attn.rescale_output_factor | |
| return hidden_states | |