app.py

import gradio as gr
import os
import sys
import argparse
import numpy as np
import trimesh
from pathlib import Path
import torch
import pytorch_lightning as pl
import spaces

sys.path.append('P3-SAM')
from demo.auto_mask import AutoMask
from demo.auto_mask_no_postprocess import AutoMask as AutoMaskNoPostProcess
sys.path.append('XPart')
from partgen.partformer_pipeline import PartFormerPipeline
from partgen.utils.misc import get_config_from_file

automask = AutoMask()
automask_no_postprocess = AutoMaskNoPostProcess(automask_instance=automask)

def _load_pipeline():
    pl.seed_everything(2026, workers=True)
    cfg_path = str(Path(__file__).parent / "XPart/partgen/config" / "infer.yaml")
    config = get_config_from_file(cfg_path)
    assert hasattr(config, "ckpt") or hasattr(
        config, "ckpt_path"
    ), "ckpt or ckpt_path must be specified in config"
    pipeline = PartFormerPipeline.from_pretrained(
        config=config,
        verbose=True,
        ignore_keys=config.get("ignore_keys", []),
    )

    device = "cuda"
    pipeline.to(device=device, dtype=torch.float32)
    return pipeline

_PIPELINE = _load_pipeline()

output_path = 'P3-SAM/results/gradio'
os.makedirs(output_path, exist_ok=True)

@spaces.GPU
def segment(mesh_path, connectivity=True, postprocess=True, postprocess_threshold=0.95, seed=42, gr_state=None):
    if mesh_path is None:
        gr.Warning("No Input Mesh")
        gr_state[0] = (None, None)
        return None, None
    mesh = trimesh.load(mesh_path, force='mesh', process=False)
    if connectivity:
        aabb, face_ids, mesh = automask.predict_aabb(mesh, seed=seed, is_parallel=False, post_process=postprocess, threshold=postprocess_threshold)
    else:
        aabb, face_ids, mesh = automask_no_postprocess.predict_aabb(mesh, seed=seed, is_parallel=False, post_process=False)
    color_map = {}
    unique_ids = np.unique(face_ids)
    for i in unique_ids:
        if i == -1:
            continue
        part_color = np.random.rand(3) * 255
        color_map[i] = part_color
    face_colors = []
    for i in face_ids:
        if i == -1:
            face_colors.append([0, 0, 0])
        else:
            face_colors.append(color_map[i])
    face_colors = np.array(face_colors).astype(np.uint8)
    mesh_save = mesh.copy()
    mesh_save.visual.face_colors = face_colors

    file_path = os.path.join(output_path, 'segment_mesh.glb')
    mesh_save.export(file_path)
    face_id_save_path = os.path.join(output_path, 'face_id.npy')
    np.save(face_id_save_path, face_ids)
    gr_state[0] = (aabb, mesh_path)
    return file_path, face_id_save_path

@spaces.GPU
def generate(mesh_path, seed=42, gr_state=None):
    if mesh_path is None:
        gr.Warning("No Input Mesh")
        gr_state[0] = (None, None)
        return None, None, None
    if gr_state[0][0] is None or mesh_path != gr_state[0][1]:
        gr.Warning("Please segment the mesh first")
        return None, None, None
    
    aabb = gr_state[0][0]
    # Ensure deterministic behavior per request
    try:
        pl.seed_everything(int(seed), workers=True)
    except Exception:
        pl.seed_everything(2026, workers=True)
    additional_params = {"output_type": "trimesh"}
    obj_mesh, (out_bbox, mesh_gt_bbox, explode_object) = _PIPELINE(
        mesh_path=mesh_path,
        aabb=aabb,
        octree_resolution=512,
        **additional_params,
    )
    # Export all results to temporary files for Gradio Model3D
    obj_path = os.path.join(output_path, 'obj_mesh.glb')
    out_bbox_path = os.path.join(output_path, 'out_bbox.glb')
    explode_path = os.path.join(output_path, 'explode.glb')
    obj_mesh.export(obj_path)
    out_bbox.export(out_bbox_path)
    explode_object.export(explode_path)
    return obj_path, out_bbox_path, explode_path

