support safetensor

P3-SAM/model.py

@@ -113,14 +113,19 @@ def load_state_dict(self,
                     ignore_seg_s2_mlp=False, 
                     ignore_iou_mlp=False):
     if ckpt_path is not None:
-        state_dict = torch.load(ckpt_path, weights_only=False, map_location="cpu")["state_dict"]
+        if ckpt_path.endswith('.pt') or ckpt_path.endswith('.ckpt'):
+            state_dict = torch.load(ckpt_path, map_location="cpu")["state_dict"]
+        elif ckpt_path.endswith('.safetensors'):
+            from safetensors.torch import load_file
+            state_dict = load_file(ckpt_path)
     elif state_dict is None:
         # download from huggingface
         print(f'trying to download model from huggingface...')
         from huggingface_hub import hf_hub_download
-        ckpt_path = hf_hub_download(repo_id="tencent/Hunyuan3D-Part", filename="p3sam.pt", local_dir='cache/P3-SAM/')
+        ckpt_path = hf_hub_download(repo_id="tencent/Hunyuan3D-Part", filename="p3sam/p3sam.safetensors", local_dir='weights')
         print(f'download model from huggingface to: {ckpt_path}')
-        state_dict = torch.load(ckpt_path, weights_only=False, map_location="cpu")["state_dict"]
+        from safetensors.torch import load_file
+        state_dict = load_file(ckpt_path)
 
     local_state_dict = self.state_dict()
     seen_keys = {k: False for k in local_state_dict.keys()}

XPart/partgen/bbox_estimator/auto_mask_api.py

@@ -33,24 +33,22 @@ class YSAM(nn.Module):
         super().__init__()
         build_P3SAM(self)
 
-    def load_state_dict(
-        self,
-        state_dict=None,
-        strict=True,
-        assign=False,
-        ignore_seg_mlp=False,
-        ignore_seg_s2_mlp=False,
-        ignore_iou_mlp=False,
-    ):
-        load_state_dict(
-            self,
-            state_dict=state_dict,
-            strict=strict,
-            assign=assign,
-            ignore_seg_mlp=ignore_seg_mlp,
-            ignore_seg_s2_mlp=ignore_seg_s2_mlp,
-            ignore_iou_mlp=ignore_iou_mlp,
-        )
+    def load_state_dict(self, 
+                        ckpt_path=None, 
+                        state_dict=None, 
+                        strict=True, 
+                        assign=False, 
+                        ignore_seg_mlp=False, 
+                        ignore_seg_s2_mlp=False, 
+                        ignore_iou_mlp=False):
+        load_state_dict(self, 
+                        ckpt_path=ckpt_path, 
+                        state_dict=state_dict, 
+                        strict=strict, 
+                        assign=assign, 
+                        ignore_seg_mlp=ignore_seg_mlp, 
+                        ignore_seg_s2_mlp=ignore_seg_s2_mlp, 
+                        ignore_iou_mlp=ignore_iou_mlp)
 
     def forward(self, feats, points, point_prompt, iter=1):
         """
@@ -1362,9 +1360,7 @@ class AutoMask:
         post_process: bool, 是否后处理
         """
         self.model = YSAM()
-        self.model.load_state_dict(
-            state_dict=torch.load(ckpt_path, weights_only=False, map_location="cpu")["state_dict"]
-        )
+        self.model.load_state_dict(ckpt_path)
         self.model.eval()
         # self.model_parallel = torch.nn.DataParallel(self.model)
         self.model_parallel = self.model

XPart/partgen/partformer_pipeline.py

@@ -8,6 +8,7 @@ from tqdm import tqdm
 import copy
 from typing import List, Optional, Union
 import os
