git init

.gitattributes

@@ -1,35 +1,2 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.glb filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,12 +1,66 @@
----
-title: SegviGen
-emoji: ⚡
-colorFrom: gray
-colorTo: gray
-sdk: gradio
-sdk_version: 6.9.0
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# SegviGen: Repurposing 3D Generative Model for Part Segmentation
+
+![teaser](assets/teaser.png)
+
+***SegviGen*** is a framework for 3D part segmentation that leverages the rich 3D structural and textural knowledge encoded in large-scale 3D generative models. 
+It learns to predict part-indicative colors while reconstructing geometry, and unifies three settings in one architecture: **interactive part segmentation**, **full segmentation**, and **2D segmentation map–guided full segmentation** with arbitrary granularity.
+
+
+## 🌟 Features
+- **Repurposed 3D Generative Priors for Data Efficiency**: By reusing the rich structural and textural knowledge encoded in large-scale native 3D generative models, ***SegviGen*** learns 3D part segmentation with minimal task-specific supervision, requiring only **0.32%** training data.
+- **Unified and Flexible Segmentation Settings**: Supports **interactive part segmentation**, **full segmentation**, and **2D segmentation map–guided full segmentation** with arbitrary part granularity under a single architecture.
+- **State-of-the-Art Accuracy**: Consistently surpasses P3-SAM, delivering a **40%** gain in IoU@1 for single-click interaction on PartObjaverse-Tiny and PartNeXT, and a **15%** improvement in overall IoU for unguided full segmentation averaged across datasets.
+
+
+## 🔨Installation
+
+### Prerequisites
+- **System**: Linux
+- **GPU**: A NVIDIA GPU with at least 24GB of memory is necessary
+- **Python**: 3.10
+
+### Installation Steps
+1. Create the environment of [TRELLIS.2](https://github.com/microsoft/TRELLIS.2)
+    ```sh
+    git clone -b main https://github.com/microsoft/TRELLIS.2.git --recursive
+    cd TRELLIS.2
+    ./setup.sh --new-env --basic --flash-attn --nvdiffrast --nvdiffrec --cumesh --o-voxel --flexgemm
+    ```
+
+2. Install the rest of requirements
+    ```sh
+    pip install mathutils
+    pip install transformers==4.57.6 # https://github.com/microsoft/TRELLIS.2/issues/101
+    pip install bpy==4.0.0 --extra-index-url https://download.blender.org/pypi/
+    sudo apt-get install -y libsm6 libxrender1 libxext6
+    pip install --upgrade Pillow
+    ```
+
+3. If want to train
+    ```sh
+    pip install pytorch_lightning
+    ```
+
+### Pretrained Weights
+
+The checkpoints of **Interactive part-segmentation**, **Full segmentation** and **Full segmentation with 2D guidance** are available on [Hugging Face](https://huggingface.co/Nelipot/tmp).
+
+## 📒Usage
+
+ - `inference_interactive.py`: **Interactive part-segmentation**
+ - `inference_full.py`: **Full segmentation** or **Full segmentation with 2D guidance**
+ - `inference_unified.py`: Unified model
+
+## Training
+
+### Data preparation
+
+ - `data_toolkit/example_interactive_seg.py`: **Interactive part-segmentation**
+ - `data_toolkit/example_full_seg.py`: **Full segmentation**
+ - `data_toolkit/example_full_seg_w_2d_map.py`: **Full segmentation with 2D guidance**
+
+### Running training
+
+ - `train_interactive.py`: **Interactive part-segmentation**
+ - `train_full.py`: **Full segmentation** or **Full segmentation with 2D guidance**
+ - `train_unified.py`: Unified model

app.py

@@ -0,0 +1,418 @@
+import os
+import urllib.request
+
+os.makedirs("pretrained_model", exist_ok=True)
+
+CKPT_FULL_SEG = "pretrained_model/full_seg.ckpt"
+CKPT_W_2D_MAP = "pretrained_model/full_seg_w_2d_map.ckpt"
+
+if not os.path.exists(CKPT_FULL_SEG):
+    urllib.request.urlretrieve(
+        "https://huggingface.co/fenghora/SegviGen/resolve/main/full_seg.ckpt",
+        CKPT_FULL_SEG,
+    )
+
+if not os.path.exists(CKPT_W_2D_MAP):
+    urllib.request.urlretrieve(
+        "https://huggingface.co/fenghora/SegviGen/resolve/main/full_seg_w_2d_map.ckpt",
+        CKPT_W_2D_MAP,
+    )
+
+import shutil
+import traceback
+from datetime import datetime
+from pathlib import Path
+from typing import List
+import inference_full as inf 
+import split as splitter
+
+
+TRANSFORMS_JSON = "./data_toolkit/transforms.json"
+
+ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
+TMP_DIR = os.path.join(ROOT_DIR, "_tmp_gradio_seg")
+EXAMPLES_CACHE_DIR = os.path.join(TMP_DIR, "examples_cache")
+os.makedirs(TMP_DIR, exist_ok=True)
+os.makedirs(EXAMPLES_CACHE_DIR, exist_ok=True)
+
+os.environ["GRADIO_TEMP_DIR"] = TMP_DIR
+os.environ["GRADIO_EXAMPLES_CACHE"] = EXAMPLES_CACHE_DIR
+
+import gradio as gr
+
+EXAMPLES_DIR = "examples"
+
+
+def _ensure_dir(p: str):
+    os.makedirs(p, exist_ok=True)
+
+
+def _normalize_path(x):
+    """
+    Compatible with different Gradio versions: File/Model3D might be str / dict / object
+    """
+    if x is None:
+        return None
+    if isinstance(x, str):
+        return x
+    if isinstance(x, dict):
+        return x.get("name") or x.get("path") or x.get("data")
+    return getattr(x, "name", None) or getattr(x, "path", None) or None
+
+
+def _raise_user_error(msg: str):
+    if hasattr(gr, "Error"):
+        raise gr.Error(msg)
+    raise RuntimeError(msg)
+
+
+def _collect_examples(example_dir: str) -> List[List[str]]:
+    """
+    Scan example_dir for pairs: <name>.glb + <name>.png
+    Return a list of examples: [[glb_path, png_path], ...]
+    """
+    d = Path(example_dir)
+    if not d.is_dir():
+        return []
+
+    examples: List[List[str]] = []
+
+    # Search recursively in case you add subfolders later
+    glb_files = sorted(d.rglob("*.glb"))
+    for glb_path in glb_files:
+        png_path = glb_path.with_suffix(".png")
+        if png_path.is_file():
+            examples.append([str(glb_path), str(png_path)])
+        # If png is missing, skip to keep examples consistent (2 inputs required)
+
+    return examples
+
+
+# Build examples once at startup
+FULL_SEG_EXAMPLES = _collect_examples(EXAMPLES_DIR)
+
+
+def run_seg(glb_in, img_in):
+    """
+    Segment button: generates whole segmented GLB and displays in the second box.
+    Returns: segmented_glb_path, segmented_glb_path(state)
+    """
+    try:
+        glb_path = _normalize_path(glb_in)
+        img_path = _normalize_path(img_in)
+
+        if glb_path is None or (not os.path.isfile(glb_path)):
+            _raise_user_error("Please upload a valid .glb file.")
+
+        _ensure_dir(TMP_DIR)
+        run_id = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
+        workdir = os.path.join(TMP_DIR, run_id)
+        _ensure_dir(workdir)
+
+        in_glb = os.path.join(workdir, "input.glb")
+        shutil.copy(glb_path, in_glb)
+
+        out_glb = os.path.join(workdir, "segmented.glb")
+        in_vxz = os.path.join(workdir, "input.vxz")
+
+        # If image is provided -> 2d_map=True; otherwise full segmentation (render_from_transforms)
+        if img_path is not None and os.path.isfile(img_path):
+            ckpt = CKPT_W_2D_MAP
+            in_img = os.path.join(workdir, "2d_map.png")
+            shutil.copy(img_path, in_img)
+            item = {
+                "2d_map": True,
+                "glb": in_glb,
+                "input_vxz": in_vxz,
+                "img": in_img,
+                "export_glb": out_glb,
+            }
+        else:
+            ckpt = CKPT_FULL_SEG
+            render_img = os.path.join(workdir, "render.png")
+            item = {
+                "2d_map": False,
+                "glb": in_glb,
+                "input_vxz": in_vxz,
+                "transforms": TRANSFORMS_JSON,
+                "img": render_img,
+                "export_glb": out_glb,
+            }
+
+        inf.inference_with_loaded_models(ckpt, item)
+
+        if not os.path.isfile(out_glb):
+            _raise_user_error("Export failed: output glb not found.")
+
+        # Apply X90 rotation for whole segmented output
+        # _apply_root_x90_rotation_glb(out_glb)
+
+        return out_glb, out_glb
+
+    except Exception as e:
+        err = "".join(traceback.format_exception(type(e), e, e.__traceback__))
+        print(err)
+        raise
+
+
+def run_refine_segmentation(
+    seg_glb_path_state,
+    color_quant_step,
+    palette_sample_pixels,
+    palette_min_pixels,
+    palette_max_colors,
+    palette_merge_dist,
+    samples_per_face,
+    flip_v,
+    uv_wrap_repeat,
+    transition_conf_thresh,
+    transition_prop_iters,
+    transition_neighbor_min,
+    small_component_action,
+    small_component_min_faces,
+    postprocess_iters,
+    min_faces_per_part,
+    bake_transforms,
+):
+    """
+    Refine Segmentation button: splits the segmented GLB into smaller parts GLB and displays in the fourth box.
+    """
+    try:
+        seg_glb_path = seg_glb_path_state if isinstance(seg_glb_path_state, str) else None
+        if (seg_glb_path is None) or (not os.path.isfile(seg_glb_path)):
+            _raise_user_error("Please run Segmentation first (the segmented GLB is missing).")
+
+        out_dir = os.path.dirname(seg_glb_path)
+        out_parts_glb = os.path.join(out_dir, "segmented_parts.glb")
+
+        splitter.split_glb_by_texture_palette_rgb(
+            in_glb_path=seg_glb_path,
+            out_glb_path=out_parts_glb,
+            min_faces_per_part=min_faces_per_part,
+            bake_transforms=bool(bake_transforms),
+            color_quant_step=color_quant_step,
+            palette_sample_pixels=palette_sample_pixels,
+            palette_min_pixels=palette_min_pixels,
+            palette_max_colors=palette_max_colors,
+            palette_merge_dist=palette_merge_dist,
+            samples_per_face=samples_per_face,
+            flip_v=flip_v,
+            uv_wrap_repeat=uv_wrap_repeat,
+            transition_conf_thresh=transition_conf_thresh,
+            transition_prop_iters=transition_prop_iters,
+            transition_neighbor_min=transition_neighbor_min,
+            small_component_action=small_component_action,
+            small_component_min_faces=small_component_min_faces,
+            postprocess_iters=postprocess_iters,
+            debug_print=True,
+        )
+
+        if not os.path.isfile(out_parts_glb):
+            _raise_user_error("Split failed: output parts glb not found.")
+
+        # If bake_transforms=False, split output will not have the wrapper transform baked, so we need to apply X90 rotation fix
+        # if (not bool(bake_transforms)) and APPLY_OUTPUT_X90_FIX:
+        #     _apply_root_x90_rotation_glb(out_parts_glb)
+
+        return out_parts_glb
+
+    except Exception as e:
+        err = "".join(traceback.format_exception(type(e), e, e.__traceback__))
+        print(err)
+        raise
+
+
+CSS_TEXT = """
+<style>
+#in_glb  { height: 520px !important; }
+#seg_glb { height: 520px !important; }
+#part_glb{ height: 520px !important; }
+#img     { height: 520px !important; }
+</style>
+"""
+
+with gr.Blocks() as demo:
+    gr.HTML(CSS_TEXT)
+    gr.Markdown(
+        """
+# SegviGen: Repurposing 3D Generative Model for Part Segmentation
+"""
+    )
+
+    # ---------------- 2x2 Layout ----------------
+    with gr.Row():
+        with gr.Column(scale=1, min_width=260):
+            in_glb = gr.Model3D(label="Input GLB", elem_id="in_glb")
+        with gr.Column(scale=1, min_width=260):
+            seg_glb = gr.Model3D(label="Processed GLB", elem_id="seg_glb")
+
+    with gr.Row():
+        with gr.Column(scale=1, min_width=260):
+            with gr.Accordion("2D Segmentation Map (Optional)", open=False):
+                in_img = gr.Image(label="2D Segmentation Map", type="filepath", elem_id="img")
+
+            seg_btn = gr.Button("Process", variant="primary")
+
+            # ✅ Examples directly under the Process button
+            if FULL_SEG_EXAMPLES:
+                gr.Examples(
+                    examples=FULL_SEG_EXAMPLES,
+                    inputs=[in_glb, in_img],
+                    label="Examples",
+                    examples_per_page=3,
+                    cache_examples=False,
+                )
+            else:
+                gr.Markdown(f"**No examples found** in: `{EXAMPLES_DIR}` (expected: `*.glb` + same-name `*.png`).")
+
+            with gr.Accordion("Advanced segmentation options", open=False):
+                def _g(name, default):
+                    return getattr(splitter, name, default)
+
+                color_quant_step = gr.Slider(
+                    1, 64, value=_g("COLOR_QUANT_STEP", 16), step=1, label="COLOR_QUANT_STEP"
+                )
+                gr.Markdown(
+                    "*COLOR_QUANT_STEP controls the RGB quantization step, where a larger value merges similar colors more aggressively and a smaller value preserves finer color differences.*"
+                )
+
+                palette_sample_pixels = gr.Number(
+                    value=_g("PALETTE_SAMPLE_PIXELS", 2_000_000), precision=0, label="PALETTE_SAMPLE_PIXELS"
+                )
+                gr.Markdown(
+                    "*PALETTE_SAMPLE_PIXELS sets the maximum number of sampled pixels used to estimate the palette, where more samples improve stability but increase runtime.*"
+                )
+
+                palette_min_pixels = gr.Number(
+                    value=_g("PALETTE_MIN_PIXELS", 500), precision=0, label="PALETTE_MIN_PIXELS"
+                )
+                gr.Markdown(
+                    "*PALETTE_MIN_PIXELS specifies the minimum pixel count required to keep a color in the palette, where a higher threshold suppresses noise but may discard small parts.*"
+                )
+
+                palette_max_colors = gr.Number(
+                    value=_g("PALETTE_MAX_COLORS", 256), precision=0, label="PALETTE_MAX_COLORS"
+                )
+                gr.Markdown(
+                    "*PALETTE_MAX_COLORS limits the maximum number of colors retained in the palette, where a larger limit yields finer partitions and a smaller limit enforces stronger merging.*"
+                )
+
+                palette_merge_dist = gr.Number(
+                    value=_g("PALETTE_MERGE_DIST", 32), precision=0, label="PALETTE_MERGE_DIST"
+                )
+                gr.Markdown(
+                    "*PALETTE_MERGE_DIST defines the distance threshold for merging nearby palette colors in RGB space, where a larger threshold merges near duplicates more often and a smaller threshold keeps colors distinct.*"
+                )
+
+                samples_per_face = gr.Dropdown(
+                    choices=[1, 4], value=_g("SAMPLES_PER_FACE", 4), label="SAMPLES_PER_FACE"
+                )
+                gr.Markdown(
+                    "*SAMPLES_PER_FACE sets the number of UV samples per triangle used for label voting, where more samples improve robustness near boundaries but increase computation.*"
+                )
+
+                flip_v = gr.Checkbox(value=_g("FLIP_V", True), label="FLIP_V")
+                gr.Markdown(
+                    "*FLIP_V toggles whether the V coordinate is flipped to match common glTF texture conventions, and you should disable it only if the texture appears vertically inverted.*"
+                )
+
+                uv_wrap_repeat = gr.Checkbox(value=_g("UV_WRAP_REPEAT", True), label="UV_WRAP_REPEAT")
+                gr.Markdown(
+                    "*UV_WRAP_REPEAT selects how out of range UVs are handled by either repeating via modulo or clamping to the unit interval, and repeating is typically preferred for tiled textures.*"
+                )
+
+                transition_conf_thresh = gr.Slider(
+                    0.25, 1.0, value=float(_g("TRANSITION_CONF_THRESH", 1.0)), step=0.25, label="TRANSITION_CONF_THRESH"
+                )
+                gr.Markdown(
+                    "*TRANSITION_CONF_THRESH sets the confidence threshold for transition handling, where a higher value makes refinement more conservative and a lower value enables more aggressive smoothing.*"
+                )
+
+                transition_prop_iters = gr.Number(
+                    value=_g("TRANSITION_PROP_ITERS", 6), precision=0, label="TRANSITION_PROP_ITERS"
+                )
+                gr.Markdown(
+                    "*TRANSITION_PROP_ITERS specifies the number of propagation iterations used in transition refinement, where more iterations strengthen diffusion effects but increase runtime.*"
+                )
+
+                transition_neighbor_min = gr.Number(
+                    value=_g("TRANSITION_NEIGHBOR_MIN", 1), precision=0, label="TRANSITION_NEIGHBOR_MIN"
+                )
+                gr.Markdown(
+                    "*TRANSITION_NEIGHBOR_MIN requires a minimum number of supporting neighbors to propagate a label, where a higher requirement is more conservative and a lower requirement is more permissive.*"
+                )
+
+                small_component_action = gr.Dropdown(
+                    choices=["reassign", "drop"], value=_g("SMALL_COMPONENT_ACTION", "reassign"), label="SMALL_COMPONENT_ACTION"
+                )
+                gr.Markdown(
+                    "*SMALL_COMPONENT_ACTION determines how small connected components are handled by either reassigning them to neighboring labels or dropping them entirely.*"
+                )
+
+                small_component_min_faces = gr.Number(
+                    value=_g("SMALL_COMPONENT_MIN_FACES", 50), precision=0, label="SMALL_COMPONENT_MIN_FACES"
+                )
+                gr.Markdown(
+                    "*SMALL_COMPONENT_MIN_FACES defines the face count threshold used to classify a component as small, where a higher threshold merges or removes more fragments and a lower threshold preserves more small parts.*"
+                )
+
+                postprocess_iters = gr.Number(
+                    value=_g("POSTPROCESS_ITERS", 3), precision=0, label="POSTPROCESS_ITERS"
+                )
+                gr.Markdown(
+                    "*POSTPROCESS_ITERS sets the number of post processing iterations, where more iterations produce stronger cleanup at the cost of additional computation.*"
+                )
+
+                min_faces_per_part = gr.Number(
+                    value=_g("MIN_FACES_PER_PART", 1), precision=0, label="MIN_FACES_PER_PART"
+                )
+                gr.Markdown(
+                    "*MIN_FACES_PER_PART enforces a minimum number of faces per exported part, where a larger value filters tiny outputs and a smaller value retains fine components.*"
+                )
+
+                bake_transforms = gr.Checkbox(value=_g("BAKE_TRANSFORMS", True), label="BAKE_TRANSFORMS")
+                gr.Markdown(
+                    "*BAKE_TRANSFORMS controls whether scene graph transforms are baked into geometry before splitting, where enabling it improves consistency in world space and disabling it preserves node transforms.*"
+                )
+
+        with gr.Column(scale=1, min_width=260):
+            refine_btn = gr.Button("Segment", variant="secondary")
+            part_glb = gr.Model3D(label="Segmented GLB", elem_id="part_glb")
+
+    # Hidden states
+    seg_glb_state = gr.State(None)
+
+    # ---------------- wiring ----------------
+    seg_btn.click(
+        fn=run_seg,
+        inputs=[in_glb, in_img],
+        outputs=[seg_glb, seg_glb_state],
+    )
+
+    refine_btn.click(
+        fn=run_refine_segmentation,
+        inputs=[
+            seg_glb_state,
+            color_quant_step,
+            palette_sample_pixels,
+            palette_min_pixels,
+            palette_max_colors,
+            palette_merge_dist,
+            samples_per_face,
+            flip_v,
+            uv_wrap_repeat,
+            transition_conf_thresh,
+            transition_prop_iters,
+            transition_neighbor_min,
+            small_component_action,
+            small_component_min_faces,
+            postprocess_iters,
+            min_faces_per_part,
+            bake_transforms
+        ],
+        outputs=[part_glb],
+    )
+
+if __name__ == "__main__":
+    inf.PIPE.load_all_models()
+    demo.launch(server_name="0.0.0.0", server_port=8012, share=False)

color_report.json

@@ -0,0 +1,30 @@
+{
+  "input_glb": "/media/nfs/tmp_data/fenghr/SegviGen/data_toolkit/assets/output.glb",
+  "color_quant_step": 8,
+  "max_image_samples": 2000000,
+  "nodes": [
+    {
+      "node": "geometry_0",
+      "geom": "geometry_0",
+      "n_faces": 92502,
+      "n_verts": 65312,
+      "visual_type": "TextureVisuals",
+      "face_colors": null,
+      "vertex_colors": null,
+      "textures": [],
+      "material": {
+        "base_color_factor_or_main_color": [
+          255.0,
+          255.0,
+          255.0,
+          255.0
+        ]
+      }
+    }
+  ],
+  "summary": {
+    "total_face_color_entries": 0,
+    "total_vertex_color_entries": 0,
+    "total_texture_pixels_sampled": 0
+  }
+}

data_toolkit/bpy_render.py

@@ -0,0 +1,623 @@
+import bpy
+import json
+import math
+import mathutils
+import numpy as np
+
+
+class BpyRenderer:
+    def __init__(self, resolution=512, engine="BLENDER_EEVEE", geo_mode=False, split_normal=False):
+        """
+        engine:
+          - "CYCLES"
+          - "BLENDER_EEVEE"   (Blender 3.x common)
+          - "BLENDER_EEVEE_NEXT" (Blender 4.x common)
+          - "EEVEE" / "EEVEE_NEXT" (aliases, optional)
+        """
+        self.resolution = resolution
+        self.engine = engine
+        self.geo_mode = geo_mode
+        self.split_normal = split_normal
+        self.import_functions = self._setup_import_functions()
+
+    def _setup_import_functions(self):
+        import_functions = {
+            "obj": bpy.ops.wm.obj_import,
+            "glb": bpy.ops.import_scene.gltf,
+            "gltf": bpy.ops.import_scene.gltf,
+            "usd": bpy.ops.import_scene.usd,
+            "fbx": bpy.ops.import_scene.fbx,
+            "stl": bpy.ops.import_mesh.stl,
+            "usda": bpy.ops.import_scene.usda,
+            "dae": bpy.ops.wm.collada_import,
+            "ply": bpy.ops.wm.ply_import,
+            "abc": bpy.ops.wm.alembic_import,
+            "blend": bpy.ops.wm.append,
+        }
+        return import_functions
+
+    # -------------------------
+    # Engine helpers
+    # -------------------------
+    def _resolve_render_engine(self, requested: str) -> str:
+        """
+        Robustly set render engine across Blender versions.
+        Blender 4.x may not accept "BLENDER_EEVEE" and instead uses "BLENDER_EEVEE_NEXT".
+        """
+        req = (requested or "").upper()
+
+        if req in {"EEVEE", "BLENDER_EEVEE"}:
+            candidates = ["BLENDER_EEVEE", "BLENDER_EEVEE_NEXT"]
+        elif req in {"EEVEE_NEXT", "BLENDER_EEVEE_NEXT"}:
+            candidates = ["BLENDER_EEVEE_NEXT", "BLENDER_EEVEE"]
+        elif req in {"CYCLES"}:
+            candidates = ["CYCLES"]
+        elif req in {"WORKBENCH", "BLENDER_WORKBENCH"}:
+            candidates = ["BLENDER_WORKBENCH"]
+        else:
+            candidates = [requested]
+
+        last_err = None
+        for eng in candidates:
+            try:
+                bpy.context.scene.render.engine = eng
+                return eng
+            except Exception as e:
+                last_err = e
+                continue
+
+        raise ValueError(f"Failed to set render engine from {candidates}. Last error: {last_err}")
+
+    def _init_eevee_settings(self, render_samples: int = 64):
+        """
+        EEVEE / EEVEE Next settings (close to huanngzh/bpy-renderer defaults).
+        """
+        scene = bpy.context.scene
+
+        # Render basics
+        scene.render.image_settings.file_format = "PNG"
+        scene.render.image_settings.color_mode = "RGBA"
+        scene.render.film_transparent = True
+
+        # EEVEE quality knobs
+        # In Blender, eevee settings live under scene.eevee.
+        # These fields are used by many scripts including bpy-renderer. :contentReference[oaicite:2]{index=2}
+        if hasattr(scene, "eevee"):
+            try:
+                scene.eevee.taa_render_samples = int(render_samples)
+            except Exception:
+                pass
+            # These flags may not exist in every minor version; guard them.
+            for name, val in [
+                ("use_gtao", True),
+                ("use_ssr", True),
+                ("use_bloom", True),
+            ]:
+                if hasattr(scene.eevee, name):
+                    try:
+                        setattr(scene.eevee, name, val)
+                    except Exception:
+                        pass
+
+        # Normals quality (also in bpy-renderer init) :contentReference[oaicite:3]{index=3}
+        if hasattr(scene.render, "use_high_quality_normals"):
+            try:
+                scene.render.use_high_quality_normals = True
+            except Exception:
+                pass
+
+    def _init_cycles_settings(self, render_samples: int = 128):
+        scene = bpy.context.scene
+
+        scene.render.image_settings.file_format = "PNG"
+        scene.render.image_settings.color_mode = "RGBA"
+        scene.render.film_transparent = True
+
+        scene.cycles.samples = int(render_samples)
+        scene.cycles.filter_type = "BOX"
+        scene.cycles.filter_width = 1
+        scene.cycles.diffuse_bounces = 1
+        scene.cycles.glossy_bounces = 1
+        scene.cycles.transparent_max_bounces = (3 if not self.geo_mode else 0)
+        scene.cycles.transmission_bounces = (3 if not self.geo_mode else 1)
+        scene.cycles.use_denoising = True
+
+        # GPU (best-effort)
+        try:
+            scene.cycles.device = "GPU"
+            bpy.context.preferences.addons["cycles"].preferences.get_devices()
+            bpy.context.preferences.addons["cycles"].preferences.compute_device_type = "CUDA"
+        except Exception:
+            pass
+
+    # -------------------------
+    # Public init
+    # -------------------------
+    def init_render_settings(self):
+        # Resolution
+        bpy.context.scene.render.resolution_x = self.resolution
+        bpy.context.scene.render.resolution_y = self.resolution
+        bpy.context.scene.render.resolution_percentage = 100
+
+        # Pick engine robustly (EEVEE vs EEVEE_NEXT etc.)
+        actual_engine = self._resolve_render_engine(self.engine)
+
+        # Samples:
+        # - For geo_mode: keep minimal samples for speed
+        # - For RGB: moderate samples
+        if actual_engine == "CYCLES":
+            samples = 128 if not self.geo_mode else 1
+            self._init_cycles_settings(render_samples=samples)
+        else:
+            # EEVEE family
+            samples = 64 if not self.geo_mode else 1
+            self._init_eevee_settings(render_samples=samples)
+
+    def init_scene(self):
+        for obj in bpy.data.objects:
+            bpy.data.objects.remove(obj, do_unlink=True)
+        for material in bpy.data.materials:
+            bpy.data.materials.remove(material, do_unlink=True)
+        for texture in bpy.data.textures:
+            bpy.data.textures.remove(texture, do_unlink=True)
+        for image in bpy.data.images:
+            bpy.data.images.remove(image, do_unlink=True)
+
+    def init_camera(self):
+        cam = bpy.data.objects.new("Camera", bpy.data.cameras.new("Camera"))
+        bpy.context.collection.objects.link(cam)
+        bpy.context.scene.camera = cam
+        cam.data.sensor_height = cam.data.sensor_width = 32
+        cam_constraint = cam.constraints.new(type="TRACK_TO")
+        cam_constraint.track_axis = "TRACK_NEGATIVE_Z"
+        cam_constraint.up_axis = "UP_Y"
+        cam_empty = bpy.data.objects.new("Empty", None)
+        cam_empty.location = (0, 0, 0)
+        bpy.context.scene.collection.objects.link(cam_empty)
+        cam_constraint.target = cam_empty
+        return cam
+
+    def init_lighting(self):
+        bpy.ops.object.select_all(action="DESELECT")
+        bpy.ops.object.select_by_type(type="LIGHT")
+        bpy.ops.object.delete()
+
+        default_light = bpy.data.objects.new("Default_Light", bpy.data.lights.new("Default_Light", type="POINT"))
+        bpy.context.collection.objects.link(default_light)
+        default_light.data.energy = 1000
+        default_light.location = (4, 1, 6)
+        default_light.rotation_euler = (0, 0, 0)
+
+        top_light = bpy.data.objects.new("Top_Light", bpy.data.lights.new("Top_Light", type="AREA"))
+        bpy.context.collection.objects.link(top_light)
+        top_light.data.energy = 10000
+        top_light.location = (0, 0, 10)
+        top_light.scale = (100, 100, 100)
+
+        bottom_light = bpy.data.objects.new("Bottom_Light", bpy.data.lights.new("Bottom_Light", type="AREA"))
+        bpy.context.collection.objects.link(bottom_light)
+        bottom_light.data.energy = 1000
+        bottom_light.location = (0, 0, -10)
+        bottom_light.rotation_euler = (0, 0, 0)
+        return {"default_light": default_light, "top_light": top_light, "bottom_light": bottom_light}
+
+    def load_object(self, object_path):
+        file_extension = object_path.split(".")[-1].lower()
+        if file_extension not in self.import_functions:
+            raise ValueError(f"Unsupported file type: {file_extension}")
+        import_function = self.import_functions[file_extension]
+        print(f"Loading object from {object_path}")
+        if file_extension == "blend":
+            import_function(directory=object_path, link=False)
+        elif file_extension in {"glb", "gltf"}:
+            import_function(filepath=object_path, merge_vertices=True, import_shading="NORMALS")
+        else:
+            import_function(filepath=object_path)
+
+    def delete_invisible_objects(self):
+        bpy.ops.object.select_all(action="DESELECT")
+        for obj in bpy.context.scene.objects:
+            if obj.hide_viewport or obj.hide_render:
+                obj.hide_viewport = False
+                obj.hide_render = False
+                obj.hide_select = False
+                obj.select_set(True)
+        bpy.ops.object.delete()
+        invisible_collections = [col for col in bpy.data.collections if col.hide_viewport]
+        for col in invisible_collections:
+            bpy.data.collections.remove(col)
+
+    def split_mesh_normal(self):
+        bpy.ops.object.select_all(action="DESELECT")
+        objs = [obj for obj in bpy.context.scene.objects if obj.type == "MESH"]
+        bpy.context.view_layer.objects.active = objs[0]
+        for obj in objs:
+            obj.select_set(True)
+        bpy.ops.object.mode_set(mode="EDIT")
+        bpy.ops.mesh.select_all(action="SELECT")
+        bpy.ops.mesh.split_normals()
+        bpy.ops.object.mode_set(mode="OBJECT")
+        bpy.ops.object.select_all(action="DESELECT")
+
+    def override_material(self):
+        new_mat = bpy.data.materials.new(name="Override0123456789")
+        new_mat.use_nodes = True
+        new_mat.node_tree.nodes.clear()
+        bsdf = new_mat.node_tree.nodes.new("ShaderNodeBsdfDiffuse")
+        bsdf.inputs[0].default_value = (0.5, 0.5, 0.5, 1)
+        bsdf.inputs[1].default_value = 1
+        output = new_mat.node_tree.nodes.new("ShaderNodeOutputMaterial")
+        new_mat.node_tree.links.new(bsdf.outputs["BSDF"], output.inputs["Surface"])
+        bpy.context.scene.view_layers["View Layer"].material_override = new_mat
+
+    def scene_bbox(self):
+        bbox_min = (math.inf,) * 3
+        bbox_max = (-math.inf,) * 3
+        found = False
+        scene_meshes = [obj for obj in bpy.context.scene.objects.values() if isinstance(obj.data, bpy.types.Mesh)]
+        for obj in scene_meshes:
+            found = True
+            for coord in obj.bound_box:
+                coord = mathutils.Vector(coord)
+                coord = obj.matrix_world @ coord
+                bbox_min = tuple(min(x, y) for x, y in zip(bbox_min, coord))
+                bbox_max = tuple(max(x, y) for x, y in zip(bbox_max, coord))
+        if not found:
+            raise RuntimeError("no objects in scene to compute bounding box for")
+        return mathutils.Vector(bbox_min), mathutils.Vector(bbox_max)
+
+    def normalize_scene(self):
+        scene_root_objects = [obj for obj in bpy.context.scene.objects.values() if not obj.parent]
+        if len(scene_root_objects) > 1:
+            scene = bpy.data.objects.new("ParentEmpty", None)
+            bpy.context.scene.collection.objects.link(scene)
+            for obj in scene_root_objects:
+                obj.parent = scene
+        else:
+            scene = scene_root_objects[0]
+
+        bbox_min, bbox_max = self.scene_bbox()
+        print(f"[INFO] Bounding box: {bbox_min}, {bbox_max}")
+        scale = 1 / max(bbox_max - bbox_min)
+        scene.scale = scene.scale * scale
+        bpy.context.view_layer.update()
+        bbox_min, bbox_max = self.scene_bbox()
+        offset = -(bbox_min + bbox_max) / 2
+        scene.matrix_world.translation += offset
+        bpy.ops.object.select_all(action="DESELECT")
+        return scale, offset
+
+    def set_camera_from_matrix(self, cam, transform_matrix):
+        matrix = mathutils.Matrix(transform_matrix)
+        cam.matrix_world = matrix
+        bpy.context.view_layer.update()
+
+    def render_from_transforms(self, file_path, transforms_json_path, output_path):
+        with open(transforms_json_path, "r") as f:
+            transforms_data = json.load(f)
+
+        self.init_render_settings()
+
+        # Load scene
+        if file_path.endswith(".blend"):
+            self.delete_invisible_objects()
+        else:
+            self.init_scene()
+            self.load_object(file_path)
+            if self.split_normal:
+                self.split_mesh_normal()
+        print("[INFO] Scene initialized.")
+
+        scale, offset = self.normalize_scene()
+        print(f"[INFO] Scene normalized with auto scale: {scale}, offset: {offset}")
+
+        cam = self.init_camera()
+        self.init_lighting()
+        print("[INFO] Camera and lighting initialized.")
+        if self.geo_mode:
+            self.override_material()
+
+        # NOTE: your transforms_json format seems like a list-of-dicts.
+        transform_matrix = transforms_data[0]["transform_matrix"]
+        camera_angle_x = transforms_data[0].get("camera_angle_x", None)
+
+        self.set_camera_from_matrix(cam, transform_matrix)
+        if camera_angle_x is not None:
+            cam.data.lens = 16 / np.tan(camera_angle_x / 2)
+
+        bpy.context.scene.render.filepath = output_path
+        bpy.ops.render.render(write_still=True)
+        bpy.context.view_layer.update()
+
+
+def render_from_transforms(
+    file_path,
+    transforms_json_path,
+    output_path,
+    resolution=512,
+    engine="BLENDER_EEVEE",
+    geo_mode=False,
+    split_normal=False,
+):
+    renderer = BpyRenderer(resolution=resolution, engine=engine, geo_mode=geo_mode, split_normal=split_normal)
+    return renderer.render_from_transforms(file_path, transforms_json_path, output_path)
+
+
+if __name__ == "__main__":
+    file_path = "./assets/example.glb"
+    transforms_json_path = "transforms.json"
+    output_path = "./assets/img.png"
+
+    # Recommended:
+    # - engine="BLENDER_EEVEE" for Blender 3.x
+    # - engine="BLENDER_EEVEE_NEXT" for Blender 4.x
+    # This script auto-fallbacks between them.
+    render_from_transforms(
+        file_path=file_path,
+        transforms_json_path=transforms_json_path,
+        output_path=output_path,
+        resolution=512,
+        engine="BLENDER_EEVEE",
+    )
+
+# import bpy
+# import json
+# import math
+# import mathutils
+# import numpy as np
+
+# class BpyRenderer:
+#     def __init__(self, resolution=512, engine="CYCLES", geo_mode=False, split_normal=False):
+#         self.resolution = resolution
+#         self.engine = engine
+#         self.geo_mode = geo_mode
+#         self.split_normal = split_normal
+#         self.import_functions = self._setup_import_functions()
+
+#     def _setup_import_functions(self):
+#         import_functions = {
+#             "obj": bpy.ops.wm.obj_import,
+#             "glb": bpy.ops.import_scene.gltf,
+#             "gltf": bpy.ops.import_scene.gltf,
+#             "usd": bpy.ops.import_scene.usd,
+#             "fbx": bpy.ops.import_scene.fbx,
+#             "stl": bpy.ops.import_mesh.stl,
+#             "usda": bpy.ops.import_scene.usda,
+#             "dae": bpy.ops.wm.collada_import,
+#             "ply": bpy.ops.wm.ply_import,
+#             "abc": bpy.ops.wm.alembic_import,
+#             "blend": bpy.ops.wm.append,
+#         }
+#         return import_functions
+
+#     def init_render_settings(self):
+#         bpy.context.scene.render.engine = self.engine
+#         bpy.context.scene.render.resolution_x = self.resolution
+#         bpy.context.scene.render.resolution_y = self.resolution
+#         bpy.context.scene.render.resolution_percentage = 100
+#         bpy.context.scene.render.image_settings.file_format = "PNG"
+#         bpy.context.scene.render.image_settings.color_mode = "RGBA"
+#         bpy.context.scene.render.film_transparent = True
+#         if self.engine == "CYCLES":
+#             bpy.context.scene.render.engine = "CYCLES"
+#             bpy.context.scene.cycles.samples = 128 if not self.geo_mode else 1
+#             bpy.context.scene.cycles.filter_type = "BOX"
+#             bpy.context.scene.cycles.filter_width = 1
+#             bpy.context.scene.cycles.diffuse_bounces = 1
+#             bpy.context.scene.cycles.glossy_bounces = 1
+#             bpy.context.scene.cycles.transparent_max_bounces = (3 if not self.geo_mode else 0)
+#             bpy.context.scene.cycles.transmission_bounces = (3 if not self.geo_mode else 1)
+#             bpy.context.scene.cycles.use_denoising = True
+#             try:
+#                 bpy.context.scene.cycles.device = "GPU"
+#                 bpy.context.preferences.addons["cycles"].preferences.get_devices()
+#                 bpy.context.preferences.addons["cycles"].preferences.compute_device_type = "CUDA"
+#             except:
+#                 pass
+
+#     def init_scene(self):
+#         for obj in bpy.data.objects:
+#             bpy.data.objects.remove(obj, do_unlink=True)
+#         for material in bpy.data.materials:
+#             bpy.data.materials.remove(material, do_unlink=True)
+#         for texture in bpy.data.textures:
+#             bpy.data.textures.remove(texture, do_unlink=True)
+#         for image in bpy.data.images:
+#             bpy.data.images.remove(image, do_unlink=True)
+
+#     def init_camera(self):
+#         cam = bpy.data.objects.new("Camera", bpy.data.cameras.new("Camera"))
+#         bpy.context.collection.objects.link(cam)
+#         bpy.context.scene.camera = cam
+#         cam.data.sensor_height = cam.data.sensor_width = 32
+#         cam_constraint = cam.constraints.new(type="TRACK_TO")
+#         cam_constraint.track_axis = "TRACK_NEGATIVE_Z"
+#         cam_constraint.up_axis = "UP_Y"
+#         cam_empty = bpy.data.objects.new("Empty", None)
+#         cam_empty.location = (0, 0, 0)
+#         bpy.context.scene.collection.objects.link(cam_empty)
+#         cam_constraint.target = cam_empty
+#         return cam
+
+#     def init_lighting(self):
+#         bpy.ops.object.select_all(action="DESELECT")
+#         bpy.ops.object.select_by_type(type="LIGHT")
+#         bpy.ops.object.delete()
+
+#         default_light = bpy.data.objects.new("Default_Light", bpy.data.lights.new("Default_Light", type="POINT"))
+#         bpy.context.collection.objects.link(default_light)
+#         default_light.data.energy = 1000
+#         default_light.location = (4, 1, 6)
+#         default_light.rotation_euler = (0, 0, 0)
+
+#         top_light = bpy.data.objects.new("Top_Light", bpy.data.lights.new("Top_Light", type="AREA"))
+#         bpy.context.collection.objects.link(top_light)
+#         top_light.data.energy = 10000
+#         top_light.location = (0, 0, 10)
+#         top_light.scale = (100, 100, 100)
+
+#         bottom_light = bpy.data.objects.new("Bottom_Light", bpy.data.lights.new("Bottom_Light", type="AREA"))
+#         bpy.context.collection.objects.link(bottom_light)
+#         bottom_light.data.energy = 1000
+#         bottom_light.location = (0, 0, -10)
+#         bottom_light.rotation_euler = (0, 0, 0)
+#         return {"default_light": default_light, "top_light": top_light, "bottom_light": bottom_light}
+
+#     def load_object(self, object_path):
+#         file_extension = object_path.split(".")[-1].lower()
+#         if file_extension not in self.import_functions:
+#             raise ValueError(f"Unsupported file type: {file_extension}")
+#         import_function = self.import_functions[file_extension]
+#         print(f"Loading object from {object_path}")
+#         if file_extension == "blend":
+#             import_function(directory=object_path, link=False)
+#         elif file_extension in {"glb", "gltf"}:
+#             import_function(filepath=object_path, merge_vertices=True, import_shading="NORMALS")
+#         else:
+#             import_function(filepath=object_path)
+
+#     def delete_invisible_objects(self):
+#         bpy.ops.object.select_all(action="DESELECT")
+#         for obj in bpy.context.scene.objects:
+#             if obj.hide_viewport or obj.hide_render:
+#                 obj.hide_viewport = False
+#                 obj.hide_render = False
+#                 obj.hide_select = False
+#                 obj.select_set(True)
+#         bpy.ops.object.delete()
+#         invisible_collections = [col for col in bpy.data.collections if col.hide_viewport]
+#         for col in invisible_collections:
+#             bpy.data.collections.remove(col)
+
+#     def unhide_all_objects(self):
+#         for obj in bpy.context.scene.objects:
+#             obj.hide_set(False)
+
+#     def convert_to_meshes(self):
+#         bpy.ops.object.select_all(action="DESELECT")
+#         bpy.context.view_layer.objects.active = [obj for obj in bpy.context.scene.objects if obj.type == "MESH"][0]
+#         for obj in bpy.context.scene.objects:
+#             obj.select_set(True)
+#         bpy.ops.object.convert(target="MESH")
+
+#     def triangulate_meshes(self):
+#         bpy.ops.object.select_all(action="DESELECT")
+#         objs = [obj for obj in bpy.context.scene.objects if obj.type == "MESH"]
+#         bpy.context.view_layer.objects.active = objs[0]
+#         for obj in objs:
+#             obj.select_set(True)
+#         bpy.ops.object.mode_set(mode="EDIT")
+#         bpy.ops.mesh.reveal()
+#         bpy.ops.mesh.select_all(action="SELECT")
+#         bpy.ops.mesh.quads_convert_to_tris(quad_method="BEAUTY", ngon_method="BEAUTY")
+#         bpy.ops.object.mode_set(mode="OBJECT")
+#         bpy.ops.object.select_all(action="DESELECT")
+
+#     def split_mesh_normal(self):
+#         bpy.ops.object.select_all(action="DESELECT")
+#         objs = [obj for obj in bpy.context.scene.objects if obj.type == "MESH"]
+#         bpy.context.view_layer.objects.active = objs[0]
+#         for obj in objs:
+#             obj.select_set(True)
+#         bpy.ops.object.mode_set(mode="EDIT")
+#         bpy.ops.mesh.select_all(action="SELECT")
+#         bpy.ops.mesh.split_normals()
+#         bpy.ops.object.mode_set(mode="OBJECT")
+#         bpy.ops.object.select_all(action="DESELECT")
+
+#     def override_material(self):
+#         new_mat = bpy.data.materials.new(name="Override0123456789")
+#         new_mat.use_nodes = True
+#         new_mat.node_tree.nodes.clear()
+#         bsdf = new_mat.node_tree.nodes.new("ShaderNodeBsdfDiffuse")
+#         bsdf.inputs[0].default_value = (0.5, 0.5, 0.5, 1)
+#         bsdf.inputs[1].default_value = 1
+#         output = new_mat.node_tree.nodes.new("ShaderNodeOutputMaterial")
+#         new_mat.node_tree.links.new(bsdf.outputs["BSDF"], output.inputs["Surface"])
+#         bpy.context.scene.view_layers["View Layer"].material_override = new_mat
+
+#     def scene_bbox(self):
+#         bbox_min = (math.inf,) * 3
+#         bbox_max = (-math.inf,) * 3
+#         found = False
+#         scene_meshes = [obj for obj in bpy.context.scene.objects.values() if isinstance(obj.data, bpy.types.Mesh)]
+#         for obj in scene_meshes:
+#             found = True
+#             for coord in obj.bound_box:
+#                 coord = mathutils.Vector(coord)
+#                 coord = obj.matrix_world @ coord
+#                 bbox_min = tuple(min(x, y) for x, y in zip(bbox_min, coord))
+#                 bbox_max = tuple(max(x, y) for x, y in zip(bbox_max, coord))
+#         if not found:
+#             raise RuntimeError("no objects in scene to compute bounding box for")
+#         return mathutils.Vector(bbox_min), mathutils.Vector(bbox_max)
+
+#     def normalize_scene(self):
+#         scene_root_objects = [obj for obj in bpy.context.scene.objects.values() if not obj.parent]
+#         if len(scene_root_objects) > 1:
+#             scene = bpy.data.objects.new("ParentEmpty", None)
+#             bpy.context.scene.collection.objects.link(scene)
+#             for obj in scene_root_objects:
+#                 obj.parent = scene
+#         else:
+#             scene = scene_root_objects[0]
+
+#         bbox_min, bbox_max = self.scene_bbox()
+#         print(f"[INFO] Bounding box: {bbox_min}, {bbox_max}")
+#         scale = 1 / max(bbox_max - bbox_min)
+#         scene.scale = scene.scale * scale
+#         bpy.context.view_layer.update()
+#         bbox_min, bbox_max = self.scene_bbox()
+#         offset = -(bbox_min + bbox_max) / 2
+#         scene.matrix_world.translation += offset
+#         bpy.ops.object.select_all(action="DESELECT")
+#         return scale, offset
+
+#     def set_camera_from_matrix(self, cam, transform_matrix):
+#         matrix = mathutils.Matrix(transform_matrix)
+#         cam.matrix_world = matrix
+#         bpy.context.view_layer.update()
+
+#     def render_from_transforms(self, file_path, transforms_json_path, output_path):
+#         with open(transforms_json_path, 'r') as f:
+#             transforms_data = json.load(f)
+        
+#         self.init_render_settings()
+        
+#         if file_path.endswith(".blend"):
+#             self.delete_invisible_objects()
+#         else:
+#             self.init_scene()
+#             self.load_object(file_path)
+#             if self.split_normal:
+#                 self.split_mesh_normal()
+#         print("[INFO] Scene initialized.")
+        
+#         scale, offset = self.normalize_scene()
+#         print(f"[INFO] Scene normalized with auto scale: {scale}, offset: {offset}")
+        
+#         cam = self.init_camera()
+#         self.init_lighting()
+#         print("[INFO] Camera and lighting initialized.")
+#         if self.geo_mode:
+#             self.override_material()
+
+#         transform_matrix = transforms_data[0]["transform_matrix"]
+#         camera_angle_x = transforms_data[0]["camera_angle_x"]
+#         self.set_camera_from_matrix(cam, transform_matrix)
+#         if camera_angle_x is not None:
+#             cam.data.lens = 16 / np.tan(camera_angle_x / 2)
+        
+#         bpy.context.scene.render.filepath = output_path
+#         bpy.ops.render.render(write_still=True)
+#         bpy.context.view_layer.update()
+
+# def render_from_transforms(file_path, transforms_json_path, output_path, resolution=512, engine="CYCLES", geo_mode=False, split_normal=False):
+#     renderer = BpyRenderer(resolution=resolution, engine=engine, geo_mode=geo_mode, split_normal=split_normal)
+#     return renderer.render_from_transforms(file_path, transforms_json_path, output_path)
+
+# if __name__ == "__main__":
+#     file_path = "./assets/example.glb"
+#     transforms_json_path = "transforms.json"
+#     output_path = "./assets/img.png"
+#     render_from_transforms(file_path=file_path, transforms_json_path=transforms_json_path, output_path=output_path)

data_toolkit/color_glb.py

@@ -0,0 +1,84 @@
+import os
+import json
+import random
+import trimesh
+import numpy as np
+
+from trimesh.visual.material import PBRMaterial
+
+def _load_as_single_mesh(part_path):
+    obj = trimesh.load(part_path, force="scene")
+    if isinstance(obj, trimesh.Scene):
+        dumped = obj.dump()
+        meshes = [m for m in dumped if isinstance(m, trimesh.Trimesh) and len(m.vertices) > 0]
+        return trimesh.util.concatenate(meshes)
+    if isinstance(obj, trimesh.Trimesh):
+        return obj
+
+def set_mesh_solid_pbr(mesh, rgba_uint8=(255, 255, 255, 255), emissive=True):
+    rgb = np.array(rgba_uint8[:3], dtype=np.float32) / 255.0
+    a = float(rgba_uint8[3]) / 255.0
+    colors = np.tile(np.array(rgba_uint8, dtype=np.uint8), (len(mesh.vertices), 1))
+    mesh.visual = trimesh.visual.ColorVisuals(mesh=mesh, vertex_colors=colors)
+    mat_kwargs = dict(
+        baseColorFactor=[float(rgb[0]), float(rgb[1]), float(rgb[2]), a],
+        metallicFactor=0.0,
+        roughnessFactor=1.0,
+    )
+    if emissive:
+        mat_kwargs["emissiveFactor"] = [float(rgb[0]), float(rgb[1]), float(rgb[2])]
+    mesh.visual.material = PBRMaterial(**mat_kwargs)
+    return mesh
+
+def color_glb(parts_path, output_path, interactive):
+    part_meshes = []
+    for part_name in sorted(os.listdir(parts_path)):
+        part_path = os.path.join(parts_path, part_name)
+        part_meshes.append(_load_as_single_mesh(part_path))
+
+    if interactive:
+        for i in range(len(part_meshes)):
+            colors = [(255, 255, 255, 255) if j == i else (0, 0, 0, 255) for j in range(len(part_meshes))]
+            scene = trimesh.Scene()
+            for j, m in enumerate(part_meshes):
+                mc = m.copy()
+                set_mesh_solid_pbr(mc, rgba_uint8=colors[j], emissive=True)
+                scene.add_geometry(mc, node_name=f"part_{j}", geom_name=f"geom_{j}")
+                os.makedirs(os.path.join(output_path, f"{i}"), exist_ok=True)
+                scene.export(os.path.join(output_path, f"{i}", "output.glb"))
+    else:
+        colors = []
+        for i in range(len(part_meshes)):
+            while True:
+                rgb = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255), 255)
+                if rgb not in colors:
+                    colors.append(rgb)
+                    break
+
+        # colors_base = [           # Static colors
+        #     (0, 0, 0, 255),
+        #     (0, 0, 255, 255),
+        #     (0, 255, 0, 255),
+        #     (0, 255, 255, 255),
+        #     (255, 0, 0, 255),
+        #     (255, 0, 255, 255),
+        #     (255, 255, 0, 255),
+        #     (255, 255, 255, 255)
+        # ]
+        # colors = random.sample(colors_base, len(part_meshes))
+        
+        with open(os.path.join(output_path, "colors.json"), "w", encoding="utf-8") as f:
+            json.dump([list(c) for c in colors], f, ensure_ascii=False, indent=4)
+
+        scene = trimesh.Scene()
+        for i, m in enumerate(part_meshes):
+            mc = m.copy()
+            set_mesh_solid_pbr(mc, rgba_uint8=colors[i], emissive=True)
+            scene.add_geometry(mc, node_name=f"part_{i}", geom_name=f"geom_{i}")
+        scene.export(os.path.join(output_path, "output.glb"))
+
+if __name__ == "__main__":
+    parts_path = "./assets/parts"
+    output_path = "./assets/interactive_seg"
+    interactive = True
+    color_glb(parts_path, output_path, interactive)

