import torch import numpy as np from vggt.models.vggt import VGGT from training.loss import MultitaskLoss from custom.utils.config import Config from vidar.utils.viz import viz_depth, viz_normals, viz_camera, viz_photo from custom.wandb import prep_video, prep_image, prep_points from custom.utils.geometry import calculate_normals from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy from custom.networks.cosmos_vggt_loss import calc_loss from torch.amp import autocast class CosmosVGGT(torch.nn.Module): """ VGGT wrapper to allow the use of the official repository as part of the predict2 pipeline Serves as template to other networks as well """ def __init__(self, **kwargs): super().__init__() # Depth loss configuration depth_loss_cfg = Config( weight=0.1, gradient_loss_fn="grad", valid_range=0.98, ) # Point loss configuration point_loss_cfg = Config( weight=1.0, gradient_loss_fn="normal", valid_range=0.98, ) # Instantiate VGGT self.model = VGGT.from_pretrained("facebook/VGGT-1B") # Instantiate VGGT losses self.loss_fn = MultitaskLoss( depth=depth_loss_cfg, point=point_loss_cfg, ) def fully_shard(self, mesh, enc_dt=torch.bfloat16): """ Use to convert the wrapper model to FSDP Inspired by cosmos_predict2/networks/text2image_dit.py (line 1543) and the VGGT repo (e.g., they do not use mixed precision in the decoder heads) """ self.precision = enc_dt # bf16 policy for encoder blocks aggregator_mp = MixedPrecisionPolicy( param_dtype=enc_dt, reduce_dtype=enc_dt, output_dtype=None, # keep FSDP from force-casting outputs cast_forward_inputs=False, # incoming tensors into each encoder block -> bf16 ) # fp32 policy for heads: also cast incoming tensors to FP32 head_mp = MixedPrecisionPolicy( param_dtype=torch.float32, reduce_dtype=torch.bfloat16, output_dtype=torch.float32, cast_forward_inputs=True, # critical: upcast the bf16 inputs from encoder to fp32 ) fully_shard(self.model.aggregator.patch_embed, mesh=mesh, mp_policy=head_mp, reshard_after_forward=True) for i in range(len(self.model.aggregator.frame_blocks)): fully_shard(self.model.aggregator.frame_blocks[i], mesh=mesh, mp_policy=aggregator_mp, reshard_after_forward=True) for i in range(len(self.model.aggregator.global_blocks)): fully_shard(self.model.aggregator.global_blocks[i], mesh=mesh, mp_policy=aggregator_mp, reshard_after_forward=True) fully_shard(self.model.depth_head, mesh=mesh, mp_policy=head_mp, reshard_after_forward=True) fully_shard(self.model.point_head, mesh=mesh, mp_policy=head_mp, reshard_after_forward=True) ########################### # fully_shard(self.model.aggregator.rope, mesh=mesh, reshard_after_forward=True) # fully_shard(self.model.aggregator.patch_embed, mesh=mesh, reshard_after_forward=True) # for i in range(len(self.model.aggregator.frame_blocks)): # fully_shard(self.model.aggregator.frame_blocks[i], mesh=mesh, reshard_after_forward=True) # for i in range(len(self.model.aggregator.global_blocks)): # fully_shard(self.model.aggregator.global_blocks[i], mesh=mesh, reshard_after_forward=True) # fully_shard(self.model.aggregator, mesh=mesh, reshard_after_forward=True) # fully_shard(self.model.point_head, mesh=mesh, reshard_after_forward=True) # fully_shard(self.model.depth_head, mesh=mesh, reshard_after_forward=True) ########################### def forward(self, data_batch): # User vidar batch batch = data_batch['anydata'] # Prepare data for model rolled = self.roll(batch) # Run model input_data = rolled['data']['rgb_stacked'] with autocast(device_type="cuda", dtype=self.precision, enabled=True): output = self.model(input_data) if self.training: return self.loss(rolled, output) # Calculate and return loss else: output['unrolled'] = self.unroll(rolled, output) # Predictions output['wandb_logs'] = self.visuals(batch, output['unrolled']) # Visuals return output, None # Return output (loss is None) def loss(self, rolled, output): """ Calculate model loss """ # Get GT data gt_depth_stacked = rolled['data']['depth_stacked'] gt_points_stacked = rolled['data']['points_stacked'] # Prepare GT dictionary gtruth = { 'depths': gt_depth_stacked.squeeze(2), 'world_points': gt_points_stacked.permute(0, 1, 3, 4, 2), 'point_masks': gt_depth_stacked.squeeze(2) > 0, } # Calculate loss loss_dict = self.loss_fn(output, gtruth) loss = loss_dict['objective'] # Get the primary loss # Copy losses to output dictionary for key, val in loss_dict.items(): output[key] = val output['loss'] = loss # Primary loss is named "loss" return output, loss def roll(self, batch): """ Prepare vidar data for model ingestion (input) and prepare ground-truth for logging and metrics (gt) """ # Store all keys keys = list(batch['rgb'].keys()) gt_rgb = batch['rgb'] gt_depth = batch['depth'] gt_cams = batch['cams'] # Stack time_cam rgb and depth information gt_rgb_stacked = torch.stack([gt_rgb[key] for key in keys], 1) gt_depth_stacked = torch.stack([gt_depth[key] for key in keys], 1) # Get pointmaps gt_points = {key: gt_cams[key].reconstruct_depth_map(gt_depth[key], to_world=True) for key in keys} gt_points_stacked = torch.stack([gt_points[key] for key in keys], 1) gt_origin = {key: gt_cams[key].get_origin() for key in keys} # Get depth/ray information from points gt_depth_points = {key: ((gt_points[key] - gt_origin[key]) ** 2).sum(1, keepdim=True).sqrt() for key in keys} gt_rays_points = {key: (gt_points[key] - gt_origin[key]) / torch.norm(gt_points[key] - gt_origin[key], dim=1).unsqueeze(1) for key in keys} return { 'data': { # Processed data 'rgb': gt_rgb, 'cams': gt_cams, 'rgb_stacked': gt_rgb_stacked, 'depth_stacked': gt_depth_stacked, 'points_stacked': gt_points_stacked, 'origin': gt_origin, 'keys': keys, }, 'gt': { # Ground-truth information 'depth_mono': gt_depth, 'depth_points': gt_depth_points, 'rays_points': gt_rays_points, 'points': gt_points, }, } def unroll(self, rolled, output): """ Unrolled model output back to vidar format """ # Get useful rolled information keys = rolled['data']['keys'] gt_origin = rolled['data']['origin'] # Unroll depth and points pred_depth_mono = {key: output['depth'][:, i].permute(0, 3, 1, 2) for i, key in enumerate(keys)} pred_points = {key: output['world_points'][:, i].permute(0, 3, 1, 2) for i, key in enumerate(keys)} # Unroll confidence pred_conf_depth = {key: output['depth_conf'][:, [i]] for i, key in enumerate(keys)} pred_conf_points = {key: output['world_points_conf'][:, [i]] for i, key in enumerate(keys)} # Get depth/ray information from points pred_depth_points = {key: ((pred_points[key] - gt_origin[key]) ** 2).sum(1, keepdim=True).sqrt() for key in keys} pred_rays_points = {key: (pred_points[key] - gt_origin[key]) / torch.norm(pred_points[key] - gt_origin[key], dim=1).unsqueeze(1) for key in keys} # Output predictions return { **rolled, 'pred': { 'depth_mono': pred_depth_mono, 'depth_points': pred_depth_points, 'rays_points': pred_rays_points, 'points': pred_points, 'conf_depth': pred_conf_depth, 'conf_points': pred_conf_points, }, } def visuals(self, batch, unrolled): """Create visuals for wandb logging""" keys = unrolled['data']['keys'] visuals = [self.visuals_single(unrolled, keys[i]) for i in range(len(keys))] suffix = f'{batch["tag"][0]}--{batch["idx"][0]:05d}' merge = np.concatenate([val['merge'] for val in visuals], axis=0) key_merge = f"VGGT--{suffix}" # gt_rgb = np.concatenate([val['gt_rgb'] for val in visuals], axis=0) # gt_points = np.concatenate([val['gt_points'] for val in visuals], axis=0) # pred_points = np.concatenate([val['pred_points'] for val in visuals], axis=0) # gt_points = np.concatenate([gt_points, gt_rgb], axis=1) # key_gt_points = f"GT_POINTS--{suffix}" # pred_points = np.concatenate([pred_points, gt_rgb], axis=1) # key_pred_points = f"PRED_POINTS--{suffix}" return { **prep_image(key_merge, merge), # **prep_points(key_gt_points, gt_points), # **prep_points(key_pred_points, pred_points), } def visuals_single(self, unrolled, key): """Create specific visuals (see above)""" gt_rgb = unrolled['data']['rgb'][key] gt_camera = unrolled['data']['cams'][key] gt_depth_mono = unrolled['gt']['depth_mono'][key] gt_depth_points = unrolled['gt']['depth_points'][key] gt_rays_points = unrolled['gt']['rays_points'][key] gt_points = unrolled['gt']['points'][key] pred_depth_mono = unrolled['pred']['depth_mono'][key] pred_depth_points = unrolled['pred']['depth_points'][key] pred_rays_points = unrolled['pred']['rays_points'][key] pred_points = unrolled['pred']['points'][key] pred_depth_conf = unrolled['pred']['conf_depth'][key] pred_points_conf = unrolled['pred']['conf_points'][key] gt_normals_mono, _ = calculate_normals(gt_depth_mono, gt_camera, to_world=False) pred_normals_mono, _ = calculate_normals(pred_depth_mono, gt_camera, to_world=False) gt_normals_points, _ = calculate_normals(gt_depth_points, gt_camera, to_world=True) pred_normals_points, _ = calculate_normals(pred_depth_points, gt_camera, to_world=True) viz_gt_rgb = gt_rgb[0].permute(1, 2, 0).cpu().numpy() zeros = np.zeros_like(viz_gt_rgb) viz_gt_depth_mono = viz_depth(gt_depth_mono[0]) viz_pred_depth_mono = viz_depth(pred_depth_mono[0]) viz_gt_depth_points = viz_depth(gt_depth_points[0]) viz_pred_depth_points = viz_depth(pred_depth_points[0]) viz_gt_rays_points = viz_camera(gt_rays_points[0]) viz_pred_rays_points = viz_camera(pred_rays_points[0]) viz_gt_normals = viz_normals(gt_normals_mono[0,...,0]) viz_pred_normals = viz_normals(pred_normals_mono[0,...,0]) viz_gt_normals_points = viz_normals(gt_normals_points[0,...,0]) viz_pred_normals_points = viz_normals(pred_normals_points[0,...,0]) viz_pred_depth_conf = viz_photo(- pred_depth_conf[0], colormap='jet', normalize=True) viz_pred_points_conf = viz_photo(- pred_points_conf[0], colormap='jet', normalize=True) col0 = np.concatenate([viz_gt_rgb, zeros ], axis=0) col1 = np.concatenate([viz_gt_depth_mono, viz_pred_depth_mono ], axis=0) col2 = np.concatenate([viz_gt_normals, viz_pred_normals ], axis=0) col3 = np.concatenate([viz_gt_depth_points, viz_pred_depth_points ], axis=0) col4 = np.concatenate([viz_gt_rays_points, viz_pred_rays_points ], axis=0) col5 = np.concatenate([viz_gt_normals_points, viz_pred_normals_points ], axis=0) col6 = np.concatenate([viz_pred_depth_conf, viz_pred_points_conf ], axis=0) merge = np.concatenate([col0,col1,col2,col3,col5,col4,col6], axis=1) merge[np.isnan(merge)] = 0.0 merge[np.isinf(merge)] = 0.0 return { 'merge': merge, 'gt_rgb': gt_rgb[0].view(3, -1).permute(1, 0)[::5].detach().to(dtype=torch.float32, device='cpu').numpy() * 255, 'gt_points': gt_points[0].view(3, -1).permute(1, 0)[::5].detach().to(dtype=torch.float32, device='cpu').numpy(), 'pred_points': pred_points[0].view(3, -1).permute(1, 0)[::5].detach().to(dtype=torch.float32, device='cpu').numpy(), }