753 lines
29 KiB
Python
753 lines
29 KiB
Python
import os
|
|
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
import trimesh
|
|
import folder_paths
|
|
from . import inference_full as inf
|
|
from . import split as splitter
|
|
|
|
# --- Native Trellis2 Integration Setup ---
|
|
import sys
|
|
import os
|
|
|
|
# Ensure Trellis2 is in the path for subclassing
|
|
# This must happen before defining SegviGenTrellis2Pipeline
|
|
custom_nodes = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
trellis2_dir = os.path.join(custom_nodes, "ComfyUI-Trellis2-GGUF")
|
|
if trellis2_dir not in sys.path:
|
|
sys.path.append(trellis2_dir)
|
|
|
|
try:
|
|
from trellis2_gguf.pipelines.trellis2_image_to_3d import Trellis2ImageTo3DPipeline
|
|
except ImportError:
|
|
print("[SegviGen] Warning: ComfyUI-Trellis2-GGUF not found. Native loader will not work.")
|
|
Trellis2ImageTo3DPipeline = object
|
|
|
|
REMOTE_CHECKPOINTS = {
|
|
"SegviGen/full_seg.safetensors": ("Aero-Ex/SegviGen", "full_seg.safetensors"),
|
|
"SegviGen/full_seg_w_2d_map.safetensors": ("Aero-Ex/SegviGen", "full_seg_w_2d_map.safetensors"),
|
|
}
|
|
|
|
def resolve_full_path(path):
|
|
if not path or not isinstance(path, str):
|
|
return path
|
|
if os.path.isabs(path):
|
|
return path
|
|
|
|
# Try output, input, and temp directories
|
|
for folder in [folder_paths.get_output_directory(), folder_paths.get_input_directory(), folder_paths.get_temp_directory()]:
|
|
if folder is None: continue
|
|
full = os.path.abspath(os.path.join(folder, path))
|
|
if os.path.exists(full):
|
|
# print(f"[SegviGen Res] Found: {full}")
|
|
return full
|
|
|
|
print(f"[SegviGen Res] FAILED to resolve relative path: {path}")
|
|
print(f" Checked directories based on folder_paths: output={folder_paths.get_output_directory()}, input={folder_paths.get_input_directory()}, temp={folder_paths.get_temp_directory()}")
|
|
# Fallback to check relative to this file
|
|
local_path = os.path.abspath(os.path.join(os.path.dirname(__file__), path))
|
|
if os.path.exists(local_path):
|
|
return local_path
|
|
|
|
return path
|
|
|
|
def extract_glb_path(mesh, temp_prefix="segvigen_tmp"):
|
|
if isinstance(mesh, list) and len(mesh) > 0:
|
|
mesh = mesh[0]
|
|
|
|
if isinstance(mesh, dict):
|
|
# Some nodes might pass a dict with 'mesh' or 'glb_path'
|
|
mesh = mesh.get("mesh") or mesh.get("glb_path") or mesh
|
|
|
|
if isinstance(mesh, str):
|
|
return resolve_full_path(mesh)
|
|
|
|
# Check for File3D specifically (from comfy_api)
|
|
if type(mesh).__name__ == "File3D":
|
|
if hasattr(mesh, "get_source"):
|
|
source = mesh.get_source()
|
|
if isinstance(source, str):
|
|
return resolve_full_path(source)
|
|
if hasattr(mesh, "save_to"):
|
|
temp_dir = folder_paths.get_temp_directory()
|
|
glb_path = os.path.join(temp_dir, f"{temp_prefix}_{os.urandom(4).hex()}.glb")
|
|
return mesh.save_to(glb_path)
|
|
if hasattr(mesh, "_source") and isinstance(mesh._source, str):
|
|
return resolve_full_path(mesh._source)
|
|
|
|
# Check for common path attributes
|
|
for attr in ["source", "path", "_path", "full_path", "filename", "abs_path"]:
|
|
if hasattr(mesh, attr):
|
|
val = getattr(mesh, attr)
|
|
if isinstance(val, str):
|
|
return resolve_full_path(val)
|
|
|
|
# Check for export method (trimesh)
|
|
if hasattr(mesh, "export"):
|
|
temp_dir = folder_paths.get_temp_directory()
|
|
glb_path = os.path.join(temp_dir, f"{temp_prefix}_{os.urandom(4).hex()}.glb")
|
|
mesh.export(glb_path)
|
|
return glb_path
|
|
|
|
# Check for trellis2 internal mesh
|
|
if type(mesh).__name__ == "Mesh" and hasattr(mesh, "vertices") and hasattr(mesh, "faces"):
|
|
import trimesh
|
|
tm = trimesh.Trimesh(vertices=mesh.vertices.cpu().numpy(), faces=mesh.faces.cpu().numpy())
|
|
temp_dir = folder_paths.get_temp_directory()
|
|
glb_path = os.path.join(temp_dir, f"{temp_prefix}_trellis_{os.urandom(4).hex()}.glb")
|
|
tm.export(glb_path)
|
|
return glb_path
|
|
|
|
# Fallback for File3D if attributes didn't match, try to parse source from repr
|
|
if type(mesh).__name__ == "File3D":
|
|
m_repr = str(mesh)
|
|
if "source='" in m_repr:
|
|
src = m_repr.split("source='")[1].split("'")[0]
|
|
return resolve_full_path(src)
|
|
|
|
raise ValueError(f"Unsupported mesh type: {type(mesh)}. Available attributes: {dir(mesh)}")
|
|
|
|
# SegviGen Model Loader
|
|
class SegviGenModelLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
ckpts = folder_paths.get_filename_list("checkpoints")
|
|
# Add remote checkpoints to the list if not already present locally
|
|
for remote_name in REMOTE_CHECKPOINTS.keys():
|
|
if remote_name not in ckpts:
|
|
ckpts.append(remote_name)
|
|
return {
|
|
"required": {
|
|
"ckpt_name": (ckpts,),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("SEG_MODEL",)
|
|
FUNCTION = "load_model"
|
|
CATEGORY = "SegviGen"
|
|
|
|
def load_model(self, ckpt_name):
|
|
ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
|
|
|
|
if ckpt_path is None or not os.path.exists(ckpt_path):
|
|
if ckpt_name in REMOTE_CHECKPOINTS:
|
|
repo_id, filename = REMOTE_CHECKPOINTS[ckpt_name]
|
|
# Try to find the best local path
|
|
base_ckpt_dir = folder_paths.get_folder_paths("checkpoints")[0]
|
|
|
|
# If name doesn't have prefix, we still prefer downloading to SegviGen/ subfolder
|
|
if "/" not in ckpt_name:
|
|
target_path = os.path.join(base_ckpt_dir, "SegviGen", ckpt_name)
|
|
else:
|
|
target_path = os.path.join(base_ckpt_dir, ckpt_name)
|
|
|
|
if not os.path.exists(target_path):
|
|
from huggingface_hub import hf_hub_download
|
|
print(f"[SegviGen] Downloading remote checkpoint: {ckpt_name} from {repo_id}")
|
|
os.makedirs(os.path.dirname(target_path), exist_ok=True)
|
|
hf_hub_download(
|
|
repo_id=repo_id,
|
|
filename=filename,
|
|
local_dir=os.path.dirname(target_path)
|
|
)
|
|
ckpt_path = target_path
|
|
else:
|
|
raise FileNotFoundError(f"Checkpoint {ckpt_name} not found and no remote mapping exists.")
|
|
|
|
inf.PIPE.load_ckpt_if_needed(ckpt_path)
|
|
return ({"ckpt_path": ckpt_path},)
|
|
|
|
# --- Native Trellis2 Integration ---
|
|
|
|
class SegviGenTrellis2Pipeline(Trellis2ImageTo3DPipeline):
|
|
"""
|
|
Subclass of Trellis2ImageTo3DPipeline to support SegviGen weights and output mapping.
|
|
"""
|
|
def __init__(self, *args, segvigen_ckpt=None, **kwargs):
|
|
# We handle initialization via from_pretrained usually
|
|
super().__init__(*args, **kwargs)
|
|
self.segvigen_ckpt = segvigen_ckpt
|
|
|
|
def _apply_segvigen_weights(self, model, ckpt_path):
|
|
from safetensors.torch import load_file
|
|
from collections import OrderedDict
|
|
print(f"[SegviGen] Loading weights from: {ckpt_path}")
|
|
state_dict = load_file(ckpt_path)
|
|
# SegviGen weights have "gen3dseg." prefix for internal models
|
|
# Also handle potential "flow_model." and "tex_slat_flow_model." prefixes
|
|
mapped_sd = OrderedDict()
|
|
for k, v in state_dict.items():
|
|
new_k = k
|
|
# Strip multiple layers of potential prefixes
|
|
for prefix in ["gen3dseg.", "tex_slat_flow_model.", "flow_model.", "tex_slat_decoder."]:
|
|
if new_k.startswith(prefix):
|
|
new_k = new_k[len(prefix):]
|
|
mapped_sd[new_k] = v
|
|
|
|
res = model.load_state_dict(mapped_sd, strict=False)
|
|
matched = len(mapped_sd) - len(res.unexpected_keys)
|
|
|
|
# Only log if we actually matched something
|
|
if matched > 0:
|
|
print(f"[SegviGen] Weights applied to {model.__class__.__name__}: {matched} keys matched.")
|
|
if res.missing_keys:
|
|
print(f"[SegviGen] Note: {len(res.missing_keys)} missing keys (standard for partial load)")
|
|
if res.unexpected_keys:
|
|
print(f"[SegviGen] Note: {len(res.unexpected_keys)} unexpected keys found in checkpoint")
|
|
|
|
def load_tex_slat_flow_model_1024(self):
|
|
super().load_tex_slat_flow_model_1024()
|
|
if self.segvigen_ckpt:
|
|
self._apply_segvigen_weights(self.models['tex_slat_flow_model_1024'], self.segvigen_ckpt)
|
|
|
|
def load_tex_slat_flow_model_512(self):
|
|
super().load_tex_slat_flow_model_512()
|
|
if self.segvigen_ckpt:
|
|
self._apply_segvigen_weights(self.models['tex_slat_flow_model_512'], self.segvigen_ckpt)
|
|
|
|
def load_tex_slat_decoder(self):
|
|
super().load_tex_slat_decoder()
|
|
if self.segvigen_ckpt:
|
|
self._apply_segvigen_weights(self.models['tex_slat_decoder'], self.segvigen_ckpt)
|
|
|
|
def postprocess_mesh(self, mesh, pbr_voxel, *args, **kwargs):
|
|
# Handle SegviGen's 3-channel output for Trellis2 compatibility
|
|
target_voxel = pbr_voxel[0] if isinstance(pbr_voxel, (list, tuple)) else pbr_voxel
|
|
|
|
if hasattr(target_voxel, 'feats') and target_voxel.feats.shape[1] == 3:
|
|
print(f"[SegviGen] Adapting 3-channel labels ({target_voxel.feats.shape[0]} voxels) to 6-channel PBR for Trellis2")
|
|
device = target_voxel.feats.device
|
|
dummy_metallic = torch.zeros(target_voxel.feats.shape[0], 1, device=device)
|
|
dummy_roughness = torch.ones(target_voxel.feats.shape[0], 1, device=device) * 0.5
|
|
dummy_alpha = torch.ones(target_voxel.feats.shape[0], 1, device=device)
|
|
|
|
new_feats = torch.cat([target_voxel.feats, dummy_metallic, dummy_roughness, dummy_alpha], dim=1)
|
|
|
|
new_voxel = target_voxel.__class__(new_feats, target_voxel.coords, target_voxel.spatial_shape)
|
|
if isinstance(pbr_voxel, list):
|
|
pbr_voxel[0] = new_voxel
|
|
else:
|
|
pbr_voxel = new_voxel
|
|
|
|
return super().postprocess_mesh(mesh, pbr_voxel, *args, **kwargs)
|
|
|
|
class Trellis2_SegviGenLoadModel:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
ckpts = folder_paths.get_filename_list("checkpoints")
|
|
for remote_name in REMOTE_CHECKPOINTS.keys():
|
|
if remote_name not in ckpts:
|
|
ckpts.append(remote_name)
|
|
return {
|
|
"required": {
|
|
"ckpt_name": (ckpts, {"default": "SegviGen/full_seg_w_2d_map.safetensors"}),
|
|
"backend": (["flash_attn", "xformers", "sdpa", "flash_attn_3"], {"default": "xformers"}),
|
|
"device": (["cpu","cuda"],{"default":"cuda"}),
|
|
"low_vram": ("BOOLEAN",{"default":True}),
|
|
"keep_models_loaded": ("BOOLEAN", {"default":True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("TRELLIS2PIPELINE", )
|
|
RETURN_NAMES = ("pipeline", )
|
|
FUNCTION = "process"
|
|
CATEGORY = "Trellis2Wrapper (GGUF)"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, ckpt_name, backend, device, low_vram, keep_models_loaded):
|
|
import sys
|
|
import os
|
|
|
|
# 1. Resolve SegviGen Checkpoint
|
|
ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
|
|
if ckpt_path is None or not os.path.exists(ckpt_path):
|
|
if ckpt_name in REMOTE_CHECKPOINTS:
|
|
repo_id, filename = REMOTE_CHECKPOINTS[ckpt_name]
|
|
base_ckpt_dir = folder_paths.get_folder_paths("checkpoints")[0]
|
|
target_path = os.path.join(base_ckpt_dir, ckpt_name if "/" in ckpt_name else os.path.join("SegviGen", ckpt_name))
|
|
if not os.path.exists(target_path):
|
|
from huggingface_hub import hf_hub_download
|
|
print(f"[SegviGen] Downloading: {ckpt_name}")
|
|
os.makedirs(os.path.dirname(target_path), exist_ok=True)
|
|
hf_hub_download(repo_id=repo_id, filename=filename, local_dir=os.path.dirname(target_path))
|
|
ckpt_path = target_path
|
|
|
|
# 2. Instantiate SegviGen pipeline
|
|
try:
|
|
from trellis2_gguf_model_manager import get_models_dir
|
|
model_path = get_models_dir()
|
|
except ImportError:
|
|
model_path = os.path.join(folder_paths.models_dir, "Trellis2")
|
|
|
|
# Instantiate correctly as SegviGenTrellis2Pipeline
|
|
pipeline = SegviGenTrellis2Pipeline.from_pretrained(
|
|
model_path,
|
|
keep_models_loaded=keep_models_loaded,
|
|
enable_gguf=False,
|
|
precision="bf16",
|
|
)
|
|
pipeline.segvigen_ckpt = ckpt_path
|
|
pipeline.low_vram = low_vram
|
|
if device == "cuda":
|
|
pipeline.to(device)
|
|
|
|
return (pipeline,)
|
|
|
|
# SegviGen Mesh Voxelizer
|
|
class SegviGenMeshVoxelizer:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"mesh": ("*",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("VXZ_DATA",)
|
|
FUNCTION = "voxelize"
|
|
CATEGORY = "SegviGen"
|
|
|
|
def voxelize(self, mesh):
|
|
glb_path = extract_glb_path(mesh, temp_prefix="segvigen_input")
|
|
|
|
vxz_path = glb_path.replace(".glb", ".vxz")
|
|
inf.process_glb_to_vxz(glb_path, vxz_path)
|
|
return ({"vxz_path": vxz_path, "glb_path": glb_path},)
|
|
|
|
# SegviGen Latent Encoder
|
|
class SegviGenLatentEncoder:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"vxz_data": ("VXZ_DATA",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("SHAPE_SLAT", "TEX_SLAT", "MESHES", "SUBS")
|
|
FUNCTION = "encode"
|
|
CATEGORY = "SegviGen"
|
|
|
|
def encode(self, vxz_data):
|
|
inf.PIPE.load_all_models()
|
|
shape_enc, tex_enc, shape_dec = inf.PIPE.get_encoders_decoder()
|
|
shape_enc.cuda()
|
|
tex_enc.cuda()
|
|
shape_dec.cuda()
|
|
|
|
shape_slat, meshes, subs, tex_slat = inf.vxz_to_latent_slat(
|
|
shape_enc, shape_dec, tex_enc, vxz_data["vxz_path"]
|
|
)
|
|
|
|
# Offload models and move latents to CPU to save VRAM
|
|
inf.PIPE.unload('shape_encoder', 'tex_encoder', 'shape_decoder')
|
|
|
|
return (
|
|
{"feats": shape_slat.feats.cpu(), "coords": shape_slat.coords.cpu()},
|
|
{"feats": tex_slat.feats.cpu(), "coords": tex_slat.coords.cpu()},
|
|
meshes,
|
|
subs
|
|
)
|
|
|
|
# SegviGen Image Conditioner
|
|
class SegviGenImageConditioner:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("CONDITION",)
|
|
FUNCTION = "condition"
|
|
CATEGORY = "SegviGen"
|
|
|
|
def condition(self, image):
|
|
# ComfyUI [B, H, W, C] -> PIL
|
|
img_np = (image[0].cpu().numpy() * 255).astype(np.uint8)
|
|
pil_img = Image.fromarray(img_np)
|
|
|
|
inf.PIPE.load_all_models()
|
|
rembg_model = inf.PIPE.get_rembg()
|
|
rembg_model.cuda()
|
|
pil_img = inf.preprocess_image(rembg_model, pil_img)
|
|
inf.PIPE.unload('rembg_model')
|
|
|
|
cond_model = inf.PIPE.get_cond_model()
|
|
cond_model.cuda()
|
|
cond = inf.get_cond(cond_model, [pil_img])
|
|
# Offload to CPU
|
|
cond = {k: v.cpu() for k, v in cond.items()}
|
|
inf.PIPE.unload('image_cond_model')
|
|
|
|
return (cond,)
|
|
|
|
# SegviGen Sampler
|
|
class SegviGenSampler:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"seg_model": ("SEG_MODEL",),
|
|
"shape_slat": ("SHAPE_SLAT",),
|
|
"tex_slat": ("TEX_SLAT",),
|
|
"condition": ("CONDITION",),
|
|
"steps": ("INT", {"default": 50, "min": 1, "max": 200}),
|
|
"guidance_scale": ("FLOAT", {"default": 7.5, "min": 0.0, "max": 20.0}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT_SLAT_OUT",)
|
|
FUNCTION = "sample"
|
|
CATEGORY = "SegviGen"
|
|
|
|
def sample(self, seg_model, shape_slat, tex_slat, condition, steps, guidance_scale, seed):
|
|
inf.PIPE.load_all_models()
|
|
inf.PIPE.load_ckpt_if_needed(seg_model["ckpt_path"])
|
|
gen3dseg = inf.PIPE.get_gen3dseg()
|
|
gen3dseg.cuda()
|
|
|
|
# We need to reconstruct SparseTensor from dict
|
|
import trellis2.modules.sparse as sp
|
|
shape_slat_sp = sp.SparseTensor(shape_slat["feats"].cuda(), shape_slat["coords"].cuda())
|
|
tex_slat_sp = sp.SparseTensor(tex_slat["feats"].cuda(), tex_slat["coords"].cuda())
|
|
cond_gpu = {k: v.cuda() for k, v in condition.items()}
|
|
|
|
output_tex_slat = inf.tex_slat_sample_single(
|
|
gen3dseg, inf.PIPE.sampler, inf.PIPE.pipeline_args,
|
|
shape_slat_sp, tex_slat_sp, cond_gpu,
|
|
steps=steps, cfg_scale=guidance_scale, seed=seed
|
|
)
|
|
res = {"feats": output_tex_slat.feats.cpu(), "coords": output_tex_slat.coords.cpu()}
|
|
inf.PIPE.unload('gen3dseg', 'tex_slat_flow_model')
|
|
return (res,)
|
|
|
|
# SegviGen Texture Decoder
|
|
class SegviGenTextureDecoder:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"latent_slat": ("LATENT_SLAT_OUT",),
|
|
"subs": ("SUBS",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("TEXTURE_VOXELS",)
|
|
FUNCTION = "decode"
|
|
CATEGORY = "SegviGen"
|
|
|
|
def decode(self, latent_slat, subs):
|
|
import trellis2.modules.sparse as sp
|
|
inf.PIPE.load_all_models()
|
|
tex_decoder = inf.PIPE.get_tex_decoder()
|
|
tex_decoder.cuda()
|
|
|
|
latent_slat_sp = sp.SparseTensor(latent_slat["feats"].cuda(), latent_slat["coords"].cuda())
|
|
subs_gpu = [s.cuda() if isinstance(s, torch.Tensor) else s for s in subs]
|
|
|
|
with torch.no_grad():
|
|
tex_voxels = tex_decoder(latent_slat_sp, guide_subs=subs_gpu) * 0.5 + 0.5
|
|
tex_voxels = [v.cpu() for v in tex_voxels]
|
|
|
|
inf.PIPE.unload('tex_decoder')
|
|
return (tex_voxels,)
|
|
|
|
# SegviGen Mesh Baker
|
|
class SegviGenMeshBaker:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"mesh": ("*",),
|
|
"texture_voxels": ("TEXTURE_VOXELS",),
|
|
"resolution": ("INT", {"default": 512, "min": 64, "max": 1024}),
|
|
"texture_size": ("INT", {"default": 2048, "min": 512, "max": 4096}),
|
|
"generate_uv": ("BOOLEAN", {"default": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("GLB_PATH",)
|
|
FUNCTION = "bake"
|
|
CATEGORY = "SegviGen"
|
|
OUTPUT_NODE = True
|
|
|
|
def bake(self, mesh, texture_voxels, resolution, texture_size, generate_uv):
|
|
glb_path = extract_glb_path(mesh, temp_prefix="segvigen_bake_src")
|
|
|
|
|
|
output_path = os.path.join(folder_paths.get_output_directory(), f"segvigen_baked_{os.urandom(4).hex()}.glb")
|
|
|
|
inf.bake_to_mesh(
|
|
glb_path,
|
|
texture_voxels,
|
|
output_path,
|
|
resolution=resolution,
|
|
texture_size=texture_size,
|
|
generate_uv=generate_uv
|
|
)
|
|
return (output_path,)
|
|
|
|
# SegviGen Mesh Exporter
|
|
class SegviGenMeshExporter:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"meshes": ("MESHES",),
|
|
"texture_voxels": ("TEXTURE_VOXELS",),
|
|
"resolution": ("INT", {"default": 512, "min": 64, "max": 1024}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("GLB_PATH",)
|
|
FUNCTION = "export_mesh"
|
|
CATEGORY = "SegviGen"
|
|
OUTPUT_NODE = True
|
|
|
|
def export_mesh(self, meshes, texture_voxels, resolution):
|
|
glb = inf.slat_to_glb(meshes, texture_voxels, resolution=resolution)
|
|
|
|
# Apply the same Y-up fix as in original app.py
|
|
T = np.eye(4, dtype=np.float64)
|
|
T[:3, :3] = np.array([[1, 0, 0], [0, 0, -1], [0, 1, 0]], dtype=np.float64)
|
|
|
|
output_path = os.path.join(folder_paths.get_output_directory(), f"segvigen_exported_{os.urandom(4).hex()}.glb")
|
|
|
|
if hasattr(glb, "apply_transform"):
|
|
glb.apply_transform(T)
|
|
glb.export(output_path)
|
|
else:
|
|
glb.export(output_path)
|
|
scene_or_mesh = trimesh.load(output_path, force="scene")
|
|
scene_or_mesh.apply_transform(T)
|
|
scene_or_mesh.export(output_path)
|
|
|
|
return (output_path,)
|
|
|
|
# SegviGen Split Refine
|
|
class SegviGenSplitRefine:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"glb_path": ("STRING",),
|
|
"min_faces_per_part": ("INT", {"default": 50, "min": 1, "max": 1000}),
|
|
"bake_transforms": ("BOOLEAN", {"default": True}),
|
|
"color_quant_step": ("INT", {"default": 16, "min": 1, "max": 64}),
|
|
"palette_sample_pixels": ("INT", {"default": 2000000}),
|
|
"palette_min_pixels": ("INT", {"default": 500}),
|
|
"palette_max_colors": ("INT", {"default": 256}),
|
|
"palette_merge_dist": ("INT", {"default": 32}),
|
|
"samples_per_face": ([1, 4], {"default": 4}),
|
|
"flip_v": ("BOOLEAN", {"default": True}),
|
|
"uv_wrap_repeat": ("BOOLEAN", {"default": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("PARTS_GLB_PATH",)
|
|
FUNCTION = "split"
|
|
CATEGORY = "SegviGen"
|
|
OUTPUT_NODE = True
|
|
|
|
def split(self, glb_path, **kwargs):
|
|
glb_path = resolve_full_path(glb_path)
|
|
out_parts_glb = os.path.join(folder_paths.get_output_directory(), f"segvigen_parts_{os.urandom(4).hex()}.glb")
|
|
splitter.split_glb_by_texture_palette_rgb(
|
|
in_glb_path=glb_path,
|
|
out_glb_path=out_parts_glb,
|
|
debug_print=True,
|
|
**kwargs
|
|
)
|
|
return (out_parts_glb,)
|
|
|
|
# SegviGen Mesh Simplify
|
|
class SegviGenMeshSimplify:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"glb_path": ("STRING",),
|
|
"target_faces": ("INT", {"default": 100000, "min": 1000, "max": 1000000}),
|
|
"aggression": ("INT", {"default": 7, "min": 1, "max": 20}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("GLB_PATH",)
|
|
FUNCTION = "simplify"
|
|
CATEGORY = "SegviGen"
|
|
OUTPUT_NODE = True
|
|
|
|
def simplify(self, glb_path, target_faces, aggression):
|
|
glb_path = resolve_full_path(glb_path)
|
|
out_glb = os.path.join(folder_paths.get_output_directory(), f"segvigen_simplified_{os.urandom(4).hex()}.glb")
|
|
inf.build_simplified_work_glb(glb_path, out_glb, target_faces=target_faces, aggression=aggression)
|
|
return (out_glb,)
|
|
|
|
# SegviGen Material Transfer
|
|
class SegviGenMaterialTransfer:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"source_mesh": ("*",),
|
|
"split_mesh": ("*",),
|
|
"flip_uv": ("BOOLEAN", {"default": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("GLB_PATH",)
|
|
FUNCTION = "transfer"
|
|
CATEGORY = "SegviGen"
|
|
OUTPUT_NODE = True
|
|
|
|
def transfer(self, source_mesh, split_mesh, flip_uv):
|
|
source_path = extract_glb_path(source_mesh, temp_prefix="segvigen_mat_src")
|
|
split_path = extract_glb_path(split_mesh, temp_prefix="segvigen_mat_split")
|
|
|
|
out_glb = os.path.join(folder_paths.get_output_directory(), f"segvigen_transferred_{os.urandom(4).hex()}.glb")
|
|
|
|
import trimesh
|
|
# Load geometries
|
|
source_scene = trimesh.load(source_path, force="scene", process=False)
|
|
target_scene = trimesh.load(split_path, force="scene", process=False)
|
|
|
|
# Build a mapping of geom_name -> material from the source scene
|
|
source_mats = {}
|
|
for node_name in source_scene.graph.nodes_geometry:
|
|
geom_name = source_scene.graph[node_name][1]
|
|
geom = source_scene.geometry.get(geom_name)
|
|
if geom and hasattr(geom.visual, 'material'):
|
|
source_mats[geom_name] = geom.visual.material
|
|
|
|
# If there's only one source geometry & material, we just use it for everything
|
|
default_mat = None
|
|
if len(source_mats) == 1:
|
|
default_mat = list(source_mats.values())[0]
|
|
|
|
# Apply to target geometries
|
|
for node_name in target_scene.graph.nodes_geometry:
|
|
geom_name = target_scene.graph[node_name][1]
|
|
geom = target_scene.geometry.get(geom_name)
|
|
if not geom:
|
|
continue
|
|
|
|
# The split mesh geom_name usually contains the original geom_name
|
|
mat_to_apply = default_mat
|
|
|
|
# Try to match the exact source geometry name if there are multiple
|
|
if len(source_mats) > 1:
|
|
# Sort by length descending to match the most specific name first
|
|
for src_name in sorted(source_mats.keys(), key=len, reverse=True):
|
|
if src_name in geom_name:
|
|
mat_to_apply = source_mats[src_name]
|
|
break
|
|
|
|
if mat_to_apply is not None:
|
|
# The split output retains the UVs in geom.visual.uv
|
|
if hasattr(geom.visual, 'uv') and geom.visual.uv is not None:
|
|
uvs = geom.visual.uv.copy()
|
|
if flip_uv:
|
|
uvs[:, 1] = 1.0 - uvs[:, 1]
|
|
geom.visual = trimesh.visual.TextureVisuals(uv=uvs, material=mat_to_apply)
|
|
else:
|
|
geom.visual.material = mat_to_apply
|
|
|
|
target_scene.export(out_glb)
|
|
print(f"[SegviGen] MaterialTransfer: Saved to {out_glb}")
|
|
return (out_glb,)
|
|
|
|
|
|
# SegviGen Monolithic Segmentation
|
|
class SegviGenMonolithicSegmentation:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"mesh": ("*",),
|
|
"seg_model": ("SEG_MODEL",),
|
|
"bake_mode": ("BOOLEAN", {"default": False}),
|
|
"generate_uv": ("BOOLEAN", {"default": False}),
|
|
},
|
|
"optional": {
|
|
"image": ("IMAGE",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("GLB_PATH",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "SegviGen"
|
|
|
|
def process(self, mesh, seg_model, bake_mode, generate_uv, image=None):
|
|
glb_path = extract_glb_path(mesh, temp_prefix="segvigen_mono")
|
|
|
|
# 2. Temp files
|
|
workdir = os.path.join(folder_paths.get_temp_directory(), f"seg_mono_{os.urandom(4).hex()}")
|
|
os.makedirs(workdir, exist_ok=True)
|
|
in_vxz = os.path.join(workdir, "input.vxz")
|
|
export_glb = os.path.join(folder_paths.get_output_directory(), f"segvigen_mono_out_{os.urandom(4).hex()}.glb")
|
|
|
|
item = {
|
|
"glb": glb_path,
|
|
"input_vxz": in_vxz,
|
|
"export_glb": export_glb,
|
|
"bake": bake_mode,
|
|
"generate_uv": generate_uv,
|
|
}
|
|
|
|
if image is not None:
|
|
img_np = (image[0].cpu().numpy() * 255).astype(np.uint8)
|
|
pil_img = Image.fromarray(img_np)
|
|
out_img = os.path.join(workdir, "input_map.png")
|
|
pil_img.save(out_img)
|
|
item["2d_map"] = True
|
|
item["img"] = out_img
|
|
else:
|
|
item["2d_map"] = False
|
|
item["transforms"] = os.path.join(os.path.dirname(__file__), "data_toolkit", "transforms.json")
|
|
item["img"] = os.path.join(workdir, "render.png")
|
|
|
|
inf.inference_with_loaded_models(seg_model["ckpt_path"], item)
|
|
return (export_glb,)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"SegviGenModelLoader": SegviGenModelLoader,
|
|
"SegviGenMeshVoxelizer": SegviGenMeshVoxelizer,
|
|
"SegviGenLatentEncoder": SegviGenLatentEncoder,
|
|
"SegviGenImageConditioner": SegviGenImageConditioner,
|
|
"SegviGenSampler": SegviGenSampler,
|
|
"SegviGenTextureDecoder": SegviGenTextureDecoder,
|
|
"SegviGenMeshBaker": SegviGenMeshBaker,
|
|
"SegviGenMeshExporter": SegviGenMeshExporter,
|
|
"SegviGenSplitRefine": SegviGenSplitRefine,
|
|
"SegviGenMeshSimplify": SegviGenMeshSimplify,
|
|
"SegviGenMaterialTransfer": SegviGenMaterialTransfer,
|
|
"SegviGenMonolithicSegmentation": SegviGenMonolithicSegmentation,
|
|
"Trellis2_SegviGenLoadModel": Trellis2_SegviGenLoadModel,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"SegviGenModelLoader": "SegviGen Model Loader",
|
|
"SegviGenMeshVoxelizer": "SegviGen Mesh Voxelizer",
|
|
"SegviGenLatentEncoder": "SegviGen Latent Encoder",
|
|
"SegviGenImageConditioner": "SegviGen Image Conditioner",
|
|
"SegviGenSampler": "SegviGen Sampler",
|
|
"SegviGenTextureDecoder": "SegviGen Texture Decoder",
|
|
"SegviGenMeshBaker": "SegviGen Mesh Baker",
|
|
"SegviGenMeshExporter": "SegviGen Mesh Exporter",
|
|
"SegviGenSplitRefine": "SegviGen Split & Refine",
|
|
"SegviGenMeshSimplify": "SegviGen Mesh Simplify",
|
|
"SegviGenMaterialTransfer": "SegviGen Material Transfer",
|
|
"SegviGenMonolithicSegmentation": "SegviGen Monolithic Segmentation",
|
|
"Trellis2_SegviGenLoadModel": "Trellis2 - Load SegviGen Model",
|
|
}
|