app.py

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)
    try:
        device = torch.device(device)
        model.to(device)  # Move the model to the specified device (CPU or GPU)
    except Exception as e:
        print(f"[ERROR] Failed to set device: {e}")
        raise
    
    # 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, padding_value, resize_height, resize_width):
        print("[DEBUG] Preprocessing image...")
        # Resize the image to the desired height and width specified in the user input
        image = TTF.resize(
            image, (resize_height, resize_width), 
            interpolation=TT.InterpolationMode.BICUBIC
        )
        # Apply padding to the image
        pad_border_fn = TT.Pad((padding_value, padding_value))
        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, num_gauss, scale_factor):
        """
        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} with scale factor {scale_factor}...")
        save_ply(outputs, ply_out_path, num_gauss=num_gauss, scale_factor=scale_factor)
        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
            """
        )
        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():
                    # Sliders for configurable parameters
                    num_gauss = gr.Slider(minimum=1, maximum=20, step=1, label="Number of Gaussians per Pixel", value=10)
                    scale_factor = gr.Slider(minimum=0.5, maximum=5.0, step=0.1, label="Scale Factor for Model Size", value=1.5, info="Test this range for stability, as extreme values may cause visual distortions or unexpected outputs.")
                    padding_value = gr.Slider(minimum=0, maximum=128, step=8, label="Padding Amount for Output Processing", value=32)
                    resize_height = gr.Slider(minimum=256, maximum=1024, step=64, label="Resize Height for Image", value=cfg.dataset.height)
                    resize_width = gr.Slider(minimum=256, maximum=1024, step=64, label="Resize Width for Image", value=cfg.dataset.width)
                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, padding_value, resize_height, resize_width],
            outputs=[processed_image],
        ).success(
            fn=reconstruct_and_export,
            inputs=[processed_image, num_gauss, scale_factor],
            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()