temp = torch.zeros((xyz.shape[0], 1)).to('cuda')

temp[:, :] = 2 * t - 1

temp = torch.cat([xyz, temp], dim=1)

tb_writer = prepare_output_and_logger(dataset)

gaussians = GaussianModel(dataset.sh_degree)

deform = DeformModel(dataset.is_blender, dataset.is_6dof)

deform.train_setting(opt)

n_layer = dataset.n_layer

weight_reg = dataset.weight_reg

weight_t = dataset.weight_t

batch_timeformer = dataset.batch

timeFormer = TimeFormer(input_dims=4, multires=4, nhead=4, hidden_dims=36, num_layer=n_layer).to('cuda')

timeFormer.train_setting(opt)

scene = Scene(dataset, gaussians)

gaussians.training_setup(opt)

bg_color = [1, 1, 1] if dataset.white_background else [0, 0, 0]

background = torch.tensor(bg_color, dtype=torch.float32, device="cuda")

iter_start = torch.cuda.Event(enable_timing=True)

iter_end = torch.cuda.Event(enable_timing=True)

viewpoint_stack = None

ema_loss_for_log = 0.0

best_psnr = 0.0

best_iteration = 0

progress_bar = tqdm(range(opt.iterations), desc="Training progress")

smooth_term = get_linear_noise_func(lr_init=0.1, lr_final=1e-15, lr_delay_mult=0.01, max_steps=20000)

for iteration in range(1, opt.iterations + 1):

if network_gui.conn == None:

network_gui.try_connect()

while network_gui.conn != None:

try:

net_image_bytes = None

custom_cam, do_training, pipe.do_shs_python, pipe.do_cov_python, keep_alive, scaling_modifer = network_gui.receive()

if custom_cam != None:

net_image = render(custom_cam, gaussians, pipe, background, scaling_modifer)["render"]

net_image_bytes = memoryview((torch.clamp(net_image, min=0, max=1.0) * 255).byte().permute(1, 2,

0).contiguous().cpu().numpy())

network_gui.send(net_image_bytes, dataset.source_path)

if do_training and ((iteration < int(opt.iterations)) or not keep_alive):

break

except Exception as e:

network_gui.conn = None

iter_start.record()

# Every 1000 its we increase the levels of SH up to a maximum degree

if iteration % 1000 == 0:

gaussians.oneupSHdegree()

# Pick a random Camera

if not viewpoint_stack or len(viewpoint_stack) < batch_timeformer:

viewpoint_stack = scene.getTrainCameras().copy()

total_frame = len(viewpoint_stack)

time_interval = 1 / total_frame

viewpoint_cams = []

for _ in range(batch_timeformer) :

viewpoint_cam = viewpoint_stack.pop(randint(0,len(viewpoint_stack)-1))

viewpoint_cams.append(viewpoint_cam)

# breakpoint()

time_ori = torch.cat([trans(gaussians.get_xyz, viewpoint_cam.fid).unsqueeze(0) for viewpoint_cam in viewpoint_cams], dim=0)

time_offset = timeFormer(time_ori)

images_gt = []

images = []

images_t = []

for i in range(len(viewpoint_cams)):

viewpoint_cam = viewpoint_cams[i]

# viewpoint_cam = viewpoint_stack.pop(randint(0, len(viewpoint_stack) - 1))

if dataset.load2gpu_on_the_fly:

viewpoint_cam.load2device()

fid = viewpoint_cam.fid

# w/ time offset

if iteration < opt.warm_up:

d_xyz, d_rotation, d_scaling = 0.0, 0.0, 0.0

else:

N = gaussians.get_xyz.shape[0]

time_input = fid.unsqueeze(0).expand(N, -1)

ast_noise = 0 if dataset.is_blender else torch.randn(1, 1, device='cuda').expand(N, -1) * time_interval * smooth_term(iteration)

d_xyz, d_rotation, d_scaling = deform.step(gaussians.get_xyz.detach()+time_offset[i, :, :], time_input + ast_noise)

render_pkg_re_t = render(viewpoint_cam, gaussians, pipe, background, d_xyz, d_rotation, d_scaling, dataset.is_6dof)

image_t, _, _, _ = render_pkg_re_t["render"], render_pkg_re_t[

"viewspace_points"], render_pkg_re_t["visibility_filter"], render_pkg_re_t["radii"]

