46 lines
1.8 KiB
Python
46 lines
1.8 KiB
Python
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) |