import os
import shutil 

import gradio as gr 
from gradio_litmodel3d import LitModel3D 
TMP_DIR = os.path.join(
    os.path.dirname(os.path.abspath(__file__)), "sessions/imageto3d"
)
os.makedirs(TMP_DIR, exist_ok=True)

# RBG_REMOVER = RembgRemover()
# SAM_PREDICTOR = SAMPredictor(model_type="vit_h")
# DELIGHT = DelightingModel()
# IMAGESR_MODEL = ImageRealESRGAN(outscale=4)
# PIPELINE = TrellisImageTo3DPipeline.from_pretrained(
#     "JeffreyXiang/TRELLIS-image-large"
# )
# # PIPELINE.cuda()

# IMAGE_BUFFER = {}
# SEG_CHECKER = ImageSegChecker(GPT_CLIENT)
# GEO_CHECKER = MeshGeoChecker(GPT_CLIENT)
# AESTHETIC_CHECKER = ImageAestheticChecker()
# CHECKERS = [GEO_CHECKER, SEG_CHECKER, AESTHETIC_CHECKER]
# URDF_CONVERTOR = URDFGenerator(GPT_CLIENT, render_view_num=4)


def start_session(req: gr.Request) -> None:
    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
    os.makedirs(user_dir, exist_ok=True)


def end_session(req: gr.Request) -> None:
    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
    if os.path.exists(user_dir):
        shutil.rmtree(user_dir)


with gr.Blocks(
    delete_cache=(43200, 43200)
) as demo:
    with gr.Column(scale=1):
        video_output = gr.Video(
            label="Generated 3D Asset",
            autoplay=True,
            loop=True,
            height=300,
        )
        model_output_gs = gr.Model3D(
            label="Gaussian Representation", height=300, interactive=False
        )
        aligned_gs = gr.Textbox(visible=False)
    
        model_output_mesh = LitModel3D(
            label="Mesh Representation",
            # height=300,
            exposure=10,
            interactive=False
        )
        gr.Markdown(
            """ The rendering of `Gaussian Representation` takes additional 10s. """  # noqa
        )

    demo.load(start_session)
    demo.unload(end_session)



if __name__ == "__main__":
    demo.launch()