Update app.py

Exposed a few parameters :)

app.py

@@ -25,11 +25,9 @@ def main():
 
     # Download model configuration and weights from Hugging Face Hub
     print("[INFO] Downloading model configuration...")
-    model_cfg_path = hf_hub_download(repo_id="einsafutdinov/flash3d", 
-                                     filename="config_re10k_v1.yaml")
+    model_cfg_path = hf_hub_download(repo_id="einsafutdinov/flash3d", filename="config_re10k_v1.yaml")
     print("[INFO] Downloading model weights...")
-    model_path = hf_hub_download(repo_id="einsafutdinov/flash3d", 
-                                 filename="model_re10k_v1.pth")
+    model_path = hf_hub_download(repo_id="einsafutdinov/flash3d", filename="model_re10k_v1.pth")
 
     # Load model configuration using OmegaConf
     print("[INFO] Loading model configuration...")
@@ -61,10 +59,7 @@ def main():
     def preprocess(image):
         print("[DEBUG] Preprocessing image...")
         # Resize the image to the desired height and width specified in the configuration
-        image = TTF.resize(
-            image, (cfg.dataset.height, cfg.dataset.width), 
-            interpolation=TT.InterpolationMode.BICUBIC
-        )
+        image = TTF.resize(image, (cfg.dataset.height, cfg.dataset.width), interpolation=TT.InterpolationMode.BICUBIC)
         # Apply padding to the image
         image = pad_border_fn(image)
         print("[INFO] Image preprocessing complete.")
@@ -72,16 +67,15 @@ def main():
 
     # Function to reconstruct the 3D model from the input image 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(image):
-        """
-        Passes image through model, outputs reconstruction in form of a dict of tensors.
-        """
+    def reconstruct_and_export(image, num_gauss, batch_size, num_iterations):
         print("[DEBUG] Starting reconstruction and export...")
         # Convert the preprocessed image to a tensor and move it to the specified device
         image = to_tensor(image).to(device).unsqueeze(0)
-        inputs = {
-            ("color_aug", 0, 0): image,
-        }
+        inputs = {("color_aug", 0, 0): image}
+
+        # Set the batch size and number of iterations in the model configuration
+        model.cfg.dataset.batch_size = batch_size
+        model.cfg.training.num_iterations = num_iterations
 
         # Pass the image through the model to get the output
         print("[INFO] Passing image through the model...")
@@ -89,11 +83,11 @@ def main():
 
         # Export the reconstruction to a PLY file
         print(f"[INFO] Saving output to {ply_out_path}...")
-        save_ply(outputs, ply_out_path, num_gauss=2)
+        save_ply(outputs, ply_out_path, num_gauss=num_gauss)
         print("[INFO] Reconstruction and export complete.")
 
         return ply_out_path
-    
+
     # Path to save the output PLY file
     ply_out_path = f'./mesh.ply'
 
@@ -107,26 +101,15 @@ def main():
 
     # Create the Gradio user interface
     with gr.Blocks(css=css) as demo:
-        gr.Markdown(
-            """
-            # Flash3D
-            """
-        )
+        gr.Markdown("# Flash3D")
         with gr.Row(variant="panel"):
             with gr.Column(scale=1):
                 with gr.Row():
                     # Input image component for the user to upload an image
-                    input_image = gr.Image(
-                        label="Input Image",
-                        image_mode="RGBA",
-                        sources="upload",
-                        type="pil",
-                        elem_id="content_image",
-                    )
+                    input_image = gr.Image(label="Input Image", image_mode="RGBA", sources="upload", type="pil", elem_id="content_image")
                 with gr.Row():
                     # Button to trigger the generation process
                     submit = gr.Button("Generate", elem_id="generate", variant="primary")
-
                 with gr.Row(variant="panel"): 
                     # Examples panel to provide sample images for users
                     gr.Examples(
@@ -143,20 +126,18 @@ def main():
                         label="Examples",
                         examples_per_page=20,
                     )
-
                 with gr.Row():
                     # Display the preprocessed image (after resizing and padding)
                     processed_image = gr.Image(label="Processed Image", interactive=False)
-
             with gr.Column(scale=2):
                 with gr.Row():
                     with gr.Tab("Reconstruction"):
                         # 3D model viewer to display the reconstructed model
-                        output_model = gr.Model3D(
-                            height=512,
-                            label="Output Model",
-                            interactive=False
-                        )
+                        output_model = gr.Model3D(height=512, label="Output Model", interactive=False)
+                with gr.Row():
+                    num_gauss = gr.Slider(minimum=1, maximum=10, step=1, label="Number of Gaussian Components", value=2)
+                    batch_size = gr.Slider(minimum=1, maximum=32, step=1, label="Batch Size", value=1)
+                    num_iterations = gr.Slider(minimum=1, maximum=1000, step=10, label="Number of Iterations", value=100)
 
         # Define the workflow for the Generate button
         submit.click(fn=check_input_image, inputs=[input_image]).success(
@@ -165,7 +146,7 @@ def main():
             outputs=[processed_image],
         ).success(
             fn=reconstruct_and_export,
-            inputs=[processed_image],
+            inputs=[processed_image, num_gauss, batch_size, num_iterations],
             outputs=[output_model],
         )