def core_fn(image: str, source_cam_dist: float, working_dir):
    dump_video_path = os.path.join(working_dir.name, "output.mp4")
    dump_mesh_path = os.path.join(working_dir.name, "output.ply")
    infer_impl(
        image_path=image,
        source_cam_dist=source_cam_dist,
        export_video=True,
        export_mesh=True,
        dump_video_path=dump_video_path,
        dump_mesh_path=dump_mesh_path,
    )
    return dump_video_path
    return dump_mesh_path

def example_fn(image: np.ndarray):
    from gradio.utils import get_cache_folder
    working_dir = get_cache_folder()
    image = preprocess_fn(
        image_in=image,
        remove_bg=True,
        recenter=True,
        working_dir=working_dir,
    )
    video = core_fn(
        image=image,
        source_cam_dist=2.0,
        working_dir=working_dir,
    )
    mesh = core_fn(
        mesh=mesh,
        working_dir=working_dir,
    )
    return image, video, mesh


_TITLE = '''OpenLRM: Open-Source Large Reconstruction Models'''

_DESCRIPTION = '''
    <div>
        <a style="display:inline-block" href='https://github.com/3DTopia/OpenLRM'><img src='https://img.shields.io/github/stars/3DTopia/OpenLRM?style=social'/></a>
        <a style="display:inline-block; margin-left: .5em" href="https://huggingface.co/zxhezexin"><img src='https://img.shields.io/badge/Model-Weights-blue'/></a>
    </div>
    OpenLRM is an open-source implementation of Large Reconstruction Models.

    <strong>Image-to-3D in 10 seconds with A100!</strong>

    <strong>Disclaimer:</strong> This demo uses `openlrm-mix-base-1.1` model with 288x288 rendering resolution here for a quick demonstration.
'''

with gr.Blocks(analytics_enabled=False) as demo:

    # HEADERS
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown('# ' + _TITLE)
    with gr.Row():
        gr.Markdown(_DESCRIPTION)

    # DISPLAY
    with gr.Row():

        with gr.Column(variant='panel', scale=1):
            with gr.Tabs(elem_id="openlrm_input_image"):
                with gr.TabItem('Input Image'):
                    with gr.Row():
                        input_image = gr.Image(label="Input Image", image_mode="RGBA", width="auto", sources="upload", type="numpy", elem_id="content_image")

        with gr.Column(variant='panel', scale=1):
            with gr.Tabs(elem_id="openlrm_processed_image"):
                with gr.TabItem('Processed Image'):
                    with gr.Row():
                        processed_image = gr.Image(label="Processed Image", image_mode="RGBA", type="filepath", elem_id="processed_image", width="auto", interactive=False)

        with gr.Column(variant='panel', scale=1):
            with gr.Tabs(elem_id="openlrm_render_video"):
                with gr.TabItem('Rendered Video'):
                    with gr.Row():
                        output_video = gr.Video(label="Rendered Video", format="mp4", width="auto", autoplay=True)

        with gr.Column(variant='panel', scale=1):
            with gr.Tabs(elem_id="openlrm_render_mesh"):
                with gr.TabItem('Rendered Mesh'):
                    with gr.Row():
                        output_mesh = gr.Mesh(label="Rendered Mesh", format="ply", width="auto", autoplay=True)

    # SETTING
    with gr.Row():
        with gr.Column(variant='panel', scale=1):
            with gr.Tabs(elem_id="openlrm_attrs"):
                with gr.TabItem('Settings'):
                    with gr.Column(variant='panel'):
                        gr.Markdown(
                            """
                            <strong>Best Practice</strong>:
                                Centered objects in reasonable sizes. Try adjusting source camera distances.
                            """
                        )
                        checkbox_rembg = gr.Checkbox(True, label='Remove background')
                        checkbox_recenter = gr.Checkbox(True, label='Recenter the object')
                        slider_cam_dist = gr.Slider(1.0, 3.5, value=2.0, step=0.1, label="Source Camera Distance")
                        submit = gr.Button('Generate', elem_id="openlrm_generate", variant='primary')

    # EXAMPLES
    with gr.Row():
        examples = [
            ['assets/sample_input/owl.png'],
            ['assets/sample_input/building.png'],
            ['assets/sample_input/mailbox.png'],
            ['assets/sample_input/fire.png'],
            ['assets/sample_input/girl.png'],
            ['assets/sample_input/lamp.png'],
            ['assets/sample_input/hydrant.png'],
            ['assets/sample_input/hotdogs.png'],
            ['assets/sample_input/traffic.png'],
            ['assets/sample_input/ceramic.png'],
        ]
        gr.Examples(
            examples=examples,
            inputs=[input_image],
            outputs=[processed_image, output_video],
            fn=example_fn,
            cache_examples=bool(os.getenv('SPACE_ID')),
            examples_per_page=20,
        )

    working_dir = gr.State()
    submit.click(
        fn=assert_input_image,
        inputs=[input_image],
        queue=False,
    ).success(
        fn=prepare_working_dir,
        outputs=[working_dir],
        queue=False,
    ).success(
        fn=preprocess_fn,
        inputs=[input_image, checkbox_rembg, checkbox_recenter, working_dir],
        outputs=[processed_image],
    ).success(
        fn=core_fn,
        inputs=[processed_image, slider_cam_dist, working_dir],
        outputs=[output_video],
    ).success(
        fn=core_fn,
        inputs=[processed_mesh, working_dir],
        outputs=[output_mesh],
    )

    demo.queue()
    demo.launch(share = True)

os.environ.update({
    "APP_ENABLED": "1",
    "APP_MODEL_NAME": "zxhezexin/openlrm-mix-base-1.1",
    "APP_INFER": "./configs/infer-gradio.yaml",
    "APP_TYPE": "infer.lrm",
    "NUMBA_THREADING_LAYER": 'omp',
})

from openlrm.runners import REGISTRY_RUNNERS
from openlrm.runners.infer.base_inferrer import Inferrer
InferrerClass : Inferrer = REGISTRY_RUNNERS[os.getenv("APP_TYPE")]
with InferrerClass() as inferrer:
    init_preprocessor()
    if not bool(os.getenv('SPACE_ID')):
        from openlrm.utils.proxy import no_proxy
        demo = no_proxy(demo_openlrm)
    else:
        demo = demo_openlrm
    demo(infer_impl=inferrer.infer_single)

launch_gradio_app()