with gr.Blocks() as demo:
    gr.Markdown(
'''
# ☯️ Hunyuan3D Part:P3-SAM&XPart
This demo allows you to generate parts given a 3D model using Hunyuan3D-Part.
First segment the 3D model using P3-SAM and then generate parts using XPart.
'''
    )
    with gr.Row():
        with gr.Column():
            # P3-SAM
            gr.Markdown(
'''
## P3-SAM: Native 3D Part Segmentation

[Paper](https://arxiv.org/abs/2509.06784) | [Project Page](https://murcherful.github.io/P3-SAM/) | [Code](https://github.com/Tencent-Hunyuan/Hunyuan3D-Part/P3-SAM/) | [Model](https://huggingface.co/tencent/Hunyuan3D-Part)

This is a demo of P3-SAM, a native 3D part segmentation method that can segment a mesh into different parts.
Input a mesh and push the "Segment" button to get the segmentation results.
'''
            )
            p3sam_button = gr.Button("Segment")
            p3sam_input = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], label="Input Mesh")
            p3sam_output = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], label="Segmentation Result")
            p3sam_face_id_output = gr.File(label='Face ID')
            p3sam_conectivity = gr.Checkbox(value=True, label="Connectivity")
            p3sam_postprocess = gr.Checkbox(value=True, label="Post-processing")
            p3sam_postprocess_threshold = gr.Number(value=0.95, label="Post-processing Threshold")
            p3sam_seed = gr.Number(value=42, label="Random Seed")
            gr.Markdown(
'''
P3-SAM will clean your mesh. To get face-aligned labels, you can download the "Segmentation Result" and "Face ID".
You can also use the "Connectivity" and "Post-processing" options to control the behavior of the algorithm.
The "Post-processing" will merge the small parts according to the threshold. The smaller the threshold, the more parts will be merged.
'''
            )
            gr.Examples(examples=[
                'P3-SAM/demo/assets/1.glb',
                'P3-SAM/demo/assets/2.glb',
                'P3-SAM/demo/assets/4.glb',
                'XPart/data/000.glb',
                'XPart/data/001.glb',
                'XPart/data/002.glb',
                'XPart/data/003.glb',
                'XPart/data/004.glb',
                ],
                inputs = [p3sam_input],
                example_labels=[
                    'Female Warrior',
                    'Suspended Island',
                    'Beetle Car',
                    'Koi Fish',
                    'Motorcycle',
                    'Gundam',
                    'Computer Desk',
                    'Coffee Machine'
                ]
            )
        with gr.Column():
            # XPart
            gr.Markdown(
'''
## XPart: High-fidelity and Structure-coherent Shapede Composition  

[Paper](https://arxiv.org/abs/2509.08643) | [Project Page](https://yanxinhao.github.io/Projects/X-Part/) | [Code](https://github.com/Tencent-Hunyuan/Hunyuan3D-Part/XPart/) | [Model](https://huggingface.co/tencent/Hunyuan3D-Part)

This is a demo of XPart, a high-fidelity and structure-coherent shape-decomposition method that can generate parts from a 3D model.
Input a mesh, segment it using P3-SAM on the left, and push the "Generate" button to get the generated parts.
'''         )
            xpart_button = gr.Button("Generate")
            xpart_output = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], label="Generated Parts")
            xpart_output_bbox = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], label="Gnerated Parts with BBox")
            xpart_output_exploded = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], label="Exploded Object")
            xpart_seed = gr.Number(value=42, label="Random Seed")
        gr_state = gr.State(value=[(None, None)])
    p3sam_button.click(segment, inputs=[p3sam_input, p3sam_conectivity, p3sam_postprocess, p3sam_postprocess_threshold, p3sam_seed, gr_state], outputs=[p3sam_output, p3sam_face_id_output])
    xpart_button.click(generate, inputs=[p3sam_input, xpart_seed, gr_state], outputs=[xpart_output, xpart_output_bbox, xpart_output_exploded])


if __name__ == '__main__':
    demo.launch(server_name='0.0.0.0', server_port=8080)