+from safetensors.torch import load_file
 from .utils.mesh_utils import (
     SampleMesh,
     load_surface_points,
@@ -21,7 +22,9 @@ from .utils.misc import (
     get_config_from_file,
     smart_load_model,
 )
+from easydict import EasyDict
 
+import json
 from diffusers.utils.torch_utils import randn_tensor
 from pathlib import Path
 
@@ -190,7 +193,7 @@ class PartFormerPipeline(TokenAllocMixin):
                 ckpt[model_name][new_key] = value
         else:
             # ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=True)
-            ckpt = torch.load(ckpt_path, weights_only=False, map_location="cpu")
+            ckpt = torch.load(ckpt_path, map_location="cpu")
         # load model
         model = instantiate_from_config(config["model"])
         # model.load_state_dict(ckpt["model"])
@@ -222,43 +225,42 @@ class PartFormerPipeline(TokenAllocMixin):
     @classmethod
     def from_pretrained(
         cls,
-        config=None,
+        model_path="tencent/Hunyuan3D-Part",
         dtype=torch.float32,
-        ignore_keys=(),
         device="cuda",
         **kwargs,
     ):
-        if config is None:
-            config = get_config_from_file(
-                str(
-                    Path(__file__).parent.parent
-                    / "config"
-                    / "partformer_full_pipeline_512_with_sonata.yaml"
-                )
-            )
-        ckpt_path = smart_load_model(
-            model_path="tencent/Hunyuan3D-Part",
+        model_dir = smart_load_model(
+            model_path=model_path,
         )
-        ckpt = torch.load(os.path.join(ckpt_path, "xpart.pt"), weights_only=False, map_location="cpu")
-        # load model
-        model = instantiate_from_config(config["model"])
-        # model.load_state_dict(ckpt["model"])
-        init_from_ckpt(model, ckpt, prefix="model", ignore_keys=ignore_keys)
-        vae = instantiate_from_config(config["shapevae"])
-        # vae.load_state_dict(ckpt["shapevae"], strict=False)
-        init_from_ckpt(vae, ckpt, prefix="shapevae", ignore_keys=ignore_keys)
-        if config.get("conditioner", None) is not None:
-            conditioner = instantiate_from_config(config["conditioner"])
-            init_from_ckpt(
-                conditioner, ckpt, prefix="conditioner", ignore_keys=ignore_keys
-            )
-        else:
-            conditioner = vae
-        scheduler = instantiate_from_config(config["scheduler"])
-        config["bbox_predictor"]["params"]["ckpt_path"] = os.path.join(
-            ckpt_path, "p3sam.pt"
+        model_ckpt = load_file(os.path.join(model_dir, "model/model.safetensors"))
+        conditioner_ckpt = load_file(
+            os.path.join(model_dir, "conditioner/conditioner.safetensors")
         )
-        bbox_predictor = instantiate_from_config(config.get("bbox_predictor", None))
+        shapevae_ckpt = load_file(
+            os.path.join(model_dir, "shapevae/shapevae.safetensors")
+        )
+        p3sam_path = os.path.join(model_dir, "p3sam/p3sam.safetensors")
+        with open(os.path.join(model_dir, "model/config.json"), "r") as f:
+            model_config = EasyDict(json.load(f))
+        with open(os.path.join(model_dir, "conditioner/config.json"), "r") as f:
+            conditioner_config = EasyDict(json.load(f))
+        with open(os.path.join(model_dir, "shapevae/config.json"), "r") as f:
+            shapevae_config = EasyDict(json.load(f))
+        with open(os.path.join(model_dir, "scheduler/config.json"), "r") as f:
+            scheduler_config = EasyDict(json.load(f))
+        with open(os.path.join(model_dir, "p3sam/config.json"), "r") as f:
+            bbox_predictor_config = EasyDict(json.load(f))
+            bbox_predictor_config["params"]["ckpt_path"] = p3sam_path
+        # load model
+        model = instantiate_from_config(model_config)
+        model.load_state_dict(model_ckpt)
+        vae = instantiate_from_config(shapevae_config)
+        vae.load_state_dict(shapevae_ckpt)
+        conditioner = instantiate_from_config(conditioner_config)
+        conditioner.load_state_dict(conditioner_ckpt)
+        scheduler = instantiate_from_config(scheduler_config)
+        bbox_predictor = instantiate_from_config(bbox_predictor_config)
         model_kwargs = dict(
             vae=vae,
             model=model,

XPart/partgen/utils/misc.py

@@ -150,7 +150,7 @@ def smart_load_model(
     original_model_path = model_path
     # try local path
     base_dir = os.environ.get("HY3DGEN_MODELS", "~/.cache/xpart")
-    model_fld = os.path.expanduser(os.path.join(base_dir, model_path))
+    model_path = os.path.expanduser(os.path.join(base_dir, model_path))
     logger.info(f"Try to load model from local path: {model_path}")
     if not os.path.exists(model_path):
         logger.info("Model path not exists, try to download from huggingface")
@@ -161,7 +161,7 @@ def smart_load_model(
             path = snapshot_download(
                 repo_id=original_model_path,
                 # allow_patterns=[f"{subfolder}/*"],  # 关键修改：模式匹配子文件夹
-                local_dir=model_fld,
+                local_dir=model_path,
             )
             model_path = path  # os.path.join(path, subfolder)  # 保持路径拼接逻辑不变
         except ImportError:

app.py

@@ -25,9 +25,8 @@ def _load_pipeline():
         config, "ckpt_path"
     ), "ckpt or ckpt_path must be specified in config"
     pipeline = PartFormerPipeline.from_pretrained(
-        config=config,
+        model_path="tencent/Hunyuan3D-Part",
         verbose=True,
-        ignore_keys=config.get("ignore_keys", []),
     )
 
     device = "cuda"

requirements.txt

@@ -41,6 +41,7 @@ fpsample
 scikit-learn
 addict
 scikit-image
+easydict
 
 # sonata
 spconv-cu126