data_toolkit/color_img.py

@@ -0,0 +1,129 @@
+import os
+import json
+
+import numpy as np
+import trimesh
+import torch
+import nvdiffrast.torch as nr
+from PIL import Image
+
+def build_projection_matrix(fov, width, height, z_near=0.01, z_far=100.0):
+    aspect = float(width) / float(height)
+    fov_y = 2.0 * np.arctan(np.tan(fov / 2.0) / aspect)
+    f = 1.0 / np.tan(fov_y / 2.0)
+    P = np.array([
+        [f / aspect, 0.0,  0.0,                                 0.0],
+        [0.0,        f,    0.0,                                 0.0],
+        [0.0,        0.0,  (z_far + z_near) / (z_near - z_far), (2.0 * z_far * z_near) / (z_near - z_far)],
+        [0.0,        0.0, -1.0,                                 0.0],
+    ], dtype=np.float32)
+    return P
+
+def compute_bbox_center_and_scale_like_blender(vertices):
+    bbox_min = vertices.min(axis=0)
+    bbox_max = vertices.max(axis=0)
+    bbox_extents = bbox_max - bbox_min
+    scale = 1.0 / np.max(bbox_extents)
+    offset = -(bbox_min + bbox_max) / 2.0
+    return offset, scale
+
+def _load_as_single_mesh(part_path):
+    obj = trimesh.load(part_path, force="scene")
+    if isinstance(obj, trimesh.Scene):
+        dumped = obj.dump()
+        meshes = [m for m in dumped if isinstance(m, trimesh.Trimesh) and len(m.vertices) > 0]
+        return trimesh.util.concatenate(meshes)
+    if isinstance(obj, trimesh.Trimesh):
+        return obj
+
+def load_parts_from_directory(object_path):
+    per_part_vertices = []
+    per_part_faces = []
+    part_names = []
+    vertices_counts = []
+    vertex_offset = 0
+    for part_name in sorted(os.listdir(object_path)):
+        part_path = os.path.join(object_path, part_name)
+        mesh = _load_as_single_mesh(part_path)
+        v = mesh.vertices.astype(np.float32)
+        f = mesh.faces.astype(np.int32)
+        per_part_vertices.append(v)
+        per_part_faces.append(f + vertex_offset)
+        part_names.append(part_name)
+        vertices_counts.append(v.shape[0])
+        vertex_offset += v.shape[0]
+    return per_part_vertices, per_part_faces, part_names, vertices_counts
+
+def save_png(path, array_uint8):
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    Image.fromarray(array_uint8, mode="RGB").save(path)
+
+def render_views(glctx, V, F, C, output_path, transforms_path):
+    width = 512
+    height = 512
+    fov = 40.0*np.pi/180.0
+    V_t = torch.from_numpy(V).to(torch.float32).cuda()
+    F_t = torch.from_numpy(F).to(torch.int32).cuda()
+    C_t = torch.from_numpy(C).to(torch.float32).cuda()
+    theta = np.pi / 2.0
+    Gx = np.array([
+        [1.0, 0.0,             0.0,            0.0],
+        [0.0, np.cos(theta),  -np.sin(theta),  0.0],
+        [0.0, np.sin(theta),   np.cos(theta),  0.0],
+        [0.0, 0.0,             0.0,            1.0],
+    ], dtype=np.float32)
+    Gx_t = torch.from_numpy(Gx).to(torch.float32).cuda()
+
+    with open(transforms_path, "r") as f:
+        transforms = json.load(f)
+    cam_to_world = np.array(transforms[0]["transform_matrix"], dtype=np.float32)
+    world_to_cam = np.linalg.inv(cam_to_world)
+    P = build_projection_matrix(fov, width, height)
+
+    V_mat = torch.from_numpy(world_to_cam).to(torch.float32).cuda()
+    P_mat = torch.from_numpy(P).to(torch.float32).cuda()
+    M_t = torch.eye(4, dtype=torch.float32).cuda()
+    pos_h = torch.cat([V_t, torch.ones((V_t.shape[0], 1), dtype=torch.float32).cuda()], dim=1)
+    pos_clip = (P_mat @ V_mat @ M_t @ Gx_t) @ pos_h.t()
+    pos_clip = pos_clip.t().contiguous().unsqueeze(0)
+
+    rast, _ = nr.rasterize(glctx, pos_clip, F_t, resolution=[height, width])
+    feat, _ = nr.interpolate(C_t.unsqueeze(0), rast, F_t)
+    cov = rast[..., 3:4]
+    img = feat.clamp(0.0, 1.0)
+    bg = torch.ones_like(img)
+    out = img * (cov > 0) + bg * (cov <= 0)
+    out_np = (out[0].cpu().numpy() * 255.0).astype(np.uint8)
+    out_np = out_np[::-1, :, :]
+    save_png(output_path, out_np)
+
+def color_img(object_path, output_path, transforms, colors_path):
+    per_part_vertices, per_part_faces, part_names, vertices_counts = load_parts_from_directory(object_path)
+    V = np.concatenate(per_part_vertices, axis=0).astype(np.float32)
+    F = np.concatenate(per_part_faces, axis=0).astype(np.int32)
+    offset, scale = compute_bbox_center_and_scale_like_blender(V)
+    V_scaled = V * scale
+    V_norm = V_scaled + offset[None, :]
+    V = V_norm
+
+    with open(colors_path, "r") as f:
+        external_colors = json.load(f)
+    color_map = {}
+    colors = []
+    for idx, part_name in enumerate(part_names):
+        rgb = external_colors[idx][:3]
+        color_map[part_name] = [int(rgb[0]), int(rgb[1]), int(rgb[2])]
+        num_v = vertices_counts[idx]
+        col = (np.array(rgb, dtype=np.float32) / 255.0)[None, :]
+        colors.append(np.repeat(col, repeats=num_v, axis=0))
+
+    C = np.concatenate(colors, axis=0).astype(np.float32)
+    glctx = nr.RasterizeCudaContext()
+    render_views(glctx, V, F, C, output_path, transforms)
+
+if __name__ == "__main__":
+    object_path = "./assets/parts"
+    output_path = "./assets/full_seg_w_2d_map/2d_map.png"
+    transforms = "transforms.json"
+    colors_path = "./assets/full_seg_w_2d_map/colors.json"
+    color_img(object_path, output_path, transforms, colors_path)

