Set GPU duration to maximum allowed value for ZeroGPU spaces

- Updated the `@spaces.GPU` decorator to allocate GPU for the maximum allowed duration of 120 seconds.
- This change was made based on findings regarding the limitations of ZeroGPU spaces, where the maximum supported duration is 120 seconds for function execution.
- Ensured that the GPU allocation is optimal to prevent errors and align with Hugging Face ZeroGPU's current constraints.

This update ensures that the GPU resource is utilized effectively for the given time without exceeding the allocation limits, which might lead to failures or unfulfilled GPU requests during extended operations.

app.py

@@ -1,6 +1,6 @@
 import sys
 import spaces
-sys.path.append("flash3d")
+sys.path.append("flash3d")  # Add the flash3d directory to the system path for importing local modules
 
 from omegaconf import OmegaConf
 import gradio as gr
@@ -8,61 +8,295 @@ 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...")
+    # Determine if CUDA (GPU) is available and set the device accordingly
     if torch.cuda.is_available():
         device = "cuda:0"
+        print("[INFO] CUDA is available. Using GPU device.")
     else:
         device = "cpu"
+        print("[INFO] CUDA is not available. Using CPU device.")
 
+    # 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")
+    print("[INFO] Downloading model weights...")
     model_path = hf_hub_download(repo_id="einsafutdinov/flash3d", 
                                  filename="model_re10k_v1.pth")
 
+    # Load model configuration using OmegaConf
+    print("[INFO] Loading model configuration...")
     cfg = OmegaConf.load(model_cfg_path)
+    
+    # Initialize the GaussianPredictor model with the loaded configuration
+    print("[INFO] Initializing GaussianPredictor model...")
     model = GaussianPredictor(cfg)
     device = torch.device(device)
-    model.to(device)
+    model.to(device)  # Move the model to the specified device (CPU or GPU)
+    
+    # Load the pre-trained model weights
+    print("[INFO] Loading model weights...")
     model.load_model(model_path)
 
-    pad_border_fn = TT.Pad((cfg.dataset.pad_border_aug, cfg.dataset.pad_border_aug))
-    to_tensor = TT.ToTensor()
+    # Define transformation functions for image preprocessing
+    pad_border_fn = TT.Pad((cfg.dataset.pad_border_aug, cfg.dataset.pad_border_aug))  # Padding to augment the image borders
+    to_tensor = TT.ToTensor()  # Convert image to tensor
 
+    # Function to check if an image is uploaded by the user
     def check_input_image(input_image):
+        print("[DEBUG] Checking input image...")
         if input_image is None:
+            print("[ERROR] No image uploaded!")
             raise gr.Error("No image uploaded!")
+        print("[INFO] Input image is valid.")
 
+    # Function to preprocess the input image before passing it to the model
     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
         )
+        # Apply padding to the image
         image = pad_border_fn(image)
+        print("[INFO] Image preprocessing complete.")
         return image
 
-    @spaces.GPU(duration=120)
+    # Function to reconstruct the 3D model from the input image and export it as a PLY file
+    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...")
+    # Determine if CUDA (GPU) is available and set the device accordingly
+    if torch.cuda.is_available():
+        device = "cuda:0"
+        print("[INFO] CUDA is available. Using GPU device.")
+    else:
+        device = "cpu"
+        print("[INFO] CUDA is not available. Using CPU device.")
+
+    # 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")
+    print("[INFO] Downloading model weights...")
+    model_path = hf_hub_download(repo_id="einsafutdinov/flash3d", 
+                                 filename="model_re10k_v1.pth")
+
+    # Load model configuration using OmegaConf
+    print("[INFO] Loading model configuration...")
+    cfg = OmegaConf.load(model_cfg_path)
+    
+    # Initialize the GaussianPredictor model with the loaded configuration
+    print("[INFO] Initializing GaussianPredictor model...")
+    model = GaussianPredictor(cfg)
+    device = torch.device(device)
+    model.to(device)  # Move the model to the specified device (CPU or GPU)
+    
+    # Load the pre-trained model weights
+    print("[INFO] Loading model weights...")
+    model.load_model(model_path)
+
+    # Define transformation functions for image preprocessing
+    pad_border_fn = TT.Pad((cfg.dataset.pad_border_aug, cfg.dataset.pad_border_aug))  # Padding to augment the image borders
+    to_tensor = TT.ToTensor()  # Convert image to tensor
+
+    # Function to check if an image is uploaded by the user
+    def check_input_image(input_image):
+        print("[DEBUG] Checking input image...")
+        if input_image is None:
+            print("[ERROR] No image uploaded!")
+            raise gr.Error("No image uploaded!")
+        print("[INFO] Input image is valid.")
+
+    # Function to preprocess the input image before passing it to the model
+    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
+        )
+        # Apply padding to the image
+        image = pad_border_fn(image)
+        print("[INFO] Image preprocessing complete.")
+        return image
+
+    # 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.
+        """
+        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,
+        }
+
+        # Pass the image through the model to get the output
+        print("[INFO] Passing image through the model...")
+        outputs = model(inputs)
+
+        # 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)
+        print("[INFO] Reconstruction and export complete.")
+
+        return ply_out_path
+    
+    # Path to save the output PLY file
+    ply_out_path = f'./mesh.ply'
+
+    # CSS styling for the Gradio interface
+    css = """
+        h1 {
+            text-align: center;
+            display:block;
+        }
+        """
+
+    # Create the Gradio user interface
+    with gr.Blocks(css=css) as demo:
+        gr.Markdown(
+            """
+            # Flash3D
+            """
+        )
+        # Comments about the app's behavior and known limitations
+        gr.Markdown(
+            """
+            ## Comments:
+            1. If you run the demo online, the first example you upload should take about 4.5 seconds (with preprocessing, saving and overhead), the following take about 1.5s.
+            2. The 3D viewer shows a .ply mesh extracted from a mix of 3D Gaussians. This is only an approximation and artifacts might show.
+            3. Known limitations include:
+            - A black dot appearing on the model from some viewpoints.
+            - See-through parts of objects, especially on the back: this is due to the model performing less well on more complicated shapes.
+            - Back of objects are blurry: this is a model limitation due to it being deterministic.
+            4. Our model is of comparable quality to state-of-the-art methods, and is **much** cheaper to train and run.
+            ## How does it work?
+            Splatter Image formulates 3D reconstruction as an image-to-image translation task. It maps the input image to another image, 
+            in which every pixel represents one 3D Gaussian and the channels of the output represent parameters of these Gaussians, including their shapes, colours, and locations.
+            The resulting image thus represents a set of Gaussians (almost like a point cloud) which reconstruct the shape and colour of the object.
+            The method is very cheap: the reconstruction amounts to a single forward pass of a neural network with only 2D operators (2D convolutions and attention).
+            The rendering is also very fast, due to using Gaussian Splatting.
+            Combined, this results in very cheap training and high-quality results.
+            For more results see the [project page](https://szymanowiczs.github.io/splatter-image) and the [CVPR article](https://arxiv.org/abs/2312.13150).
+            """
+        )
+        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",
+                    )
+                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(
+                        examples=[
+                            './demo_examples/bedroom_01.png',
+                            './demo_examples/kitti_02.png',
+                            './demo_examples/kitti_03.png',
+                            './demo_examples/re10k_04.jpg',
+                            './demo_examples/re10k_05.jpg',
+                            './demo_examples/re10k_06.jpg',
+                        ],
+                        inputs=[input_image],
+                        cache_examples=False,
+                        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
+                        )
+
+        # Define the workflow for the Generate button
+        submit.click(fn=check_input_image, inputs=[input_image]).success(
+            fn=preprocess,
+            inputs=[input_image],
+            outputs=[processed_image],
+        ).success(
+            fn=reconstruct_and_export,
+            inputs=[processed_image],
+            outputs=[output_model],
+        )
+
+    # Queue the requests to handle them sequentially (to avoid GPU resource conflicts)
+    demo.queue(max_size=1)
+    print("[INFO] Launching Gradio demo...")
+    demo.launch(share=True)  # Launch the Gradio interface and allow public sharing
+
+if __name__ == "__main__":
+    print("[INFO] Running application...")
+    main()  # 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.
         """
+        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,
         }
 
