from zoedepth.utils.config import get_config
from zoedepth.models.builder import build_model

import gradio as gr
import torch

from depth import depth_interface
from mesh import mesh_interface


css = """
#img-display-container {
    max-height: 50vh;
    }
#img-display-input {
    max-height: 40vh;
    }
#img-display-output {
    max-height: 40vh;
    }
    
"""
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
model = torch.hub.load('isl-org/ZoeDepth', "ZoeD_N", pretrained=True).to("cpu").eval()

# config_mode="infer"
# pretrained_resource = f"local::C:/Users/Charl/.cache/torch/hub/checkpoints/ZoeD_M12_N.pt"
# config = get_config("zoedepth", config_mode, pretrained_resource=pretrained_resource)
# model = build_model(config).to(DEVICE).eval()

# title = "# ZoeDepth"
# description = """Official demo for **ZoeDepth: Zero-shot Transfer by Combining Relative and Metric Depth**."""

with gr.Blocks(css=css) as API:
    # gr.Markdown(title)
    # gr.Markdown(description)
    with gr.Tab("Depth Prediction"):
        depth_interface(model, DEVICE)
    with gr.Tab("Image to 3D"):
        mesh_interface(model, DEVICE)
    # with gr.Tab("360 Panorama to 3D"):
    #     create_pano_to_3d_demo(model)

if __name__ == '__main__':
    API.launch()