import shlex
import subprocess
import os
import sys
from huggingface_hub import snapshot_download
import torch
import fire
import gradio as gr
from gradio_app.gradio_3dgen import create_ui as create_3d_ui
from gradio_app.all_models import model_zoo

# Install required packages
def setup_dependencies():
    subprocess.run(shlex.split("pip install pip==24.0"), check=True)
    subprocess.run(
        shlex.split(
            "pip install package/onnxruntime_gpu-1.17.0-cp310-cp310-manylinux_2_28_x86_64.whl --force-reinstall --no-deps"
        ),
        check=True
    )
    subprocess.run(
        shlex.split(
            "pip install package/nvdiffrast-0.3.1.torch-cp310-cp310-linux_x86_64.whl --force-reinstall --no-deps"
        ),
        check=True
    )

# Download model checkpoints
def setup_model():
    snapshot_download("public-data/Unique3D", repo_type="model", local_dir="./ckpt")
    
    # Configure PyTorch settings
    torch.set_float32_matmul_precision('medium')
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.set_grad_enabled(False)

# Application title
_TITLE = 'Text to 3D'

def launch():
    # Initialize models
    model_zoo.init_models()
    
    # Create Gradio interface
    with gr.Blocks(title=_TITLE) as demo:
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown('# ' + _TITLE)
                create_3d_ui("wkl")
        
        demo.queue().launch(share=True)

if __name__ == '__main__':
    setup_dependencies()
    setup_model()
    sys.path.append(os.curdir)
    fire.Fire(launch)