Update app.py

app.py

@@ -25,9 +25,11 @@ 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...")
@@ -59,7 +61,10 @@ 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.")
@@ -67,15 +72,16 @@ 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, num_gauss, batch_size, num_iterations):
+    def reconstruct_and_export(image):
+        """
+        Passes image through model, outputs reconstruction in form of a dict of tensors.
+        """
         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}
-
-        # 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
+        inputs = {
+            ("color_aug", 0, 0): image,
+        }
 
         # Pass the image through the model to get the output
         print("[INFO] Passing image through the model...")
@@ -83,11 +89,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=num_gauss)
+        save_ply(outputs, ply_out_path, num_gauss=2)
         print("[INFO] Reconstruction and export complete.")
 
         return ply_out_path
-
+    
     # Path to save the output PLY file
     ply_out_path = f'./mesh.ply'
 
@@ -101,15 +107,26 @@ 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(
@@ -126,18 +143,20 @@ 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)
-                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)
+                        output_model = gr.Model3D(
+                            height=512,
+                            label="Output Model",
+                            interactive=False
+                        )
 
         # Define the workflow for the Generate button
         submit.click(fn=check_input_image, inputs=[input_image]).success(
@@ -146,7 +165,7 @@ def main():
             outputs=[processed_image],
         ).success(
             fn=reconstruct_and_export,
-            inputs=[processed_image, num_gauss, batch_size, num_iterations],
+            inputs=[processed_image],
             outputs=[output_model],
         )