import torch from custom.dataloader.utils import prepare_key, random_value, random_choice def prepare_dense(key, data, data_lat, pixel_condition_frames, avail_times, dst_cam_idx, latent_downsampling_factors, device, tensor_kwargs): if isinstance(dst_cam_idx, list): a4d_dicts = [prepare_dense( 'rgb', data, data_lat, pixel_condition_frames, avail_times, cam, latent_downsampling_factors, device, tensor_kwargs, ) for cam in dst_cam_idx] a4d_dict = {key: torch.cat([val[key] for val in a4d_dicts], 0) for key in a4d_dicts[0]} for key, val in a4d_dicts[0].items(): print('aaa', key, val.shape) for key, val in a4d_dict.items(): print('bbb', key, val.shape) return a4d_dict # Prepare data if data is not None: dst_list = [data[(t, dst_cam_idx)] for t in avail_times] dst = torch.stack(dst_list, dim=2).to(**tensor_kwargs) elif data_lat is not None: b, c, tt, hh, ww = data_lat[(avail_times[0], dst_cam_idx)].shape t = (tt - 1) * latent_downsampling_factors[0] + 1 h = hh * latent_downsampling_factors[1] w = ww * latent_downsampling_factors[2] dst = torch.zeros((b, 3, t, h, w), device=device, dtype=tensor_kwargs['dtype']) else: raise ValueError(f'Insufficient data for key {key}') # Prepare masks dst_input_mask = torch.zeros_like(dst[:, :1]) dst_input_mask[:, :, :pixel_condition_frames] = 1.0 dst_output_mask = 1.0 - dst_input_mask # Prepare dictionary a4d_dict = { f'{key}_dst': dst, f'{key}_dst_input_mask': dst_input_mask, f'{key}_dst_output_mask': dst_output_mask, f'{key}_dst_supervise_mask': dst_output_mask, } return a4d_dict def anydata_to_a4d_fn(data_dict, phase=None, config=None, generator=None, **leftover): # Valid phases assert phase in ['train', 'val', 'test'] # Default parameters latent_downsampling_factors = (4, 8, 8) tensor_kwargs = dict(dtype=torch.bfloat16) # Get vidar dictionary vidar_dict = data_dict['vidar'] # Construct the dictionaries. a4d_dict = dict() load_rgb = True load_depth = False # True load_points = False # True load_cams = False # True load_language = True load_action = False # True # Extract relevant raw / pixel modalities. vidar_rgb = vidar_dict.get('rgb', None) if load_rgb else None # ^ dict mapping (time, camera) to (B, 3, H, W) tensor of float in [0, 1]. vidar_depth = vidar_dict.get('depth', None) if load_depth else None # ^ dict mapping (time, camera) to (B, 3, H, W) tensor of float in [0, +inf). vidar_points = vidar_dict.get('points', None) if load_points else None # ^ dict mapping (time, camera) to (B, 3, H, W) tensor of float in (-inf, +inf). vidar_cams = vidar_dict.get('cams', None) if load_cams else None # ^ dict mapping (time, camera) to CameraPinhole object. vidar_language = vidar_dict.get('language', None) if load_language else None # ^ optional dict mapping (0, 0) to {'prompt': [str]} vidar_action = vidar_dict.get('action', None) if load_action else None # ^ optional dict mapping (time, 0) to {'action', 'proprio': (B, 16, 20) tensor of float}. # Extract relevant latent modalities. latent_rgb = vidar_dict.get('lat_rgb', None) if load_rgb else None # ^ dict mapping (time, camera) to (B, 16, H, W) tensor of float. latent_depth = vidar_dict.get('lat_depth', None) if load_depth else None # ^ dict mapping (time, camera) to (B, 16, H, W) tensor of float. latent_points = vidar_dict.get('lat_points', None) if load_points else None # ^ dict mapping (time, camera) to (B, 16, H, W) tensor of float. latent_cams = vidar_dict.get('lat_cams', None) if load_cams else None # ^ dict mapping (time, camera) to (B, 32/48, H, W) tensor of float. latent_language = vidar_dict.get('lat_language', None) if load_language else None # ^ optional dict mapping (0, 0) to {'prompt': [str]} # Obtain available (time, camera) keys & device to use first_valid = (vidar_rgb or vidar_depth or vidar_points or vidar_cams or \ latent_rgb or latent_depth or latent_points or latent_cams) assert first_valid is not None, 'Insufficient input information' keys = first_valid.keys() device = first_valid[keys[0]].device # Available timesteps and cameras avail_times = sorted(list(set([key[0] for key in keys]))) avail_cams = sorted(list(set([key[1] for key in keys]))) # T = len(avail_times) # Destination camera dst_cam_idx = avail_cams[0] # src_cam_idx = avail_cams[1] # Number of conditioning frames if phase == 'train': pixel_condition_frames = 9 # random_choice([1,9,17], device, generator) elif phase == 'val': pixel_condition_frames = 9 else: raise ValueError(f'Invalid phase {phase}') if vidar_rgb is not None or latent_rgb is not None: # a4d_dict.update(prepare_dense( # 'rgb', vidar_rgb, latent_rgb, pixel_condition_frames, # avail_times, dst_cam_idx, latent_downsampling_factors, device, tensor_kwargs, # )) # Prepare RGB data if vidar_rgb is not None: rgb_dst_list = [vidar_rgb[(t, dst_cam_idx)] for t in avail_times] rgb_dst = torch.stack(rgb_dst_list, dim=2).to(**tensor_kwargs) # (3, T, H, W) tensor of float in [0, 1]. else: b, c, tt, hh, ww = latent_rgb[keys[0]].shape t = (tt - 1) * latent_downsampling_factors[0] + 1 h = hh * latent_downsampling_factors[1] w = ww * latent_downsampling_factors[2] rgb_dst = torch.zeros((b, 3, t, h, w), device=device, dtype=tensor_kwargs['dtype']) # Prepare masks rgb_dst_input_mask = torch.zeros_like(rgb_dst[:, :1]) rgb_dst_input_mask[:, :, :pixel_condition_frames] = 1.0 rgb_dst_output_mask = 1.0 - rgb_dst_input_mask # Prepare dictionary a4d_dict['rgb_dst'] = rgb_dst # (3, T, H, W) tensor of bfloat16 in [-1, 1]. a4d_dict['rgb_dst_input_mask'] = rgb_dst_input_mask # (1, T, H, W) tensor of bfloat16. a4d_dict['rgb_dst_output_mask'] = rgb_dst_output_mask # (1, T, H, W) tensor of bfloat16. a4d_dict['rgb_dst_supervise_mask'] = rgb_dst_output_mask # (1, T, H, W) tensor of bfloat16. # # Prepare RGB data # rgb_src_list = [vidar_rgb[(t, src_cam_idx)] for t in avail_times] # rgb_src = torch.stack(rgb_src_list, dim=2) # (3, T, H, W) tensor of float in [0, 1]. # rgb_src = torch.tensor(rgb_src, **tensor_kwargs) # (3, T, H, W) tensor of bfloat16 in [-1, 1]. # # Prepare masks # rgb_src_input_mask = torch.zeros_like(rgb_src[:, :1]) # rgb_src_input_mask[:, :, :pixel_condition_frames] = 1.0 # rgb_src_output_mask = 1.0 - rgb_src_input_mask # # Prepare dictionary # a4d_dict['rgb_src'] = rgb_src # (3, T, H, W) tensor of bfloat16 in [-1, 1]. # a4d_dict['rgb_src_input_mask'] = rgb_src_input_mask # (1, T, H, W) tensor of bfloat16. # a4d_dict['rgb_src_output_mask'] = rgb_src_output_mask # (1, T, H, W) tensor of bfloat16. # # a4d_dict['rgb_dst_supervise_mask'] = rgb_dst_output_mask # (1, T, H, W) tensor of bfloat16. if vidar_depth is not None: # Prepare Depth data depth_list = [vidar_depth[(t, dst_cam_idx)] for t in avail_times] depth_dst = torch.stack(depth_list, dim=2) # (1, T, H, W) tensor of float in [0, inf). depth_dst = torch.tensor(depth_dst, **tensor_kwargs) # (1, T, H, W) tensor of bfloat16 in [0, inf). # Prepare masks depth_dst_input_mask = torch.zeros_like(depth_dst[:, :1]) depth_dst_input_mask[:, :, :pixel_condition_frames] = 1.0 depth_dst_output_mask = 1.0 - depth_dst_input_mask # Prepare dictionary a4d_dict['depth_dst'] = depth_dst # (1, T, H, W) tensor of bfloat16. a4d_dict['depth_dst_input_mask'] = depth_dst_input_mask # (1, T, H, W) tensor of bfloat16. a4d_dict['depth_dst_output_mask'] = depth_dst_output_mask # (1, T, H, W) tensor of bfloat16. if vidar_points is not None: # Prepare Points data points_list = [vidar_points[(t, dst_cam_idx)] for t in avail_times] points_dst = torch.stack(points_list, dim=2) # (1, T, H, W) tensor of float in [0, inf). points_dst = torch.tensor(points_dst, **tensor_kwargs) # (1, T, H, W) tensor of bfloat16 in [0, inf). # Prepare masks points_dst_input_mask = torch.zeros_like(points_dst[:, :1]) points_dst_input_mask[:, :, :pixel_condition_frames] = 1.0 points_dst_output_mask = 1.0 - points_dst_input_mask # Prepare dictionary a4d_dict['points_dst'] = points_dst # (3, T, H, W) tensor of bfloat16. a4d_dict['points_dst_input_mask'] = points_dst_input_mask # (1, T, H, W) tensor of bfloat16. a4d_dict['points_dst_output_mask'] = points_dst_output_mask # (1, T, H, W) tensor of bfloat16. if vidar_cams is not None: # Prepare Camera data cams_list = [vidar_cams[(t, dst_cam_idx)].get_plucker() for t in avail_times] cams_dst = torch.stack(cams_list, dim=2) # (1, T, H, W) tensor of float in [0, inf). cams_dst = torch.tensor(cams_dst, **tensor_kwargs) # (1, T, H, W) tensor of bfloat16 in [0, inf). # Prepare masks cams_dst_input_mask = torch.zeros_like(cams_dst[:, :1]) cams_dst_input_mask[:, :, :pixel_condition_frames] = 1.0 cams_dst_output_mask = 1.0 - cams_dst_input_mask # Prepare dictionary a4d_dict['cams_dst'] = cams_dst # (9, T, H, W) tensor of bfloat16. a4d_dict['cams_dst_input_mask'] = cams_dst_input_mask # (1, T, H, W) tensor of bfloat16. a4d_dict['cams_dst_output_mask'] = cams_dst_output_mask # (1, T, H, W) tensor of bfloat16. if vidar_action is not None: # Prepare Acion data action_dst = vidar_action[(0,0)]['action'] # Prepare dictionary action_dst_input_mask = torch.zeros_like(action_dst[:, :1]) action_dst_input_mask[:, :1] = 1.0 action_dst_output_mask = 1.0 - torch.zeros_like(action_dst[:, :1]) # Prepare dictionary a4d_dict['action_dst'] = action_dst # (9, T, H, W) tensor of bfloat16. a4d_dict['action_dst_input_mask'] = action_dst_input_mask # (1, T, H, W) tensor of bfloat16. a4d_dict['action_dst_output_mask'] = action_dst_output_mask # (1, T, H, W) tensor of bfloat16. # Get relevant pixel dimensions (Tpl, Hpl, Wpl) = latent_downsampling_factors (Tpm, Hpm, Wpm) = a4d_dict['rgb_dst'].shape[2:] # Get relevant latent dimensions Tl = (Tpm - 1) // Tpl + 1 Hl = Hpm // Hpl Wl = Wpm // Wpl prompt = vidar_language[(0,0)]['prompt'] a4d_dict['prompt'] = prompt[0] # string a4d_dict['num_frames'] = torch.tensor([Tpm], **tensor_kwargs) # in pixel space # This will be processed by Matrix4DVAE later: data_dict['a4d'] = a4d_dict # Apply more Cosmos / Matrix4D related parameters. num_low_sattn_tokens = 0 latent_vid_seq_len = Tl * Hl * Wl latent_total_seq_len = latent_vid_seq_len + num_low_sattn_tokens latent_loss_mask = torch.ones(latent_total_seq_len, **tensor_kwargs) # Specify relevant parameters data_dict['vid_seq_len'] = latent_vid_seq_len data_dict['low_sattn_seq_len'] = num_low_sattn_tokens data_dict['total_seq_len'] = latent_total_seq_len data_dict['loss_mask'] = latent_loss_mask # Specify keys of interest for visualization: data_dict['pkeys_input'] = ['rgb_dst,rgb_src'] data_dict['pkeys_output'] = ['rgb_dst,rgb_src'] # Necessary for vanilla cosmos data_dict['video'] = a4d_dict['rgb_dst'] data_dict['num_conditional_frames'] = 3 return data_dict