fix req and state

app.py

@@ -8,50 +8,41 @@ from pathlib import Path
 import torch
 import pytorch_lightning as pl
 import spaces
-# from spaces import zero
-# zero.startup()
-# import torch.multiprocessing as mp
-# mp.set_start_method('spawn')
-# print('using torch spawm')
-# print('zero gpu startup')
 
 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
-print('no automask no postprocess')
+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
+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", []),
+    )
 
-# _PIPELINE = _load_pipeline()
+    device = "cuda"
+    pipeline.to(device=device, dtype=torch.float32)
+    return pipeline
 
-print('no xpart 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):
+def segment(mesh_path, connectivity=True, postprocess=True, postprocess_threshold=0.95, seed=42):
     if mesh_path is None:
         gr.Warning("No Input Mesh")
         gr_state[0] = (None, None)
@@ -82,8 +73,8 @@ def segment(mesh_path, connectivity=True, postprocess=True, postprocess_threshol
     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
+    gr_state = [(aabb, mesh_path)]
+    return file_path, face_id_save_path, gr_state
 
 @spaces.GPU
 def generate(mesh_path, seed=42, gr_state=None):
@@ -192,7 +183,7 @@ Input a mesh, segment it using P3-SAM on the left, and push the "Generate" butto
             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])
+    p3sam_button.click(segment, inputs=[p3sam_input, p3sam_conectivity, p3sam_postprocess, p3sam_postprocess_threshold, p3sam_seed], outputs=[p3sam_output, p3sam_face_id_output, gr_state])
     xpart_button.click(generate, inputs=[p3sam_input, xpart_seed, gr_state], outputs=[xpart_output, xpart_output_bbox, xpart_output_exploded])
 
 

requirements.txt

@@ -43,5 +43,6 @@ scikit-image
 
 # sonata
 spconv-cu126
-torch-scatter -f https://data.pyg.org/whl/torch-2.8.0+cu128.html
-git+https://github.com/Dao-AILab/flash-attention.git
+torch-scatter -f https://data.pyg.org/whl/torch-2.8.0+cu129.html
+git+https://github.com/Dao-AILab/flash-attention.git
+torch-cluster -f https://data.pyg.org/whl/torch-2.8.0+cu129.html