# w/o time offset

if iteration < opt.warm_up:

d_xyz, d_rotation, d_scaling = 0.0, 0.0, 0.0

else:

N = gaussians.get_xyz.shape[0]

time_input = fid.unsqueeze(0).expand(N, -1)

ast_noise = 0 if dataset.is_blender else torch.randn(1, 1, device='cuda').expand(N, -1) * time_interval * smooth_term(iteration)

d_xyz, d_rotation, d_scaling = deform.step(gaussians.get_xyz.detach(), time_input + ast_noise)

render_pkg_re = render(viewpoint_cam, gaussians, pipe, background, d_xyz, d_rotation, d_scaling, dataset.is_6dof)

image, viewspace_point_tensor, visibility_filter, radii = render_pkg_re["render"], render_pkg_re[

"viewspace_points"], render_pkg_re["visibility_filter"], render_pkg_re["radii"]

# Loss

gt_image = viewpoint_cam.original_image.cuda()

images.append(image.unsqueeze(0))

images_t.append(image_t.unsqueeze(0))

images_gt.append(gt_image.unsqueeze(0))

image_tensor = torch.cat(images, 0)

image_tensor_t = torch.cat(images_t, 0)

gt_image_tensor = torch.cat(images_gt, 0)

# breakpoint()

Ll1 = l1_loss(image_tensor, gt_image_tensor[:,:3,:,:])

Ll1_t = l1_loss(image_tensor_t, gt_image_tensor[:, :3, :, :])

l_reg = l1_loss(time_offset, 0)

loss = Ll1 + weight_t * Ll1_t + weight_reg * (l_reg if iteration >= opt.warm_up else 0.0)

# loss += (1.0 - opt.lambda_dssim) * Ll1 + opt.lambda_dssim * (1.0 - ssim(image, gt_image))

loss.backward()

iter_end.record()

if dataset.load2gpu_on_the_fly:

viewpoint_cam.load2device('cpu')

with torch.no_grad():

# Progress bar

ema_loss_for_log = 0.4 * loss.item() + 0.6 * ema_loss_for_log

if iteration % 10 == 0:

progress_bar.set_postfix({"Loss": f"{ema_loss_for_log:.{7}f}"})

progress_bar.update(10)

if iteration == opt.iterations:

progress_bar.close()

# Keep track of max radii in image-space for pruning

gaussians.max_radii2D[visibility_filter] = torch.max(gaussians.max_radii2D[visibility_filter],

radii[visibility_filter])

# Log and save

if iteration % 10 == 0:

cur_psnr = training_report(tb_writer, iteration, Ll1, loss, l1_loss, iter_start.elapsed_time(iter_end),

testing_iterations, scene, render, (pipe, background), deform,

dataset.load2gpu_on_the_fly, dataset.is_6dof, Ll1_t, l_reg)

if iteration in testing_iterations:

if cur_psnr.item() > best_psnr:

best_psnr = cur_psnr.item()

best_iteration = iteration

if iteration in saving_iterations:

print("\n[ITER {}] Saving Gaussians".format(iteration))

scene.save(iteration)

deform.save_weights(args.model_path, iteration)

# Densification

if iteration < opt.densify_until_iter:

viewspace_point_tensor_densify = render_pkg_re["viewspace_points_densify"]

gaussians.add_densification_stats(viewspace_point_tensor_densify, visibility_filter)

if iteration > opt.densify_from_iter and iteration % opt.densification_interval == 0:

size_threshold = 20 if iteration > opt.opacity_reset_interval else None

gaussians.densify_and_prune(opt.densify_grad_threshold, 0.005, scene.cameras_extent, size_threshold)

if iteration % opt.opacity_reset_interval == 0 or (

dataset.white_background and iteration == opt.densify_from_iter):

# gaussians.reset_opacity()

pass

# Optimizer step

if iteration < opt.iterations:

gaussians.optimizer.step()

gaussians.update_learning_rate(iteration)

deform.optimizer.step()

gaussians.optimizer.zero_grad(set_to_none=True)

deform.optimizer.zero_grad()

deform.update_learning_rate(iteration)

timeFormer.optimizer.step()