data_toolkit/example_full_seg.py

@@ -0,0 +1,45 @@
+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="briaai/RMBG-2.0")
+rembg_model.cuda()
+image_cond_model = DinoV3FeatureExtractor(model_name="facebook/dinov3-vitl16-pretrain-lvd1689m")
+image_cond_model.cuda()
+
+shape_encoder = models.from_pretrained("microsoft/TRELLIS.2-4B/ckpts/shape_enc_next_dc_f16c32_fp16").cuda().eval()
+tex_encoder = models.from_pretrained("microsoft/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"
+full_seg_path = "./assets/full_seg"
+interactive = False
+
+transforms = "transforms.json"
+img = "./assets/img.png"
+cond = "./assets/cond.pth"
+
+output_glb_path = os.path.join(full_seg_path, "output.glb")
+output_vxz_path = os.path.join(full_seg_path, "output.vxz")
+
+glb_to_vxz(glb, input_vxz)
+glb_to_parts(glb, parts_path)
+color_glb(parts_path, full_seg_path, interactive)
+
+render_from_transforms(glb, transforms, img)
+img_to_cond(rembg_model, image_cond_model, img, cond)
+glb_to_vxz(output_glb_path, output_vxz_path)
+vxz_to_slat(shape_encoder, tex_encoder, input_vxz, output_vxz_path, full_seg_path, interactive)

data_toolkit/example_full_seg_w_2d_map.py

@@ -0,0 +1,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 color_img import color_img
+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 trellis2.pipelines.rembg import BiRefNet
+from trellis2.modules.image_feature_extractor import DinoV3FeatureExtractor
+
+rembg_model = BiRefNet(model_name="briaai/RMBG-2.0")
+rembg_model.cuda()
+image_cond_model = DinoV3FeatureExtractor(model_name="facebook/dinov3-vitl16-pretrain-lvd1689m")
+image_cond_model.cuda()
+
+shape_encoder = models.from_pretrained("microsoft/TRELLIS.2-4B/ckpts/shape_enc_next_dc_f16c32_fp16").cuda().eval()
+tex_encoder = models.from_pretrained("microsoft/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"
+full_seg_w_2d_map_path = "./assets/full_seg_w_2d_map"
+interactive = False
+
+transforms = "transforms.json"
+img = "./assets/full_seg_w_2d_map/2d_map.png"
+cond = "./assets/full_seg_w_2d_map/cond.pth"
+
+colors = os.path.join(full_seg_w_2d_map_path, "colors.json")
+output_glb_path = os.path.join(full_seg_w_2d_map_path, "output.glb")
+output_vxz_path = os.path.join(full_seg_w_2d_map_path, "output.vxz")
+
+glb_to_vxz(glb, input_vxz)
+glb_to_parts(glb, parts_path)
+color_glb(parts_path, full_seg_w_2d_map_path, interactive)
+
+color_img(parts_path, img, transforms, colors)
+img_to_cond(rembg_model, image_cond_model, img, cond)
+glb_to_vxz(output_glb_path, output_vxz_path)
+vxz_to_slat(shape_encoder, tex_encoder, input_vxz, output_vxz_path, full_seg_w_2d_map_path, interactive)

data_toolkit/example_interactive_seg.py

@@ -0,0 +1,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)

data_toolkit/glb_to_parts.py

@@ -0,0 +1,17 @@
+import os
+import trimesh
+
+def glb_to_parts(glb_path, output_dir):
+    scene = trimesh.load(glb_path, force='scene')
+    os.makedirs(output_dir, exist_ok=True)
+    geometries = list(scene.geometry.values())
+    for idx, geometry in enumerate(geometries):
+        part_scene = trimesh.Scene()
+        part_scene.add_geometry(geometry)
+        output_path = os.path.join(output_dir, f"{idx}.glb")
+        part_scene.export(output_path)
+
+if __name__ == "__main__":
+    glb_path = "./assets/example.glb"
+    output_dir = "./assets/parts"
+    glb_to_parts(glb_path, output_dir)

data_toolkit/glb_to_vxz.py

@@ -0,0 +1,87 @@
+import torch
+import trimesh
+import o_voxel
+
+from PIL import Image
+
+def make_texture_square_pow2(img: Image.Image, target_size=None):
+    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, 2048)
+    return img.resize((pow2, pow2), Image.BILINEAR)
+
+def preprocess_scene_textures(asset):
+    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))
+            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)
+        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)
+    return asset
+
+def glb_to_vxz(glb_path, vxz_path):
+    asset = trimesh.load(glb_path, force='scene')
+    asset = preprocess_scene_textures(asset)
+    aabb = asset.bounding_box.bounds
+    center = (aabb[0] + aabb[1]) / 2
+    scale = 0.99999 / (aabb[1] - aabb[0]).max()
+    asset.apply_translation(-center)
+    asset.apply_scale(scale)
+    mesh = asset.to_mesh()
+    vertices = torch.from_numpy(mesh.vertices).float()
+    faces = torch.from_numpy(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]
+
+    voxel_indices_mat, attributes = o_voxel.convert.textured_mesh_to_volumetric_attr(
+        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()}
+
+    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['dual_vertices'] = dual_vertices
+    attributes['intersected'] = intersected
+    o_voxel.io.write(vxz_path, voxel_indices, attributes)
+
+if __name__ == "__main__":
+    glb_path = "./assets/example.glb"
+    vxz_path = "./assets/input.vxz"
+    glb_to_vxz(glb_path, vxz_path)

data_toolkit/img_to_cond.py

@@ -0,0 +1,66 @@
+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)
+
+import torch
+import numpy as np
+
+from PIL import Image
+from trellis2.pipelines.rembg import BiRefNet
+from trellis2.modules.image_feature_extractor import DinoV3FeatureExtractor
+
+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)  # type: ignore
+    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.cpu(), 'neg_cond': neg_cond.cpu()}
+
+def img_to_cond(rembg_model, image_cond_model, image_path, save_path):
+    image = Image.open(image_path)
+    image = preprocess_image(rembg_model, image)
+    cond_dict = get_cond(image_cond_model, [image])
+    torch.save(cond_dict, save_path)
+
+if __name__ == "__main__":
+    image_path = "./assets/img.png"
+    save_path = "./assets/cond.pth"
+    rembg_model = BiRefNet(model_name="briaai/RMBG-2.0")
+    rembg_model.cuda()
+    image_cond_model = DinoV3FeatureExtractor(model_name="facebook/dinov3-vitl16-pretrain-lvd1689m")
+    image_cond_model.cuda()
+    img_to_cond(rembg_model, image_cond_model, image_path, save_path)

data_toolkit/texturing_pipeline.json

@@ -0,0 +1,64 @@
+{
+    "name": "Trellis2TexturingPipeline",
+    "args": {
+        "models": {
+            "shape_slat_encoder": "ckpts/shape_enc_next_dc_f16c32_fp16",
+            "tex_slat_decoder": "ckpts/tex_dec_next_dc_f16c32_fp16",
+            "tex_slat_flow_model_512": "ckpts/slat_flow_imgshape2tex_dit_1_3B_512_bf16",
+            "tex_slat_flow_model_1024": "ckpts/slat_flow_imgshape2tex_dit_1_3B_1024_bf16"
+        },
+        "shape_slat_normalization": {
+            "mean": [
+                0.781296, 0.018091, -0.495192, -0.558457, 1.060530, 0.093252, 1.518149, -0.933218,
+                -0.732996, 2.604095, -0.118341, -2.143904, 0.495076, -2.179512, -2.130751, -0.996944,
+                0.261421, -2.217463, 1.260067, -0.150213, 3.790713, 1.481266, -1.046058, -1.523667,
+                -0.059621, 2.220780, 1.621212, 0.877230, 0.567247, -3.175944, -3.186688, 1.578665
+            ],
+            "std": [
+                5.972266, 4.706852, 5.445010, 5.209927, 5.320220, 4.547237, 5.020802, 5.444004,
+                5.226681, 5.683095, 4.831436, 5.286469, 5.652043, 5.367606, 5.525084, 4.730578,
+                4.805265, 5.124013, 5.530808, 5.619001, 5.103930, 5.417670, 5.269677, 5.547194,
+                5.634698, 5.235274, 6.110351, 5.511298, 6.237273, 4.879207, 5.347008, 5.405691
+            ]
+        },
+        "tex_slat_sampler": {
+            "name": "FlowEulerGuidanceIntervalSampler",
+            "args": {
+                "sigma_min": 1e-5
+            },
+            "params": {
+                "steps": 12,
+                "guidance_strength": 1.0,
+                "guidance_rescale": 0.0,
+                "guidance_interval": [0.6, 0.9],
+                "rescale_t": 3.0
+            }
+        },
+        "tex_slat_normalization": {
+            "mean": [
+                3.501659, 2.212398, 2.226094, 0.251093, -0.026248, -0.687364, 0.439898, -0.928075,
+                0.029398, -0.339596, -0.869527, 1.038479, -0.972385, 0.126042, -1.129303, 0.455149,
+                -1.209521, 2.069067, 0.544735, 2.569128, -0.323407, 2.293000, -1.925608, -1.217717,
+                1.213905, 0.971588, -0.023631, 0.106750, 2.021786, 0.250524, -0.662387, -0.768862
+            ],
+            "std": [
+                2.665652, 2.743913, 2.765121, 2.595319, 3.037293, 2.291316, 2.144656, 2.911822,
+                2.969419, 2.501689, 2.154811, 3.163343, 2.621215, 2.381943, 3.186697, 3.021588,
+                2.295916, 3.234985, 3.233086, 2.260140, 2.874801, 2.810596, 3.292720, 2.674999,
+                2.680878, 2.372054, 2.451546, 2.353556, 2.995195, 2.379849, 2.786195, 2.775190
+            ]
+        },
+        "image_cond_model": {
+            "name": "DinoV3FeatureExtractor",
+            "args": {
+                "model_name": "facebook/dinov3-vitl16-pretrain-lvd1689m"
+            }
+        },
+        "rembg_model": {
+            "name": "BiRefNet",
+            "args": {
+                "model_name": "briaai/RMBG-2.0"
+            }
+        }
+    }
+}

data_toolkit/transforms.json

@@ -0,0 +1,31 @@
+[
+    {
+        "camera_angle_x": 0.6981317007977318,
+        "transform_matrix": [
+            [
+                0.8819212913513184,
+                0.06494797021150589,
+                -0.46690112352371216,
+                -0.9338021874427795
+            ],
+            [
+                -0.4713967740535736,
+                0.12150891125202179,
+                -0.8735106587409973,
+                -1.7470210790634155
+            ],
+            [
+                -4.881157167346828e-08,
+                0.9904633164405823,
+                0.13777753710746765,
+                0.2755555510520935
+            ],
+            [
+                0,
+                0,
+                0,
+                1
+            ]
+        ]
+    }
+]

data_toolkit/vxz_to_slat.py

