import os os.environ["OPENCV_IO_ENABLE_OPENEXR"] = '1' # os.environ["ATTN_BACKEND"] = "flash_attn_3" # Removed to allow fallbacks import json import math import torch import trimesh import o_voxel import numpy as np import torch.nn as nn import trellis2.modules.sparse as sp from PIL import Image from tqdm import tqdm from trellis2 import models from collections import OrderedDict from transformers import Mistral3ForConditionalGeneration from trellis2.pipelines.rembg import BiRefNet from trellis2.representations import MeshWithVoxel from trellis2.modules.image_feature_extractor import DinoV3FeatureExtractor try: import cv2 import nvdiffrast.torch as nr from flex_gemm.ops.grid_sample import grid_sample_3d BAKE_ENABLED = True except ImportError as e: print(f"[SegviGen] Warning: Texture baking dependencies NOT found ({e}). Skipping bake mode.") BAKE_ENABLED = False BASE_DIR = os.path.dirname(os.path.abspath(__file__)) TRELLIS_PIPELINE_JSON = os.path.join(BASE_DIR, "data_toolkit/texturing_pipeline.json") TRELLIS_TEX_FLOW = "microsoft/TRELLIS.2-4B/ckpts/slat_flow_imgshape2tex_dit_1_3B_512_bf16" TRELLIS_SHAPE_ENC = "microsoft/TRELLIS.2-4B/ckpts/shape_enc_next_dc_f16c32_fp16" TRELLIS_TEX_ENC = "microsoft/TRELLIS.2-4B/ckpts/tex_enc_next_dc_f16c32_fp16" TRELLIS_SHAPE_DEC = "microsoft/TRELLIS.2-4B/ckpts/shape_dec_next_dc_f16c32_fp16" TRELLIS_TEX_DEC = "microsoft/TRELLIS.2-4B/ckpts/tex_dec_next_dc_f16c32_fp16" DINO_PATH = "Aero-Ex/Dinov3" import folder_paths COMFY_MODELS_DIR = folder_paths.models_dir EARLY_SIMPLIFY_ENABLED = True EARLY_SIMPLIFY_TARGET_FACES = 120000 EARLY_SIMPLIFY_AGGRESSION = 2 MAX_PREPROCESS_TEX_SIZE = 1024 EXPORT_TEXTURE_SIZE = 2048 def _scene_to_single_mesh(asset): if isinstance(asset, trimesh.Scene): mesh = asset.to_mesh() elif isinstance(asset, trimesh.Trimesh): mesh = asset else: raise TypeError(f"Unsupported asset type: {type(asset)}") if mesh is None or (hasattr(mesh, 'faces') and len(mesh.faces) == 0): # Fallback for empty geometries return mesh return mesh def make_texture_square_pow2(img: Image.Image, target_size=None, max_size=1024): w, h = img.size max_side = max(w, h) pow2 = 1 while pow2 < max_side: pow2 *= 2 if target_size is not None: pow2 = target_size pow2 = min(pow2, max_size) return img.resize((pow2, pow2), Image.BILINEAR) def preprocess_scene_textures(asset, max_texture_size=1024): if not isinstance(asset, trimesh.Scene): return asset tex_keys = ["baseColorTexture", "normalTexture", "metallicRoughnessTexture", "emissiveTexture", "occlusionTexture"] for geom in asset.geometry.values(): visual = getattr(geom, "visual", None) mat = getattr(visual, "material", None) if mat is None: continue for key in tex_keys: if not hasattr(mat, key): continue tex = getattr(mat, key) if tex is None: continue if isinstance(tex, Image.Image): setattr(mat, key, make_texture_square_pow2(tex, max_size=max_texture_size)) elif hasattr(tex, "image") and tex.image is not None: img = tex.image if not isinstance(img, Image.Image): img = Image.fromarray(img) tex.image = make_texture_square_pow2(img, max_size=max_texture_size) if hasattr(mat, "image") and mat.image is not None: img = mat.image if not isinstance(img, Image.Image): img = Image.fromarray(img) mat.image = make_texture_square_pow2(img, max_size=max_texture_size) return asset def _apply_neutral_visual(mesh: trimesh.Trimesh) -> trimesh.Trimesh: if mesh.visual is not None: return mesh face_colors = np.tile( np.array([[200, 200, 200, 255]], dtype=np.uint8), (len(mesh.faces), 1), ) mesh.visual = trimesh.visual.color.ColorVisuals(mesh=mesh, face_colors=face_colors) return mesh def get_work_glb_path(item): base, _ = os.path.splitext(item["input_vxz"]) return f"{base}_work.glb" def build_simplified_work_glb( input_glb_path, output_glb_path, target_faces=EARLY_SIMPLIFY_TARGET_FACES, aggression=EARLY_SIMPLIFY_AGGRESSION, ): """ Build a simplified geometry-only GLB for early rendering and shape processing. This keeps the original GLB untouched for texture/attribute extraction. """ os.makedirs(os.path.dirname(output_glb_path), exist_ok=True) asset = trimesh.load(input_glb_path, force="scene", process=False) mesh = _scene_to_single_mesh(asset) src_faces = int(len(mesh.faces)) if src_faces <= target_faces: mesh = _apply_neutral_visual(mesh.copy()) mesh.export(output_glb_path) print(f"[Simplify] Skip: faces={src_faces} <= target={target_faces}") return output_glb_path, src_faces, src_faces try: simplified = mesh.simplify_quadric_decimation( face_count=target_faces, aggression=aggression, ) if simplified is None or len(simplified.faces) == 0: raise RuntimeError("simplify_quadric_decimation returned empty mesh") simplified = _apply_neutral_visual(simplified) simplified.export(output_glb_path) dst_faces = int(len(simplified.faces)) print(f"[Simplify] faces: {src_faces} -> {dst_faces}") return output_glb_path, src_faces, dst_faces except Exception as e: print(f"[Simplify] Failed, fallback to original mesh: {e}") mesh = _apply_neutral_visual(mesh.copy()) def generate_2d_map_from_glb(glb_path, transforms_path, out_img_path, render_img_path=None): """ Render the GLB first, then generate a 2D segmentation map with FLUX2. """ PIPE.load_all_models() if render_img_path is None: base, _ = os.path.splitext(out_img_path) render_img_path = f"{base}_render.png" render_from_transforms(glb_path, transforms_path, render_img_path) prompt = "Apply distinct colors to different regions of this image" render_img = Image.open(render_img_path).convert("RGB") if max(render_img.size) > 768: scale = 768 / max(render_img.size) render_img = render_img.resize( (int(render_img.width * scale), int(render_img.height * scale)), Image.Resampling.LANCZOS, ) torch.cuda.empty_cache() # Flux generation removed due to missing pipeline # image = PIPE.flux2( # prompt=prompt, # image=render_img, # num_inference_steps=4, # ).images[0] # image.save(out_img_path) print("[Warning] 2D map generation with FLUX is currently disabled.") return None def _colorvisuals_to_texturevisuals(mesh: trimesh.Trimesh) -> trimesh.Trimesh: """ Convert ColorVisuals to TextureVisuals by baking per-face colors into a tiny atlas and generating per-face UVs. Ensure the resulting material is PBRMaterial to satisfy downstream GLTF/PBR-only pipelines. """ if mesh.visual is None: return mesh if isinstance(mesh.visual, trimesh.visual.texture.TextureVisuals): mat = getattr(mesh.visual, "material", None) if isinstance(mat, trimesh.visual.material.SimpleMaterial): mesh = mesh.copy() try: mesh.visual.material = mat.to_pbr() except Exception: mesh.visual.material = trimesh.visual.material.PBRMaterial( baseColorTexture=mat.image ) return mesh if not isinstance(mesh.visual, trimesh.visual.color.ColorVisuals): return mesh F = int(len(mesh.faces)) if F <= 0: return mesh face_rgba = None if hasattr(mesh.visual, "face_colors") and mesh.visual.face_colors is not None: fc = np.asarray(mesh.visual.face_colors) if fc.ndim == 2 and fc.shape[0] == F: face_rgba = fc[:, :4].astype(np.uint8) if face_rgba is None and hasattr(mesh.visual, "vertex_colors") and mesh.visual.vertex_colors is not None: vc = np.asarray(mesh.visual.vertex_colors) if vc.ndim == 2 and vc.shape[0] == len(mesh.vertices): tri = mesh.faces vcol = vc[tri] face_rgba = np.rint(vcol.mean(axis=1)).astype(np.uint8) if face_rgba is None: face_rgba = np.tile(np.array([[255, 255, 255, 255]], dtype=np.uint8), (F, 1)) grid = int(math.ceil(math.sqrt(F))) img = np.zeros((grid, grid, 4), dtype=np.uint8) for i in range(F): x = i % grid y = i // grid if y >= grid: break img[y, x, :] = face_rgba[i] pil_img = Image.fromarray(img, mode="RGBA") v_new = mesh.vertices[mesh.faces].reshape(-1, 3) f_new = np.arange(F * 3, dtype=np.int64).reshape(F, 3) uv_new = np.zeros((F * 3, 2), dtype=np.float32) for i in range(F): x = i % grid y = i // grid u = (x + 0.5) / float(grid) v = (y + 0.5) / float(grid) uv_new[i * 3 : i * 3 + 3, 0] = u uv_new[i * 3 : i * 3 + 3, 1] = v pbr = trimesh.visual.material.PBRMaterial( baseColorTexture=pil_img, metallicFactor=0.0, roughnessFactor=1.0, doubleSided=True, alphaMode="BLEND", ) visual = trimesh.visual.texture.TextureVisuals(uv=uv_new, material=pbr) out = trimesh.Trimesh(vertices=v_new, faces=f_new, visual=visual, process=False) return out def ensure_texture_visuals(asset): """ Ensure all geometries in a Scene (or a single Trimesh) use TextureVisuals. For ColorVisuals, we bake them into a synthetic atlas. """ if isinstance(asset, trimesh.Scene): for geom_name, g in list(asset.geometry.items()): if isinstance(g, trimesh.Trimesh): asset.geometry[geom_name] = _colorvisuals_to_texturevisuals(g) return asset if isinstance(asset, trimesh.Trimesh): return _colorvisuals_to_texturevisuals(asset) return asset class Sampler: def _inference_model(self, model, x_t, tex_slat, shape_slat, coords_len_list, t, cond): t = torch.tensor([t * 1000] * x_t.shape[0], dtype=torch.float32).cuda() return model(x_t, tex_slat, shape_slat, t, cond, coords_len_list) def guidance_inference_model(self, model, x_t, tex_slat, shape_slat, coords_len_list, t, cond_dict, guidance_strength, guidance_rescale=0.0): if guidance_strength == 1: return self._inference_model(model, x_t, tex_slat, shape_slat, coords_len_list, t, cond_dict['cond']) elif guidance_strength == 0: return self._inference_model(model, x_t, tex_slat, shape_slat, coords_len_list, t, cond_dict['neg_cond']) else: pred_pos = self._inference_model(model, x_t, tex_slat, shape_slat, coords_len_list, t, cond_dict['cond']) pred_neg = self._inference_model(model, x_t, tex_slat, shape_slat, coords_len_list, t, cond_dict['neg_cond']) pred = guidance_strength * pred_pos + (1 - guidance_strength) * pred_neg if guidance_rescale > 0: x_0_pos = self._pred_to_xstart(x_t, t, pred_pos) x_0_cfg = self._pred_to_xstart(x_t, t, pred) std_pos = x_0_pos.std(dim=list(range(1, x_0_pos.ndim)), keepdim=True) std_cfg = x_0_cfg.std(dim=list(range(1, x_0_cfg.ndim)), keepdim=True) x_0_rescaled = x_0_cfg * (std_pos / std_cfg) x_0 = guidance_rescale * x_0_rescaled + (1 - guidance_rescale) * x_0_cfg pred = self._xstart_to_pred(x_t, t, x_0) return pred def interval_inference_model(self, model, x_t, tex_slat, shape_slat, coords_len_list, t, cond_dict, sampler_params): guidance_strength = sampler_params['guidance_strength'] guidance_interval = sampler_params['guidance_interval'] guidance_rescale = sampler_params['guidance_rescale'] if guidance_interval[0] <= t <= guidance_interval[1]: return self.guidance_inference_model(model, x_t, tex_slat, shape_slat, coords_len_list, t, cond_dict, guidance_strength, guidance_rescale) else: return self.guidance_inference_model(model, x_t, tex_slat, shape_slat, coords_len_list, t, cond_dict, 1, guidance_rescale) @torch.no_grad() def sample_once(self, model, x_t, tex_slat, shape_slat, coords_len_list, t, t_prev, cond_dict, sampler_params): pred_v = self.interval_inference_model(model, x_t, tex_slat, shape_slat, coords_len_list, t, cond_dict, sampler_params) pred_x_prev = x_t - (t - t_prev) * pred_v return pred_x_prev @torch.no_grad() def sample(self, model, noise, tex_slat, shape_slat, coords_len_list, cond_dict, sampler_params): sample = noise steps = sampler_params['steps'] rescale_t = sampler_params['rescale_t'] t_seq = np.linspace(1, 0, steps + 1) t_seq = rescale_t * t_seq / (1 + (rescale_t - 1) * t_seq) t_seq = t_seq.tolist() t_pairs = list((t_seq[i], t_seq[i + 1]) for i in range(steps)) for t, t_prev in tqdm(t_pairs, desc="Sampling"): sample = self.sample_once(model, sample, tex_slat, shape_slat, coords_len_list, t, t_prev, cond_dict, sampler_params) return sample class Gen3DSeg(nn.Module): def __init__(self, flow_model): super().__init__() self.flow_model = flow_model def forward(self, x_t, tex_slats, shape_slats, t, cond, coords_len_list): input_tex_feats_list = [] input_tex_coords_list = [] shape_feats_list = [] shape_coords_list = [] begin = 0 for coords_len in coords_len_list: end = begin + coords_len input_tex_feats_list.append(x_t.feats[begin:end]) input_tex_feats_list.append(tex_slats.feats[begin:end]) input_tex_coords_list.append(x_t.coords[begin:end]) input_tex_coords_list.append(tex_slats.coords[begin:end]) shape_feats_list.append(shape_slats.feats[begin:end]) shape_feats_list.append(shape_slats.feats[begin:end]) shape_coords_list.append(shape_slats.coords[begin:end]) shape_coords_list.append(shape_slats.coords[begin:end]) begin = end x_t = sp.SparseTensor(torch.cat(input_tex_feats_list), torch.cat(input_tex_coords_list)) shape_slats = sp.SparseTensor(torch.cat(shape_feats_list), torch.cat(shape_coords_list)) output_tex_slats = self.flow_model(x_t, t, cond, shape_slats) output_tex_feats_list = [] output_tex_coords_list = [] begin = 0 for coords_len in coords_len_list: end = begin + coords_len output_tex_feats_list.append(output_tex_slats.feats[begin:end]) output_tex_coords_list.append(output_tex_slats.coords[begin:end]) begin = begin + 2 * coords_len output_tex_slat = sp.SparseTensor(torch.cat(output_tex_feats_list), torch.cat(output_tex_coords_list)) return output_tex_slat def make_texture_square_pow2(img: Image.Image, target_size=None, max_size=MAX_PREPROCESS_TEX_SIZE): w, h = img.size max_side = max(w, h) pow2 = 1 while pow2 < max_side: pow2 *= 2 if target_size is not None: pow2 = target_size pow2 = min(pow2, max_size) return img.resize((pow2, pow2), Image.BILINEAR) def preprocess_scene_textures(asset, max_texture_size=MAX_PREPROCESS_TEX_SIZE): if not isinstance(asset, trimesh.Scene): return asset tex_keys = ["baseColorTexture", "normalTexture", "metallicRoughnessTexture", "emissiveTexture", "occlusionTexture"] for geom in asset.geometry.values(): visual = getattr(geom, "visual", None) mat = getattr(visual, "material", None) if mat is None: continue for key in tex_keys: if not hasattr(mat, key): continue tex = getattr(mat, key) if tex is None: continue if isinstance(tex, Image.Image): setattr(mat, key, make_texture_square_pow2(tex, max_size=max_texture_size)) elif hasattr(tex, "image") and tex.image is not None: img = tex.image if not isinstance(img, Image.Image): img = Image.fromarray(img) tex.image = make_texture_square_pow2(img, max_size=max_texture_size) if hasattr(mat, "image") and mat.image is not None: img = mat.image if not isinstance(img, Image.Image): img = Image.fromarray(img) mat.image = make_texture_square_pow2(img, max_size=max_texture_size) return asset def process_glb_to_vxz(glb_path, vxz_path, shape_glb_path=None): """ Use the original GLB for texture/material attributes, and an optional simplified GLB for the geometry-heavy shape branch. """ tex_asset = trimesh.load(glb_path, force='scene') tex_asset = ensure_texture_visuals(tex_asset) tex_asset = preprocess_scene_textures(tex_asset, max_texture_size=MAX_PREPROCESS_TEX_SIZE) if shape_glb_path is None: shape_asset = trimesh.load(glb_path, force='scene') else: shape_asset = trimesh.load(shape_glb_path, force='scene') aabb = tex_asset.bounding_box.bounds center = (aabb[0] + aabb[1]) / 2 max_side = (aabb[1] - aabb[0]).max() scale = 0.99999 / max_side print(f"[SegviGen] Vxz Calibration: AABB_MIN={aabb[0].tolist()}, AABB_MAX={aabb[1].tolist()}, Center={center.tolist()}, Scale={scale:.12f}, MaxSide={max_side:.12f}") tex_asset.apply_translation(-center) tex_asset.apply_scale(scale) shape_asset.apply_translation(-center) shape_asset.apply_scale(scale) shape_mesh = _scene_to_single_mesh(shape_asset) vertices = torch.from_numpy(shape_mesh.vertices).float() faces = torch.from_numpy(shape_mesh.faces).long() voxel_indices, dual_vertices, intersected = o_voxel.convert.mesh_to_flexible_dual_grid( vertices, faces, grid_size=512, aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]], face_weight=1.0, boundary_weight=0.2, regularization_weight=1e-2, timing=False, ) vid = o_voxel.serialize.encode_seq(voxel_indices) mapping = torch.argsort(vid) voxel_indices = voxel_indices[mapping] dual_vertices = dual_vertices[mapping] intersected = intersected[mapping] # Material Voxelization (using full scene asset for correct attribute mapping) voxel_indices_mat, attributes = o_voxel.convert.textured_mesh_to_volumetric_attr( tex_asset, grid_size=512, aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]], timing=False ) vid_mat = o_voxel.serialize.encode_seq(voxel_indices_mat) mapping_mat = torch.argsort(vid_mat) attributes = {k: v[mapping_mat] for k, v in attributes.items()} # Quantize dual_vertices and intersection flags dual_vertices = dual_vertices * 512 - voxel_indices dual_vertices = (torch.clamp(dual_vertices, 0, 1) * 255).type(torch.uint8) intersected = (intersected[:, 0:1] + 2 * intersected[:, 1:2] + 4 * intersected[:, 2:3]).type(torch.uint8) attributes.update({'dual_vertices': dual_vertices}) attributes.update({'intersected': intersected}) os.makedirs(os.path.dirname(vxz_path), exist_ok=True) o_voxel.io.write(vxz_path, voxel_indices, attributes) print(f"[SegviGen] Vxz: done") def vxz_to_latent_slat(shape_encoder, shape_decoder, tex_encoder, vxz_path): device = next(shape_encoder.parameters()).device coords, data = o_voxel.io.read(vxz_path) coords = torch.cat([torch.zeros(coords.shape[0], 1, dtype=torch.int32), coords], dim=1).to(device) vertices = (data['dual_vertices'].to(device) / 255) intersected = torch.cat([data['intersected'] % 2, data['intersected'] // 2 % 2, data['intersected'] // 4 % 2], dim=-1).bool().to(device) vertices_sparse = sp.SparseTensor(vertices, coords) intersected_sparse = sp.SparseTensor(intersected.float(), coords) with torch.no_grad(): shape_slat = shape_encoder(vertices_sparse, intersected_sparse) shape_slat = sp.SparseTensor(shape_slat.feats.to(device), shape_slat.coords.to(device)) shape_decoder.set_resolution(512) meshes, subs = shape_decoder(shape_slat, return_subs=True) try: for i, m in enumerate(meshes): v_min = m.vertices.min(dim=0).values.tolist() v_max = m.vertices.max(dim=0).values.tolist() print(f"[Log - Decoder Output] Component {i} bounds: {v_min} to {v_max}") except Exception: pass base_color = (data['base_color'] / 255) metallic = (data['metallic'] / 255) roughness = (data['roughness'] / 255) alpha = (data['alpha'] / 255) attr = torch.cat([base_color, metallic, roughness, alpha], dim=-1).float().to(device) * 2 - 1 with torch.no_grad(): tex_slat = tex_encoder(sp.SparseTensor(attr, coords)) return shape_slat, meshes, subs, tex_slat def preprocess_image(rembg_model, input): if input.mode != "RGB": bg = Image.new("RGB", input.size, (255, 255, 255)) bg.paste(input, mask=input.split()[3]) input = bg has_alpha = False if input.mode == 'RGBA': alpha = np.array(input)[:, :, 3] if not np.all(alpha == 255): has_alpha = True max_size = max(input.size) scale = min(1, 1024 / max_size) if scale < 1: input = input.resize((int(input.width * scale), int(input.height * scale)), Image.Resampling.LANCZOS) if has_alpha: output = input else: input = input.convert('RGB') output = rembg_model(input) output_np = np.array(output) alpha = output_np[:, :, 3] bbox = np.argwhere(alpha > 0.8 * 255) bbox = np.min(bbox[:, 1]), np.min(bbox[:, 0]), np.max(bbox[:, 1]), np.max(bbox[:, 0]) center = (bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2 size = max(bbox[2] - bbox[0], bbox[3] - bbox[1]) size = int(size * 1) bbox = center[0] - size // 2, center[1] - size // 2, center[0] + size // 2, center[1] + size // 2 output = output.crop(bbox) output = np.array(output).astype(np.float32) / 255 output = output[:, :, :3] * output[:, :, 3:4] output = Image.fromarray((output * 255).astype(np.uint8)) return output def get_cond(image_cond_model, image): image_cond_model.image_size = 512 cond = image_cond_model(image) neg_cond = torch.zeros_like(cond) return {'cond': cond, 'neg_cond': neg_cond} def tex_slat_sample_single(gen3dseg, sampler, pipeline_args, shape_slat, input_tex_slat, cond_dict, steps=None, cfg_scale=None, seed=None): device = shape_slat.feats.device shape_std = torch.tensor(pipeline_args['shape_slat_normalization']['std'])[None].to(device) shape_mean = torch.tensor(pipeline_args['shape_slat_normalization']['mean'])[None].to(device) tex_std = torch.tensor(pipeline_args['tex_slat_normalization']['std'])[None].to(device) tex_mean = torch.tensor(pipeline_args['tex_slat_normalization']['mean'])[None].to(device) shape_slat = ((shape_slat - shape_mean) / shape_std) input_tex_slat = ((input_tex_slat - tex_mean) / tex_std) coords_len_list = [shape_slat.coords.shape[0]] # Handle seed with generator if seed is not None: generator = torch.Generator(device=device).manual_seed(seed) noise_feats = torch.randn(input_tex_slat.feats.shape, generator=generator, device=device) else: noise_feats = torch.randn_like(input_tex_slat.feats) noise = sp.SparseTensor(noise_feats, shape_slat.coords) # Update sampler params if provided params = pipeline_args['tex_slat_sampler']['params'].copy() if steps is not None: params['steps'] = steps if cfg_scale is not None: params['cfg_scale'] = cfg_scale # Note: the sampler's .sample method might not take a 'seed' key in its internal params, # but the noise already handles it. We pass it anyway if it is used internally. if seed is not None: params['seed'] = seed output_tex_slat = sampler.sample(gen3dseg, noise, input_tex_slat, shape_slat, coords_len_list, cond_dict, params) output_tex_slat = output_tex_slat * tex_std + tex_mean return output_tex_slat def make_texture_square_pow2(img: Image.Image, target_size=None, max_size=1024): w, h = img.size max_side = max(w, h) pow2 = 1 while pow2 < max_side: pow2 *= 2 if target_size is not None: pow2 = target_size pow2 = min(pow2, max_size) return img.resize((pow2, pow2), Image.BILINEAR) def bake_to_mesh(glb_path, tex_voxels, output_path, resolution=512, texture_size=2048, generate_uv=False): """ Bake texture from voxels onto an existing GLB mesh using the same logic as to_glb. """ if not BAKE_ENABLED: raise RuntimeError("Texture baking is not enabled (missing dependencies).") print(f"[SegviGen] Bake: baking onto {glb_path} -> {output_path}") asset = trimesh.load(glb_path, force='scene') device = torch.device("cuda") # Prepare the attribute volume (sparse) as to_glb expects if isinstance(tex_voxels, list): attr_volume = torch.cat([vox.feats for vox in tex_voxels]).to(device) attr_coords = torch.cat([vox.coords for vox in tex_voxels]).to(device)[:, -3:] # <--- Take only (Z, Y, X) else: attr_volume = tex_voxels.feats.to(device) attr_coords = tex_voxels.coords.to(device)[:, -3:] print(f"[SegviGen] Bake Debug: attr_volume shape={attr_volume.shape}, coords bounds={attr_coords.min(dim=0)[0].tolist()} to {attr_coords.max(dim=0)[0].tolist()}") pbr_attr_layout = { 'base_color': slice(0, 3), 'metallic': slice(3, 4), 'roughness': slice(4, 5), 'alpha': slice(5, 6), } # [AI Voxel Data Logging] num_voxels = attr_volume.shape[0] print(f"\n[AI Voxel Data] Found {num_voxels} voxels generated by AI.") # Log a sample of 10 voxels (Position -> Color) print("[AI Voxel Data] Sample Voxels (Grid Pos -> RGB):") sample_indices = torch.linspace(0, num_voxels - 1, 10).long() for idx in sample_indices: pos = attr_coords[idx].tolist() # Extract RGB from base_color slice (multiplied by 255) color = (attr_volume[idx, pbr_attr_layout['base_color']] * 255).int().tolist() print(f" - Pos {pos} -> RGB {color}") print("-" * 50) # Scale voxels to match the slat to glb Unit Cube AABB aabb = [[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]] # Calculate normalization (must match process_glb_to_vxz) full_aabb = asset.bounding_box.bounds center = (full_aabb[0] + full_aabb[1]) / 2 max_side = (full_aabb[1] - full_aabb[0]).max() scale = 0.99999 / max_side # Iterate over nodes to handle scene graph transforms for node_name in asset.graph.nodes_geometry: transform, geom_name = asset.graph[node_name] geom = asset.geometry[geom_name] # Robust UV check has_uv = False if hasattr(geom.visual, 'uv') and geom.visual.uv is not None: has_uv = True elif isinstance(geom.visual, trimesh.visual.color.ColorVisuals): # If it's color visuals, it definitely has no UVs. # We must generate them or skip. pass elif hasattr(geom, 'texture_visuals') or hasattr(geom.visual, 'material'): # Sometimes UVs are hidden in material or other attributes in some trimesh versions try: geom.visual = geom.visual.to_texture() has_uv = hasattr(geom.visual, 'uv') and geom.visual.uv is not None except Exception: pass if not generate_uv and not has_uv: print(f"[SegviGen] Bake Warning: Skipping {node_name} - No UVs found (Visual type: {type(geom.visual)}).") continue print(f"[SegviGen] Bake: sampling {node_name}...") # Apply normalization to vertices for sampling (must use World Space) # 1. Transform local vertices to world space world_vertices = trimesh.transformations.transform_points(geom.vertices, transform) # 2. Normalize world vertices to fit AI unit cube [-0.5, 0.5] norm_vertices = (world_vertices - center) * scale # LOGS FOR USER: v_orig_min = world_vertices.min(axis=0).tolist() v_orig_max = world_vertices.max(axis=0).tolist() v_norm_min = norm_vertices.min(axis=0).tolist() v_norm_max = norm_vertices.max(axis=0).tolist() print(f"[Log - Bake Input] {node_name}: Original Bounds={v_orig_min} to {v_orig_max}") print(f"[Log - Bake Input] {node_name}: Normalized Bounds={v_norm_min} to {v_norm_max} (Target: [-0.5, 0.5])") if not generate_uv and has_uv: # --- UV Preservation Mode: Exact replication of to_glb sampling --- # Key steps copied from o_voxel/postprocess.py to_glb(): # 1. Rasterize in UV space (NO V-flip - to_glb has none) # 2. Interpolate 3D positions # 3. BVH-snap positions onto original mesh surface # 4. Sample voxel attributes using proper aabb/voxel_size print(f"[SegviGen] Bake: Preserving original UVs for {node_name}...") verts = torch.from_numpy(norm_vertices).float().cuda() faces_t = torch.from_numpy(geom.faces).int().cuda() uvs = torch.from_numpy(geom.visual.uv).float().cuda() # Compute per-axis voxel_size from aabb and actual grid dims aabb_t = torch.tensor(aabb, dtype=torch.float32, device=device) grid_size_t = torch.tensor([resolution, resolution, resolution], dtype=torch.float32, device=device) voxel_size_t = (aabb_t[1] - aabb_t[0]) / grid_size_t # (3,) # Build BVH for snapping (same as to_glb: uses original normalized vertices) # cumesh may or may not be importable in this context - try lazily try: import cumesh as _cumesh bvh = _cumesh.cuBVH(verts, faces_t) bvh_ok = True except Exception: bvh_ok = False # Setup rasterizer ctx = nr.RasterizeCudaContext() # NO V-flip — matching to_glb line 254: "out_uvs * 2 - 1" uvs_rast = torch.cat([uvs * 2 - 1, torch.zeros_like(uvs[:, :1]), torch.ones_like(uvs[:, :1])], dim=-1).unsqueeze(0) rast = torch.zeros((1, texture_size, texture_size, 4), device='cuda', dtype=torch.float32) # Rasterize in chunks (to_glb does the same +i trick for accumulation) for i in range(0, faces_t.shape[0], 100000): rast_chunk, _ = nr.rasterize(ctx, uvs_rast, faces_t[i:i+100000], resolution=[texture_size, texture_size]) mask_chunk = rast_chunk[..., 3:4] > 0 rast_chunk[..., 3:4] += i rast = torch.where(mask_chunk, rast_chunk, rast) mask = rast[0, ..., 3] > 0 # Interpolate 3D positions in UV space (produces approx positions) pos = nr.interpolate(verts.unsqueeze(0), rast, faces_t)[0][0] valid_pos = pos[mask] # BVH snap: project valid_pos exactly onto original mesh surface # This is the critical step from to_glb line 276-278! if bvh_ok: try: _, face_id, uvw = bvh.unsigned_distance(valid_pos, return_uvw=True) orig_tri_verts = verts[faces_t[face_id.long()]] # (N, 3, 3) valid_pos = (orig_tri_verts * uvw.unsqueeze(-1)).sum(dim=1) except Exception as e: print(f"[SegviGen] Bake: BVH snap failed ({e}), using raw rasterized positions") # Sample attribute volume: grid = (pos - aabb[0]) / voxel_size (index space) sampled_attrs = torch.zeros(texture_size, texture_size, attr_volume.shape[1], device='cuda') sampled_attrs[mask] = grid_sample_3d( attr_volume, torch.cat([torch.zeros_like(attr_coords[:, :1]), attr_coords], dim=-1), shape=torch.Size([1, attr_volume.shape[1], resolution, resolution, resolution]), grid=((valid_pos - aabb_t[0]) / voxel_size_t).reshape(1, -1, 3), mode='trilinear', ) # Construct material (matching to_glb lines 302-325) import cv2 mask_np = mask.cpu().numpy() mask_inv = (~mask_np).astype(np.uint8) base_color = np.clip(sampled_attrs[..., pbr_attr_layout['base_color']].cpu().numpy() * 255, 0, 255).astype(np.uint8) metallic = np.clip(sampled_attrs[..., pbr_attr_layout['metallic']].cpu().numpy() * 255, 0, 255).astype(np.uint8) roughness = np.clip(sampled_attrs[..., pbr_attr_layout['roughness']].cpu().numpy() * 255, 0, 255).astype(np.uint8) alpha = np.clip(sampled_attrs[..., pbr_attr_layout['alpha']].cpu().numpy() * 255, 0, 255).astype(np.uint8) base_color = cv2.inpaint(base_color, mask_inv, 3, cv2.INPAINT_TELEA) metallic = cv2.inpaint(metallic, mask_inv, 1, cv2.INPAINT_TELEA)[..., None] roughness = cv2.inpaint(roughness, mask_inv, 1, cv2.INPAINT_TELEA)[..., None] alpha = cv2.inpaint(alpha, mask_inv, 1, cv2.INPAINT_TELEA)[..., None] material = trimesh.visual.material.PBRMaterial( baseColorTexture=Image.fromarray(np.concatenate([base_color, alpha], axis=-1)), baseColorFactor=np.array([255, 255, 255, 255], dtype=np.uint8), metallicRoughnessTexture=Image.fromarray(np.concatenate([np.zeros_like(metallic), roughness, metallic], axis=-1)), metallicFactor=1.0, roughnessFactor=1.0, doubleSided=True, ) # Flip UV V for the saved mesh (to_glb does uvs_np[:, 1] = 1 - uvs_np[:, 1]) original_uv = geom.visual.uv.copy() flipped_uv = original_uv.copy() flipped_uv[:, 1] = 1 - flipped_uv[:, 1] geom.visual = trimesh.visual.TextureVisuals(uv=flipped_uv, material=material) baked_mesh = geom world_baked_vertices = world_vertices else: # --- Standard Pipeline (to_glb with new UVs) --- # Use to_glb with sparse input baked_mesh = o_voxel.postprocess.to_glb( vertices=torch.from_numpy(norm_vertices).float().to(device).contiguous(), faces=torch.from_numpy(geom.faces).int().to(device).contiguous(), attr_volume=attr_volume, coords=attr_coords, attr_layout=pbr_attr_layout, aabb=aabb, voxel_size=1.0 / resolution, texture_size=texture_size, remesh=False, verbose=False ) if isinstance(baked_mesh, trimesh.Scene): baked_mesh = list(baked_mesh.geometry.values())[0] # 1. Reverse the axis swap/invert from to_glb (Y->Z, Z->-Y) # Internal to_glb: v[:, 1], v[:, 2] = v[:, 2], -v[:, 1] y_glb = baked_mesh.vertices[:, 1].copy() z_glb = baked_mesh.vertices[:, 2].copy() baked_mesh.vertices[:, 1] = -z_glb # restore original Y baked_mesh.vertices[:, 2] = y_glb # restore original Z # 2. Denormalize output vertices to match original GLB coordinates (Back to World Space) world_baked_vertices = (baked_mesh.vertices / scale) + center # 3. Transform back to LOCAL space for replacement in the original node inv_transform = np.linalg.inv(transform) baked_mesh.vertices = trimesh.transformations.transform_points(world_baked_vertices, inv_transform) # [Comparison: AI Voxel vs Mesh Vertex] print(f"[Comparison] Picking 3 sample vertices from {node_name} to verify alignment:") v_indices_sample = torch.linspace(0, len(geom.vertices) - 1, 3).long() v_tensor_sample = torch.from_numpy(norm_vertices).float().to(device) # --- Manual Trilinear Sampling (Exact Match) --- # 1. Prepare original coords (Z, Y, X) for axis detection orig_attr_coords_3 = torch.cat([vox.coords for vox in tex_voxels]).to(device)[:, -3:] # 2. Detect axis mapping by comparing ranges mesh_ranges = norm_vertices.max(axis=0) - norm_vertices.min(axis=0) vox_ranges = (orig_attr_coords_3.max(dim=0)[0] - orig_attr_coords_3.min(dim=0)[0]).float().cpu().numpy() m_sort = np.argsort(mesh_ranges) v_sort = np.argsort(vox_ranges) axis_map = {m_idx: v_idx for m_idx, v_idx in zip(m_sort, v_sort)} for v_idx_sample in v_indices_sample: v_idx_sample = v_idx_sample.item() v_pos_sample = v_tensor_sample[v_idx_sample].cpu().numpy() # [3] (X, Y, Z normalized) # Map mesh units to float grid positions grid_pos = np.zeros(3) for m_i, v_i in axis_map.items(): grid_pos[v_i] = (v_pos_sample[m_i] + 0.5) * (resolution - 1) # Manual Trilinear Interpolation # Get 8 neighbors (floor/ceil) p0 = np.floor(grid_pos).astype(int) p1 = np.clip(p0 + 1, 0, resolution - 1) d = grid_pos - p0 # fractional parts # Sparse Lookup for 8 corners def get_vox_color(pos): pos_t = torch.from_numpy(pos).to(device) mask = torch.all(torch.eq(orig_attr_coords_3, pos_t), dim=1) if not torch.any(mask): return torch.zeros(3).to(device) idx = torch.where(mask)[0][0] return attr_volume[idx, pbr_attr_layout['base_color']] # Sample 8 corners [z, y, x] c000 = get_vox_color(p0) c001 = get_vox_color(np.array([p0[0], p0[1], p1[2]])) c010 = get_vox_color(np.array([p0[0], p1[1], p0[2]])) c011 = get_vox_color(np.array([p0[0], p1[1], p1[2]])) c100 = get_vox_color(np.array([p1[0], p0[1], p0[2]])) c101 = get_vox_color(np.array([p1[0], p0[1], p1[2]])) c110 = get_vox_color(np.array([p1[0], p1[1], p0[2]])) c111 = get_vox_color(p1) # Lerp across 3 axes # x-axis c00 = c000 * (1 - d[2]) + c001 * d[2] c01 = c010 * (1 - d[2]) + c011 * d[2] c10 = c100 * (1 - d[2]) + c101 * d[2] c11 = c110 * (1 - d[2]) + c111 * d[2] # y-axis c0 = c00 * (1 - d[1]) + c01 * d[1] c1 = c10 * (1 - d[1]) + c11 * d[1] # z-axis final_c = c0 * (1 - d[0]) + c1 * d[0] near_color = (final_c * 255).int().tolist() # --- Mesh Side: Find nearest vertex on baked mesh (Topology-Aware) --- v_orig = world_vertices[v_idx_sample] # Original world pos # baked_mesh.vertices is LOCAL here, so use world_baked_vertices for comparison dist_baked_all = np.linalg.norm(world_baked_vertices - v_orig, axis=1) v_idx_baked = np.argmin(dist_baked_all) match_dist = dist_baked_all[v_idx_baked] # Sample the Baked Color from the finished mesh baked_color = "N/A" try: img = None vis = baked_mesh.visual def find_image(obj, depth=0): if depth > 2: return None if hasattr(obj, 'image') and obj.image is not None: return obj.image if hasattr(obj, 'diffuse') and obj.diffuse is not None: return obj.diffuse for attr in ['baseColorTexture', 'base_color_texture']: if hasattr(obj, attr): sub = getattr(obj, attr) if hasattr(sub, 'image'): return sub.image if isinstance(sub, Image.Image): return sub return None img = find_image(vis) if img is None and hasattr(vis, 'material'): img = find_image(vis.material) if img is not None: u, v_coord = baked_mesh.visual.uv[v_idx_baked] w, h = img.size px, py = int(u * (w-1)), int((1.0 - v_coord) * (h-1)) baked_color = list(img.getpixel((px, py))[:3]) except Exception: pass print(f" - Vertex {v_idx_sample} (Match Dist: {match_dist:.4f})") print(f" -> AI Expected Color: RGB {near_color}") print(f" -> Baked Color on Mesh: RGB {baked_color}") print("-" * 50) v_final_max = world_baked_vertices.max(axis=0).tolist() v_final_min = world_baked_vertices.min(axis=0).tolist() print(f"[Log - Bake Output] {node_name}: Final Denormalized (World) Bounds={v_final_min} to {v_final_max}") # Create a unique geometry for this node if we want to preserve instances with different bakes # However, for simple replacement: asset.geometry[geom_name] = baked_mesh asset.export(output_path) print(f"[SegviGen] Bake: saved to {output_path}") return output_path def slat_to_glb(meshes, tex_voxels, resolution=512): pbr_attr_layout = { 'base_color': slice(0, 3), 'metallic': slice(3, 4), 'roughness': slice(4, 5), 'alpha': slice(5, 6), } out_mesh = [] for i, (m, v) in enumerate(zip(meshes, tex_voxels)): try: v_min = m.vertices.min(dim=0).values.tolist() v_max = m.vertices.max(dim=0).values.tolist() print(f"[Log - Slat Input] Component {i}: vertices={len(m.vertices)}, bounds={v_min} to {v_max}") except Exception: pass out_mesh.append( MeshWithVoxel( m.vertices, m.faces, origin=[-0.5, -0.5, -0.5], voxel_size=1 / resolution, coords=v.coords[:, 1:], attrs=v.feats, voxel_shape=torch.Size([*v.shape, *v.spatial_shape]), layout=pbr_attr_layout, ) ) mesh = out_mesh[0] try: mesh.simplify(200000) except Exception as e: print(f"[Export] mesh.simplify skipped: {e}") glb = o_voxel.postprocess.to_glb( vertices=mesh.vertices, faces=mesh.faces, attr_volume=mesh.attrs.cuda(), coords=mesh.coords.cuda(), attr_layout=mesh.layout, voxel_size=mesh.voxel_size, aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]], decimation_target=50000, texture_size=EXPORT_TEXTURE_SIZE, remesh=True, remesh_band=1, remesh_project=0, verbose=False, ) return glb class _LoadedPipeline: def __init__(self): self.loaded = False self.current_ckpt = None self.pipeline_args = None self.tex_slat_flow_model = None self.gen3dseg = None self.sampler = None self.shape_encoder = None self.tex_encoder = None self.shape_decoder = None self.tex_decoder = None self.rembg_model = None self.image_cond_model = None def load_all_models(self): if self.loaded: return print("-" * 100) print("[Init] Loading pipeline config .....................") with open(TRELLIS_PIPELINE_JSON, "r") as f: pipeline_config = json.load(f) self.pipeline_args = pipeline_config['args'] self.sampler = Sampler() self.loaded = True print("[Init] Config loaded. Models will be loaded on-demand and unloaded to save RAM.") def _download_if_missing(self, repo_id, local_root, subfolder=None): # Target dir is where we expect the files to actually exist target_dir = os.path.join(local_root, subfolder) if subfolder else local_root # Check if target dir exists and is not empty if os.path.exists(target_dir) and os.listdir(target_dir): # Already exists locally, skip download return target_dir from huggingface_hub import snapshot_download parts = repo_id.split("/") actual_repo_id = "/".join(parts[:2]) # If repo_id had more than 2 parts, the rest is the subfolder pattern if subfolder is None and len(parts) > 2: subfolder = "/".join(parts[2:]) print(f"[SegviGen] Downloading {actual_repo_id} (subfolder: {subfolder}) to {local_root}...") os.makedirs(local_root, exist_ok=True) snapshot_download( repo_id=actual_repo_id, local_dir=local_root, allow_patterns=[f"{subfolder}/*"] if subfolder else None ) return target_dir def get_rembg(self): if self.rembg_model is None: print(f"[OnDemand] Loading RMBG-2.0") target_path = os.path.join(COMFY_MODELS_DIR, "Aero-Ex", "RMBG-2.0") self._download_if_missing("Aero-Ex/RMBG-2.0", target_path) self.rembg_model = BiRefNet(model_name=target_path) self.rembg_model.eval() return self.rembg_model def get_cond_model(self): if self.image_cond_model is None: print(f"[OnDemand] Loading DinoV3 conditioners") # repo Aero-Ex/Dinov3 has subfolder facebook/dinov3-vitl16-pretrain-lvd1689m base_dir = COMFY_MODELS_DIR sub = "facebook/dinov3-vitl16-pretrain-lvd1689m" self._download_if_missing(DINO_PATH, base_dir, subfolder=sub) # The files will be in base_dir/sub actual_local_dir = os.path.join(base_dir, sub) self.image_cond_model = DinoV3FeatureExtractor(DINO_PATH, local_dir=actual_local_dir, subfolder=sub) self.image_cond_model.eval() return self.image_cond_model def get_encoders_decoder(self): # Load encoders and shape decoder together as they are used in vxz_to_latent_slat if self.shape_encoder is None: print(f"[OnDemand] Loading Encoders and Shape Decoder") # Download to COMFY_MODELS_DIR root; snapshot_download keeps subfolder structure enc_path = self._download_if_missing(TRELLIS_SHAPE_ENC, COMFY_MODELS_DIR) tex_enc_path = self._download_if_missing(TRELLIS_TEX_ENC, COMFY_MODELS_DIR) dec_path = self._download_if_missing(TRELLIS_SHAPE_DEC, COMFY_MODELS_DIR) self.shape_encoder = models.from_pretrained(TRELLIS_SHAPE_ENC, local_dir=enc_path).eval() self.tex_encoder = models.from_pretrained(TRELLIS_TEX_ENC, local_dir=tex_enc_path).eval() self.shape_decoder = models.from_pretrained(TRELLIS_SHAPE_DEC, local_dir=dec_path).eval() return self.shape_encoder, self.tex_encoder, self.shape_decoder def get_gen3dseg(self): if self.gen3dseg is None: print(f"[OnDemand] Loading Backbone (Gen3DSeg)") flow_path = self._download_if_missing(TRELLIS_TEX_FLOW, COMFY_MODELS_DIR) self.tex_slat_flow_model = models.from_pretrained(TRELLIS_TEX_FLOW, local_dir=flow_path) self.gen3dseg = Gen3DSeg(self.tex_slat_flow_model) local_ckpt = self.current_ckpt if local_ckpt is not None: print(f"[OnDemand] Applying deferred checkpoint: {local_ckpt}") if local_ckpt.endswith(".safetensors"): from safetensors.torch import load_file state_dict = load_file(local_ckpt) else: # Use weights_only=False for legacy .ckpt files if on newer torch try: state_dict = torch.load(local_ckpt, weights_only=False) except TypeError: state_dict = torch.load(local_ckpt) if 'state_dict' in state_dict: state_dict = state_dict['state_dict'] state_dict = OrderedDict([(k.replace("gen3dseg.", ""), v) for k, v in state_dict.items()]) self.gen3dseg.load_state_dict(state_dict) self.gen3dseg.eval() return self.gen3dseg def get_tex_decoder(self): if self.tex_decoder is None: print(f"[OnDemand] Loading Texture Decoder") dec_path = self._download_if_missing(TRELLIS_TEX_DEC, COMFY_MODELS_DIR) self.tex_decoder = models.from_pretrained(TRELLIS_TEX_DEC, local_dir=dec_path).eval() return self.tex_decoder def unload(self, *attr_names): import gc for name in attr_names: val = getattr(self, name, None) if val is not None: if hasattr(val, 'cpu'): val.cpu() setattr(self, name, None) gc.collect() torch.cuda.empty_cache() def load_ckpt_if_needed(self, ckpt_path: str): if self.current_ckpt == ckpt_path: return self.current_ckpt = ckpt_path print(f"[OnDemand] Checkpoint set to: {ckpt_path}. Will be applied when model is loaded.") PIPE = _LoadedPipeline() def inference_with_loaded_models(ckpt_path, item): PIPE.load_all_models() PIPE.load_ckpt_if_needed(ckpt_path) work_glb = item["glb"] if EARLY_SIMPLIFY_ENABLED: work_glb = get_work_glb_path(item) build_simplified_work_glb( input_glb_path=item["glb"], output_glb_path=work_glb, target_faces=EARLY_SIMPLIFY_TARGET_FACES, aggression=EARLY_SIMPLIFY_AGGRESSION, ) if not item["2d_map"]: generate_2d_map_from_glb( glb_path=work_glb, transforms_path=item["transforms"], out_img_path=item["img"], ) process_glb_to_vxz( glb_path=item["glb"], vxz_path=item["input_vxz"], ) image = Image.open(item["img"]) # 1. Background removal rembg_model = PIPE.get_rembg() rembg_model.cuda() image = preprocess_image(rembg_model, image) PIPE.unload('rembg_model') # 2. Condition generation cond_model = PIPE.get_cond_model() cond_model.cuda() cond = get_cond(cond_model, [image]) # Offload cond to CPU cond = {k: v.cpu() for k, v in cond.items()} PIPE.unload('image_cond_model') # 3. VXZ to Latent SLat (Encoders + Shape Decoder) shape_enc, tex_enc, shape_dec = PIPE.get_encoders_decoder() shape_enc.cuda() tex_enc.cuda() shape_dec.cuda() shape_slat, meshes, subs, tex_slat = vxz_to_latent_slat( shape_enc, shape_dec, tex_enc, item["input_vxz"], ) # Offload products to CPU shape_slat = sp.SparseTensor(shape_slat.feats.cpu(), shape_slat.coords.cpu()) tex_slat = sp.SparseTensor(tex_slat.feats.cpu(), tex_slat.coords.cpu()) # subs are usually small indices or similar, keep them as is PIPE.unload('shape_encoder', 'tex_encoder', 'shape_decoder') # 4. Sampling (Backbone) gen3dseg = PIPE.get_gen3dseg() gen3dseg.cuda() # Move inputs to GPU for sampling shape_slat_gpu = sp.SparseTensor(shape_slat.feats.cuda(), shape_slat.coords.cuda()) tex_slat_gpu = sp.SparseTensor(tex_slat.feats.cuda(), tex_slat.coords.cuda()) cond_gpu = {k: v.cuda() for k, v in cond.items()} output_tex_slat = tex_slat_sample_single( gen3dseg, PIPE.sampler, PIPE.pipeline_args, shape_slat_gpu, tex_slat_gpu, cond_gpu ) # Offload result to CPU output_tex_slat_cpu = sp.SparseTensor(output_tex_slat.feats.cpu(), output_tex_slat.coords.cpu()) # Cleanup GPU inputs del shape_slat_gpu, tex_slat_gpu, cond_gpu, output_tex_slat PIPE.unload('gen3dseg', 'tex_slat_flow_model') # 5. Texture Decoding tex_decoder = PIPE.get_tex_decoder() tex_decoder.cuda() output_tex_slat_gpu = sp.SparseTensor(output_tex_slat_cpu.feats.cuda(), output_tex_slat_cpu.coords.cuda()) with torch.no_grad(): subs_gpu = [s.cuda() if isinstance(s, torch.Tensor) else s for s in subs] tex_voxels = tex_decoder(output_tex_slat_gpu, guide_subs=subs_gpu) * 0.5 + 0.5 # Move result to CPU immediately tex_voxels = [v.cpu() for v in tex_voxels] PIPE.unload('tex_decoder') if item.get("bake", False): bake_to_mesh( item["glb"], tex_voxels, item["export_glb"], resolution=512, texture_size=2048, generate_uv=item.get("generate_uv", False) ) else: glb = slat_to_glb(meshes, tex_voxels) T = np.eye(4, dtype=np.float64) T[:3, :3] = np.array( [ [1, 0, 0], [0, 0, -1], [0, 1, 0], ], dtype=np.float64, ) if hasattr(glb, "apply_transform") and callable(getattr(glb, "apply_transform")): glb.apply_transform(T) glb.export(item["export_glb"]) else: glb.export(item["export_glb"]) scene_or_mesh = trimesh.load(item["export_glb"], force="scene") scene_or_mesh.apply_transform(T) scene_or_mesh.export(item["export_glb"])