+        # Pass the image through the model to get the output
+        print("[INFO] Passing image through the model...")
         outputs = model(inputs)
 
-        # export reconstruction to ply
+        # 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)
+        print("[INFO] Reconstruction and export complete.")
 
         return ply_out_path
     
+    # Path to save the output PLY file
     ply_out_path = f'./mesh.ply'
 
+    # CSS styling for the Gradio interface
     css = """
         h1 {
             text-align: center;
@@ -70,15 +304,38 @@ def main():
         }
         """
 
+    # Create the Gradio user interface
     with gr.Blocks(css=css) as demo:
         gr.Markdown(
             """
             # Flash3D
             """
-            )
+        )
+        # Comments about the app's behavior and known limitations
+        gr.Markdown(
+            """
+            ## Comments:
+            1. If you run the demo online, the first example you upload should take about 4.5 seconds (with preprocessing, saving and overhead), the following take about 1.5s.
+            2. The 3D viewer shows a .ply mesh extracted from a mix of 3D Gaussians. This is only an approximation and artifacts might show.
+            3. Known limitations include:
+            - A black dot appearing on the model from some viewpoints.
+            - See-through parts of objects, especially on the back: this is due to the model performing less well on more complicated shapes.
+            - Back of objects are blurry: this is a model limitation due to it being deterministic.
+            4. Our model is of comparable quality to state-of-the-art methods, and is **much** cheaper to train and run.
+            ## How does it work?
+            Splatter Image formulates 3D reconstruction as an image-to-image translation task. It maps the input image to another image, 
+            in which every pixel represents one 3D Gaussian and the channels of the output represent parameters of these Gaussians, including their shapes, colours, and locations.
+            The resulting image thus represents a set of Gaussians (almost like a point cloud) which reconstruct the shape and colour of the object.
+            The method is very cheap: the reconstruction amounts to a single forward pass of a neural network with only 2D operators (2D convolutions and attention).
+            The rendering is also very fast, due to using Gaussian Splatting.
+            Combined, this results in very cheap training and high-quality results.
+            For more results see the [project page](https://szymanowiczs.github.io/splatter-image) and the [CVPR article](https://arxiv.org/abs/2312.13150).
+            """
+        )
         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",
@@ -87,9 +344,11 @@ def main():
                         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(
                         examples=[
                             './demo_examples/bedroom_01.png',
@@ -106,17 +365,20 @@ def main():
                     )
 
                 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
                         )
 
+        # Define the workflow for the Generate button
         submit.click(fn=check_input_image, inputs=[input_image]).success(
             fn=preprocess,
             inputs=[input_image],
@@ -127,8 +389,11 @@ def main():
             outputs=[output_model],
         )
 
+    # Queue the requests to handle them sequentially (to avoid GPU resource conflicts)
     demo.queue(max_size=1)
-    demo.launch(share=True)
+    print("[INFO] Launching Gradio demo...")
+    demo.launch(share=True)  # Launch the Gradio interface and allow public sharing
 
 if __name__ == "__main__":
+    print("[INFO] Running application...")
     main()