Update app.py

Pseudocode
Add sliders for resolution and num_gauss to the Gradio interface.
Modify the preprocess function to accept resolution as a parameter.
Modify the reconstruct_and_export function to accept num_gauss as a parameter.
Update the Gradio interface to include the new sliders and pass their values to the respective functions

app.py

@@ -1,18 +1,3 @@
-import sys
-import spaces
-sys.path.append("flash3d")  # Add the flash3d directory to the system path for importing local modules
-
-from omegaconf import OmegaConf
-import gradio as gr
-import torch
-import torchvision.transforms as TT
-import torchvision.transforms.functional as TTF
-from huggingface_hub import hf_hub_download
-import numpy as np
-
-from networks.gaussian_predictor import GaussianPredictor
-from util.vis3d import save_ply
-
 def main():
     print("[INFO] Starting main function...")
     if torch.cuda.is_available():
@@ -32,12 +17,8 @@ def main():
     
     print("[INFO] Initializing GaussianPredictor model...")
     model = GaussianPredictor(cfg)
-    try:
-        device = torch.device(device)
-        model.to(device)
-    except Exception as e:
-        print(f"[ERROR] Failed to set device: {e}")
-        raise
+    device = torch.device(device)
+    model.to(device)
     
     print("[INFO] Loading model weights...")
     model.load_model(model_path)
@@ -52,31 +33,23 @@ def main():
             raise gr.Error("No image uploaded!")
         print("[INFO] Input image is valid.")
 
-    def preprocess(image, padding_value):
+    def preprocess(image, resolution):
         print("[DEBUG] Preprocessing image...")
-        image = TTF.resize(image, (cfg.dataset.height, cfg.dataset.width), interpolation=TT.InterpolationMode.BICUBIC)
-        pad_border_fn = TT.Pad((padding_value, padding_value))
+        image = TTF.resize(image, (resolution, resolution), interpolation=TT.InterpolationMode.BICUBIC)
         image = pad_border_fn(image)
         print("[INFO] Image preprocessing complete.")
         return image
 
     @spaces.GPU(duration=120)
-    def reconstruct_and_export(image, num_gauss, max_sh_degree, scaling_modifier):
+    def reconstruct_and_export(image, num_gauss):
         print("[DEBUG] Starting reconstruction and export...")
         image = to_tensor(image).to(device).unsqueeze(0)
         inputs = {("color_aug", 0, 0): image}
-
         print("[INFO] Passing image through the model...")
         outputs = model(inputs)
-
-        gauss_means = outputs[('gauss_means',0, 0)]
-        if gauss_means.shape[0] % num_gauss != 0:
-            raise ValueError(f"Shape mismatch: cannot divide axis of length {gauss_means.shape[0]} into chunks of {num_gauss}")
-
         print(f"[INFO] Saving output to {ply_out_path}...")
-        save_ply(outputs, ply_out_path, num_gauss=num_gauss, max_sh_degree=max_sh_degree, scaling_modifier=scaling_modifier)
+        save_ply(outputs, ply_out_path, num_gauss=num_gauss)
         print("[INFO] Reconstruction and export complete.")
-
         return ply_out_path
     
     ply_out_path = f'./mesh.ply'
@@ -94,15 +67,9 @@ def main():
             with gr.Column(scale=1):
                 with gr.Row():
                     input_image = gr.Image(label="Input Image", image_mode="RGBA", sources="upload", type="pil", elem_id="content_image")
-                with gr.Row():
-                    num_gauss = gr.Slider(minimum=1, maximum=20, step=1, label="Number of Gaussians per Pixel", value=10)
-                    padding_value = gr.Slider(minimum=0, maximum=128, step=8, label="Padding Amount for Output Processing", value=32)
-                    max_sh_degree = gr.Slider(minimum=1, maximum=10, step=1, label="Max SH Degree", value=1)
-                    scaling_modifier = gr.Slider(minimum=0.1, maximum=2.0, step=0.1, label="Scaling Modifier", value=1.0)
                 with gr.Row():
                     submit = gr.Button("Generate", elem_id="generate", variant="primary")
-
-                with gr.Row(variant="panel"):
+                with gr.Row(variant="panel"): 
                     gr.Examples(
                         examples=[
                             './demo_examples/bedroom_01.png',
@@ -117,22 +84,23 @@ def main():
                         label="Examples",
                         examples_per_page=20,
                     )
-
                 with gr.Row():
                     processed_image = gr.Image(label="Processed Image", interactive=False)
-
             with gr.Column(scale=2):
                 with gr.Row():
                     with gr.Tab("Reconstruction"):
                         output_model = gr.Model3D(height=512, label="Output Model", interactive=False)
+                with gr.Row():
+                    resolution = gr.Slider(minimum=256, maximum=1024, step=64, label="Image Resolution", value=cfg.dataset.height)
+                    num_gauss = gr.Slider(minimum=1, maximum=10, step=1, label="Number of Gaussian Components", value=2)
 
         submit.click(fn=check_input_image, inputs=[input_image]).success(
             fn=preprocess,
-            inputs=[input_image, padding_value],
+            inputs=[input_image, resolution],
             outputs=[processed_image],
         ).success(
             fn=reconstruct_and_export,
-            inputs=[processed_image, num_gauss, max_sh_degree, scaling_modifier],
+            inputs=[processed_image, num_gauss],
             outputs=[output_model],
         )