git init

app.py

@@ -1,9 +1,17 @@
 import os
-os.environ["OPENCV_IO_ENABLE_OPENEXR"] = '1'
+os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
 os.environ["ATTN_BACKEND"] = "flash_attn_3"
 
 import urllib.request
+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
 
 os.makedirs("pretrained_model", exist_ok=True)
 
@@ -22,17 +30,6 @@ if not os.path.exists(CKPT_W_2D_MAP):
         CKPT_W_2D_MAP,
     )
 
-CKPT_FULL_SEG = 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__))
@@ -83,25 +80,36 @@ def _collect_examples(example_dir: str) -> List[List[str]]:
 
     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):
+def _toggle_map_input(mode: str):
+    """
+    Show upload image input only when Upload mode is selected.
     """
-    Segment button: generates whole segmented GLB and displays in the second box.
-    Returns: segmented_glb_path, segmented_glb_path(state)
+    return gr.update(visible=(mode == "Upload"))
+
+
+def run_seg(glb_in, map_mode, img_in):
+    """
+    Process button:
+    - Upload mode: use the uploaded 2D map directly
+    - Generate mode: generate a 2D map with FLUX2, show it to the user,
+      and use it as if it were the uploaded map
+
+    Returns:
+        segmented_glb_path,
+        segmented_glb_path(state),
+        used_2d_map_path
     """
     try:
         glb_path = _normalize_path(glb_in)
@@ -121,39 +129,49 @@ def run_seg(glb_in, img_in):
         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
+        # Always build an item that uses a 2D map in the end.
+        # If the user chooses Generate, we generate the map first.
+        ckpt = CKPT_W_2D_MAP
+
+        if map_mode == "Upload":
+            if img_path is None or (not os.path.isfile(img_path)):
+                _raise_user_error("Please upload a valid 2D segmentation map, or switch to Generate mode.")
+
+            used_img = os.path.join(workdir, "2d_map_uploaded.png")
+            shutil.copy(img_path, used_img)
+
+        elif map_mode == "Generate":
             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,
