207 lines
8.1 KiB
Python
207 lines
8.1 KiB
Python
# !/usr/bin/env python
|
|
# -*- coding: UTF-8 -*-
|
|
import numpy as np
|
|
import gc
|
|
import os
|
|
import torch
|
|
import uuid
|
|
from .app import image_to_3d
|
|
from .trellis.pipelines import TrellisImageTo3DPipeline
|
|
from .utils import glb2obj_,obj2fbx_,tensor2imglist,pre_img
|
|
import folder_paths
|
|
|
|
MAX_SEED = np.iinfo(np.int32).max
|
|
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
|
|
current_path = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
weigths_dinov2_current_path = os.path.join(folder_paths.models_dir, "dinov2")
|
|
if not os.path.exists(weigths_dinov2_current_path):
|
|
os.makedirs(weigths_dinov2_current_path)
|
|
|
|
try:
|
|
folder_paths.add_model_folder_path("dinov2", weigths_dinov2_current_path, False)
|
|
except:
|
|
folder_paths.add_model_folder_path("dinov2", weigths_dinov2_current_path)
|
|
|
|
class Trellis_LoadModel:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"repo": ("STRING", {"default": "JeffreyXiang/TRELLIS-image-large"}),
|
|
"dino": (["none"] + folder_paths.get_filename_list("dinov2"),),
|
|
"attn_backend":(["xformers","flash-attn"],),
|
|
"spconv_algo":(["auto","flash-native"],),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL_TRELLIS", )
|
|
RETURN_NAMES = ("model",)
|
|
FUNCTION = "main_loader"
|
|
CATEGORY = "Trellis"
|
|
|
|
def main_loader(self, repo,dino,attn_backend,spconv_algo):
|
|
if attn_backend=="xformers":
|
|
os.environ['ATTN_BACKEND'] = 'xformers'
|
|
else:
|
|
os.environ['ATTN_BACKEND'] = 'flash-attn'
|
|
if spconv_algo=="auto":
|
|
os.environ['SPCONV_ALGO'] = 'auto'
|
|
else:
|
|
os.environ['SPCONV_ALGO'] = 'native'
|
|
|
|
if dino=="none":
|
|
raise "need choice dinov2 checkpoint"
|
|
|
|
TrellisImageTo3DPipeline.dino=folder_paths.get_full_path("dinov2", dino)
|
|
TrellisImageTo3DPipeline.dino_moudel=os.path.join(current_path,"facebookresearch/dinov2")
|
|
if repo:
|
|
model=TrellisImageTo3DPipeline.from_pretrained(repo)
|
|
else:
|
|
raise "need fill repo"
|
|
return (model,)
|
|
|
|
|
|
class Trellis_Sampler:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",), # [B,H,W,C], C=3
|
|
"model": ("MODEL_TRELLIS",),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": MAX_SEED}),
|
|
"cfg": ("FLOAT", {"default": 7.5, "min": 0.0, "max": 10.0, "step": 0.1, "round": 0.01}),
|
|
"steps": ("INT", {"default": 12, "min": 1, "max": 50}),
|
|
"slat_cfg": ("FLOAT", {"default": 3.0, "min": 0.0, "max": 10.0, "step": 0.1, "round": 0.01}),
|
|
"slat_steps": ("INT", {"default": 12, "min": 1, "max": 50}),
|
|
"preprocess_image": ("BOOLEAN", {"default": False},),
|
|
"texture_size": ("INT", {"default": 512, "min": 512, "max": 2048, "step": 512, "display": "number"}),
|
|
"mesh_simplify": ("FLOAT", {"default": 0.95, "min": 0.0, "max": 0.98, "step": 0.01}),
|
|
"mode":(["fast","opt"],),
|
|
"multi_image": ("BOOLEAN", {"default": False},),
|
|
"multiimage_algo":(["multidiffusion", "stochastic"],),
|
|
"gaussians2ply": ("BOOLEAN", {"default": False},),
|
|
"covert2video": ("BOOLEAN", {"default": False},),
|
|
"glb2obj": ("BOOLEAN", {"default": False},),
|
|
"glb2fbx": ("BOOLEAN", {"default": False},),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", )
|
|
RETURN_NAMES = ("model_path",)
|
|
FUNCTION = "sampler_main"
|
|
CATEGORY = "Trellis"
|
|
|
|
def sampler_main(self, image, model, seed, cfg, steps,slat_cfg, slat_steps,preprocess_image,texture_size,mesh_simplify,mode,multi_image,multiimage_algo,gaussians2ply,covert2video,glb2obj,glb2fbx):
|
|
|
|
image_list,image_batch=tensor2imglist(image) #pil_list,batch
|
|
|
|
if multi_image and image_batch % 3 == 0:
|
|
print("********infer multi image,like Three views ******")
|
|
image_list=[image_list[i:i + 3] for i in range(0, len(image_list), 3)] #三等分列表
|
|
is_multiimage=True
|
|
else:
|
|
is_multiimage = False
|
|
trial_id = str(uuid.uuid4())
|
|
|
|
output_path = []
|
|
for i,img in enumerate(image_list):
|
|
model.cuda()
|
|
glb=image_to_3d(model,img,preprocess_image,covert2video,trial_id,seed,cfg,steps,slat_cfg,slat_steps,mesh_simplify,texture_size,mode,is_multiimage,gaussians2ply,multiimage_algo)
|
|
glb_path = f"{folder_paths.get_output_directory()}/{trial_id}_{i}.glb"
|
|
glb.export(glb_path)
|
|
output_path.append(glb_path)
|
|
model.cpu()
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
print(f"glb save in {glb_path} ")
|
|
|
|
if glb2obj:
|
|
obj_paths=[]
|
|
for path in output_path:
|
|
obj_path=os.path.join(os.path.split(path)[0],os.path.split(path)[1].replace(".glb",".obj"))
|
|
glb2obj_(path, obj_path)
|
|
obj_paths.append(obj_path)
|
|
if glb2fbx:
|
|
fbx_paths=[]
|
|
for i in obj_paths:
|
|
fbx_path = os.path.join(os.path.split(i)[0], os.path.split(i)[1].replace(".obj", ".fbx"))
|
|
obj2fbx_(i, fbx_path)
|
|
fbx_paths.append(fbx_path)
|
|
output_path=fbx_paths
|
|
else:
|
|
output_path=obj_paths
|
|
else:
|
|
if glb2fbx:
|
|
obj_paths = []
|
|
fbx_paths = []
|
|
for path in output_path:
|
|
obj_path = os.path.join(os.path.split(path)[0], os.path.split(path)[1].replace(".glb", ".obj"))
|
|
glb2obj_(path, obj_path)
|
|
obj_paths.append(obj_path)
|
|
for i in obj_paths:
|
|
fbx_path = os.path.join(os.path.split(i)[0], os.path.split(i)[1].replace(".obj", ".fbx"))
|
|
obj2fbx_(i, fbx_path)
|
|
fbx_paths.append(fbx_path)
|
|
output_path = obj_paths
|
|
model_path = '\n'.join(output_path)
|
|
return (model_path,)
|
|
|
|
class Trellis_multiimage_loader:
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"image_a": ("IMAGE",),
|
|
},
|
|
"optional": {"image_b": ("IMAGE",),
|
|
"image_c": ("IMAGE",),}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
ETURN_NAMES = ("image",)
|
|
FUNCTION = "main_batch"
|
|
CATEGORY = "Trellis"
|
|
|
|
def main_batch(self, image_a, **kwargs):
|
|
image_b = kwargs.get("image_b")
|
|
image_c = kwargs.get("image_c")
|
|
_,height_a,_,_ = image_a.shape
|
|
if isinstance(image_b, torch.Tensor) and isinstance(image_c, torch.Tensor):
|
|
_, height_b, _, _ = image_b.shape
|
|
_, height_c, _, _ = image_c.shape
|
|
height = max(height_a, height_b, height_c)
|
|
img_list=[pre_img(image_a, height),pre_img(image_b, height),pre_img(image_c, height)]
|
|
image = torch.cat(img_list, dim=0)
|
|
elif isinstance(image_b, torch.Tensor) and not isinstance(image_c, torch.Tensor):
|
|
_, height_b, _, _ = image_b.shape
|
|
height = max(height_a, height_b,)
|
|
img_list = [pre_img(image_a, height), pre_img(image_b, height)]
|
|
image = torch.cat(img_list, dim=0)
|
|
elif not isinstance(image_b, torch.Tensor) and isinstance(image_c, torch.Tensor):
|
|
_, height_c, _, _ = image_c.shape
|
|
height = max(height_a, height_c, )
|
|
img_list = [pre_img(image_a, height), pre_img(image_b, height)]
|
|
image = torch.cat(img_list, dim=0)
|
|
else:
|
|
image=image_a
|
|
|
|
return (image,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"Trellis_LoadModel": Trellis_LoadModel,
|
|
"Trellis_Sampler": Trellis_Sampler,
|
|
"Trellis_multiimage_loader":Trellis_multiimage_loader
|
|
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"Trellis_LoadModel": "Trellis_LoadModel",
|
|
"Trellis_Sampler": "Trellis_Sampler",
|
|
"Trellis_multiimage_loader":"Trellis_multiimage_loader"
|
|
}
|
|
|