@@ -0,0 +1,123 @@
+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)
+
+import torch
+import o_voxel
+import trellis2.modules.sparse as sp
+
+from trellis2 import models
+
+def vxz_to_latent_slat(shape_encoder, tex_encoder, vxz_path, return_foreground=False):
+    coords, data = o_voxel.io.read(vxz_path)
+    coords = torch.cat([torch.zeros(coords.shape[0], 1, dtype=torch.int32), coords], dim=1).cuda()
+    vertices = (data['dual_vertices'].cuda() / 255)
+    intersected = torch.cat([data['intersected'] % 2, data['intersected'] // 2 % 2, data['intersected'] // 4 % 2], dim=-1).bool().cuda()
+    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.cuda(), shape_slat.coords.cuda())
+    
+    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().cuda() * 2 - 1
+    with torch.no_grad():
+        tex_slat = tex_encoder(sp.SparseTensor(attr, coords))
+        if return_foreground:
+            mask = ((base_color == 1.0).sum(dim=1) == 3)
+            neg_mask = ((base_color != 1.0).sum(dim=1) == 3)
+            tex_slat_foreground = tex_encoder(sp.SparseTensor(attr[mask], coords[mask]))
+            tex_slat_background = tex_encoder(sp.SparseTensor(attr[neg_mask], coords[neg_mask]))
+            foreground_coords = torch.unique(tex_slat_foreground.coords, dim=0)
+            background_coords = torch.unique(tex_slat_background.coords, dim=0)
+            N = background_coords.shape[0]
+            all_coords = torch.cat([background_coords, foreground_coords], dim=0)
+            _, inv = torch.unique(all_coords, dim=0, return_inverse=True)
+            inv_background = inv[:N]
+            inv_foreground = inv[N:]
+            keep = ~torch.isin(inv_foreground, inv_background)
+            foreground_coords = foreground_coords[keep]
+    if return_foreground:
+        return shape_slat, tex_slat, foreground_coords
+    else:
+        return shape_slat, tex_slat
+
+def get_common_coords(slat1, slat2, slat3, slat4, foreground_coords_origin=None):
+    coords_list = [slat1.coords, slat2.coords, slat3.coords, slat4.coords]
+    xs = [torch.unique(x, dim=0) for x in coords_list]
+    all_coords = torch.cat(xs, dim=0)
+    uniq_coords, counts = torch.unique(all_coords, dim=0, return_counts=True)
+    common_coords = uniq_coords[counts == len(coords_list)].cuda()
+
+    if foreground_coords_origin is not None:
+        xs_foreground = [torch.unique(x, dim=0) for x in (common_coords, foreground_coords_origin)]
+        all_coords_foreground = torch.cat(xs_foreground, dim=0)
+        uniq_coords_foreground, counts_foreground = torch.unique(all_coords_foreground, dim=0, return_counts=True)
+        foreground_coords = uniq_coords_foreground[counts_foreground == 2].cuda()
+        return common_coords, foreground_coords
+    else:
+        return common_coords
+
+def get_slat_by_common_coords(slat_origin, common_coords):
+    N = slat_origin.coords.shape[0]
+    all_coords = torch.cat([slat_origin.coords, common_coords], dim=0)
+    uniq_coords, inv = torch.unique(all_coords, dim=0, return_inverse=True)
+    inv_slat = inv[:N].cuda()
+    inv_common = inv[N:].cuda()
+    device = slat_origin.coords.device
+    idx_map = torch.full((uniq_coords.shape[0],), -1, dtype=torch.int32, device=device)
+    slat_idx = torch.arange(N, dtype=torch.int32, device=device)
+    idx_map.scatter_reduce_(0, inv_slat, slat_idx, reduce='amin', include_self=False)
+    idx_in_slat = idx_map[inv_common]
+    feats = slat_origin.feats[idx_in_slat]
+    return sp.SparseTensor(feats, common_coords)
+
+def get_point(point_num, common_coords, foreground_coords):
+    device = common_coords.device
+    point_feats_coords = torch.zeros((10, 4), dtype=torch.int32, device=device)
+    point_labels = torch.zeros((10, 1), dtype=torch.int32, device=device)
+    foreground_idx = torch.randperm(foreground_coords.shape[0], device=device)[:point_num]
+    point_foreground = foreground_coords[foreground_idx]
+    if point_foreground.shape[0] != point_num:
+        return None
+    point_feats_coords[:point_num] = point_foreground
+    point_labels[:point_num] = 1
+    return {'point_feats': point_feats_coords.cpu(), 'point_labels': point_labels.cpu()}
+
+def vxz_to_slat(shape_encoder, tex_encoder, input_vxz_path, output_vxz_path, save_dir, interactive):
+    input_shape_slat_origin, input_tex_slat_origin = vxz_to_latent_slat(shape_encoder, tex_encoder, input_vxz_path)
+    if interactive:
+        output_shape_slat_origin, output_tex_slat_origin, foreground_coords_origin = vxz_to_latent_slat(shape_encoder, tex_encoder, output_vxz_path, return_foreground=interactive)
+        common_coords, foreground_coords = get_common_coords(input_shape_slat_origin, input_tex_slat_origin, output_shape_slat_origin, output_tex_slat_origin, foreground_coords_origin)
+    else:
+        output_shape_slat_origin, output_tex_slat_origin = vxz_to_latent_slat(shape_encoder, tex_encoder, output_vxz_path)
+        common_coords = get_common_coords(input_shape_slat_origin, input_tex_slat_origin, output_shape_slat_origin, output_tex_slat_origin)
+
+    os.makedirs(save_dir, exist_ok=True)
+    shape_slat = get_slat_by_common_coords(input_shape_slat_origin, common_coords)
+    torch.save({"feats": shape_slat.feats.cpu(), "coords": shape_slat.coords.cpu()}, os.path.join(save_dir, "shape_slat.pth"))
+    input_tex_slat = get_slat_by_common_coords(input_tex_slat_origin, common_coords)
+    torch.save({"feats": input_tex_slat.feats.cpu(), "coords": input_tex_slat.coords.cpu()}, os.path.join(save_dir, "input_tex_slat.pth"))
+    output_tex_slat_gt = get_slat_by_common_coords(output_tex_slat_origin, common_coords)
+    torch.save({"feats": output_tex_slat_gt.feats.cpu(), "coords": output_tex_slat_gt.coords.cpu()}, os.path.join(save_dir, "output_tex_slat.pth"))
+    
+    if interactive:
+        for point_num in range(1, 11):
+            input_points = get_point(point_num, common_coords, foreground_coords)
+            if input_points is None:
+                continue
+            torch.save(input_points, os.path.join(save_dir, f"point_{point_num}.pth"))
+
+if __name__ == "__main__":
+    input_vxz_path = "./assets/input.vxz"
+    output_vxz_path = "./assets/interactive_seg/0/output.vxz"
+    save_dir = "./assets/interactive_seg/0"
+    interactive = True
+    shape_encoder = models.from_pretrained("microsoft/TRELLIS.2-4B/ckpts/shape_enc_next_dc_f16c32_fp16").cuda().eval()
+    tex_encoder = models.from_pretrained("microsoft/TRELLIS.2-4B/ckpts/tex_enc_next_dc_f16c32_fp16").cuda().eval()
+    vxz_to_slat(shape_encoder, tex_encoder, input_vxz_path, output_vxz_path, save_dir, interactive)

examples/00aee5c2fef743d69421bb642d446a5b.glb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f2d10a086262ccdf1109d5eb2fa385efb005da7503df9d39e5ab5cce0aa7d236
+size 1585072

examples/01b8043112e74366a21256d5e64398fb.glb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:da1ad73c6eeb2338ed7d982bb5307daaa2f5faab91e0f883de214960c739be19
+size 467796

examples/0c070001a3904cd6809a31345475e930.glb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:57cd98fb3ca54128942744a8090fd040487159107cc810acdd1c20594ceac9bb
+size 1524860

examples/0c3ca2b32545416f8f1e6f0e87def1a6.glb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:984c76bea31761c1f600240e6212b51c407b4cb69c19af30352d11879274a3b9
+size 649004

examples/1b3e8b99913442308aa989e3f87680b3.glb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fa19b47445f18cbd75d812078103d0ae9bbc0b1ad404eeac12e78c9ed6ab41a1
+size 2320124

examples/1c33b2e86c023a72905a5bea4ae713d0.glb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7f0589018f02feaafa00770b8e1fe434009ed5bd419127452edd87d5534fa51b
+size 1164628

examples/1ca8ea337fbc4bcfbeb3c633bc4c43f0.glb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5417c20a8f8b7fad7df93f2c656d1a6d22c44c284512e2bbd1ad0f498b9f7287
+size 1755800

examples/2260799ee4e342398b64ab4ce8af1559.glb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:69914b1a07a6e3866b1d88bbb6f8d33fc13d492315626165455de902be818cb2
+size 515608

examples/2ae5cf2990c34e7db704f677de8de74c.glb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:902905bc4cc0563f242a7cc5f96ee92165a4912527fda52f2beaa43ab6e67367
+size 468560

examples/2ceb6778ac114101833e4c531544ada8.glb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4d0ad0049486e6f46ea3ec4df8ea0acf2d846a0f5cac22032a1b26cf2050dc3f
+size 225560

examples/4b57e73e82ab400aa307adac36ea0e5e.glb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:89dabad77dffd5cfc26fd4e6f977fb9ed686fa165a950d0c35df0693262a9352
+size 3001252

inference_full.py

@@ -0,0 +1,553 @@
+import os
+os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1'
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+
+import json
+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 trellis2.pipelines.rembg import BiRefNet
+from trellis2.representations import MeshWithVoxel
+from data_toolkit.bpy_render import render_from_transforms
+from trellis2.modules.image_feature_extractor import DinoV3FeatureExtractor
+
+
+TRELLIS_PIPELINE_JSON = "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 = "facebook/dinov3-vitl16-pretrain-lvd1689m"
+# DINO_PATH = "/bj-mlp-buaa-prod/pretrained_weights/pretrained_weights/dinov3"
+
+
+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 already textured, just ensure PBR material
+    if isinstance(mesh.visual, trimesh.visual.texture.TextureVisuals):
+        mat = getattr(mesh.visual, "material", None)
+        if isinstance(mat, trimesh.visual.material.SimpleMaterial):
+            # Avoid side-effects if the mesh is shared elsewhere
+            mesh = mesh.copy()
+            try:
+                mesh.visual.material = mat.to_pbr()
+            except Exception:
+                # Fallback: construct a minimal PBR material from the image
+                mesh.visual.material = trimesh.visual.material.PBRMaterial(
+                    baseColorTexture=mat.image
+                )
+        return mesh
+
+    # Only handle ColorVisuals here
+    if not isinstance(mesh.visual, trimesh.visual.color.ColorVisuals):
+        return mesh
+
+    F = int(len(mesh.faces))
+    if F <= 0:
+        return mesh
+
+    # ---- Get per-face RGBA (uint8) ----
+    face_rgba = None
+
+    # Prefer face colors if present
+    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)
+
+    # Fallback: average vertex colors per face
+    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]  # (F,3,4)
+            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):
+        # Replace geometry objects in-place; graph nodes still refer to geometry names
+        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):
+    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, 2048)
+    return img.resize((pow2, pow2), Image.BILINEAR)
+
+
+def preprocess_scene_textures(asset):
+    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))
+            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)
+        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)
+    return asset
+
+
+def process_glb_to_vxz(glb_path, vxz_path):
+    asset = trimesh.load(glb_path, force='scene')
+    asset = preprocess_scene_textures(asset)
+
+    asset = ensure_texture_visuals(asset)
+    
+    aabb = asset.bounding_box.bounds
+    center = (aabb[0] + aabb[1]) / 2
+    scale = 0.99999 / (aabb[1] - aabb[0]).max()
+    asset.apply_translation(-center)
+    asset.apply_scale(scale)
+    mesh = asset.to_mesh()
+    vertices = torch.from_numpy(mesh.vertices).float()
+    faces = torch.from_numpy(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]
+
+    voxel_indices_mat, attributes = o_voxel.convert.textured_mesh_to_volumetric_attr(
+        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()}
+
+    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['dual_vertices'] = dual_vertices
+    attributes['intersected'] = intersected
+    o_voxel.io.write(vxz_path, voxel_indices, attributes)
+
+
+def vxz_to_latent_slat(shape_encoder, shape_decoder, tex_encoder, vxz_path):
+    coords, data = o_voxel.io.read(vxz_path)
+    coords = torch.cat([torch.zeros(coords.shape[0], 1, dtype=torch.int32), coords], dim=1).cuda()
+    vertices = (data['dual_vertices'].cuda() / 255)
+    intersected = torch.cat([data['intersected'] % 2, data['intersected'] // 2 % 2, data['intersected'] // 4 % 2], dim=-1).bool().cuda()
+    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.cuda(), shape_slat.coords.cuda())
+        shape_decoder.set_resolution(512)
+        meshes, subs = shape_decoder(shape_slat, return_subs=True)
+
+    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().cuda() * 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)  # type: ignore
+    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):
+    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]]
+    noise = sp.SparseTensor(torch.randn_like(input_tex_slat.feats), shape_slat.coords)
+    output_tex_slat = sampler.sample(gen3dseg, noise, input_tex_slat, shape_slat, coords_len_list, cond_dict, pipeline_args['tex_slat_sampler']['params'])
+    output_tex_slat = output_tex_slat * tex_std + tex_mean
+    return output_tex_slat
+
+
+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 m, v in zip(meshes, tex_voxels):
+        m.fill_holes()
+        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]
+    mesh.simplify(10000000)
+    glb = o_voxel.postprocess.to_glb(
+        vertices            =   mesh.vertices,
+        faces               =   mesh.faces,
+        attr_volume         =   mesh.attrs,
+        coords              =   mesh.coords,
+        attr_layout         =   mesh.layout,
+        voxel_size          =   mesh.voxel_size,
+        aabb                =   [[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
+        decimation_target   =   100000,
+        texture_size        =   4096,
+        remesh              =   True,
+        remesh_band         =   1,
+        remesh_project      =   0,
+        verbose             =   True
+    )
+    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']
+
+        print("-" * 100)
+        print("[Init] Loading TRELLIS backbone ............")
+        self.tex_slat_flow_model = models.from_pretrained(TRELLIS_TEX_FLOW)
+
+        self.gen3dseg = Gen3DSeg(self.tex_slat_flow_model)
+        self.gen3dseg.eval()
+        self.gen3dseg.cuda()
+
+        self.sampler = Sampler()
+
+        self.shape_encoder = models.from_pretrained(TRELLIS_SHAPE_ENC).cuda().eval()
+        self.tex_encoder = models.from_pretrained(TRELLIS_TEX_ENC).cuda().eval()
+        self.shape_decoder = models.from_pretrained(TRELLIS_SHAPE_DEC).cuda().eval()
+        self.tex_decoder = models.from_pretrained(TRELLIS_TEX_DEC).cuda().eval()
+
+        print("-" * 100)
+        print("[Init] Loading conditioners ............")
+
+        self.rembg_model = BiRefNet(model_name="briaai/RMBG-2.0")
+        self.rembg_model.cuda()
+
+        self.image_cond_model = DinoV3FeatureExtractor(DINO_PATH)
+        self.image_cond_model.cuda()
+
+        self.loaded = True
+        print("[Init] Done.")
+
+    def load_ckpt_if_needed(self, ckpt_path: str):
+        if self.current_ckpt == ckpt_path:
+            return
+
+        print("-" * 100)
+        print(f"[CKPT] Loading ckpt: {ckpt_path}")
+        state_dict = torch.load(ckpt_path)['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()
+        self.gen3dseg.cuda()
+        self.current_ckpt = ckpt_path
+
+
+PIPE = _LoadedPipeline()
+
+
+def inference_with_loaded_models(ckpt_path, item):
+    PIPE.load_all_models()
+    PIPE.load_ckpt_if_needed(ckpt_path)
+
+    if PIPE.rembg_model is None:
+        raise RuntimeError("PIPE.rembg_model is None. Check BiRefNet loading and .cuda() usage.")
+    if PIPE.image_cond_model is None:
+        raise RuntimeError("PIPE.image_cond_model is None. Check DinoV3FeatureExtractor loading and .cuda() usage.")
+
+    process_glb_to_vxz(item['glb'], item['input_vxz'])
+    shape_slat, meshes, subs, tex_slat = vxz_to_latent_slat(
+        PIPE.shape_encoder, PIPE.shape_decoder, PIPE.tex_encoder, item['input_vxz']
+    )
+
+    if not item['2d_map']:
+        render_from_transforms(item['glb'], item['transforms'], item['img'])
+
+    image = Image.open(item['img'])
+    image = preprocess_image(PIPE.rembg_model, image)
+    cond = get_cond(PIPE.image_cond_model, [image])
+
+    output_tex_slat = tex_slat_sample_single(
+        PIPE.gen3dseg, PIPE.sampler, PIPE.pipeline_args, shape_slat, tex_slat, cond
+    )
+    with torch.no_grad():
+        tex_voxels = PIPE.tex_decoder(output_tex_slat, guide_subs=subs) * 0.5 + 0.5
+
+    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"])

inference_full_ori.py

@@ -0,0 +1,383 @@
+import os
+os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1'
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+
+import json
+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 trellis2.pipelines.rembg import BiRefNet
+from trellis2.representations import MeshWithVoxel
+from data_toolkit.bpy_render import render_from_transforms
+from trellis2.modules.image_feature_extractor import DinoV3FeatureExtractor
+
+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):
+    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, 2048)
+    return img.resize((pow2, pow2), Image.BILINEAR)
+
+def preprocess_scene_textures(asset):
+    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))
+            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)
+        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)
+    return asset
+
+def process_glb_to_vxz(glb_path, vxz_path):
+    asset = trimesh.load(glb_path, force='scene')
+    asset = preprocess_scene_textures(asset)
+    aabb = asset.bounding_box.bounds
+    center = (aabb[0] + aabb[1]) / 2
+    scale = 0.99999 / (aabb[1] - aabb[0]).max()
+    asset.apply_translation(-center)
+    asset.apply_scale(scale)
+    mesh = asset.to_mesh()
+    vertices = torch.from_numpy(mesh.vertices).float()
+    faces = torch.from_numpy(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]
+
+    voxel_indices_mat, attributes = o_voxel.convert.textured_mesh_to_volumetric_attr(
+        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()}
+
+    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['dual_vertices'] = dual_vertices
+    attributes['intersected'] = intersected
+    o_voxel.io.write(vxz_path, voxel_indices, attributes)
+
+def vxz_to_latent_slat(shape_encoder, shape_decoder, tex_encoder, vxz_path):
+    coords, data = o_voxel.io.read(vxz_path)
+    coords = torch.cat([torch.zeros(coords.shape[0], 1, dtype=torch.int32), coords], dim=1).cuda()
+    vertices = (data['dual_vertices'].cuda() / 255)
+    intersected = torch.cat([data['intersected'] % 2, data['intersected'] // 2 % 2, data['intersected'] // 4 % 2], dim=-1).bool().cuda()
+    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.cuda(), shape_slat.coords.cuda())
+        shape_decoder.set_resolution(512)
+        meshes, subs = shape_decoder(shape_slat, return_subs=True)
+    
+    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().cuda() * 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)  # type: ignore
+    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):
+    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]]
+    noise = sp.SparseTensor(torch.randn_like(input_tex_slat.feats), shape_slat.coords)
+    output_tex_slat = sampler.sample(gen3dseg, noise, input_tex_slat, shape_slat, coords_len_list, cond_dict, pipeline_args['tex_slat_sampler']['params'])
+    output_tex_slat = output_tex_slat * tex_std + tex_mean
+    return output_tex_slat
+
+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 m, v in zip(meshes, tex_voxels):
+        m.fill_holes()
+        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]
+    mesh.simplify(10000000)
+    # mesh.simplify(16777216) # nvdiffrast limit
+    glb = o_voxel.postprocess.to_glb(
+        vertices            =   mesh.vertices,
+        faces               =   mesh.faces,
+        attr_volume         =   mesh.attrs,
+        coords              =   mesh.coords,
+        attr_layout         =   mesh.layout,
+        voxel_size          =   mesh.voxel_size,
+        aabb                =   [[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
+        decimation_target   =   100000, # 1000000
+        texture_size        =   4096,
+        remesh              =   True,
+        # remesh              =   False,
+        remesh_band         =   1,
+        remesh_project      =   0,
+        verbose             =   True
+    )
+    return glb
+
+def inference(ckpt_path, item):
+    print("-"*100)
+    print("Loading model ............")
+
+    with open("/media/nfs/tmp_data/fenghr/download/TRELLIS.2-4B/texturing_pipeline.json", "r") as f:
+        pipeline_config = json.load(f)
+    pipeline_args = pipeline_config['args']
+    tex_slat_flow_model = models.from_pretrained(
+        "/media/nfs/tmp_data/fenghr/download/TRELLIS.2-4B/ckpts/slat_flow_imgshape2tex_dit_1_3B_512_bf16")
+
+    gen3dseg = Gen3DSeg(tex_slat_flow_model)
+    state_dict = torch.load(ckpt_path)['state_dict']
+    state_dict = OrderedDict([(k.replace("gen3dseg.", ""), v) for k, v in state_dict.items()])
+    gen3dseg.load_state_dict(state_dict)
+    gen3dseg.eval()
+    gen3dseg.cuda()
+    sampler = Sampler()
+
+    shape_encoder = models.from_pretrained(
+        "/media/nfs/tmp_data/fenghr/download/TRELLIS.2-4B/ckpts/shape_enc_next_dc_f16c32_fp16").cuda().eval()
+    tex_encoder = models.from_pretrained(
+        "/media/nfs/tmp_data/fenghr/download/TRELLIS.2-4B/ckpts/tex_enc_next_dc_f16c32_fp16").cuda().eval()
+    shape_decoder = models.from_pretrained(
+        "/media/nfs/tmp_data/fenghr/download/TRELLIS.2-4B/ckpts/shape_dec_next_dc_f16c32_fp16").cuda().eval()
+    tex_decoder = models.from_pretrained(
+        "/media/nfs/tmp_data/fenghr/download/TRELLIS.2-4B/ckpts/tex_dec_next_dc_f16c32_fp16").cuda().eval()
+
+    rembg_model = BiRefNet(model_name="briaai/RMBG-2.0")
+    rembg_model.cuda()
+    # image_cond_model = DinoV3FeatureExtractor(
+    #     model_name="facebook/dinov3-vitl16-pretrain-lvd1689m")
+    image_cond_model = DinoV3FeatureExtractor("/media/nfs/tmp_data/fenghr/download/dinov3")
+    image_cond_model.cuda()
+
+    process_glb_to_vxz(item['glb'], item['input_vxz'])
+    shape_slat, meshes, subs, tex_slat = vxz_to_latent_slat(shape_encoder, shape_decoder, tex_encoder, item['input_vxz'])
+
+    print("-"*100)
+    print("Getting cond ............")
+    if not item['2d_map']:
+        render_from_transforms(item['glb'], item['transforms'], item['img'])
+    image = Image.open(item['img'])
+    image = preprocess_image(rembg_model, image)
+    cond = get_cond(image_cond_model, [image])
+
+    print("-"*100)
+    print("Sampling .................")
+    output_tex_slat = tex_slat_sample_single(gen3dseg, sampler, pipeline_args, shape_slat, tex_slat, cond)
+    with torch.no_grad():
+        tex_voxels = tex_decoder(output_tex_slat, guide_subs=subs) * 0.5 + 0.5
+
+    print("-"*100)
+    print("Exporting glb ............")
+    glb = slat_to_glb(meshes, tex_voxels)
+    glb.export(item['export_glb'])
+
+if __name__ == "__main__":
+    _2d_map = False
+    if _2d_map:
+        ckpt_path = "/media/nfs/tmp_data/fenghr/SegviGen/pretrained_models/full_seg_w_2d_map.ckpt"
+        item = {
+            "2d_map": True,
+            "glb": "./data_toolkit/assets/example.glb",
+            "input_vxz": "./data_toolkit/assets/input.vxz",
+            "img": "./data_toolkit/assets/full_seg_w_2d_map/2d_map.png",
+            "export_glb": "./data_toolkit/assets/output.glb"
+        }
+    else:
+        ckpt_path = "/media/nfs/tmp_data/fenghr/SegviGen/pretrained_models/full_seg.ckpt"
+        item = {
+            "2d_map": False,
+            "glb": "./data_toolkit/assets/example.glb",
+            "input_vxz": "./data_toolkit/assets/input.vxz",
+            "transforms": "./data_toolkit/transforms.json",
+            "img": "./data_toolkit/assets/img.png",
+            "export_glb": "./data_toolkit/assets/output.glb"
+        }
+    inference(ckpt_path, item)

inference_interactive.py

@@ -0,0 +1,435 @@
+import os
+os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1'
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+
+import json
+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 types import MethodType
+from collections import OrderedDict
+from torch.nn import functional as F
+from trellis2.pipelines.rembg import BiRefNet
+from trellis2.modules.utils import manual_cast
+from trellis2.representations import MeshWithVoxel
+from data_toolkit.bpy_render import render_from_transforms
+from trellis2.modules.image_feature_extractor import DinoV3FeatureExtractor
+
+class Sampler:
+    def _inference_model(self, model, x_t, tex_slat, shape_slat, input_points, 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, input_points, coords_len_list)
+
+    def guidance_inference_model(self, model, x_t, tex_slat, shape_slat, input_points, 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, input_points, coords_len_list, t, cond_dict['cond'])
+        elif guidance_strength == 0:
+            return self._inference_model(model, x_t, tex_slat, shape_slat, input_points, coords_len_list, t, cond_dict['neg_cond'])
+        else:
+            pred_pos = self._inference_model(model, x_t, tex_slat, shape_slat, input_points, coords_len_list, t, cond_dict['cond'])
+            pred_neg = self._inference_model(model, x_t, tex_slat, shape_slat, input_points, 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, input_points, 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, input_points, coords_len_list, t, cond_dict, guidance_strength, guidance_rescale)
+        else:
+            return self.guidance_inference_model(model, x_t, tex_slat, shape_slat, input_points, coords_len_list, t, cond_dict, 1, guidance_rescale)
+
+    @torch.no_grad()
+    def sample_once(self, model, x_t, tex_slat, shape_slat, input_points, coords_len_list, t, t_prev, cond_dict, sampler_params):
+        pred_v = self.interval_inference_model(model, x_t, tex_slat, shape_slat, input_points, 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, input_points, 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, input_points, coords_len_list, t, t_prev, cond_dict, sampler_params)
+        return sample
+
+def flow_forward(self, x, t, cond, concat_cond, point_embeds, coords_len_list):
+    # x.feats: [N, 32]
+    x = sp.sparse_cat([x, concat_cond], dim=-1)
+    if isinstance(cond, list):
+        cond = sp.VarLenTensor.from_tensor_list(cond)
+    # x.feats: [N, 64]
+    h = self.input_layer(x)
+    # h.feats: [N, 1536]
+    h = manual_cast(h, self.dtype)
+    t_emb = self.t_embedder(t)
+    t_emb = self.adaLN_modulation(t_emb)
+    t_emb = manual_cast(t_emb, self.dtype)
+    cond = manual_cast(cond, self.dtype)
+    point_embeds = manual_cast(point_embeds, self.dtype)
+
+    h_feats_list = []
+    h_coords_list = []
+    begin = 0
+    for i, coords_len in enumerate(coords_len_list):
+        end = begin + 2 * coords_len
+        h_feats_list.append(h.feats[begin:end])
+        h_coords_list.append(h.coords[begin:end])
+        h_feats_list.append(point_embeds.feats[i*10:(i+1)*10])
+        h_coords_list.append(point_embeds.coords[i*10:(i+1)*10])
+        begin = end + 10
+    h = sp.SparseTensor(torch.cat(h_feats_list), torch.cat(h_coords_list))
+
+    for block in self.blocks:
+        h = block(h, t_emb, cond)
+
+    h_feats_list = []
+    h_coords_list = []
+    begin = 0
+    for i, coords_len in enumerate(coords_len_list):
+        end = begin + 2 * coords_len
+        h_feats_list.append(h.feats[begin:end])
+        h_coords_list.append(h.coords[begin:end])
+        begin = end
+    h = sp.SparseTensor(torch.cat(h_feats_list), torch.cat(h_coords_list))
+
+    h = manual_cast(h, x.dtype)
+    h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:]))
+    # h.feats: [N, 1536]
+    h = self.out_layer(h)
+    # h.feats: [N, 32]
+    return h
+
+class Gen3DSeg(nn.Module):
+    def __init__(self, flow_model):
+        super().__init__()
+        self.flow_model = flow_model
+        self.seg_embeddings = nn.Embedding(1, 1536)
+        
+    def get_positional_encoding(self, input_points):
+        point_feats_embed = torch.zeros((10, 1536), dtype=torch.float32).to(input_points['point_slats'].feats.device)
+        labels = input_points['point_labels'].squeeze(-1)
+        point_feats_embed[labels == 1] = self.seg_embeddings.weight
+        return sp.SparseTensor(point_feats_embed, input_points['point_slats'].coords)
+
+    def forward(self, x_t, tex_slats, shape_slats, t, cond, input_points, 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))
+
+        point_embeds = self.get_positional_encoding(input_points)
+        output_tex_slats = self.flow_model(x_t, t, cond, shape_slats, point_embeds, coords_len_list)
+        
+        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):
+    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, 2048)
+    return img.resize((pow2, pow2), Image.BILINEAR)
+
+def preprocess_scene_textures(asset):
+    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))
+            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)
+        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)
+    return asset
+
+def process_glb_to_vxz(glb_path, vxz_path):
+    asset = trimesh.load(glb_path, force='scene')
+    asset = preprocess_scene_textures(asset)
+    aabb = asset.bounding_box.bounds
+    center = (aabb[0] + aabb[1]) / 2
+    scale = 0.99999 / (aabb[1] - aabb[0]).max()
+    asset.apply_translation(-center)
+    asset.apply_scale(scale)
+    mesh = asset.to_mesh()
+    vertices = torch.from_numpy(mesh.vertices).float()
+    faces = torch.from_numpy(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]
+
+    voxel_indices_mat, attributes = o_voxel.convert.textured_mesh_to_volumetric_attr(
+        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()}
+
+    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['dual_vertices'] = dual_vertices
+    attributes['intersected'] = intersected
+    o_voxel.io.write(vxz_path, voxel_indices, attributes)
+
+def vxz_to_latent_slat(shape_encoder, shape_decoder, tex_encoder, vxz_path):
+    coords, data = o_voxel.io.read(vxz_path)
+    coords = torch.cat([torch.zeros(coords.shape[0], 1, dtype=torch.int32), coords], dim=1).cuda()
+    vertices = (data['dual_vertices'].cuda() / 255)
+    intersected = torch.cat([data['intersected'] % 2, data['intersected'] // 2 % 2, data['intersected'] // 4 % 2], dim=-1).bool().cuda()
+    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.cuda(), shape_slat.coords.cuda())
+        shape_decoder.set_resolution(512)
+        meshes, subs = shape_decoder(shape_slat, return_subs=True)
+    
+    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().cuda() * 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)  # type: ignore
+    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, input_points):
+    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]]
+    noise = sp.SparseTensor(torch.randn_like(input_tex_slat.feats), shape_slat.coords)
+    output_tex_slat = sampler.sample(gen3dseg, noise, input_tex_slat, shape_slat, input_points, coords_len_list, cond_dict, pipeline_args['tex_slat_sampler']['params'])
+    output_tex_slat = output_tex_slat * tex_std + tex_mean
+    return output_tex_slat
+
+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 m, v in zip(meshes, tex_voxels):
+        m.fill_holes()
+        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]
+    mesh.simplify(10000000)
+    # mesh.simplify(16777216) # nvdiffrast limit
+    glb = o_voxel.postprocess.to_glb(
+        vertices            =   mesh.vertices,
+        faces               =   mesh.faces,
+        attr_volume         =   mesh.attrs,
+        coords              =   mesh.coords,
+        attr_layout         =   mesh.layout,
+        voxel_size          =   mesh.voxel_size,
+        aabb                =   [[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
+        decimation_target   =   100000, # 1000000
+        texture_size        =   4096,
+        remesh              =   True,
+        # remesh              =   False,
+        remesh_band         =   1,
+        remesh_project      =   0,
+        verbose             =   True
+    )
+    return glb
+
+def inference(ckpt_path, item, input_vxz_points_list):
+    print("-"*100)
+    print("Loading model ............")
+    with open("microsoft/TRELLIS.2-4B/pipeline.json", "r") as f:
+        pipeline_config = json.load(f)
+    pipeline_args = pipeline_config['args']
+    tex_slat_flow_model = models.from_pretrained("microsoft/TRELLIS.2-4B/ckpts/slat_flow_imgshape2tex_dit_1_3B_512_bf16")
+    tex_slat_flow_model.forward = MethodType(flow_forward, tex_slat_flow_model)
+
+    gen3dseg = Gen3DSeg(tex_slat_flow_model)
+    state_dict = torch.load(ckpt_path)['state_dict']
+    state_dict = OrderedDict([(k.replace("gen3dseg.", ""), v) for k, v in state_dict.items()])
+    gen3dseg.load_state_dict(state_dict)
+    gen3dseg.eval()
+    gen3dseg.cuda()
+    sampler = Sampler()
+
+    shape_encoder = models.from_pretrained("microsoft/TRELLIS.2-4B/ckpts/shape_enc_next_dc_f16c32_fp16").cuda().eval()
+    tex_encoder = models.from_pretrained("microsoft/TRELLIS.2-4B/ckpts/tex_enc_next_dc_f16c32_fp16").cuda().eval()
+    shape_decoder = models.from_pretrained("microsoft/TRELLIS.2-4B/ckpts/shape_dec_next_dc_f16c32_fp16").cuda().eval()
+    tex_decoder = models.from_pretrained("microsoft/TRELLIS.2-4B/ckpts/tex_dec_next_dc_f16c32_fp16").cuda().eval()
+
+    rembg_model = BiRefNet(model_name="briaai/RMBG-2.0")
+    rembg_model.cuda()
+    image_cond_model = DinoV3FeatureExtractor(model_name="facebook/dinov3-vitl16-pretrain-lvd1689m")
+    image_cond_model.cuda()
+
+    process_glb_to_vxz(item['glb'], item['input_vxz'])
+    shape_slat, meshes, subs, tex_slat = vxz_to_latent_slat(shape_encoder, shape_decoder, tex_encoder, item['input_vxz'])
+
+    print("-"*100)
+    print("Getting cond ............")
+    render_from_transforms(item['glb'], item['transforms'], item['img'])
+    image = Image.open(item['img'])
+    image = preprocess_image(rembg_model, image)
+    cond = get_cond(image_cond_model, [image])
+
+    print("-"*100)
+    print("Sampling .................")
+    vxz_points_coords = torch.tensor(input_vxz_points_list, dtype=torch.int32).cuda()
+    vxz_points_coords = torch.cat([torch.zeros((vxz_points_coords.shape[0], 1), dtype=torch.int32).cuda(), vxz_points_coords], dim=1)
+    input_points_coords = tex_encoder(sp.SparseTensor(torch.zeros((vxz_points_coords.shape[0], 6), dtype=torch.float32).cuda(), vxz_points_coords)).coords
+    input_points_coords = torch.unique(input_points_coords, dim=0)
+    point_num = input_points_coords.shape[0]
+    if point_num >= 10:
+        input_points_coords = input_points_coords[:10]
+        point_labels = torch.tensor(([[1]]*10), dtype=torch.int32).cuda()
+    else:
+        input_points_coords = torch.cat([input_points_coords, torch.zeros((10 - point_num, 4), dtype=torch.int32).cuda()], dim=0)
+        point_labels = torch.tensor(([[1]]*point_num+[[0]]*(10-point_num)), dtype=torch.int32).cuda()
+    input_points = {'point_slats': sp.SparseTensor(input_points_coords, input_points_coords), 'point_labels': point_labels}
+
+    output_tex_slat = tex_slat_sample_single(gen3dseg, sampler, pipeline_args, shape_slat, tex_slat, cond, input_points)
+    with torch.no_grad():
+        tex_voxels = tex_decoder(output_tex_slat, guide_subs=subs) * 0.5 + 0.5
+
+    print("-"*100)
+    print("Exporting glb ............")
+    glb = slat_to_glb(meshes, tex_voxels)
+    glb.export(item['export_glb'])
+
+if __name__ == "__main__":
+    ckpt_path = "path/to/interactive_seg.ckpt"
+    item = {
+        "glb": "./data_toolkit/assets/example.glb",
+        "input_vxz": "./data_toolkit/assets/input.vxz",
+        "transforms": "./data_toolkit/transforms.json",
+        "img": "./data_toolkit/assets/img.png",
+        "export_glb": "./data_toolkit/assets/output.glb"
+    }
+    input_vxz_points_list = [[388, 448, 392]] # example
+    inference(ckpt_path, item, input_vxz_points_list)

inference_unified.py

@@ -0,0 +1,473 @@
+import os
+os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1'
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+
+import json
+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 types import MethodType
+from collections import OrderedDict
+from torch.nn import functional as F
+from trellis2.pipelines.rembg import BiRefNet
+from trellis2.modules.utils import manual_cast
+from trellis2.representations import MeshWithVoxel
+from data_toolkit.bpy_render import render_from_transforms
+from trellis2.modules.image_feature_extractor import DinoV3FeatureExtractor
+
+class Sampler:
+    def _inference_model(self, model, x_t, tex_slat, shape_slat, input_points, coords_len_list, t, cond, tag):
+        t = torch.tensor([t*1000] * x_t.shape[0], dtype=torch.float32).cuda()
+        tag = torch.tensor([tag] * x_t.shape[0], dtype=torch.float32).cuda()
+        return model(x_t, tex_slat, shape_slat, t, tag, cond, input_points, coords_len_list)
+
+    def guidance_inference_model(self, model, x_t, tex_slat, shape_slat, input_points, coords_len_list, t, cond_dict, tag, guidance_strength, guidance_rescale=0.0):
+        if guidance_strength == 1:
+            return self._inference_model(model, x_t, tex_slat, shape_slat, input_points, coords_len_list, t, cond_dict['cond'], tag)
+        elif guidance_strength == 0:
+            return self._inference_model(model, x_t, tex_slat, shape_slat, input_points, coords_len_list, t, cond_dict['neg_cond'], tag)
+        else:
+            pred_pos = self._inference_model(model, x_t, tex_slat, shape_slat, input_points, coords_len_list, t, cond_dict['cond'], tag)
+            pred_neg = self._inference_model(model, x_t, tex_slat, shape_slat, input_points, coords_len_list, t, cond_dict['neg_cond'], tag)
+            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, input_points, coords_len_list, t, cond_dict, tag, 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, input_points, coords_len_list, t, cond_dict, tag, guidance_strength, guidance_rescale)
+        else:
+            return self.guidance_inference_model(model, x_t, tex_slat, shape_slat, input_points, coords_len_list, t, cond_dict, tag, 1, guidance_rescale)
+
+    @torch.no_grad()
+    def sample_once(self, model, x_t, tex_slat, shape_slat, input_points, coords_len_list, t, t_prev, cond_dict, tag, sampler_params):
+        pred_v = self.interval_inference_model(model, x_t, tex_slat, shape_slat, input_points, coords_len_list, t, cond_dict, tag, 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, input_points, coords_len_list, cond_dict, tag, 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, input_points, coords_len_list, t, t_prev, cond_dict, tag, sampler_params)
+        return sample
+
+def flow_forward(self, x, t, tag_embeds, cond, concat_cond, point_embeds, coords_len_list):
+    # x.feats: [N, 32]
+    x = sp.sparse_cat([x, concat_cond], dim=-1)
+    if isinstance(cond, list):
+        cond = sp.VarLenTensor.from_tensor_list(cond)
+    # x.feats: [N, 64]
+    h = self.input_layer(x)
+    # h.feats: [N, 1536]
+    h = manual_cast(h, self.dtype)
+    t_emb = self.t_embedder(t)
+    t_emb = self.adaLN_modulation(t_emb)
+    tag_embeds = self.adaLN_modulation(tag_embeds)
+    t_emb = t_emb + tag_embeds
+    t_emb = manual_cast(t_emb, self.dtype)
+    cond = manual_cast(cond, self.dtype)
+    point_embeds = manual_cast(point_embeds, self.dtype)
+
+    h_feats_list = []
+    h_coords_list = []
+    begin = 0
+    for i, coords_len in enumerate(coords_len_list):
+        end = begin + 2 * coords_len
+        h_feats_list.append(h.feats[begin:end])
+        h_coords_list.append(h.coords[begin:end])
+        h_feats_list.append(point_embeds.feats[i*10:(i+1)*10])
+        h_coords_list.append(point_embeds.coords[i*10:(i+1)*10])
+        begin = end + 10
+    h = sp.SparseTensor(torch.cat(h_feats_list), torch.cat(h_coords_list))
+
+    for block in self.blocks:
+        h = block(h, t_emb, cond)
+
+    h_feats_list = []
+    h_coords_list = []
+    begin = 0
+    for i, coords_len in enumerate(coords_len_list):
+        end = begin + 2 * coords_len
+        h_feats_list.append(h.feats[begin:end])
+        h_coords_list.append(h.coords[begin:end])
+        begin = end
+    h = sp.SparseTensor(torch.cat(h_feats_list), torch.cat(h_coords_list))
+
+    h = manual_cast(h, x.dtype)
+    h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:]))
+    # h.feats: [N, 1536]
+    h = self.out_layer(h)
+    # h.feats: [N, 32]
+    return h
+
+class Gen3DSeg(nn.Module):
+    def __init__(self, flow_model):
+        super().__init__()
+        self.flow_model = flow_model
+        self.seg_embeddings = nn.Embedding(1, 1536)
+        self.tag_mlp = nn.Sequential(nn.Linear(256, 1536, bias=True), nn.SiLU(), nn.Linear(1536, 1536, bias=True))
+        
+    def tag_embedding(self, tag):
+        freqs = torch.exp(-np.log(10000) * torch.arange(start=0, end=128, dtype=torch.float32) / 128).to(device=tag.device)
+        args = tag[:, None].float() * freqs[None]
+        tag_freq = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        tag_embeds = self.tag_mlp(tag_freq)
+        return tag_embeds
+
+    def get_positional_encoding(self, input_points):
+        point_feats_embed = torch.zeros((10, 1536), dtype=torch.float32).to(input_points['point_slats'].feats.device)
+        labels = input_points['point_labels'].squeeze(-1)
+        point_feats_embed[labels == 1] = self.seg_embeddings.weight
+        return sp.SparseTensor(point_feats_embed, input_points['point_slats'].coords)
+
+    def forward(self, x_t, tex_slats, shape_slats, t, tags, cond, input_points, 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))
+
+        tag_embeds = self.tag_embedding(tags)
+        point_embeds = self.get_positional_encoding(input_points)
+        output_tex_slats = self.flow_model(x_t, t, tag_embeds, cond, shape_slats, point_embeds, coords_len_list)
+        
+        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):
+    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, 2048)
+    return img.resize((pow2, pow2), Image.BILINEAR)
+
+def preprocess_scene_textures(asset):
+    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))
+            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)
+        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)
+    return asset
+
+def process_glb_to_vxz(glb_path, vxz_path):
+    asset = trimesh.load(glb_path, force='scene')
+    asset = preprocess_scene_textures(asset)
+    aabb = asset.bounding_box.bounds
+    center = (aabb[0] + aabb[1]) / 2
+    scale = 0.99999 / (aabb[1] - aabb[0]).max()
+    asset.apply_translation(-center)
+    asset.apply_scale(scale)
+    mesh = asset.to_mesh()
+    vertices = torch.from_numpy(mesh.vertices).float()
+    faces = torch.from_numpy(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]
+
+    voxel_indices_mat, attributes = o_voxel.convert.textured_mesh_to_volumetric_attr(
+        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()}
+
+    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['dual_vertices'] = dual_vertices
+    attributes['intersected'] = intersected
+    o_voxel.io.write(vxz_path, voxel_indices, attributes)
+
+def vxz_to_latent_slat(shape_encoder, shape_decoder, tex_encoder, vxz_path):
+    coords, data = o_voxel.io.read(vxz_path)
+    coords = torch.cat([torch.zeros(coords.shape[0], 1, dtype=torch.int32), coords], dim=1).cuda()
+    vertices = (data['dual_vertices'].cuda() / 255)
+    intersected = torch.cat([data['intersected'] % 2, data['intersected'] // 2 % 2, data['intersected'] // 4 % 2], dim=-1).bool().cuda()
+    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.cuda(), shape_slat.coords.cuda())
+        shape_decoder.set_resolution(512)
+        meshes, subs = shape_decoder(shape_slat, return_subs=True)
+    
+    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().cuda() * 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)  # type: ignore
+    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, input_points, tag):
+    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]]
+    noise = sp.SparseTensor(torch.randn_like(input_tex_slat.feats), shape_slat.coords)
+    output_tex_slat = sampler.sample(gen3dseg, noise, input_tex_slat, shape_slat, input_points, coords_len_list, cond_dict, tag, pipeline_args['tex_slat_sampler']['params'])
+    output_tex_slat = output_tex_slat * tex_std + tex_mean
+    return output_tex_slat
+
+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 m, v in zip(meshes, tex_voxels):
+        m.fill_holes()
+        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]
+    mesh.simplify(10000000)
+    # mesh.simplify(16777216) # nvdiffrast limit
+    glb = o_voxel.postprocess.to_glb(
+        vertices            =   mesh.vertices,
+        faces               =   mesh.faces,
+        attr_volume         =   mesh.attrs,
+        coords              =   mesh.coords,
+        attr_layout         =   mesh.layout,
+        voxel_size          =   mesh.voxel_size,
+        aabb                =   [[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
+        decimation_target   =   100000, # 1000000
+        texture_size        =   4096,
+        remesh              =   True,
+        # remesh              =   False,
+        remesh_band         =   1,
+        remesh_project      =   0,
+        verbose             =   True
+    )
+    return glb
+
+def inference(ckpt_path, item, tag, input_vxz_points_list=None):
+    print("-"*100)
+    print("Loading model ............")
+    with open("microsoft/TRELLIS.2-4B/pipeline.json", "r") as f:
+        pipeline_config = json.load(f)
+    pipeline_args = pipeline_config['args']
+    tex_slat_flow_model = models.from_pretrained("microsoft/TRELLIS.2-4B/ckpts/slat_flow_imgshape2tex_dit_1_3B_512_bf16")
+    tex_slat_flow_model.forward = MethodType(flow_forward, tex_slat_flow_model)
+
+    gen3dseg = Gen3DSeg(tex_slat_flow_model)
+    state_dict = torch.load(ckpt_path)['state_dict']
+    state_dict = OrderedDict([(k.replace("gen3dseg.", ""), v) for k, v in state_dict.items()])
+    gen3dseg.load_state_dict(state_dict)
+    gen3dseg.eval()
+    gen3dseg.cuda()
+    sampler = Sampler()
+
+    shape_encoder = models.from_pretrained("microsoft/TRELLIS.2-4B/ckpts/shape_enc_next_dc_f16c32_fp16").cuda().eval()
+    tex_encoder = models.from_pretrained("microsoft/TRELLIS.2-4B/ckpts/tex_enc_next_dc_f16c32_fp16").cuda().eval()
+    shape_decoder = models.from_pretrained("microsoft/TRELLIS.2-4B/ckpts/shape_dec_next_dc_f16c32_fp16").cuda().eval()
+    tex_decoder = models.from_pretrained("microsoft/TRELLIS.2-4B/ckpts/tex_dec_next_dc_f16c32_fp16").cuda().eval()
+
+    rembg_model = BiRefNet(model_name="briaai/RMBG-2.0")
+    rembg_model.cuda()
+    image_cond_model = DinoV3FeatureExtractor(model_name="facebook/dinov3-vitl16-pretrain-lvd1689m")
+    image_cond_model.cuda()
+
+    process_glb_to_vxz(item['glb'], item['input_vxz'])
+    shape_slat, meshes, subs, tex_slat = vxz_to_latent_slat(shape_encoder, shape_decoder, tex_encoder, item['input_vxz'])
+
+    print("-"*100)
+    print("Getting cond ............")
+    if tag in [0, 1]:
+        render_from_transforms(item['glb'], item['transforms'], item['img'])
+    image = Image.open(item['img'])
+    image = preprocess_image(rembg_model, image)
+    cond = get_cond(image_cond_model, [image])
+
+    print("-"*100)
+    print("Sampling .................")
+    if tag == 0:
+        vxz_points_coords = torch.tensor(input_vxz_points_list, dtype=torch.int32).cuda()
+        vxz_points_coords = torch.cat([torch.zeros((vxz_points_coords.shape[0], 1), dtype=torch.int32).cuda(), vxz_points_coords], dim=1)
+        input_points_coords = tex_encoder(sp.SparseTensor(torch.zeros((vxz_points_coords.shape[0], 6), dtype=torch.float32).cuda(), vxz_points_coords)).coords
+        input_points_coords = torch.unique(input_points_coords, dim=0)
+        point_num = input_points_coords.shape[0]
+        if point_num >= 10:
+            input_points_coords = input_points_coords[:10]
+            point_labels = torch.tensor(([[1]]*10), dtype=torch.int32).cuda()
+        else:
+            input_points_coords = torch.cat([input_points_coords, torch.zeros((10 - point_num, 4), dtype=torch.int32).cuda()], dim=0)
+            point_labels = torch.tensor(([[1]]*point_num+[[0]]*(10-point_num)), dtype=torch.int32).cuda()
+    else:
+        input_points_coords = torch.zeros((10, 4), dtype=torch.int32).cuda()
+        point_labels = torch.tensor(([[0]]*10), dtype=torch.int32).cuda()
+    input_points = {'point_slats': sp.SparseTensor(input_points_coords, input_points_coords), 'point_labels': point_labels}
+
+    output_tex_slat = tex_slat_sample_single(gen3dseg, sampler, pipeline_args, shape_slat, tex_slat, cond, input_points, tag)
+    with torch.no_grad():
+        tex_voxels = tex_decoder(output_tex_slat, guide_subs=subs) * 0.5 + 0.5
+
+    print("-"*100)
+    print("Exporting glb ............")
+    glb = slat_to_glb(meshes, tex_voxels)
+    glb.export(item['export_glb'])
+
+if __name__ == "__main__":
+    ckpt_path = "path/to/unified.ckpt"
+    tag = 0
+    if tag == 0: # interactive seg
+        item = {
+            "glb": "./data_toolkit/assets/example.glb",
+            "input_vxz": "./data_toolkit/assets/input.vxz",
+            "transforms": "./data_toolkit/transforms.json",
+            "img": "./data_toolkit/assets/img.png",
+            "export_glb": "./data_toolkit/assets/output.glb"
+        }
+        input_vxz_points_list = [[388, 448, 392]] # example
+        inference(ckpt_path, item, tag, input_vxz_points_list)
+    elif tag == 1: # full seg
+        item = {
+            "glb": "./data_toolkit/assets/example.glb",
+            "input_vxz": "./data_toolkit/assets/input.vxz",
+            "transforms": "./data_toolkit/transforms.json",
+            "img": "./data_toolkit/assets/img.png",
+            "export_glb": "./data_toolkit/assets/output.glb"
+        }
+        inference(ckpt_path, item, tag)
+    elif tag == 2: # full seg with 2d map
+        item = {
+            "glb": "./data_toolkit/assets/example.glb",
+            "input_vxz": "./data_toolkit/assets/input.vxz",
+            "img": "./data_toolkit/assets/full_seg_w_2d_map/2d_map.png.png",
+            "export_glb": "./data_toolkit/assets/output.glb"
+        }
+        inference(ckpt_path, item, tag)
+    else:
+        raise ValueError(f"Invalid tag: {tag}")

requirements.txt

@@ -0,0 +1,23 @@
+--extra-index-url https://download.pytorch.org/whl/cu124
+
+torch==2.6.0
+torchvision==0.21.0
+triton==3.2.0
+pillow==12.0.0
+imageio==2.37.2
+imageio-ffmpeg==0.6.0
+tqdm==4.67.1
+easydict==1.13
+opencv-python-headless==4.12.0.88
+trimesh==4.10.1
+transformers==4.57.3
+zstandard==0.25.0
+kornia==0.8.2
+timm==1.0.22
+git+https://github.com/EasternJournalist/utils3d.git@9a4eb15e4021b67b12c460c7057d642626897ec8
+https://github.com/JeffreyXiang/Storages/releases/download/Space_Wheels_251210/flash_attn_3-3.0.0b1-cp39-abi3-linux_x86_64.whl
+https://github.com/JeffreyXiang/Storages/releases/download/Space_Wheels_251210/cumesh-0.0.1-cp310-cp310-linux_x86_64.whl
+https://github.com/JeffreyXiang/Storages/releases/download/Space_Wheels_251210/flex_gemm-0.0.1-cp310-cp310-linux_x86_64.whl
+https://github.com/JeffreyXiang/Storages/releases/download/Space_Wheels_251210/o_voxel-0.0.1-cp310-cp310-linux_x86_64.whl
+https://github.com/JeffreyXiang/Storages/releases/download/Space_Wheels_251210/nvdiffrast-0.4.0-cp310-cp310-linux_x86_64.whl
+https://github.com/JeffreyXiang/Storages/releases/download/Space_Wheels_251210/nvdiffrec_render-0.0.0-cp310-cp310-linux_x86_64.whl

split.py

@@ -0,0 +1,833 @@
+import os
+import struct
+from collections import defaultdict
+from typing import Optional, Tuple
+
+import numpy as np
+import trimesh
+from PIL import Image
+
+# =========================
+# 你只需要改这里
+# =========================
+# INPUT_GLB = "/mnt/pfs/users/huangzehuan/projects/SegviGen/examples/trellis2_output.glb"
+
+
+# UID = "demonic_warrior_red_bronze_armor"
+# UID = "playful_pose_white_top_portrait"
+# UID = "african_inspired_metallic_silver_ensemble_with_headwrap"
+# UID = "cyberpunk_bowser_motorcycle"
+# UID = "crimson_battle_mecha_with_spikes"
+UID = "black_lace_lingerie_ensemble"
+
+INPUT_GLB = (
+    f"/mnt/pfs/users/maxueqi/studio/datasets/dense_mesh/segvigen_bak/{UID}/output.glb"
+)
+
+# 只用 RGB（忽略透明度/alpha）
+COLOR_QUANT_STEP = 16  # RGB 量化步长：0/4/8/16（越大越“合并”）
+PALETTE_SAMPLE_PIXELS = 2_000_000
+PALETTE_MIN_PIXELS = 500  # 少于该像素数的颜色当噪声丢掉（边界抗锯齿中间色）
+PALETTE_MAX_COLORS = 256  # 最多保留多少个主颜色
+PALETTE_MERGE_DIST = 32  # ✅ 合并 palette 内近似颜色（解决“看着同色却拆两块”）
+
+SAMPLES_PER_FACE = 4  # 1 或 4（推荐 4，能明显减少边界采样误差）
+FLIP_V = True  # glTF 常见需要 flip V
+UV_WRAP_REPEAT = True  # True: repeat (mod 1)；False: clamp 到 [0,1]
+
+MIN_FACES_PER_PART = 50
+BAKE_TRANSFORMS = True
+DEBUG_PRINT = True
+# =========================
+
+
+CHUNK_TYPE_JSON = 0x4E4F534A  # b'JSON'
+CHUNK_TYPE_BIN = 0x004E4942  # b'BIN\0'
+
+
+def _default_out_path(in_path: str) -> str:
+    root, ext = os.path.splitext(in_path)
+    if ext.lower() not in [".glb", ".gltf"]:
+        ext = ".glb"
+    return f"{root}_seg.glb"
+
+
+def _quantize_rgb(rgb: np.ndarray, step: int) -> np.ndarray:
+    """
+    rgb: (...,3) uint8
+    """
+    if step is None or step <= 0:
+        return rgb
+    q = (rgb.astype(np.int32) + step // 2) // step * step
+    q = np.clip(q, 0, 255).astype(np.uint8)
+    return q
+
+
+def _load_glb_json_and_bin(glb_path: str) -> Tuple[dict, bytes]:
+    data = open(glb_path, "rb").read()
+    if len(data) < 12:
+        raise RuntimeError("Invalid GLB: too small")
+
+    magic, version, length = struct.unpack_from("<4sII", data, 0)
+    if magic != b"glTF":
+        raise RuntimeError("Not a GLB file (missing glTF header)")
+
+    offset = 12
+    gltf_json = None
+    bin_chunk = None
+
+    while offset + 8 <= len(data):
+        chunk_len, chunk_type = struct.unpack_from("<II", data, offset)
+        offset += 8
+        chunk_data = data[offset : offset + chunk_len]
+        offset += chunk_len
+
+        if chunk_type == CHUNK_TYPE_JSON:
+            gltf_json = chunk_data.decode("utf-8", errors="replace")
+        elif chunk_type == CHUNK_TYPE_BIN:
+            bin_chunk = chunk_data
+
+    if gltf_json is None:
+        raise RuntimeError("GLB missing JSON chunk")
+    if bin_chunk is None:
+        raise RuntimeError("GLB missing BIN chunk")
+
+    import json
+
+    return json.loads(gltf_json), bin_chunk
+
+
+def _extract_basecolor_texture_image(glb_path: str) -> np.ndarray:
+    """
+    从 GLB 内嵌资源里拿 baseColorTexture 的 PNG/JPG，返回 (H,W,4) uint8 RGBA
+    """
+    gltf, bin_chunk = _load_glb_json_and_bin(glb_path)
+
+    materials = gltf.get("materials", [])
+    textures = gltf.get("textures", [])
+    images = gltf.get("images", [])
+    buffer_views = gltf.get("bufferViews", [])
+
+    if not materials:
+        raise RuntimeError("No materials in GLB")
+
+    # 这里按 material[0] 取 baseColorTexture（你的 glb 只有一个材质/primitive）
+    pbr = materials[0].get("pbrMetallicRoughness", {})
+    base_tex_index = pbr.get("baseColorTexture", {}).get("index", None)
+    if base_tex_index is None:
+        raise RuntimeError("Material has no baseColorTexture")
+
+    if base_tex_index >= len(textures):
+        raise RuntimeError("baseColorTexture index out of range")
+
+    tex = textures[base_tex_index]
+    img_index = tex.get("source", None)
+    if img_index is None or img_index >= len(images):
+        raise RuntimeError("Texture has no valid image source")
+
+    img_info = images[img_index]
+    bv_index = img_info.get("bufferView", None)
+    mime = img_info.get("mimeType", None)
+    if bv_index is None:
+        uri = img_info.get("uri", None)
+        raise RuntimeError(f"Image is not embedded (bufferView missing). uri={uri}")
+
+    if bv_index >= len(buffer_views):
+        raise RuntimeError("image.bufferView out of range")
+
+    bv = buffer_views[bv_index]
+    bo = int(bv.get("byteOffset", 0))
+    bl = int(bv.get("byteLength", 0))
+    img_bytes = bin_chunk[bo : bo + bl]
+
+    if DEBUG_PRINT:
+        print(
+            f"[Texture] baseColorTextureIndex={base_tex_index}, imageIndex={img_index}, "
+            f"bufferView={bv_index}, mime={mime}, bytes={len(img_bytes)}"
+        )
+
+    pil = Image.open(trimesh.util.wrap_as_stream(img_bytes)).convert("RGBA")
+    return np.array(pil, dtype=np.uint8)
+
+
+def _merge_palette_rgb(
+    palette_rgb: np.ndarray, counts: np.ndarray, merge_dist: float
+) -> np.ndarray:
+    """
+    对 palette 内 RGB 做“近似合并”，用 counts 作为权重更新中心。
+    palette_rgb: (K,3) uint8
+    counts: (K,) int
+    """
+    if palette_rgb is None or len(palette_rgb) == 0:
+        return palette_rgb
+    if merge_dist is None or merge_dist <= 0:
+        return palette_rgb
+
+    rgb = palette_rgb.astype(np.float32)
+    counts = counts.astype(np.int64)
+
+    order = np.argsort(-counts)
+
+    centers = []
+    center_w = []
+    thr2 = float(merge_dist) * float(merge_dist)
+
+    for idx in order:
+        x = rgb[idx]
+        w = int(counts[idx])
+
+        if not centers:
+            centers.append(x.copy())
+            center_w.append(w)
+            continue
+
+        C = np.stack(centers, axis=0)  # (M,3)
+        d2 = np.sum((C - x[None, :]) ** 2, axis=1)
+        k = int(np.argmin(d2))
+
+        if float(d2[k]) <= thr2:
+            cw = center_w[k]
+            centers[k] = (centers[k] * cw + x * w) / (cw + w)
+            center_w[k] = cw + w
+        else:
+            centers.append(x.copy())
+            center_w.append(w)
+
+    merged = np.clip(np.rint(np.stack(centers, axis=0)), 0, 255).astype(np.uint8)
+
+    if DEBUG_PRINT:
+        print(
+            f"[PaletteMerge] before={len(palette_rgb)} after={len(merged)} merge_dist={merge_dist}"
+        )
+
+    return merged
+
+
+def _build_palette_rgb(tex_rgba: np.ndarray) -> np.ndarray:
+    """
+    从贴图中提取 RGB 主颜色调色板（忽略 alpha）。
+    返回: (K,3) uint8
+    """
+    rgb = tex_rgba[:, :, :3].reshape(-1, 3)
+    n = rgb.shape[0]
+
+    if n > PALETTE_SAMPLE_PIXELS:
+        rng = np.random.default_rng(0)
+        idx = rng.choice(n, size=PALETTE_SAMPLE_PIXELS, replace=False)
+        rgb = rgb[idx]
+
+    rgb = _quantize_rgb(rgb, COLOR_QUANT_STEP)
+
+    uniq, counts = np.unique(rgb, axis=0, return_counts=True)
+    order = np.argsort(-counts)
+    uniq = uniq[order]
+    counts = counts[order]
+
+    keep = counts >= PALETTE_MIN_PIXELS
+    uniq = uniq[keep]
+    counts = counts[keep]
+
+    if len(uniq) > PALETTE_MAX_COLORS:
+        uniq = uniq[:PALETTE_MAX_COLORS]
+        counts = counts[:PALETTE_MAX_COLORS]
+
+    if DEBUG_PRINT:
+        print(
+            f"[Palette] quant_step={COLOR_QUANT_STEP} palette_size(before_merge)={len(uniq)} "
+            f"min_pixels={PALETTE_MIN_PIXELS}"
+        )
+        for i in range(min(15, len(uniq))):
+            r, g, b = [int(x) for x in uniq[i]]
+            print(f"  {i:02d} rgb=({r},{g},{b}) count={int(counts[i])}")
+
+    uniq = _merge_palette_rgb(uniq.astype(np.uint8), counts, PALETTE_MERGE_DIST)
+
+    if DEBUG_PRINT:
+        print(f"[Palette] palette_size(after_merge)={len(uniq)}")
+        for i in range(min(15, len(uniq))):
+            r, g, b = [int(x) for x in uniq[i]]
+            print(f"  {i:02d} rgb=({r},{g},{b})")
+
+    return uniq.astype(np.uint8)
+
+
+def _unwrap_uv3_for_seam(uv3: np.ndarray) -> np.ndarray:
+    """
+    uv3: (F,3,2). 若跨 seam（跨度>0.5），把小于0.5的一侧 +1，避免均值跑到另一边。
+    """
+    out = uv3.copy()
+    for d in range(2):
+        v = out[:, :, d]
+        vmin = v.min(axis=1)
+        vmax = v.max(axis=1)
+        seam = (vmax - vmin) > 0.5
+        if np.any(seam):
+            vv = v[seam]
+            vv = np.where(vv < 0.5, vv + 1.0, vv)
+            out[seam, :, d] = vv
+    return out
+
+
+def _barycentric_samples(uv3: np.ndarray, samples_per_face: int) -> np.ndarray:
+    """
+    uv3: (F,3,2)
+    return: (F,S,2)
+    """
+    uv3 = _unwrap_uv3_for_seam(uv3)
+
+    if samples_per_face == 1:
+        w = np.array([1 / 3, 1 / 3, 1 / 3], dtype=np.float32)
+        uvs = uv3[:, 0, :] * w[0] + uv3[:, 1, :] * w[1] + uv3[:, 2, :] * w[2]
+        return uvs[:, None, :]
+
+    # 4 个点：中心 + 三个靠近顶点的内点（尽量远离边界抗锯齿带）
+    ws = np.array(
+        [
+            [1 / 3, 1 / 3, 1 / 3],
+            [0.80, 0.10, 0.10],
+            [0.10, 0.80, 0.10],
+            [0.10, 0.10, 0.80],
+        ],
+        dtype=np.float32,
+    )
+    uvs = (
+        uv3[:, None, 0, :] * ws[None, :, 0, None]
+        + uv3[:, None, 1, :] * ws[None, :, 1, None]
+        + uv3[:, None, 2, :] * ws[None, :, 2, None]
+    )
+    return uvs
+
+
+def _wrap_or_clamp_uv(uv: np.ndarray) -> np.ndarray:
+    if UV_WRAP_REPEAT:
+        return np.mod(uv, 1.0)
+    return np.clip(uv, 0.0, 1.0)
+
+
+def _sample_texture_nearest_rgb(tex_rgba: np.ndarray, uv: np.ndarray) -> np.ndarray:
+    """
+    tex_rgba: (H,W,4) uint8
+    uv: (N,2) float
+    return: (N,3) uint8
+    """
+    h, w = tex_rgba.shape[0], tex_rgba.shape[1]
+    uv = _wrap_or_clamp_uv(uv)
+
+    u = uv[:, 0]
+    v = uv[:, 1]
+    if FLIP_V:
+        v = 1.0 - v
+
+    x = np.rint(u * (w - 1)).astype(np.int32)
+    y = np.rint(v * (h - 1)).astype(np.int32)
+    x = np.clip(x, 0, w - 1)
+    y = np.clip(y, 0, h - 1)
+
+    return tex_rgba[y, x, :3].astype(np.uint8)
+
+
+def _map_to_palette_rgb(
+    colors_rgb: np.ndarray, palette_rgb: np.ndarray, chunk: int = 20000
+) -> Tuple[np.ndarray, np.ndarray]:
+    """
+    把采样到的 RGB 映射到最近的 palette RGB.
+    如果 palette 为空，则用 colors_rgb 的 unique 作为“临时 palette”.
+    返回:
+      labels: (N,) int
+      used_palette_rgb: (K,3) uint8
+    """
+    if palette_rgb is None or len(palette_rgb) == 0:
+        uniq, inv = np.unique(colors_rgb, axis=0, return_inverse=True)
+        return inv.astype(np.int32), uniq.astype(np.uint8)
+
+    c = colors_rgb.astype(np.float32)
+    p = palette_rgb.astype(np.float32)
+
+    out = np.empty((c.shape[0],), dtype=np.int32)
+    for i in range(0, c.shape[0], chunk):
+        cc = c[i : i + chunk]
+        d2 = ((cc[:, None, :] - p[None, :, :]) ** 2).sum(axis=2)
+        out[i : i + chunk] = np.argmin(d2, axis=1).astype(np.int32)
+
+    return out, palette_rgb
+
+
+def _face_labels_from_texture_rgb(
+    mesh: trimesh.Trimesh,
+    tex_rgba: np.ndarray,
+    palette_rgb: np.ndarray,
+) -> Optional[Tuple[np.ndarray, np.ndarray]]:
+    """
+    用 TEXCOORD_0 + baseColorTexture，为每个 face 采样 RGB，并映射到 palette label。
+    返回:
+      face_label: (F,) int
+      label_rgb: (K,3) uint8
+    """
+    uv = getattr(mesh.visual, "uv", None)
+    if uv is None:
+        return None
+
+    uv = np.asarray(uv, dtype=np.float32)
+    if uv.ndim != 2 or uv.shape[1] != 2 or uv.shape[0] != len(mesh.vertices):
+        return None
+
+    faces = mesh.faces
+    uv3 = uv[faces]  # (F,3,2)
+
+    uvs = _barycentric_samples(uv3, SAMPLES_PER_FACE)  # (F,S,2)
+    F, S = uvs.shape[0], uvs.shape[1]
+    flat_uv = uvs.reshape(-1, 2)
+
+    sampled_rgb = _sample_texture_nearest_rgb(tex_rgba, flat_uv)  # (F*S,3)
+    sampled_rgb = _quantize_rgb(sampled_rgb, COLOR_QUANT_STEP)
+
+    sample_label, used_palette = _map_to_palette_rgb(sampled_rgb, palette_rgb)
+    sample_label = sample_label.reshape(F, S)
+
+    if S == 1:
+        return sample_label[:, 0].astype(np.int32), used_palette
+
+    # 4 票投票（向量化）
+    l0, l1, l2, l3 = (
+        sample_label[:, 0],
+        sample_label[:, 1],
+        sample_label[:, 2],
+        sample_label[:, 3],
+    )
+    c0 = 1 + (l0 == l1) + (l0 == l2) + (l0 == l3)
+    c1 = 1 + (l1 == l0) + (l1 == l2) + (l1 == l3)
+    c2 = 1 + (l2 == l0) + (l2 == l1) + (l2 == l3)
+    c3 = 1 + (l3 == l0) + (l3 == l1) + (l3 == l2)
+
+    counts = np.stack([c0, c1, c2, c3], axis=1)  # (F,4)
+    vals = np.stack([l0, l1, l2, l3], axis=1)  # (F,4)
+    best = vals[np.arange(F), np.argmax(counts, axis=1)]
+    return best.astype(np.int32), used_palette
+
+
+# =========================
+# 拓扑纠错
+# =========================
+
+import numpy as np
+import trimesh
+from scipy.sparse import coo_matrix
+from scipy.sparse.csgraph import connected_components
+
+
+def _get_physical_face_adjacency(mesh: trimesh.Trimesh) -> np.ndarray:
+    """
+    忽略 UV 接缝，计算纯物理空间上的面片相邻关系。
+    """
+    # 1. 四舍五入顶点坐标（处理浮点数微小误差），找出空间中真正唯一的物理顶点
+    v_rounded = np.round(mesh.vertices, decimals=3)
+    v_unique, inv_indices = np.unique(v_rounded, axis=0, return_inverse=True)
+
+    # 2. 将原本的面片索引，映射到这些“唯一物理顶点”上
+    # 这样，跨越 UV 接缝的面片，此时它们引用的顶点索引就变成一样的了
+    physical_faces = inv_indices[mesh.faces]
+
+    # 3. 创建一个临时的“影子网格”（process=False 极其重要，防止 trimesh 内部重排面片）
+    tmp_mesh = trimesh.Trimesh(vertices=v_unique, faces=physical_faces, process=False)
+
+    # 返回影子网格的物理相邻边
+    return tmp_mesh.face_adjacency
+
+
+def smooth_face_labels_by_topology(
+    mesh: trimesh.Trimesh, face_label: np.ndarray, min_faces: int = 50
+) -> np.ndarray:
+    """
+    通过真实的 3D 物理拓扑关系过滤飞点，跨越 UV 接缝合并色块。
+
+    Phase 1: 在同色连通图上，把挨着大块的小块吞并到大块中。
+    Phase 2: 对残留小块（邻居全是小块），回退到全物理邻接，
+             按物理邻居中的多数 label 吞并。
+    Phase 3: 对完全孤立的面片（无物理邻接边），按面片质心距离
+             找最近的非孤立面片，继承其 label。
+    """
+    labels = face_label.copy()
+    edges = _get_physical_face_adjacency(mesh)
+    F = len(mesh.faces)
+
+    # ---- Phase 1: 同色连通域平滑 ----
+    for iteration in range(3):
+        same_label = labels[edges[:, 0]] == labels[edges[:, 1]]
+        sub_edges = edges[same_label]
+
+        if len(sub_edges) > 0:
+            data = np.ones(len(sub_edges), dtype=bool)
+            graph = coo_matrix((data, (sub_edges[:, 0], sub_edges[:, 1])), shape=(F, F))
+            graph = graph.maximum(graph.T)
+            n_components, comp_labels = connected_components(graph, directed=False)
+        else:
+            n_components = F
+            comp_labels = np.arange(F)
+
+        comp_sizes = np.bincount(comp_labels, minlength=n_components)
+        small_comps = np.where(comp_sizes < min_faces)[0]
+        if len(small_comps) == 0:
+            break
+
+        is_small = np.isin(comp_labels, small_comps)
+
+        mask0 = is_small[edges[:, 0]]
+        mask1 = is_small[edges[:, 1]]
+
+        boundary_edges_0 = edges[mask0 & ~mask1]
+        boundary_edges_1 = edges[mask1 & ~mask0]
+
+        b_inner = np.concatenate([boundary_edges_0[:, 0], boundary_edges_1[:, 1]])
+        b_outer = np.concatenate([boundary_edges_0[:, 1], boundary_edges_1[:, 0]])
+
+        if len(b_inner) == 0:
+            break
+
+        outer_labels = labels[b_outer]
+        inner_comps = comp_labels[b_inner]
+
+        for cid in np.unique(inner_comps):
+            cid_mask = inner_comps == cid
+            surrounding_labels = outer_labels[cid_mask]
+            if len(surrounding_labels) > 0:
+                best_label = np.bincount(surrounding_labels).argmax()
+                labels[comp_labels == cid] = best_label
+
+    # ---- Phase 2: 用全物理邻接处理残留小块 ----
+    # 重新计算同色连通域，找出还残留的小块
+    same_label = labels[edges[:, 0]] == labels[edges[:, 1]]
+    sub_edges = edges[same_label]
+    if len(sub_edges) > 0:
+        data = np.ones(len(sub_edges), dtype=bool)
+        graph = coo_matrix((data, (sub_edges[:, 0], sub_edges[:, 1])), shape=(F, F))
+        graph = graph.maximum(graph.T)
+        n_components, comp_labels = connected_components(graph, directed=False)
+    else:
+        n_components = F
+        comp_labels = np.arange(F)
+
+    comp_sizes = np.bincount(comp_labels, minlength=n_components)
+    small_comps_set = set(np.where(comp_sizes < min_faces)[0])
+
+    if small_comps_set:
+        is_small = np.array([comp_labels[i] in small_comps_set for i in range(F)])
+
+        # 构建全物理邻接查找表: face -> set of neighbor faces
+        adj = defaultdict(set)
+        for e0, e1 in edges:
+            adj[int(e0)].add(int(e1))
+            adj[int(e1)].add(int(e0))
+
+        # 迭代：每轮让小块面片从物理邻居（忽略颜色）中投票取多数 label
+        for _ in range(3):
+            changed = False
+            small_comps_now = set(
+                int(c)
+                for c in range(n_components)
+                if comp_sizes[c] < min_faces and c in small_comps_set
+            )
+            if not small_comps_now:
+                break
+
+            for cid in small_comps_now:
+                cid_faces = np.where(comp_labels == cid)[0]
+                # 收集所有物理邻居中不属于本连通域的面片的 label
+                neighbor_labels = []
+                for fi in cid_faces:
+                    for nf in adj[int(fi)]:
+                        if comp_labels[nf] != cid:
+                            neighbor_labels.append(labels[nf])
+
+                if len(neighbor_labels) > 0:
+                    best_label = int(np.bincount(neighbor_labels).argmax())
+                    labels[cid_faces] = best_label
+                    changed = True
+
+            if not changed:
+                break
+
+            # 重新计算连通域
+            same_label = labels[edges[:, 0]] == labels[edges[:, 1]]
+            sub_edges = edges[same_label]
+            if len(sub_edges) > 0:
+                data = np.ones(len(sub_edges), dtype=bool)
+                graph = coo_matrix(
+                    (data, (sub_edges[:, 0], sub_edges[:, 1])), shape=(F, F)
+                )
+                graph = graph.maximum(graph.T)
+                n_components, comp_labels = connected_components(graph, directed=False)
+            else:
+                n_components = F
+                comp_labels = np.arange(F)
+            comp_sizes = np.bincount(comp_labels, minlength=n_components)
+            small_comps_set = set(np.where(comp_sizes < min_faces)[0])
+
+    # ---- Phase 3: 完全孤立面片（无物理邻接边），按质心距离继承 label ----
+    same_label = labels[edges[:, 0]] == labels[edges[:, 1]]
+    sub_edges = edges[same_label]
+    if len(sub_edges) > 0:
+        data = np.ones(len(sub_edges), dtype=bool)
+        graph = coo_matrix((data, (sub_edges[:, 0], sub_edges[:, 1])), shape=(F, F))
+        graph = graph.maximum(graph.T)
+        _, comp_labels = connected_components(graph, directed=False)
+    else:
+        comp_labels = np.arange(F)
+    comp_sizes = np.bincount(comp_labels)
+    orphan_comps = set(np.where(comp_sizes < min_faces)[0])
+
+    if orphan_comps:
+        orphan_mask = np.array([comp_labels[i] in orphan_comps for i in range(F)])
+        non_orphan_mask = ~orphan_mask
+        if non_orphan_mask.any() and orphan_mask.any():
+            centroids = mesh.triangles_center
+            orphan_indices = np.where(orphan_mask)[0]
+            non_orphan_indices = np.where(non_orphan_mask)[0]
+            non_orphan_centroids = centroids[non_orphan_indices]
+
+            for oi in orphan_indices:
+                dists = np.linalg.norm(non_orphan_centroids - centroids[oi], axis=1)
+                nearest = non_orphan_indices[np.argmin(dists)]
+                labels[oi] = labels[nearest]
+
+        if DEBUG_PRINT:
+            n_orphan = int(orphan_mask.sum())
+            print(f"  [Phase3] Assigned {n_orphan} orphan faces by centroid proximity")
+
+    return labels
+
+
+# =========================
+# 分割主函数
+# =========================
+
+
+# def split_glb_by_texture_palette_rgb(
+#     in_glb_path: str,
+#     out_glb_path: Optional[str] = None,
+#     min_faces_per_part: int = 1,
+#     bake_transforms: bool = True,
+# ) -> str:
+#     """
+#     输入：glb（无 COLOR_0，但有 baseColorTexture + TEXCOORD_0）
+#     输出：先从贴图提取 RGB 主色 palette（忽略 alpha），再按 palette label 分割
+#     """
+#     if out_glb_path is None:
+#         out_glb_path = _default_out_path(in_glb_path)
+
+#     tex_rgba = _extract_basecolor_texture_image(in_glb_path)
+#     palette_rgb = _build_palette_rgb(tex_rgba)
+
+#     scene = trimesh.load(in_glb_path, force="scene", process=False)
+#     out_scene = trimesh.Scene()
+
+#     part_count = 0
+#     base = os.path.splitext(os.path.basename(in_glb_path))[0]
+
+#     for node_name in scene.graph.nodes_geometry:
+#         geom_name = scene.graph[node_name][1]
+#         if geom_name is None:
+#             continue
+
+#         geom = scene.geometry.get(geom_name, None)
+#         if geom is None or not isinstance(geom, trimesh.Trimesh):
+#             continue
+
+#         mesh = geom.copy()
+
+#         if bake_transforms:
+#             T, _ = scene.graph.get(node_name)
+#             if T is not None:
+#                 mesh.apply_transform(T)
+
+#         res = _face_labels_from_texture_rgb(mesh, tex_rgba, palette_rgb)
+#         if res is None:
+#             if DEBUG_PRINT:
+#                 print(f"[{node_name}] no uv / cannot sample -> keep orig")
+#             out_scene.add_geometry(mesh, geom_name=f"{base}__{node_name}__orig")
+#             continue
+
+#         face_label, label_rgb = res
+
+#         # =========================
+#         # 🔥 新增调用：进行拓扑纠错，合并飞点
+#         # =========================
+#         face_label = smooth_face_labels_by_topology(mesh, face_label, min_faces=100)
+
+#         if DEBUG_PRINT:
+#             uniq_labels, cnts = np.unique(face_label, return_counts=True)
+#             order = np.argsort(-cnts)
+#             print(
+#                 f"[{node_name}] faces={len(mesh.faces)} labels_used={len(uniq_labels)} palette_size={len(label_rgb)}"
+#             )
+#             for i in order[:10]:
+#                 lab = int(uniq_labels[i])
+#                 r, g, b = (
+#                     [int(x) for x in label_rgb[lab]]
+#                     if 0 <= lab < len(label_rgb)
+#                     else (0, 0, 0)
+#                 )
+#                 print(f"  label={lab} rgb=({r},{g},{b}) faces={int(cnts[i])}")
+
+#         groups = defaultdict(list)
+#         for fi, lab in enumerate(face_label):
+#             groups[int(lab)].append(fi)
+
+#         for lab, face_ids in groups.items():
+#             if len(face_ids) < min_faces_per_part:
+#                 continue
+
+#             sub = mesh.submesh(
+#                 [np.array(face_ids, dtype=np.int64)], append=True, repair=False
+#             )
+#             if sub is None:
+#                 continue
+#             if isinstance(sub, (list, tuple)):
+#                 if not sub:
+#                     continue
+#                 sub = sub[0]
+
+#             if 0 <= lab < len(label_rgb):
+#                 r, g, b = [int(x) for x in label_rgb[lab]]
+#                 part_name = f"{base}__{node_name}__label_{lab}__rgb_{r}_{g}_{b}"
+#             else:
+#                 part_name = f"{base}__{node_name}__label_{lab}"
+
+#             out_scene.add_geometry(sub, geom_name=part_name)
+#             part_count += 1
+
+#     if part_count == 0:
+#         if DEBUG_PRINT:
+#             print("[INFO] part_count==0, fallback to original scene export.")
+#         out_scene = scene
+
+#     out_scene.export(out_glb_path)
+#     return out_glb_path
+
+
+def split_glb_by_texture_palette_rgb(
+    in_glb_path: str,
+    out_glb_path: Optional[str] = None,
+    min_faces_per_part: int = 1,
+    bake_transforms: bool = True,
+    color_quant_step: int = 16,
+    palette_sample_pixels: int = 2_000_000,
+    palette_min_pixels: int = 500,
+    palette_max_colors: int = 256,
+    palette_merge_dist: int = 32,
+    samples_per_face: int = 4,
+    flip_v: bool = True,
+    uv_wrap_repeat: bool = True,
+    transition_conf_thresh: float = 1.0,
+    transition_prop_iters: int = 6,
+    transition_neighbor_min: int = 1,
+    small_component_action: str = "reassign",
+    small_component_min_faces: int = 50,
+    postprocess_iters: int = 3,
+    debug_print: bool = True,
+) -> str:
+    """
+    Input: GLB (no COLOR_0, but with baseColorTexture + TEXCOORD_0)
+    Output: Split based on palette labels derived from baseColorTexture
+    """
+    if out_glb_path is None:
+        out_glb_path = _default_out_path(in_glb_path)
+
+    tex_rgba = _extract_basecolor_texture_image(in_glb_path)
+    palette_rgb = _build_palette_rgb(tex_rgba)
+
+    scene = trimesh.load(in_glb_path, force="scene", process=False)
+    out_scene = trimesh.Scene()
+
+    part_count = 0
+    base = os.path.splitext(os.path.basename(in_glb_path))[0]
+
+    for node_name in scene.graph.nodes_geometry:
+        geom_name = scene.graph[node_name][1]
+        if geom_name is None:
+            continue
+
+        geom = scene.geometry.get(geom_name, None)
+        if geom is None or not isinstance(geom, trimesh.Trimesh):
+            continue
+
+        mesh = geom.copy()
+
+        if bake_transforms:
+            T, _ = scene.graph.get(node_name)
+            if T is not None:
+                mesh.apply_transform(T)
+
+        res = _face_labels_from_texture_rgb(mesh, tex_rgba, palette_rgb)
+        if res is None:
+            if debug_print:
+                print(f"[{node_name}] no uv / cannot sample -> keep orig")
+            out_scene.add_geometry(mesh, geom_name=f"{base}__{node_name}__orig")
+            continue
+
+        face_label, label_rgb = res
+
+        # =========================
+        # 🔥 New: Apply topology correction to merge small disconnected components
+        # =========================
+        face_label = smooth_face_labels_by_topology(mesh, face_label, min_faces=100)
+
+        if debug_print:
+            uniq_labels, cnts = np.unique(face_label, return_counts=True)
+            order = np.argsort(-cnts)
+            print(
+                f"[{node_name}] faces={len(mesh.faces)} labels_used={len(uniq_labels)} palette_size={len(label_rgb)}"
+            )
+            for i in order[:10]:
+                lab = int(uniq_labels[i])
+                r, g, b = (
+                    [int(x) for x in label_rgb[lab]]
+                    if 0 <= lab < len(label_rgb)
+                    else (0, 0, 0)
+                )
+                print(f"  label={lab} rgb=({r},{g},{b}) faces={int(cnts[i])}")
+
+        groups = defaultdict(list)
+        for fi, lab in enumerate(face_label):
+            groups[int(lab)].append(fi)
+
+        for lab, face_ids in groups.items():
+            if len(face_ids) < min_faces_per_part:
+                continue
+
+            sub = mesh.submesh([np.array(face_ids, dtype=np.int64)], append=True, repair=False)
+            if sub is None:
+                continue
+            if isinstance(sub, (list, tuple)):
+                if not sub:
+                    continue
+                sub = sub[0]
+
+            if 0 <= lab < len(label_rgb):
+                r, g, b = [int(x) for x in label_rgb[lab]]
+                part_name = f"{base}__{node_name}__label_{lab}__rgb_{r}_{g}_{b}"
+            else:
+                part_name = f"{base}__{node_name}__label_{lab}"
+
+            out_scene.add_geometry(sub, geom_name=part_name)
+            part_count += 1
+
+    if part_count == 0:
+        if debug_print:
+            print("[INFO] part_count==0, fallback to original scene export.")
+        out_scene = scene
+
+    out_scene.export(out_glb_path)
+    return out_glb_path
+
+def main():
+    out_path = split_glb_by_texture_palette_rgb(
+        INPUT_GLB,
+        out_glb_path=None,
+        min_faces_per_part=MIN_FACES_PER_PART,
+        bake_transforms=BAKE_TRANSFORMS,
+    )
+    print("Done. Exported:", out_path)
+
+
+if __name__ == "__main__":
+    main()

split_ori.py

@@ -0,0 +1,686 @@
+import os
+import struct
+from collections import defaultdict, Counter
+from typing import Optional, Tuple
+
+import numpy as np
+import trimesh
+from PIL import Image
+
+
+# =========================
+# 你只需要改这里
+# =========================
+INPUT_GLB = "/media/nfs/tmp_data/fenghr/SegviGen/data_toolkit/assets/output.glb"  # 输入 GLB 路径
+
+# -------------------------
+# 颜色/调色板（RGB-only）
+# -------------------------
+COLOR_QUANT_STEP = 28
+# RGB 量化步长：把颜色通道四舍五入到该步长的倍数（减少颜色抖动/过渡色）
+# - 变大：颜色更粗、更容易合并为少数类（杂点/过渡色更少），但可能把本该不同的类别合并
+# - 变小：颜色更细、更容易保留差异，但过渡/噪声颜色会增多，后处理压力更大
+
+PALETTE_SAMPLE_PIXELS = 2_000_000
+# 构建调色板时，从贴图中最多抽样多少像素用于统计颜色频次
+# - 变大：调色板统计更准（更接近真实分布），但更慢、内存更高
+# - 变小：更快，但可能漏掉小类别或统计不稳定
+
+PALETTE_MIN_PIXELS = 300
+# 调色板过滤阈值：在抽样像素中，出现次数 < 该值的颜色视为“边界过渡/噪声”，不进调色板
+# - 变大：更激进地丢掉稀有颜色（减少边界过渡色导致的新label），但可能误删小零件类别
+# - 变小：保留更多稀有颜色（小类别更容易保住），但也会引入更多噪声颜色
+
+PALETTE_MAX_COLORS = 256
+# 调色板最多保留的主颜色数量（按频次从高到低截断）
+# - 变大：允许更多类别（细分更多），但更可能包含噪声色、导致碎块
+# - 变小：类别数更少、更干净，但可能把多类合并
+
+PALETTE_MERGE_DIST = 65
+# 调色板内近似颜色合并阈值（RGB欧氏距离）：把“看起来同色”的多种RGB合并成一个代表色
+# - 变大：更容易把近似色合并（解决“肉眼同色却拆两块”），但过大可能把相邻类别合并
+# - 变小：更保留差异，但同色可能仍分裂成多个label
+
+# -------------------------
+# UV 采样
+# -------------------------
+SAMPLES_PER_FACE = 4
+# 每个三角面采样点数：1=只取中心；4=中心+靠近3个顶点（投票），更抗边界抗锯齿
+# - 用 4：显著减少“采到边界中间色”导致的错分/过渡带
+# - 用 1：最快，但更容易生成过渡带/杂点
+
+FLIP_V = True
+# UV 的 V 轴是否翻转（glTF 常见需要翻转才能和贴图坐标对齐）
+# - 若你发现颜色整体错位/全错：优先尝试把它改成 False
+
+UV_WRAP_REPEAT = True
+# UV 超出 [0,1] 时的处理：True=repeat（mod 1），False=clamp 到边界
+# - repeat：适合贴图采用重复寻址的情况
+# - clamp：适合贴图不重复、超界应贴边的情况（repeat 可能采到错误区域）
+
+# -------------------------
+# ✅ “过渡面/边界带”归并（解决你说的：多出来一整块过渡区域）
+# -------------------------
+TRANSITION_CONF_THRESH = 1.0
+# 过渡面判定阈值：置信度 = 4次采样中最多票数 / 4
+# - 1.0：只要不是 4/4 完全一致，就当过渡面（最强去过渡带，最不容易多出一整块）
+# - 0.75：只有出现 2-2 或更不稳定才算过渡面（更保守，边界更“原汁原味”，但可能残留过渡带）
+# - 变大：更多面被当过渡面，会更强力贴合到两侧，但可能“抹边”更明显
+# - 变小：更少面被当过渡面，保留边界细节，但更可能留下过渡区域
+
+TRANSITION_PROP_ITERS = 6
+# 标签传播迭代次数：把过渡面按邻居多数投票逐轮吸收到稳定区域
+# - 变大：传播更充分，过渡带更容易被“吃掉”，但边界可能被推得更远/更平滑
+# - 变小：传播更少，边界更保留，但可能仍残留部分过渡带
+
+TRANSITION_NEIGHBOR_MIN = 1
+# 过渡面更新时，邻居投票的最小票数要求（防止太少证据就改）
+# - 变大：更新更谨慎，不容易被少数邻居误导，但可能残留过渡点
+# - 变小：更容易被吸收，去噪更强，但可能稍微更“糊边”
+
+# -------------------------
+# ✅ 小连通块收尾（主要消“杂点小岛”，不是边界带）
+# -------------------------
+SMALL_COMPONENT_ACTION = "reassign"
+# 小连通块处理方式：
+# - "reassign"：把小块按空间邻接归到周围大块（通常更符合“去杂点”）
+# - "drop"：直接丢掉这些面（输出会缺面，不建议除非你能接受空洞）
+
+SMALL_COMPONENT_MIN_FACES = 50
+# 小连通块阈值：某个 label 的一个连通块 face 数 < 该值，就当杂点处理
+# - 变大：更强去杂点，但可能误伤真实小零件
+# - 变小：更保留小零件，但杂点可能更多
+
+POSTPROCESS_ITERS = 3
+# 小连通块处理迭代次数（reassign时）：
+# - 变大：更彻底清理杂点，但更可能“抹掉”小细节
+# - 变小：更保守
+
+# -------------------------
+# 导出过滤/其他
+# -------------------------
+MIN_FACES_PER_PART = 1
+# 导出时��最小面数过滤：某个 part 的 face 数 < 该值就不导出
+# - 变大：输出更干净（少碎片），但会丢失小零件
+# - 变小：保留全部（包括小碎片）
+
+BAKE_TRANSFORMS = True
+# 是否把 node 的世界变换烘焙到顶点（True 更稳，导出后位置不容易错）
+# - 一般保持 True；除非你明确想保留层级变换
+
+DEBUG_PRINT = True
+# 是否打印调试信息（palette大小、过渡面数量、迭代变化等）
+# - True：方便调参；稳定后可关掉
+# =========================
+
+
+CHUNK_TYPE_JSON = 0x4E4F534A  # b'JSON'
+CHUNK_TYPE_BIN = 0x004E4942   # b'BIN\0'
+
+
+def _default_out_path(in_path: str) -> str:
+    root, ext = os.path.splitext(in_path)
+    if ext.lower() not in [".glb", ".gltf"]:
+        ext = ".glb"
+    return f"{root}_seg.glb"
+
+
+def _quantize_rgb(rgb: np.ndarray, step: int) -> np.ndarray:
+    if step is None or step <= 0:
+        return rgb
+    q = (rgb.astype(np.int32) + step // 2) // step * step
+    q = np.clip(q, 0, 255).astype(np.uint8)
+    return q
+
+
+def _load_glb_json_and_bin(glb_path: str) -> Tuple[dict, bytes]:
+    data = open(glb_path, "rb").read()
+    if len(data) < 12:
+        raise RuntimeError("Invalid GLB: too small")
+
+    magic, version, length = struct.unpack_from("<4sII", data, 0)
+    if magic != b"glTF":
+        raise RuntimeError("Not a GLB file (missing glTF header)")
+
+    offset = 12
+    gltf_json = None
+    bin_chunk = None
+
+    while offset + 8 <= len(data):
+        chunk_len, chunk_type = struct.unpack_from("<II", data, offset)
+        offset += 8
+        chunk_data = data[offset: offset + chunk_len]
+        offset += chunk_len
+
+        if chunk_type == CHUNK_TYPE_JSON:
+            gltf_json = chunk_data.decode("utf-8", errors="replace")
+        elif chunk_type == CHUNK_TYPE_BIN:
+            bin_chunk = chunk_data
+
+    if gltf_json is None:
+        raise RuntimeError("GLB missing JSON chunk")
+    if bin_chunk is None:
+        raise RuntimeError("GLB missing BIN chunk")
+
+    import json
+    return json.loads(gltf_json), bin_chunk
+
+
+def _extract_basecolor_texture_image(glb_path: str) -> np.ndarray:
+    gltf, bin_chunk = _load_glb_json_and_bin(glb_path)
+
+    materials = gltf.get("materials", [])
+    textures = gltf.get("textures", [])
+    images = gltf.get("images", [])
+    buffer_views = gltf.get("bufferViews", [])
+
+    if not materials:
+        raise RuntimeError("No materials in GLB")
+
+    pbr = materials[0].get("pbrMetallicRoughness", {})
+    base_tex_index = pbr.get("baseColorTexture", {}).get("index", None)
+    if base_tex_index is None:
+        raise RuntimeError("Material has no baseColorTexture")
+    if base_tex_index >= len(textures):
+        raise RuntimeError("baseColorTexture index out of range")
+
+    tex = textures[base_tex_index]
+    img_index = tex.get("source", None)
+    if img_index is None or img_index >= len(images):
+        raise RuntimeError("Texture has no valid image source")
+
+    img_info = images[img_index]
+    bv_index = img_info.get("bufferView", None)
+    mime = img_info.get("mimeType", None)
+    if bv_index is None:
+        uri = img_info.get("uri", None)
+        raise RuntimeError(f"Image is not embedded (bufferView missing). uri={uri}")
+    if bv_index >= len(buffer_views):
+        raise RuntimeError("image.bufferView out of range")
+
+    bv = buffer_views[bv_index]
+    bo = int(bv.get("byteOffset", 0))
+    bl = int(bv.get("byteLength", 0))
+    img_bytes = bin_chunk[bo: bo + bl]
+
+    if DEBUG_PRINT:
+        print(
+            f"[Texture] baseColorTextureIndex={base_tex_index}, imageIndex={img_index}, "
+            f"bufferView={bv_index}, mime={mime}, bytes={len(img_bytes)}"
+        )
+
+    pil = Image.open(trimesh.util.wrap_as_stream(img_bytes)).convert("RGBA")
+    return np.array(pil, dtype=np.uint8)
+
+
+def _merge_palette_rgb(palette_rgb: np.ndarray, counts: np.ndarray, merge_dist: float) -> np.ndarray:
+    if palette_rgb is None or len(palette_rgb) == 0:
+        return palette_rgb
+    if merge_dist is None or merge_dist <= 0:
+        return palette_rgb
+
+    rgb = palette_rgb.astype(np.float32)
+    counts = counts.astype(np.int64)
+
+    order = np.argsort(-counts)
+    centers = []
+    center_w = []
+    thr2 = float(merge_dist) * float(merge_dist)
+
+    for idx in order:
+        x = rgb[idx]
+        w = int(counts[idx])
+
+        if not centers:
+            centers.append(x.copy())
+            center_w.append(w)
+            continue
+
+        C = np.stack(centers, axis=0)
+        d2 = np.sum((C - x[None, :]) ** 2, axis=1)
+        k = int(np.argmin(d2))
+
+        if float(d2[k]) <= thr2:
+            cw = center_w[k]
+            centers[k] = (centers[k] * cw + x * w) / (cw + w)
+            center_w[k] = cw + w
+        else:
+            centers.append(x.copy())
+            center_w.append(w)
+
+    merged = np.clip(np.rint(np.stack(centers, axis=0)), 0, 255).astype(np.uint8)
+
+    if DEBUG_PRINT:
+        print(f"[PaletteMerge] before={len(palette_rgb)} after={len(merged)} merge_dist={merge_dist}")
+
+    return merged
+
+
+def _build_palette_rgb(tex_rgba: np.ndarray) -> np.ndarray:
+    rgb = tex_rgba[:, :, :3].reshape(-1, 3)
+    n = rgb.shape[0]
+
+    if n > PALETTE_SAMPLE_PIXELS:
+        rng = np.random.default_rng(0)
+        idx = rng.choice(n, size=PALETTE_SAMPLE_PIXELS, replace=False)
+        rgb = rgb[idx]
+
+    rgb = _quantize_rgb(rgb, COLOR_QUANT_STEP)
+
+    uniq, counts = np.unique(rgb, axis=0, return_counts=True)
+    order = np.argsort(-counts)
+    uniq = uniq[order]
+    counts = counts[order]
+
+    keep = counts >= PALETTE_MIN_PIXELS
+    uniq = uniq[keep]
+    counts = counts[keep]
+
+    if len(uniq) > PALETTE_MAX_COLORS:
+        uniq = uniq[:PALETTE_MAX_COLORS]
+        counts = counts[:PALETTE_MAX_COLORS]
+
+    if DEBUG_PRINT:
+        print(
+            f"[Palette] quant_step={COLOR_QUANT_STEP} palette_size(before_merge)={len(uniq)} "
+            f"min_pixels={PALETTE_MIN_PIXELS}"
+        )
+        for i in range(min(15, len(uniq))):
+            r, g, b = [int(x) for x in uniq[i]]
+            print(f"  {i:02d} rgb=({r},{g},{b}) count={int(counts[i])}")
+
+    uniq = _merge_palette_rgb(uniq.astype(np.uint8), counts, PALETTE_MERGE_DIST)
+
+    if DEBUG_PRINT:
+        print(f"[Palette] palette_size(after_merge)={len(uniq)}")
+        for i in range(min(15, len(uniq))):
+            r, g, b = [int(x) for x in uniq[i]]
+            print(f"  {i:02d} rgb=({r},{g},{b})")
+
+    return uniq.astype(np.uint8)
+
+
+def _unwrap_uv3_for_seam(uv3: np.ndarray) -> np.ndarray:
+    out = uv3.copy()
+    for d in range(2):
+        v = out[:, :, d]
+        vmin = v.min(axis=1)
+        vmax = v.max(axis=1)
+        seam = (vmax - vmin) > 0.5
+        if np.any(seam):
+            vv = v[seam]
+            vv = np.where(vv < 0.5, vv + 1.0, vv)
+            out[seam, :, d] = vv
+    return out
+
+
+def _barycentric_samples(uv3: np.ndarray, samples_per_face: int) -> np.ndarray:
+    uv3 = _unwrap_uv3_for_seam(uv3)
+
+    if samples_per_face == 1:
+        w = np.array([1 / 3, 1 / 3, 1 / 3], dtype=np.float32)
+        uvs = uv3[:, 0, :] * w[0] + uv3[:, 1, :] * w[1] + uv3[:, 2, :] * w[2]
+        return uvs[:, None, :]
+
+    ws = np.array(
+        [
+            [1 / 3, 1 / 3, 1 / 3],
+            [0.80, 0.10, 0.10],
+            [0.10, 0.80, 0.10],
+            [0.10, 0.10, 0.80],
+        ],
+        dtype=np.float32,
+    )
+    uvs = (
+        uv3[:, None, 0, :] * ws[None, :, 0, None]
+        + uv3[:, None, 1, :] * ws[None, :, 1, None]
+        + uv3[:, None, 2, :] * ws[None, :, 2, None]
+    )
+    return uvs
+
+
+def _wrap_or_clamp_uv(uv: np.ndarray) -> np.ndarray:
+    if UV_WRAP_REPEAT:
+        return np.mod(uv, 1.0)
+    return np.clip(uv, 0.0, 1.0)
+
+
+def _sample_texture_nearest_rgb(tex_rgba: np.ndarray, uv: np.ndarray) -> np.ndarray:
+    h, w = tex_rgba.shape[0], tex_rgba.shape[1]
+    uv = _wrap_or_clamp_uv(uv)
+
+    u = uv[:, 0]
+    v = uv[:, 1]
+    if FLIP_V:
+        v = 1.0 - v
+
+    x = np.rint(u * (w - 1)).astype(np.int32)
+    y = np.rint(v * (h - 1)).astype(np.int32)
+    x = np.clip(x, 0, w - 1)
+    y = np.clip(y, 0, h - 1)
+
+    return tex_rgba[y, x, :3].astype(np.uint8)
+
+
+def _map_to_palette_rgb(colors_rgb: np.ndarray, palette_rgb: np.ndarray, chunk: int = 20000) -> Tuple[np.ndarray, np.ndarray]:
+    if palette_rgb is None or len(palette_rgb) == 0:
+        uniq, inv = np.unique(colors_rgb, axis=0, return_inverse=True)
+        return inv.astype(np.int32), uniq.astype(np.uint8)
+
+    c = colors_rgb.astype(np.float32)
+    p = palette_rgb.astype(np.float32)
+
+    out = np.empty((c.shape[0],), dtype=np.int32)
+    for i in range(0, c.shape[0], chunk):
+        cc = c[i:i + chunk]
+        d2 = ((cc[:, None, :] - p[None, :, :]) ** 2).sum(axis=2)
+        out[i:i + chunk] = np.argmin(d2, axis=1).astype(np.int32)
+
+    return out, palette_rgb
+
+
+def _face_labels_and_confidence_from_texture_rgb(
+    mesh: trimesh.Trimesh,
+    tex_rgba: np.ndarray,
+    palette_rgb: np.ndarray,
+) -> Optional[Tuple[np.ndarray, np.ndarray, np.ndarray]]:
+    uv = getattr(mesh.visual, "uv", None)
+    if uv is None:
+        return None
+
+    uv = np.asarray(uv, dtype=np.float32)
+    if uv.ndim != 2 or uv.shape[1] != 2 or uv.shape[0] != len(mesh.vertices):
+        return None
+
+    faces = mesh.faces
+    uv3 = uv[faces]
+
+    uvs = _barycentric_samples(uv3, SAMPLES_PER_FACE)
+    F, S = uvs.shape[0], uvs.shape[1]
+    flat_uv = uvs.reshape(-1, 2)
+
+    sampled_rgb = _sample_texture_nearest_rgb(tex_rgba, flat_uv)
+    sampled_rgb = _quantize_rgb(sampled_rgb, COLOR_QUANT_STEP)
+
+    sample_label, used_palette = _map_to_palette_rgb(sampled_rgb, palette_rgb)
+    sample_label = sample_label.reshape(F, S)
+
+    if S == 1:
+        face_label = sample_label[:, 0].astype(np.int32)
+        face_conf = np.ones((F,), dtype=np.float32)
+        return face_label, face_conf, used_palette
+
+    face_label = np.empty((F,), dtype=np.int32)
+    face_conf = np.empty((F,), dtype=np.float32)
+    for i in range(F):
+        row = sample_label[i].tolist()
+        c = Counter(row)
+        lab, cnt = c.most_common(1)[0]
+        face_label[i] = int(lab)
+        face_conf[i] = float(cnt) / float(S)
+
+    return face_label, face_conf, used_palette
+
+
+def _build_face_adjacency_list(mesh: trimesh.Trimesh) -> Optional[list]:
+    adj_pairs = mesh.face_adjacency
+    if adj_pairs is None or len(adj_pairs) == 0:
+        return None
+    F = len(mesh.faces)
+    adj = [[] for _ in range(F)]
+    for a, b in adj_pairs:
+        adj[a].append(b)
+        adj[b].append(a)
+    return adj
+
+
+def _reassign_transition_faces(
+    face_label: np.ndarray,
+    face_conf: np.ndarray,
+    adj: list,
+    conf_thresh: float,
+    iters: int,
+    neighbor_min: int,
+) -> np.ndarray:
+    labels = face_label.copy()
+    F = labels.shape[0]
+
+    transition = face_conf < float(conf_thresh)
+    if DEBUG_PRINT:
+        print(f"[Transition] faces={F}, transition_faces={int(transition.sum())}, conf_thresh={conf_thresh}")
+
+    if not np.any(transition):
+        return labels
+
+    for it in range(max(1, iters)):
+        changed = 0
+        for f in range(F):
+            if not transition[f]:
+                continue
+            neigh = adj[f]
+            if not neigh:
+                continue
+
+            votes = defaultdict(int)
+            for nb in neigh:
+                if transition[nb]:
+                    continue
+                votes[int(labels[nb])] += 1
+
+            if not votes:
+                for nb in neigh:
+                    votes[int(labels[nb])] += 1
+
+            if not votes:
+                continue
+
+            best_lab, best_cnt = max(votes.items(), key=lambda x: x[1])
+            if best_cnt < neighbor_min:
+                continue
+
+            if int(labels[f]) != int(best_lab):
+                labels[f] = int(best_lab)
+                changed += 1
+
+        if DEBUG_PRINT:
+            print(f"[Transition] iter={it+1}/{max(1,iters)} changed={changed}")
+        if changed == 0:
+            break
+
+    return labels
+
+
+def _postprocess_small_components(
+    mesh: trimesh.Trimesh,
+    face_label: np.ndarray,
+    min_component_faces: int,
+    action: str,
+    iters: int,
+) -> np.ndarray:
+    if min_component_faces is None or min_component_faces <= 0:
+        return face_label
+    if action not in ("drop", "reassign"):
+        raise ValueError('SMALL_COMPONENT_ACTION must be "drop" or "reassign"')
+
+    adj = _build_face_adjacency_list(mesh)
+    if adj is None:
+        return face_label
+
+    F = len(mesh.faces)
+    labels = face_label.copy()
+
+    for it in range(max(1, iters)):
+        visited = np.zeros(F, dtype=bool)
+        changed = False
+
+        for seed in range(F):
+            if visited[seed]:
+                continue
+            lab = int(labels[seed])
+            if lab < 0:
+                visited[seed] = True
+                continue
+
+            q = [seed]
+            visited[seed] = True
+            comp = [seed]
+
+            while q:
+                f = q.pop()
+                for nb in adj[f]:
+                    if not visited[nb] and int(labels[nb]) == lab:
+                        visited[nb] = True
+                        q.append(nb)
+                        comp.append(nb)
+
+            if len(comp) >= min_component_faces:
+                continue
+
+            if action == "drop":
+                labels[np.array(comp, dtype=np.int64)] = -1
+                changed = True
+                continue
+
+            neigh_counts = defaultdict(int)
+            for f in comp:
+                for nb in adj[f]:
+                    nl = int(labels[nb])
+                    if nl >= 0 and nl != lab:
+                        neigh_counts[nl] += 1
+
+            if not neigh_counts:
+                continue
+
+            new_lab = max(neigh_counts.items(), key=lambda x: x[1])[0]
+            labels[np.array(comp, dtype=np.int64)] = int(new_lab)
+            changed = True
+
+        if DEBUG_PRINT:
+            print(
+                f"[SmallComp] iter={it+1}/{max(1,iters)} action={action} "
+                f"min_comp_faces={min_component_faces} changed={changed}"
+            )
+
+        if not changed:
+            break
+
+    return labels
+
+
+def split_glb_by_texture_palette_rgb(
+    in_glb_path: str,
+    out_glb_path: Optional[str] = None,
+    bake_transforms: bool = True,
+) -> str:
+    if out_glb_path is None:
+        out_glb_path = _default_out_path(in_glb_path)
+
+    tex_rgba = _extract_basecolor_texture_image(in_glb_path)
+    palette_rgb = _build_palette_rgb(tex_rgba)
+
+    scene = trimesh.load(in_glb_path, force="scene", process=False)
+    out_scene = trimesh.Scene()
+
+    part_count = 0
+    base = os.path.splitext(os.path.basename(in_glb_path))[0]
+
+    for node_name in scene.graph.nodes_geometry:
+        geom_name = scene.graph[node_name][1]
+        if geom_name is None:
+            continue
+
+        geom = scene.geometry.get(geom_name, None)
+        if geom is None or not isinstance(geom, trimesh.Trimesh):
+            continue
+
+        mesh = geom.copy()
+
+        if bake_transforms:
+            T, _ = scene.graph.get(node_name)
+            if T is not None:
+                mesh.apply_transform(T)
+
+        res = _face_labels_and_confidence_from_texture_rgb(mesh, tex_rgba, palette_rgb)
+        if res is None:
+            if DEBUG_PRINT:
+                print(f"[{node_name}] no uv / cannot sample -> keep orig")
+            out_scene.add_geometry(mesh, geom_name=f"{base}__{node_name}__orig")
+            continue
+
+        face_label, face_conf, label_rgb = res
+
+        if DEBUG_PRINT:
+            u, _ = np.unique(face_label, return_counts=True)
+            print(f"[{node_name}] raw labels_used={len(u)} palette_size={len(label_rgb)}")
+
+        adj = _build_face_adjacency_list(mesh)
+        if adj is not None:
+            face_label = _reassign_transition_faces(
+                face_label=face_label,
+                face_conf=face_conf,
+                adj=adj,
+                conf_thresh=TRANSITION_CONF_THRESH,
+                iters=TRANSITION_PROP_ITERS,
+                neighbor_min=TRANSITION_NEIGHBOR_MIN,
+            )
+
+        face_label = _postprocess_small_components(
+            mesh=mesh,
+            face_label=face_label,
+            min_component_faces=SMALL_COMPONENT_MIN_FACES,
+            action=SMALL_COMPONENT_ACTION,
+            iters=POSTPROCESS_ITERS,
+        )
+
+        if DEBUG_PRINT:
+            u2, _ = np.unique(face_label[face_label >= 0], return_counts=True)
+            print(f"[{node_name}] after post labels_used={len(u2)}")
+
+        groups = defaultdict(list)
+        for fi, lab in enumerate(face_label):
+            lab = int(lab)
+            if lab < 0:
+                continue
+            groups[lab].append(fi)
+
+        for lab, face_ids in groups.items():
+            if len(face_ids) < MIN_FACES_PER_PART:
+                continue
+
+            sub = mesh.submesh([np.array(face_ids, dtype=np.int64)], append=True, repair=False)
+            if sub is None:
+                continue
+            if isinstance(sub, (list, tuple)):
+                if not sub:
+                    continue
+                sub = sub[0]
+
+            if 0 <= lab < len(label_rgb):
+                r, g, b = [int(x) for x in label_rgb[lab]]
+                part_name = f"{base}__{node_name}__label_{lab}__rgb_{r}_{g}_{b}"
+            else:
+                part_name = f"{base}__{node_name}__label_{lab}"
+
+            out_scene.add_geometry(sub, geom_name=part_name)
+            part_count += 1
+
+    if part_count == 0:
+        if DEBUG_PRINT:
+            print("[INFO] part_count==0, fallback to original scene export.")
+        out_scene = scene
+
+    out_scene.export(out_glb_path)
+    return out_glb_path
+
+
+def main():
+    out_path = split_glb_by_texture_palette_rgb(
+        INPUT_GLB,
+        out_glb_path=None,
+        bake_transforms=BAKE_TRANSFORMS,
+    )
+    print("Done. Exported:", out_path)
+
+
+if __name__ == "__main__":
+    main()

train_full.py

@@ -0,0 +1,227 @@
+import os
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+
+import json
+import torch
+import numpy as np
+import torch.nn as nn
+import pytorch_lightning as pl
+import trellis2.modules.sparse as sp
+
+from trellis2 import models
+from torch.nn import functional as F
+from pytorch_lightning import Trainer
+from torch.utils.data import Dataset, DataLoader
+from pytorch_lightning.callbacks import ModelCheckpoint
+
+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
+
+class Gen3DSegDataset(Dataset):
+    def __init__(self, dataset_path, indices, split="train", repeat=1):
+        super().__init__()
+        self.repeat = repeat
+        self.split = split
+        self.indices = indices
+        with open(dataset_path, "r") as f:
+            all_samples = json.load(f)
+        if self.indices == -1:
+            self.indices = [0, len(all_samples)]
+        self.all_samples = self.split_data(all_samples, split)
+
+    def split_data(self, all_samples, split):
+        repeat = self.repeat if split == "train" else 1
+        all_samples = all_samples[self.indices[0] : self.indices[1]]
+        all_samples = all_samples * repeat
+        return all_samples
+
+    def __len__(self):
+        return len(self.all_samples)
+    
+    def load_instance(self, index):
+        shape_slat = torch.load(self.all_samples[index]["shape_slat"])
+        shape_slat = sp.SparseTensor(shape_slat["feats"], shape_slat["coords"])
+        input_tex_slat = torch.load(self.all_samples[index]["input_tex_slat"])
+        input_tex_slat = sp.SparseTensor(input_tex_slat["feats"], input_tex_slat["coords"])
+        output_tex_slat_gt = torch.load(self.all_samples[index]["output_tex_slat_gt"])
+        output_tex_slat_gt = sp.SparseTensor(output_tex_slat_gt["feats"], output_tex_slat_gt["coords"])
+        cond_dict = torch.load(self.all_samples[index]["cond"])
+        return {"shape_slat": shape_slat, "input_tex_slat": input_tex_slat, "output_tex_slat_gt": output_tex_slat_gt, "cond_dict": cond_dict}
+    
+    def __getitem__(self, index):
+        try:
+            return self.load_instance(index)
+        except Exception as e:
+            print(f"Error in {self.all_samples[index]}: {e}")
+            return self.__getitem__((index + 1) % self.__len__())
+
+class DataModule(pl.LightningDataModule):
+    def __init__(self, batch_size, num_workers, dataset_path, indices, repeat, shuffle, seed):
+        super().__init__()
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.dataset_path = dataset_path
+        self.indices = indices
+        self.repeat = repeat
+        self.shuffle = shuffle
+        self.seed = seed
+
+    def setup(self, stage=None):
+        if stage in (None, "fit"):
+            self.train_dataset = Gen3DSegDataset(self.dataset_path, self.indices, "train", self.repeat)
+
+    def collate_fn(self, batch):
+        shape_slats = sp.sparse_cat([sample["shape_slat"] for sample in batch])
+        input_tex_slats = sp.sparse_cat([sample["input_tex_slat"] for sample in batch])
+        output_tex_slat_gts = sp.sparse_cat([sample["output_tex_slat_gt"] for sample in batch])
+        cond_dicts = [sample["cond_dict"] for sample in batch]
+        coords_len_list = [sample["shape_slat"].coords.shape[0] for sample in batch]
+        return {"shape_slats": shape_slats, "input_tex_slats": input_tex_slats, "output_tex_slat_gts": output_tex_slat_gts, "cond_dicts": cond_dicts, "coords_len_list": coords_len_list}
+
+    def train_dataloader(self):
+        distributed_sampler = None
+        if hasattr(self.trainer, "world_size") and self.trainer.world_size > 1:
+            from torch.utils.data.distributed import DistributedSampler
+            distributed_sampler = DistributedSampler(
+                self.train_dataset,
+                num_replicas=self.trainer.world_size,
+                rank=self.trainer.global_rank,
+                shuffle=self.shuffle,
+                seed=self.seed
+            )
+        return DataLoader(
+            self.train_dataset,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            collate_fn=self.collate_fn,
+            sampler=distributed_sampler,
+            shuffle=False,
+        )
+
+class System(pl.LightningModule):
+    def __init__(self, gen3dseg, pipeline_args, sigma_min, p_uncond, print_every):
+        super().__init__()
+        self.gen3dseg = gen3dseg
+        self.sigma_min = sigma_min
+        self.p_uncond = p_uncond
+        self.print_every = print_every
+        self.shape_std = torch.tensor(pipeline_args['shape_slat_normalization']['std'])[None]
+        self.shape_mean = torch.tensor(pipeline_args['shape_slat_normalization']['mean'])[None]
+        self.tex_std = torch.tensor(pipeline_args['tex_slat_normalization']['std'])[None]
+        self.tex_mean = torch.tensor(pipeline_args['tex_slat_normalization']['mean'])[None]
+        for param in self.gen3dseg.parameters():
+            param.requires_grad = True
+        self.gen3dseg.train()
+
+    def forward(self, shape_slats, input_tex_slats, output_tex_slat_gts, cond_dicts, coords_len_list):
+        batch_size = len(coords_len_list)
+        device = shape_slats.feats.device
+        shape_slats = ((shape_slats - self.shape_mean.to(device)) / self.shape_std.to(device))
+        input_tex_slats = ((input_tex_slats - self.tex_mean.to(device)) / self.tex_std.to(device))
+
+        x_0 = (output_tex_slat_gts - self.tex_mean.to(device)) / self.tex_std.to(device)
+        t = torch.sigmoid(torch.randn(batch_size) * 1.0 + 1.0).to(device)
+        t_x = t.view(-1, *[1 for _ in range(len(x_0.shape) - 1)])
+        noise = sp.SparseTensor(torch.randn_like(x_0.feats), x_0.coords).to(device)
+        x_t = (1 - t_x) * x_0 + (self.sigma_min + (1 - self.sigma_min) * t_x) * noise
+
+        mask = list(np.random.rand(batch_size) < self.p_uncond)
+        cond_list = []
+        for i in range(batch_size):
+            if mask[i]:
+                cond_list.append(cond_dicts[i]['neg_cond'])
+            else:
+                cond_list.append(cond_dicts[i]['cond'])
+        cond = torch.cat(cond_list, dim=0)
+
+        pred = self.gen3dseg(x_t, input_tex_slats, shape_slats, t*1000, cond, coords_len_list)
+        
+        target = (1 - self.sigma_min) * noise - x_0
+        loss = F.mse_loss(pred.feats, target.feats)
+        return loss
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.AdamW(self.gen3dseg.parameters(), lr=1e-4, betas=(0.9, 0.999), weight_decay=0.01)
+        scheduler = {"scheduler": torch.optim.lr_scheduler.ConstantLR(optimizer, factor=1.0, total_iters=9999999), "interval": "step"}
+        return {"optimizer": optimizer, "lr_scheduler": scheduler}
+
+    def training_step(self, batch, batch_idx):
+        loss = self(batch["shape_slats"], batch["input_tex_slats"], batch["output_tex_slat_gts"], batch["cond_dicts"], batch["coords_len_list"])
+        self.log("train_loss", loss.item(), on_step=True, on_epoch=True, prog_bar=True, logger=True)
+        torch.cuda.empty_cache()
+
+        if (self.global_step + 1) % self.print_every == 0:
+            self.print(f"[step {self.global_step+1}] train_loss = {loss.item():.6f}")
+        return loss
+
+def train(dataset_path, ckpts_path):
+    pl.seed_everything(42, workers=True)
+    data_module = DataModule(1, 16, dataset_path, -1, 1, True, 42)
+
+    with open("microsoft/TRELLIS.2-4B/pipeline.json", "r") as f:
+        pipeline_config = json.load(f)
+    pipeline_args = pipeline_config['args']
+    tex_slat_flow_model = models.from_pretrained("microsoft/TRELLIS.2-4B/ckpts/slat_flow_imgshape2tex_dit_1_3B_512_bf16")
+
+    gen3dseg = Gen3DSeg(tex_slat_flow_model)
+    sigma_min = pipeline_args['tex_slat_sampler']['args']['sigma_min']
+    system = System(gen3dseg, pipeline_args, sigma_min, p_uncond=0.1, print_every=10)
+    ckpt_callback = ModelCheckpoint(
+        dirpath=ckpts_path,
+        filename="step_{step}",
+        every_n_train_steps=500,
+        save_top_k=-1
+    )
+    trainer = Trainer(
+        callbacks=[ckpt_callback],
+        accelerator="gpu",
+        devices=-1,
+        max_epochs=1,
+        gradient_clip_val=1.0
+    )
+    trainer.fit(system, datamodule=data_module)
+
+if __name__ == "__main__":
+    _2d_map = True
+    if _2d_map:
+        dataset_path = "./data_toolkit/assets/full_seg_w_2d_map/dataset.json"
+        ckpts_path = "path/to/ckpts_full_seg_w_2d_map"
+    else:
+        dataset_path = "./data_toolkit/assets/full_seg/dataset.json"
+        ckpts_path = "path/to/ckpts_full_seg"
+    train(dataset_path, ckpts_path)

train_interactive.py

@@ -0,0 +1,287 @@
+import os
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+
+import json
+import torch
+import random
+import numpy as np
+import torch.nn as nn
+import pytorch_lightning as pl
+import trellis2.modules.sparse as sp
+
+from trellis2 import models
+from types import MethodType
+from torch.nn import functional as F
+from pytorch_lightning import Trainer
+from trellis2.modules.utils import manual_cast
+from torch.utils.data import Dataset, DataLoader
+from pytorch_lightning.callbacks import ModelCheckpoint
+
+def flow_forward(self, x, t, cond, concat_cond, point_embeds, coords_len_list):
+    # x.feats: [N, 32]
+    x = sp.sparse_cat([x, concat_cond], dim=-1)
+    if isinstance(cond, list):
+        cond = sp.VarLenTensor.from_tensor_list(cond)
+    # x.feats: [N, 64]
+    h = self.input_layer(x)
+    # h.feats: [N, 1536]
+    h = manual_cast(h, self.dtype)
+    t_emb = self.t_embedder(t)
+    t_emb = self.adaLN_modulation(t_emb)
+    t_emb = manual_cast(t_emb, self.dtype)
+    cond = manual_cast(cond, self.dtype)
+    point_embeds = manual_cast(point_embeds, self.dtype)
+
+    h_feats_list = []
+    h_coords_list = []
+    begin = 0
+    for i, coords_len in enumerate(coords_len_list):
+        end = begin + 2 * coords_len
+        h_feats_list.append(h.feats[begin:end])
+        h_coords_list.append(h.coords[begin:end])
+        h_feats_list.append(point_embeds.feats[i*10:(i+1)*10])
+        h_coords_list.append(point_embeds.coords[i*10:(i+1)*10])
+        begin = end + 10
+    h = sp.SparseTensor(torch.cat(h_feats_list), torch.cat(h_coords_list))
+
+    for block in self.blocks:
+        h = block(h, t_emb, cond)
+
+    h_feats_list = []
+    h_coords_list = []
+    begin = 0
+    for i, coords_len in enumerate(coords_len_list):
+        end = begin + 2 * coords_len
+        h_feats_list.append(h.feats[begin:end])
+        h_coords_list.append(h.coords[begin:end])
+        begin = end
+    h = sp.SparseTensor(torch.cat(h_feats_list), torch.cat(h_coords_list))
+
+    h = manual_cast(h, x.dtype)
+    h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:]))
+    # h.feats: [N, 1536]
+    h = self.out_layer(h)
+    # h.feats: [N, 32]
+    return h
+
+class Gen3DSeg(nn.Module):
+    def __init__(self, flow_model):
+        super().__init__()
+        self.flow_model = flow_model
+        self.seg_embeddings = nn.Embedding(1, 1536)
+        
+    def get_positional_encoding(self, input_points):
+        point_feats_embed = torch.zeros((10, 1536), dtype=torch.float32).to(input_points['point_slats'].feats.device)
+        labels = input_points['point_labels'].squeeze(-1)
+        point_feats_embed[labels == 1] = self.seg_embeddings.weight
+        return sp.SparseTensor(point_feats_embed, input_points['point_slats'].coords)
+
+    def forward(self, x_t, tex_slats, shape_slats, t, cond, input_points, 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))
+
+        point_embeds = self.get_positional_encoding(input_points)
+        output_tex_slats = self.flow_model(x_t, t, cond, shape_slats, point_embeds, coords_len_list)
+        
+        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
+
+class Gen3DSegDataset(Dataset):
+    def __init__(self, dataset_path, indices, split="train", repeat=1):
+        super().__init__()
+        self.repeat = repeat
+        self.split = split
+        self.indices = indices
+        with open(dataset_path, "r") as f:
+            all_samples = json.load(f)
+        if self.indices == -1:
+            self.indices = [0, len(all_samples)]
+        self.all_samples = self.split_data(all_samples, split)
+
+    def split_data(self, all_samples, split):
+        repeat = self.repeat if split == "train" else 1
+        all_samples = all_samples[self.indices[0] : self.indices[1]]
+        all_samples = all_samples * repeat
+        return all_samples
+
+    def __len__(self):
+        return len(self.all_samples)
+    
+    def load_instance(self, index):
+        shape_slat = torch.load(self.all_samples[index]["shape_slat"])
+        shape_slat = sp.SparseTensor(shape_slat["feats"], shape_slat["coords"])
+        input_tex_slat = torch.load(self.all_samples[index]["input_tex_slat"])
+        input_tex_slat = sp.SparseTensor(input_tex_slat["feats"], input_tex_slat["coords"])
+        output_tex_slat_gt = torch.load(self.all_samples[index]["output_tex_slat_gt"])
+        output_tex_slat_gt = sp.SparseTensor(output_tex_slat_gt["feats"], output_tex_slat_gt["coords"])
+        cond_dict = torch.load(self.all_samples[index]["cond"])
+        max_point_num = self.all_samples[index]["max_point_num"]
+        point_num = random.randint(1, max_point_num)
+        input_points = torch.load(self.all_samples[index]["input_points"].format(point_num=point_num))
+        return {"shape_slat": shape_slat, "input_tex_slat": input_tex_slat, "output_tex_slat_gt": output_tex_slat_gt, "cond_dict": cond_dict, "input_points": input_points}
+    
+    def __getitem__(self, index):
+        try:
+            return self.load_instance(index)
+        except Exception as e:
+            print(f"Error in {self.all_samples[index]}: {e}")
+            return self.__getitem__((index + 1) % self.__len__())
+
+class DataModule(pl.LightningDataModule):
+    def __init__(self, batch_size, num_workers, dataset_path, indices, repeat, shuffle, seed):
+        super().__init__()
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.dataset_path = dataset_path
+        self.indices = indices
+        self.repeat = repeat
+        self.shuffle = shuffle
+        self.seed = seed
+
+    def setup(self, stage=None):
+        if stage in (None, "fit"):
+            self.train_dataset = Gen3DSegDataset(self.dataset_path, self.indices, "train", self.repeat)
+
+    def collate_fn(self, batch):
+        shape_slats = sp.sparse_cat([sample["shape_slat"] for sample in batch])
+        input_tex_slats = sp.sparse_cat([sample["input_tex_slat"] for sample in batch])
+        output_tex_slat_gts = sp.sparse_cat([sample["output_tex_slat_gt"] for sample in batch])
+        cond_dicts = [sample["cond_dict"] for sample in batch]
+        point_slats = sp.sparse_cat([sp.SparseTensor(sample["input_points"]["point_feats"], sample["input_points"]["point_feats"]) for sample in batch])
+        point_labels = torch.cat([sample["input_points"]["point_labels"] for sample in batch])
+        input_points = {'point_slats': point_slats, 'point_labels': point_labels}
+        coords_len_list = [sample["shape_slat"].coords.shape[0] for sample in batch]
+        return {"shape_slats": shape_slats, "input_tex_slats": input_tex_slats, "output_tex_slat_gts": output_tex_slat_gts, "cond_dicts": cond_dicts, "input_points": input_points, "coords_len_list": coords_len_list}
+
+    def train_dataloader(self):
+        distributed_sampler = None
+        if hasattr(self.trainer, "world_size") and self.trainer.world_size > 1:
+            from torch.utils.data.distributed import DistributedSampler
+            distributed_sampler = DistributedSampler(
+                self.train_dataset,
+                num_replicas=self.trainer.world_size,
+                rank=self.trainer.global_rank,
+                shuffle=self.shuffle,
+                seed=self.seed
+            )
+        return DataLoader(
+            self.train_dataset,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            collate_fn=self.collate_fn,
+            sampler=distributed_sampler,
+            shuffle=False,
+        )
+
+class System(pl.LightningModule):
+    def __init__(self, gen3dseg, pipeline_args, sigma_min, p_uncond, print_every):
+        super().__init__()
+        self.gen3dseg = gen3dseg
+        self.sigma_min = sigma_min
+        self.p_uncond = p_uncond
+        self.print_every = print_every
+        self.shape_std = torch.tensor(pipeline_args['shape_slat_normalization']['std'])[None]
+        self.shape_mean = torch.tensor(pipeline_args['shape_slat_normalization']['mean'])[None]
+        self.tex_std = torch.tensor(pipeline_args['tex_slat_normalization']['std'])[None]
+        self.tex_mean = torch.tensor(pipeline_args['tex_slat_normalization']['mean'])[None]
+        for param in self.gen3dseg.parameters():
+            param.requires_grad = True
+        self.gen3dseg.train()
+
+    def forward(self, shape_slats, input_tex_slats, output_tex_slat_gts, cond_dicts, input_points, coords_len_list):
+        batch_size = len(coords_len_list)
+        device = shape_slats.feats.device
+        shape_slats = ((shape_slats - self.shape_mean.to(device)) / self.shape_std.to(device))
+        input_tex_slats = ((input_tex_slats - self.tex_mean.to(device)) / self.tex_std.to(device))
+
+        x_0 = (output_tex_slat_gts - self.tex_mean.to(device)) / self.tex_std.to(device)
+        t = torch.sigmoid(torch.randn(batch_size) * 1.0 + 1.0).to(device)
+        t_x = t.view(-1, *[1 for _ in range(len(x_0.shape) - 1)])
+        noise = sp.SparseTensor(torch.randn_like(x_0.feats), x_0.coords).to(device)
+        x_t = (1 - t_x) * x_0 + (self.sigma_min + (1 - self.sigma_min) * t_x) * noise
+
+        mask = list(np.random.rand(batch_size) < self.p_uncond)
+        cond_list = []
+        for i in range(batch_size):
+            if mask[i]:
+                cond_list.append(cond_dicts[i]['neg_cond'])
+            else:
+                cond_list.append(cond_dicts[i]['cond'])
+        cond = torch.cat(cond_list, dim=0)
+
+        pred = self.gen3dseg(x_t, input_tex_slats, shape_slats, t*1000, cond, input_points, coords_len_list)
+        
+        target = (1 - self.sigma_min) * noise - x_0
+        loss = F.mse_loss(pred.feats, target.feats)
+        return loss
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.AdamW(self.gen3dseg.parameters(), lr=1e-4, betas=(0.9, 0.999), weight_decay=0.01)
+        scheduler = {"scheduler": torch.optim.lr_scheduler.ConstantLR(optimizer, factor=1.0, total_iters=9999999), "interval": "step"}
+        return {"optimizer": optimizer, "lr_scheduler": scheduler}
+
+    def training_step(self, batch, batch_idx):
+        loss = self(batch["shape_slats"], batch["input_tex_slats"], batch["output_tex_slat_gts"], batch["cond_dicts"], batch["input_points"], batch["coords_len_list"])
+        self.log("train_loss", loss.item(), on_step=True, on_epoch=True, prog_bar=True, logger=True)
+        torch.cuda.empty_cache()
+
+        if (self.global_step + 1) % self.print_every == 0:
+            self.print(f"[step {self.global_step+1}] train_loss = {loss.item():.6f}")
+        return loss
+
+def train(dataset_path, ckpts_path):
+    pl.seed_everything(42, workers=True)
+    data_module = DataModule(1, 16, dataset_path, -1, 1, True, 42)
+
+    with open("microsoft/TRELLIS.2-4B/pipeline.json", "r") as f:
+        pipeline_config = json.load(f)
+    pipeline_args = pipeline_config['args']
+    tex_slat_flow_model = models.from_pretrained("microsoft/TRELLIS.2-4B/ckpts/slat_flow_imgshape2tex_dit_1_3B_512_bf16")
+    tex_slat_flow_model.forward = MethodType(flow_forward, tex_slat_flow_model)
+
+    gen3dseg = Gen3DSeg(tex_slat_flow_model)
+    sigma_min = pipeline_args['tex_slat_sampler']['args']['sigma_min']
+    system = System(gen3dseg, pipeline_args, sigma_min, p_uncond=0.1, print_every=10)
+    ckpt_callback = ModelCheckpoint(
+        dirpath=ckpts_path,
+        filename="step_{step}",
+        every_n_train_steps=500,
+        save_top_k=-1
+    )
+    trainer = Trainer(
+        callbacks=[ckpt_callback],
+        accelerator="gpu",
+        devices=-1,
+        max_epochs=1,
+        gradient_clip_val=1.0
+    )
+    trainer.fit(system, datamodule=data_module)
+
+if __name__ == "__main__":
+    dataset_path = "./data_toolkit/assets/interactive_seg/dataset.json"
+    ckpts_path = "path/to/ckpts_interactive_seg"
+    train(dataset_path, ckpts_path)

