import sys
sys.path.append('wild-gaussian-splatting/mast3r/')
sys.path.append('demo/')

import os
import tempfile
import gradio as gr
from mast3r.demo import get_args_parser
from mast3r.utils.misc import hash_md5
from mast3r_demo import mast3r_demo_tab
from gs_demo import gs_demo_tab

if __name__ == '__main__':
    parser = get_args_parser()
    args = parser.parse_args()

    if args.server_name is not None:
        server_name = args.server_name
    else:
        server_name = '0.0.0.0' if args.local_network else '127.0.0.1'

    weights_path = args.weights if args.weights is not None else "naver/" + args.model_name
    chkpt_tag = hash_md5(weights_path)

    with tempfile.TemporaryDirectory(suffix='demo') as tmpdirname:
        cache_path = os.path.join(tmpdirname, chkpt_tag)
        os.makedirs(cache_path, exist_ok=True)

        with gr.Blocks() as demo:
            with gr.Tabs():
                with gr.Tab("MASt3R Demo"):
                    mast3r_demo_tab(cache_path, weights_path, args.device)
                with gr.Tab("Gaussian Splatting Demo"):
                    gs_demo_tab(cache_path)

            demo.launch(server_name=server_name, server_port=args.server_port) 

# python3 demo.py --weights "/app/mast3r/checkpoints/MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric.pth" --device "cuda" --server_port 3334 --local_network "$@"