Just giving up on this idea

demo/gs_train.py

@@ -12,8 +12,6 @@ from demo_globals import DEVICE
 import spaces
 from simple_knn._C import distCUDA2
 
-
-
 @dataclass
 class PipelineParams:
     convert_SHs_python: bool = False
@@ -169,73 +167,73 @@ def train(
     first_iter += 1
 
     point_cloud_path = ""
-    progress = gr.Progress()  # Initialize the progress bar
-    for iteration in range(first_iter, opt.iterations + 1):
-        iter_start.record()
-        gaussians.update_learning_rate(iteration)
+    # progress = gr.Progress()  # Initialize the progress bar
+    # for iteration in range(first_iter, opt.iterations + 1):
+    #     iter_start.record()
+    #     gaussians.update_learning_rate(iteration)
         
-        # Every 1000 its we increase the levels of SH up to a maximum degree
-        if iteration % 1000 == 0:
-            gaussians.oneupSHdegree()
+    #     # 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:
-            viewpoint_stack = scene.getTrainCameras().copy()
-        viewpoint_cam = viewpoint_stack.pop(randint(0, len(viewpoint_stack)-1))
+    #     # Pick a random Camera
+    #     if not viewpoint_stack:
+    #         viewpoint_stack = scene.getTrainCameras().copy()
+    #     viewpoint_cam = viewpoint_stack.pop(randint(0, len(viewpoint_stack)-1))
         
-        # Render
-        if (iteration - 1) == debug_from:
-            pipe.debug = True
-        bg = torch.rand((3), device=DEVICE) if opt.random_background else background
+    #     # Render
+    #     if (iteration - 1) == debug_from:
+    #         pipe.debug = True
+    #     bg = torch.rand((3), device=DEVICE) if opt.random_background else background
         
-        render_pkg = render(viewpoint_cam, gaussians, pipe, bg)
-        image, viewspace_point_tensor, visibility_filter, radii = render_pkg["render"], render_pkg["viewspace_points"], render_pkg["visibility_filter"], render_pkg["radii"]
-
-        # Loss
-        gt_image = viewpoint_cam.original_image.cuda()
-        Ll1 = l1_loss(image, gt_image)
-        loss = (1.0 - opt.lambda_dssim) * Ll1 + opt.lambda_dssim * (1.0 - ssim(image, gt_image))
-        loss.backward()
-        iter_end.record()
+    #     render_pkg = render(viewpoint_cam, gaussians, pipe, bg)
+    #     image, viewspace_point_tensor, visibility_filter, radii = render_pkg["render"], render_pkg["viewspace_points"], render_pkg["visibility_filter"], render_pkg["radii"]
+
+    #     # Loss
+    #     gt_image = viewpoint_cam.original_image.cuda()
+    #     Ll1 = l1_loss(image, gt_image)
+    #     loss = (1.0 - opt.lambda_dssim) * Ll1 + opt.lambda_dssim * (1.0 - ssim(image, gt_image))
+    #     loss.backward()
+    #     iter_end.record()
         
-        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)
-                progress(iteration / opt.iterations)  # Update Gradio progress bar
-            if iteration == opt.iterations:
-                progress_bar.close()
-
-            # Log and save
-            training_report(tb_writer, iteration, Ll1, loss, l1_loss, iter_start.elapsed_time(iter_end), testing_iterations, scene, render, (pipe, background))
-            if (iteration == opt.iterations):
-                point_cloud_path = os.path.join(os.path.join(dataset.model_path, "point_cloud/iteration_{}".format(iteration)), "point_cloud.ply")
-                print("\n[ITER {}] Saving Gaussians to {}".format(iteration, point_cloud_path))
-                scene.save(iteration)
-
-            # Densification
-            if iteration < opt.densify_until_iter:
-                # 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])
-                gaussians.add_densification_stats(viewspace_point_tensor, 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()
-
-                # Optimizer step
-                if iteration < opt.iterations:
-                    gaussians.optimizer.step()
-                    gaussians.optimizer.zero_grad(set_to_none = True)
-
-                if (iteration == opt.iterations):
-                    print("\n[ITER {}] Saving Checkpoint".format(iteration))
-                    torch.save((gaussians.capture(), iteration), scene.model_path + "/chkpnt" + str(iteration) + ".pth")
+    #     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)
+    #             progress(iteration / opt.iterations)  # Update Gradio progress bar
+    #         if iteration == opt.iterations:
+    #             progress_bar.close()
+
+    #         # Log and save
+    #         training_report(tb_writer, iteration, Ll1, loss, l1_loss, iter_start.elapsed_time(iter_end), testing_iterations, scene, render, (pipe, background))
+    #         if (iteration == opt.iterations):
+    #             point_cloud_path = os.path.join(os.path.join(dataset.model_path, "point_cloud/iteration_{}".format(iteration)), "point_cloud.ply")
+    #             print("\n[ITER {}] Saving Gaussians to {}".format(iteration, point_cloud_path))
+    #             scene.save(iteration)
+
+    #         # Densification
+    #         if iteration < opt.densify_until_iter:
+    #             # 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])
+    #             gaussians.add_densification_stats(viewspace_point_tensor, 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()
+
+    #             # Optimizer step
+    #             if iteration < opt.iterations:
+    #                 gaussians.optimizer.step()
+    #                 gaussians.optimizer.zero_grad(set_to_none = True)
+
+    #             if (iteration == opt.iterations):
+    #                 print("\n[ITER {}] Saving Checkpoint".format(iteration))
+    #                 torch.save((gaussians.capture(), iteration), scene.model_path + "/chkpnt" + str(iteration) + ".pth")
 
 
     from os import makedirs