-            }
+            used_img = os.path.join(workdir, "2d_map_generated.png")
+
+            # Generate the 2D map first, and then use it as the uploaded image.
+            inf.generate_2d_map_from_glb(
+                glb_path=in_glb,
+                transforms_path=TRANSFORMS_JSON,
+                out_img_path=used_img,
+                render_img_path=render_img,
+            )
+
+            if not os.path.isfile(used_img):
+                _raise_user_error("2D map generation failed: generated image not found.")
+
+        else:
+            _raise_user_error(f"Unsupported map mode: {map_mode}")
+
+        item = {
+            "2d_map": True,
+            "glb": in_glb,
+            "input_vxz": in_vxz,
+            "img": used_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
+        return out_glb, out_glb, used_img
 
     except Exception as e:
         err = "".join(traceback.format_exception(type(e), e, e.__traceback__))
@@ -216,10 +234,6 @@ def run_refine_segmentation(
         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:
@@ -230,10 +244,11 @@ def run_refine_segmentation(
 
 CSS_TEXT = """
 <style>
-#in_glb  { height: 520px !important; }
-#seg_glb { height: 520px !important; }
-#part_glb{ height: 520px !important; }
-#img     { height: 520px !important; }
+#in_glb      { height: 520px !important; }
+#seg_glb     { height: 520px !important; }
+#part_glb    { height: 520px !important; }
+#img         { height: 520px !important; }
+#used_img    { height: 520px !important; }
 </style>
 """
 
@@ -245,7 +260,6 @@ with gr.Blocks() as demo:
 """
     )
 
-    # ---------------- 2x2 Layout ----------------
     with gr.Row():
         with gr.Column(scale=1, min_width=260):
             in_glb = gr.Model3D(label="Input GLB", elem_id="in_glb")
@@ -254,12 +268,28 @@ with gr.Blocks() as demo:
 
     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")
+            map_mode = gr.Radio(
+                choices=["Upload", "Generate"],
+                value="Upload",
+                label="2D Map Source",
+            )
+
+            with gr.Accordion("2D Segmentation Map", open=True):
+                in_img = gr.Image(
+                    label="Upload 2D Segmentation Map",
+                    type="filepath",
+                    elem_id="img",
+                    visible=True,
+                )
+
+            used_img_preview = gr.Image(
+                label="Used 2D Segmentation Map",
+                type="filepath",
+                elem_id="used_img",
+            )
 
             seg_btn = gr.Button("Process", variant="primary")
 
-            # ✅ Examples directly under the Process button
             if FULL_SEG_EXAMPLES:
                 gr.Examples(
                     examples=FULL_SEG_EXAMPLES,
@@ -269,7 +299,9 @@ with gr.Blocks() as demo:
                     cache_examples=False,
                 )
             else:
-                gr.Markdown(f"**No examples found** in: `{EXAMPLES_DIR}` (expected: `*.glb` + same-name `*.png`).")
+                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):
@@ -385,14 +417,18 @@ with gr.Blocks() as demo:
             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 ----------------
+    map_mode.change(
+        fn=_toggle_map_input,
+        inputs=[map_mode],
+        outputs=[in_img],
+    )
+
     seg_btn.click(
         fn=run_seg,
-        inputs=[in_glb, in_img],
-        outputs=[seg_glb, seg_glb_state],
+        inputs=[in_glb, map_mode, in_img],
+        outputs=[seg_glb, seg_glb_state, used_img_preview],
     )
 
     refine_btn.click(
@@ -414,7 +450,7 @@ with gr.Blocks() as demo:
             small_component_min_faces,
             postprocess_iters,
             min_faces_per_part,
-            bake_transforms
+            bake_transforms,
         ],
         outputs=[part_glb],
     )

app_ori.py

@@ -0,0 +1,423 @@
+import os
+os.environ["OPENCV_IO_ENABLE_OPENEXR"] = '1'
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+os.environ["ATTN_BACKEND"] = "flash_attn_3"
+
+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()
+    inf.PIPE.load_ckpt_if_needed(CKPT_W_2D_MAP)
+    demo.launch()

inference_full.py

@@ -33,6 +33,32 @@ TRELLIS_TEX_DEC = "microsoft/TRELLIS.2-4B/ckpts/tex_dec_next_dc_f16c32_fp16"
 DINO_PATH = "fenghora/dinov3"
 
 
+def generate_2d_map_from_glb(glb_path, transforms_path, out_img_path, render_img_path=None):
+    """
+    Render the GLB first, then generate a 2D segmentation map with FLUX2.
+    """
+    PIPE.load_all_models()
+
+    if render_img_path is None:
+        base, _ = os.path.splitext(out_img_path)
+        render_img_path = f"{base}_render.png"
+
+    render_from_transforms(glb_path, transforms_path, render_img_path)
+
+    prompt = "Apply distinct colors to different regions of this image"
+    image = PIPE.flux2(
+        height=512,
+        width=512,
+        prompt=prompt,
+        image=Image.open(render_img_path),
+        num_inference_steps=28,
+        guidance_scale=4,
+    ).images[0]
+
+    image.save(out_img_path)
+    return out_img_path
+
+
 def _colorvisuals_to_texturevisuals(mesh: trimesh.Trimesh) -> trimesh.Trimesh:
     """
     Convert ColorVisuals to TextureVisuals by baking per-face colors into a tiny atlas
@@ -525,25 +551,26 @@ def inference_with_loaded_models(ckpt_path, item):
     PIPE.load_all_models()
     PIPE.load_ckpt_if_needed(ckpt_path)
 
-    if not item['2d_map']:
-        render_from_transforms(item['glb'], item['transforms'], item['img'])
+    # if not item['2d_map']:
+    #     render_from_transforms(item['glb'], item['transforms'], item['img'])
         
-        prompt = "Apply distinct colors to different regions of this image"
-        image = PIPE.flux2(
-            height=512,
-            width=512,
-            prompt=prompt,
-            image=Image.open(item['img']),
-            num_inference_steps=28,
-            guidance_scale=4,
-        ).images[0]
-        image.save(item['img'])
-
-        # import gc
-        # del flux2
-        # gc.collect()
-        # torch.cuda.empty_cache()
-        # torch.cuda.ipc_collect()
+    #     prompt = "Apply distinct colors to different regions of this image"
+    #     image = PIPE.flux2(
+    #         height=512,
+    #         width=512,
+    #         prompt=prompt,
+    #         image=Image.open(item['img']),
+    #         num_inference_steps=28,
+    #         guidance_scale=4,
+    #     ).images[0]
+    #     image.save(item['img'])
+    
+    if not item["2d_map"]:
+        generate_2d_map_from_glb(
+            glb_path=item["glb"],
+            transforms_path=item["transforms"],
+            out_img_path=item["img"],
+        )
 
     if PIPE.rembg_model is None:
         raise RuntimeError("PIPE.rembg_model is None. Check BiRefNet loading and .cuda() usage.")