train_unified.py

@@ -0,0 +1,303 @@
+import os
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+
+import json
+import torch
+import random
+import numpy as np
+import torch.nn as nn
+import pytorch_lightning as pl
+import trellis2.modules.sparse as sp
+
+from trellis2 import models
+from types import MethodType
+from torch.nn import functional as F
+from pytorch_lightning import Trainer
+from trellis2.modules.utils import manual_cast
+from torch.utils.data import Dataset, DataLoader
+from pytorch_lightning.callbacks import ModelCheckpoint
+
+def flow_forward(self, x, t, tag_embeds, cond, concat_cond, point_embeds, coords_len_list):
+    # x.feats: [N, 32]
+    x = sp.sparse_cat([x, concat_cond], dim=-1)
+    if isinstance(cond, list):
+        cond = sp.VarLenTensor.from_tensor_list(cond)
+    # x.feats: [N, 64]
+    h = self.input_layer(x)
+    # h.feats: [N, 1536]
+    h = manual_cast(h, self.dtype)
+    t_emb = self.t_embedder(t)
+    t_emb = self.adaLN_modulation(t_emb)
+    tag_embeds = self.adaLN_modulation(tag_embeds)
+    t_emb = t_emb + tag_embeds
+    t_emb = manual_cast(t_emb, self.dtype)
+    cond = manual_cast(cond, self.dtype)
+    point_embeds = manual_cast(point_embeds, self.dtype)
+
+    h_feats_list = []
+    h_coords_list = []
+    begin = 0
+    for i, coords_len in enumerate(coords_len_list):
+        end = begin + 2 * coords_len
+        h_feats_list.append(h.feats[begin:end])
+        h_coords_list.append(h.coords[begin:end])
+        h_feats_list.append(point_embeds.feats[i*10:(i+1)*10])
+        h_coords_list.append(point_embeds.coords[i*10:(i+1)*10])
+        begin = end + 10
+    h = sp.SparseTensor(torch.cat(h_feats_list), torch.cat(h_coords_list))
+
+    for block in self.blocks:
+        h = block(h, t_emb, cond)
+
+    h_feats_list = []
+    h_coords_list = []
+    begin = 0
+    for i, coords_len in enumerate(coords_len_list):
+        end = begin + 2 * coords_len
+        h_feats_list.append(h.feats[begin:end])
+        h_coords_list.append(h.coords[begin:end])
+        begin = end
+    h = sp.SparseTensor(torch.cat(h_feats_list), torch.cat(h_coords_list))
+
+    h = manual_cast(h, x.dtype)
+    h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:]))
+    # h.feats: [N, 1536]
+    h = self.out_layer(h)
+    # h.feats: [N, 32]
+    return h
+
+class Gen3DSeg(nn.Module):
+    def __init__(self, flow_model):
+        super().__init__()
+        self.flow_model = flow_model
+        self.seg_embeddings = nn.Embedding(1, 1536)
+        self.tag_mlp = nn.Sequential(nn.Linear(256, 1536, bias=True), nn.SiLU(), nn.Linear(1536, 1536, bias=True))
+        
+    def tag_embedding(self, tag):
+        freqs = torch.exp(-np.log(10000) * torch.arange(start=0, end=128, dtype=torch.float32) / 128).to(device=tag.device)
+        args = tag[:, None].float() * freqs[None]
+        tag_freq = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        tag_embeds = self.tag_mlp(tag_freq)
+        return tag_embeds
+
+    def get_positional_encoding(self, input_points):
+        point_feats_embed = torch.zeros((10, 1536), dtype=torch.float32).to(input_points['point_slats'].feats.device)
+        labels = input_points['point_labels'].squeeze(-1)
+        point_feats_embed[labels == 1] = self.seg_embeddings.weight
+        return sp.SparseTensor(point_feats_embed, input_points['point_slats'].coords)
+
+    def forward(self, x_t, tex_slats, shape_slats, t, tags, cond, input_points, 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))
+
+        tag_embeds = self.tag_embedding(tags)
+        point_embeds = self.get_positional_encoding(input_points)
+        output_tex_slats = self.flow_model(x_t, t, tag_embeds, cond, shape_slats, point_embeds, coords_len_list)
+        
+        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
+
+class Gen3DSegDataset(Dataset):
+    def __init__(self, dataset_path, indices, split="train", repeat=1):
+        super().__init__()
+        self.repeat = repeat
+        self.split = split
+        self.indices = indices
+        with open(dataset_path, "r") as f:
+            all_samples = json.load(f)
+        if self.indices == -1:
+            self.indices = [0, len(all_samples)]
+        self.all_samples = self.split_data(all_samples, split)
+
+    def split_data(self, all_samples, split):
+        repeat = self.repeat if split == "train" else 1
+        all_samples = all_samples[self.indices[0] : self.indices[1]]
+        all_samples = all_samples * repeat
+        return all_samples
+
+    def __len__(self):
+        return len(self.all_samples)
+    
+    def load_instance(self, index):
+        shape_slat = torch.load(self.all_samples[index]["shape_slat"])
+        shape_slat = sp.SparseTensor(shape_slat["feats"], shape_slat["coords"])
+        input_tex_slat = torch.load(self.all_samples[index]["input_tex_slat"])
+        input_tex_slat = sp.SparseTensor(input_tex_slat["feats"], input_tex_slat["coords"])
+        output_tex_slat_gt = torch.load(self.all_samples[index]["output_tex_slat_gt"])
+        output_tex_slat_gt = sp.SparseTensor(output_tex_slat_gt["feats"], output_tex_slat_gt["coords"])
+        cond_dict = torch.load(self.all_samples[index]["cond"])
+        tag = torch.tensor([self.all_samples[index]["tag"]])
+        max_point_num = self.all_samples[index]["max_point_num"]
+        if max_point_num == 0:
+            input_points = {"point_feats": torch.zeros(10, 4, dtype=torch.int32), "point_labels": torch.zeros(10, 1, dtype=torch.int32)}
+        else:
+            point_num = random.randint(1, max_point_num)
+            input_points = torch.load(self.all_samples[index]["input_points"].format(point_num=point_num))
+        return {"shape_slat": shape_slat, "input_tex_slat": input_tex_slat, "output_tex_slat_gt": output_tex_slat_gt, "cond_dict": cond_dict, "input_points": input_points, "tag": tag}
+    
+    def __getitem__(self, index):
+        try:
+            return self.load_instance(index)
+        except Exception as e:
+            print(f"Error in {self.all_samples[index]}: {e}")
+            return self.__getitem__((index + 1) % self.__len__())
+
+class DataModule(pl.LightningDataModule):
+    def __init__(self, batch_size, num_workers, dataset_path, indices, repeat, shuffle, seed):
+        super().__init__()
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.dataset_path = dataset_path
+        self.indices = indices
+        self.repeat = repeat
+        self.shuffle = shuffle
+        self.seed = seed
+
+    def setup(self, stage=None):
+        if stage in (None, "fit"):
+            self.train_dataset = Gen3DSegDataset(self.dataset_path, self.indices, "train", self.repeat)
+
+    def collate_fn(self, batch):
+        shape_slats = sp.sparse_cat([sample["shape_slat"] for sample in batch])
+        input_tex_slats = sp.sparse_cat([sample["input_tex_slat"] for sample in batch])
+        output_tex_slat_gts = sp.sparse_cat([sample["output_tex_slat_gt"] for sample in batch])
+        cond_dicts = [sample["cond_dict"] for sample in batch]
+        point_slats = sp.sparse_cat([sp.SparseTensor(sample["input_points"]["point_feats"], sample["input_points"]["point_feats"]) for sample in batch])
+        point_labels = torch.cat([sample["input_points"]["point_labels"] for sample in batch])
+        input_points = {'point_slats': point_slats, 'point_labels': point_labels}
+        coords_len_list = [sample["shape_slat"].coords.shape[0] for sample in batch]
+        tags = [sample["tag"] for sample in batch]
+        return {"shape_slats": shape_slats, "input_tex_slats": input_tex_slats, "output_tex_slat_gts": output_tex_slat_gts, "cond_dicts": cond_dicts, "input_points": input_points, "coords_len_list": coords_len_list, "tags": tags}
+
+    def train_dataloader(self):
+        distributed_sampler = None
+        if hasattr(self.trainer, "world_size") and self.trainer.world_size > 1:
+            from torch.utils.data.distributed import DistributedSampler
+            distributed_sampler = DistributedSampler(
+                self.train_dataset,
+                num_replicas=self.trainer.world_size,
+                rank=self.trainer.global_rank,
+                shuffle=self.shuffle,
+                seed=self.seed
+            )
+        return DataLoader(
+            self.train_dataset,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            collate_fn=self.collate_fn,
+            sampler=distributed_sampler,
+            shuffle=False,
+        )
+
+class System(pl.LightningModule):
+    def __init__(self, gen3dseg, pipeline_args, sigma_min, p_uncond, print_every):
+        super().__init__()
+        self.gen3dseg = gen3dseg
+        self.sigma_min = sigma_min
+        self.p_uncond = p_uncond
+        self.print_every = print_every
+        self.shape_std = torch.tensor(pipeline_args['shape_slat_normalization']['std'])[None]
+        self.shape_mean = torch.tensor(pipeline_args['shape_slat_normalization']['mean'])[None]
+        self.tex_std = torch.tensor(pipeline_args['tex_slat_normalization']['std'])[None]
+        self.tex_mean = torch.tensor(pipeline_args['tex_slat_normalization']['mean'])[None]
+        for param in self.gen3dseg.parameters():
+            param.requires_grad = True
+        self.gen3dseg.train()
+
+    def forward(self, shape_slats, input_tex_slats, output_tex_slat_gts, cond_dicts, input_points, coords_len_list, tags):
+        batch_size = len(coords_len_list)
+        device = shape_slats.feats.device
+        shape_slats = ((shape_slats - self.shape_mean.to(device)) / self.shape_std.to(device))
+        input_tex_slats = ((input_tex_slats - self.tex_mean.to(device)) / self.tex_std.to(device))
+
+        x_0 = (output_tex_slat_gts - self.tex_mean.to(device)) / self.tex_std.to(device)
+        t = torch.sigmoid(torch.randn(batch_size) * 1.0 + 1.0).to(device)
+        t_x = t.view(-1, *[1 for _ in range(len(x_0.shape) - 1)])
+        noise = sp.SparseTensor(torch.randn_like(x_0.feats), x_0.coords).to(device)
+        x_t = (1 - t_x) * x_0 + (self.sigma_min + (1 - self.sigma_min) * t_x) * noise
+
+        mask = list(np.random.rand(batch_size) < self.p_uncond)
+        cond_list = []
+        for i in range(batch_size):
+            if mask[i]:
+                cond_list.append(cond_dicts[i]['neg_cond'])
+            else:
+                cond_list.append(cond_dicts[i]['cond'])
+        cond = torch.cat(cond_list, dim=0)
+
+        pred = self.gen3dseg(x_t, input_tex_slats, shape_slats, t*1000, tags[0], cond, input_points, coords_len_list)
+        
+        target = (1 - self.sigma_min) * noise - x_0
+        loss = F.mse_loss(pred.feats, target.feats)
+        return loss
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.AdamW(self.gen3dseg.parameters(), lr=1e-4, betas=(0.9, 0.999), weight_decay=0.01)
+        scheduler = {"scheduler": torch.optim.lr_scheduler.ConstantLR(optimizer, factor=1.0, total_iters=9999999), "interval": "step"}
+        return {"optimizer": optimizer, "lr_scheduler": scheduler}
+
+    def training_step(self, batch, batch_idx):
+        loss = self(batch["shape_slats"], batch["input_tex_slats"], batch["output_tex_slat_gts"], batch["cond_dicts"], batch["input_points"], batch["coords_len_list"], batch["tags"])
+        self.log("train_loss", loss.item(), on_step=True, on_epoch=True, prog_bar=True, logger=True)
+        torch.cuda.empty_cache()
+
+        if (self.global_step + 1) % self.print_every == 0:
+            self.print(f"[step {self.global_step+1}] train_loss = {loss.item():.6f}")
+        return loss
+
+def train(dataset_path, ckpts_path):
+    pl.seed_everything(42, workers=True)
+    data_module = DataModule(1, 16, dataset_path, -1, 1, True, 42)
+
+    with open("microsoft/TRELLIS.2-4B/pipeline.json", "r") as f:
+        pipeline_config = json.load(f)
+    pipeline_args = pipeline_config['args']
+    tex_slat_flow_model = models.from_pretrained("microsoft/TRELLIS.2-4B/ckpts/slat_flow_imgshape2tex_dit_1_3B_512_bf16")
+    tex_slat_flow_model.forward = MethodType(flow_forward, tex_slat_flow_model)
+
+    gen3dseg = Gen3DSeg(tex_slat_flow_model)
+    sigma_min = pipeline_args['tex_slat_sampler']['args']['sigma_min']
+    system = System(gen3dseg, pipeline_args, sigma_min, p_uncond=0.1, print_every=10)
+    ckpt_callback = ModelCheckpoint(
+        dirpath=ckpts_path,
+        filename="step_{step}",
+        every_n_train_steps=500,
+        save_top_k=-1
+    )
+    trainer = Trainer(
+        callbacks=[ckpt_callback],
+        accelerator="gpu",
+        devices=-1,
+        max_epochs=1,
+        gradient_clip_val=1.0
+    )
+    trainer.fit(system, datamodule=data_module)
+
+if __name__ == "__main__":
+    dataset_path = "./data_toolkit/assets/unified/dataset.json"
+    ckpts_path = "path/to/ckpts_unified"
+    train(dataset_path, ckpts_path)

trellis2/__init__.py

@@ -0,0 +1,6 @@
+from . import models
+from . import modules
+from . import pipelines
+from . import renderers
+from . import representations
+from . import utils