import gradio as gr
import spaces
import torch
from gradio_rerun import Rerun
import rerun as rr
from pathlib import Path


from mini_dust3r.api import OptimizedResult, inferece_dust3r, log_optimized_result
from mini_dust3r.model import AsymmetricCroCo3DStereo

DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
model = AsymmetricCroCo3DStereo.from_pretrained(
        "naver/DUSt3R_ViTLarge_BaseDecoder_512_dpt"
    ).to(DEVICE)

@spaces.GPU
def predict(image_dir: str):
    rr.init("my data")
    optimized_results: OptimizedResult = inferece_dust3r(
        image_dir=image_dir,
        model=model,
        device=DEVICE,
        batch_size=1,
    )
    log_optimized_result(optimized_results, Path("world"))
    rr.save("dust3r.rrd")
    return "dust3r.rrd"

with gr.Blocks(css=""".gradio-container {margin: 0 !important; min-width: 100%};""", title="DUSt3R Demo") as demo:
    # scene state is save so that you can change conf_thr, cam_size... without rerunning the inference
    gr.HTML('<h2 style="text-align: center;">DUSt3R Demo</h2>')
    with gr.Column():
        inputfiles = gr.File(file_count="multiple")
        rerun_viewer = Rerun(height=900)

        run_btn = gr.Button("Run")
        run_btn.click(fn=predict,
                        inputs=[inputfiles],
                        outputs=[rerun_viewer])
        
demo.launch()