timeFormer.optimizer.zero_grad()

timeFormer.update_learning_rate(iteration)

# def print_model_param_values(model):

# for name, param in model.named_parameters():

# print(name, param.data)

# print_model_param_values(timeFormer.output_layer)

if not args.model_path:

if os.getenv('OAR_JOB_ID'):

unique_str = os.getenv('OAR_JOB_ID')

else:

unique_str = str(uuid.uuid4())

args.model_path = os.path.join("./output/", unique_str[0:10])

# Set up output folder

print("Output folder: {}".format(args.model_path))

os.makedirs(args.model_path, exist_ok=True)

with open(os.path.join(args.model_path, "cfg_args"), 'w') as cfg_log_f:

cfg_log_f.write(str(Namespace(**vars(args))))

# Create Tensorboard writer

tb_writer = None

if TENSORBOARD_FOUND:

tb_writer = SummaryWriter(args.model_path)

else:

print("Tensorboard not available: not logging progress")

renderArgs, deform, load2gpu_on_the_fly, is_6dof=False, loss_t=0.0, loss_reg=0.0):

if tb_writer:

tb_writer.add_scalar('train_loss_patches/l1_loss', Ll1.item(), iteration)

tb_writer.add_scalar('train_loss_patches/total_loss', loss.item(), iteration)

tb_writer.add_scalar('train_loss_patches/iter_time', elapsed, iteration)

tb_writer.add_scalar(f'train_loss_patches/l1_timeformer_loss', loss_t.item(), iteration)

tb_writer.add_scalar(f'train_loss_patches/reg_loss', loss_reg.item(), iteration)

if tb_writer:

# tb_writer.add_histogram("scene/opacity_histogram", scene.gaussians.get_opacity, iteration)

tb_writer.add_scalar('train_loss_patches/total_points', scene.gaussians.get_xyz.shape[0], iteration)

torch.cuda.empty_cache()

test_psnr = 0.0

# Report test and samples of training set

if iteration in testing_iterations:

torch.cuda.empty_cache()

validation_configs = ({'name': 'test', 'cameras': scene.getTestCameras()},

{'name': 'train',

'cameras': [scene.getTrainCameras()[idx % len(scene.getTrainCameras())] for idx in

range(5, 30, 5)]})

for config in validation_configs:

if config['cameras'] and len(config['cameras']) > 0:

images = torch.tensor([], device="cuda")

gts = torch.tensor([], device="cuda")

for idx, viewpoint in enumerate(config['cameras']):

if load2gpu_on_the_fly:

viewpoint.load2device()

fid = viewpoint.fid

xyz = scene.gaussians.get_xyz

time_input = fid.unsqueeze(0).expand(xyz.shape[0], -1)

d_xyz, d_rotation, d_scaling = deform.step(xyz.detach(), time_input)

image = torch.clamp(

renderFunc(viewpoint, scene.gaussians, *renderArgs, d_xyz, d_rotation, d_scaling, is_6dof)["render"],

0.0, 1.0)

gt_image = torch.clamp(viewpoint.original_image.to("cuda"), 0.0, 1.0)

images = torch.cat((images, image.unsqueeze(0)), dim=0)

gts = torch.cat((gts, gt_image.unsqueeze(0)), dim=0)

if load2gpu_on_the_fly:

viewpoint.load2device('cpu')

if tb_writer and (idx < 5):

tb_writer.add_images(config['name'] + "_view_{}/render".format(viewpoint.image_name),

image[None], global_step=iteration)

if iteration == testing_iterations[0]:

tb_writer.add_images(config['name'] + "_view_{}/ground_truth".format(viewpoint.image_name),

gt_image[None], global_step=iteration)

l1_test = l1_loss(images, gts)

psnr_test = psnr(images, gts).mean()

if config['name'] == 'test' or len(validation_configs[0]['cameras']) == 0:

test_psnr = psnr_test

print("\n[ITER {}] Evaluating {}: L1 {} PSNR {}".format(iteration, config['name'], l1_test, psnr_test))

if tb_writer:

tb_writer.add_scalar(config['name'] + '/loss_viewpoint - l1_loss', l1_test, iteration)

tb_writer.add_scalar(config['name'] + '/loss_viewpoint - psnr', psnr_test, iteration)