Spaces:
Running on L40S
Running on L40S
File size: 1,829 Bytes
84bb8a3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 | 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) |