Update app.py

app.py

@@ -78,6 +78,38 @@ def main():
     # Function to reconstruct the 3D model from the input images and export it as a PLY file
     @spaces.GPU(duration=120)  # Decorator to allocate a GPU for this function during execution
     def reconstruct_and_export(images, num_gauss):
+        """
+        Passes a batch of images through the model, outputs reconstruction in the form of a dict of tensors.
+        """
+        print("[DEBUG] Starting reconstruction and export...")
+        # Stack the images along a new dimension to create a batch
+        images_batch = torch.stack([to_tensor(image) for image in images]).to(device)  # Create a batch of images
+
+        # Create input dictionary expected by the model
+        inputs = {
+            ("color_aug", 0, 0): images_batch,  # Batch of input images
+        }
+
+        # Pass the batch of images through the model to get the output
+        print("[INFO] Passing batch of images through the model...")
+        outputs = model(inputs)  # Perform inference to get model outputs
+
+        # Use the first output for illustration (or modify to combine outputs as needed)
+        gauss_means = outputs[('gauss_means', 0, 0)]
+        if gauss_means.size(0) < num_gauss or gauss_means.size(0) % num_gauss != 0:
+            adjusted_num_gauss = max(1, gauss_means.size(0) // (gauss_means.size(0) // num_gauss))
+            print(f"[WARNING] Adjusting num_gauss from {num_gauss} to {adjusted_num_gauss} to avoid shape mismatch.")
+            num_gauss = adjusted_num_gauss  # Adjust num_gauss to prevent errors during tensor reshaping
+
+        # Debugging tensor shape
+        print(f"[DEBUG] gauss_means tensor shape: {gauss_means.shape}")
+
+        # Export the reconstruction to a PLY file
+        print(f"[INFO] Saving output to {ply_out_path}...")
+        save_ply(outputs, ply_out_path, num_gauss=num_gauss)  # Save the output 3D model to a PLY file
+        print("[INFO] Reconstruction and export complete.")
+
+        return ply_out_path  # Return the path to the saved PLY file
         """
         Passes images through model, outputs reconstruction in form of a dict of tensors.
         """
@@ -137,11 +169,11 @@ def main():
                     # Input images component for the user to upload multiple images
                     input_images = gr.Gallery(
                         label="Input Images",
-                        image_mode="RGBA",  # Accept RGBA images
+                        # Accept RGBA images
                         sources="upload",  # Allow users to upload images
-                        type="pil",  # The images are returned as PIL images
+                        # The images are returned as PIL images
                         elem_id="content_images",
-                        tool="editor",  # Optional, for editing images
+                        # Optional, for editing images
                         # Allow multiple image uploads
                     )
                 with gr.Row():