import os
import sys
ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
if ROOT_DIR not in sys.path:
    sys.path.insert(0, ROOT_DIR)

from trellis2 import models
from color_glb import color_glb
from glb_to_vxz import glb_to_vxz
from vxz_to_slat import vxz_to_slat
from img_to_cond import img_to_cond
from glb_to_parts import glb_to_parts
from bpy_render import render_from_transforms
from trellis2.pipelines.rembg import BiRefNet
from trellis2.modules.image_feature_extractor import DinoV3FeatureExtractor

rembg_model = BiRefNet(model_name="/root/autodl-tmp/RMBG-2.0")
rembg_model.cuda()
image_cond_model = DinoV3FeatureExtractor(model_name="/root/autodl-tmp/dinov3-vitl16-pretrain-lvd1689m")
image_cond_model.cuda()

shape_encoder = models.from_pretrained("/root/autodl-tmp/TRELLIS.2-4B/ckpts/shape_enc_next_dc_f16c32_fp16").cuda().eval()
tex_encoder = models.from_pretrained("/root/autodl-tmp/TRELLIS.2-4B/ckpts/tex_enc_next_dc_f16c32_fp16").cuda().eval()

glb = "./assets/example.glb"
input_vxz = "./assets/input.vxz"
parts_path = "./assets/parts"
interactive_seg_path = "./assets/interactive_seg"
interactive = True

transforms = "transforms.json"
img = "./assets/img.png"
cond = "./assets/cond.pth"

glb_to_vxz(glb, input_vxz)
glb_to_parts(glb, parts_path)
color_glb(parts_path, interactive_seg_path, interactive)

render_from_transforms(glb, transforms, img)
img_to_cond(rembg_model, image_cond_model, img, cond)
for part_name in sorted(os.listdir(interactive_seg_path)):
    part_path = os.path.join(interactive_seg_path, part_name)
    output_glb_path = os.path.join(part_path, "output.glb")
    output_vxz_path = os.path.join(part_path, "output.vxz")
    glb_to_vxz(output_glb_path, output_vxz_path)
    vxz_to_slat(shape_encoder, tex_encoder, input_vxz, output_vxz_path, part_path, interactive)