Spaces:
Running on Zero
Running on Zero
| # Project EmbodiedGen | |
| # | |
| # Copyright (c) 2025 Horizon Robotics. All Rights Reserved. | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or | |
| # implied. See the License for the specific language governing | |
| # permissions and limitations under the License. | |
| import importlib | |
| import os | |
| import sys | |
| import torch | |
| import torch.nn as nn | |
| def monkey_patch_p3sam(): | |
| """Patch P3-SAM model construction to use the shared Hugging Face cache.""" | |
| current_file_path = os.path.abspath(__file__) | |
| current_dir = os.path.dirname(current_file_path) | |
| hunyuan_part_root = os.path.abspath( | |
| os.path.join(current_dir, "../../../thirdparty/Hunyuan3D-Part") | |
| ) | |
| p3sam_root = os.path.join(hunyuan_part_root, "P3-SAM") | |
| partgen_root = os.path.join(hunyuan_part_root, "XPart/partgen") | |
| for path in [p3sam_root, partgen_root]: | |
| if path not in sys.path: | |
| sys.path.insert(0, path) | |
| from models import sonata | |
| p3sam_model = importlib.import_module("model") | |
| def build_P3SAM(self): | |
| self.sonata = sonata.load( | |
| "sonata", | |
| repo_id="facebook/sonata", | |
| download_root=os.path.expanduser( | |
| "~/.cache/huggingface/hub/sonata" | |
| ), | |
| ) | |
| self.mlp = nn.Sequential( | |
| nn.Linear(1232, 512), | |
| nn.GELU(), | |
| nn.Linear(512, 512), | |
| nn.GELU(), | |
| nn.Linear(512, 512), | |
| ) | |
| self.transform = sonata.transform.default() | |
| self.seg_mlp_1 = nn.Sequential( | |
| nn.Linear(512 + 3 + 3, 512), | |
| nn.GELU(), | |
| nn.Linear(512, 512), | |
| nn.GELU(), | |
| nn.Linear(512, 1), | |
| ) | |
| self.seg_mlp_2 = nn.Sequential( | |
| nn.Linear(512 + 3 + 3, 512), | |
| nn.GELU(), | |
| nn.Linear(512, 512), | |
| nn.GELU(), | |
| nn.Linear(512, 1), | |
| ) | |
| self.seg_mlp_3 = nn.Sequential( | |
| nn.Linear(512 + 3 + 3, 512), | |
| nn.GELU(), | |
| nn.Linear(512, 512), | |
| nn.GELU(), | |
| nn.Linear(512, 1), | |
| ) | |
| self.seg_s2_mlp_g = nn.Sequential( | |
| nn.Linear(512 + 3 + 3 + 3, 256), | |
| nn.GELU(), | |
| nn.Linear(256, 256), | |
| nn.GELU(), | |
| nn.Linear(256, 256), | |
| ) | |
| self.seg_s2_mlp_1 = nn.Sequential( | |
| nn.Linear(512 + 3 + 3 + 3 + 256, 256), | |
| nn.GELU(), | |
| nn.Linear(256, 256), | |
| nn.GELU(), | |
| nn.Linear(256, 1), | |
| ) | |
| self.seg_s2_mlp_2 = nn.Sequential( | |
| nn.Linear(512 + 3 + 3 + 3 + 256, 256), | |
| nn.GELU(), | |
| nn.Linear(256, 256), | |
| nn.GELU(), | |
| nn.Linear(256, 1), | |
| ) | |
| self.seg_s2_mlp_3 = nn.Sequential( | |
| nn.Linear(512 + 3 + 3 + 3 + 256, 256), | |
| nn.GELU(), | |
| nn.Linear(256, 256), | |
| nn.GELU(), | |
| nn.Linear(256, 1), | |
| ) | |
| self.iou_mlp = nn.Sequential( | |
| nn.Linear(512 + 3 + 3 + 3 + 256, 256), | |
| nn.GELU(), | |
| nn.Linear(256, 256), | |
| nn.GELU(), | |
| nn.Linear(256, 256), | |
| ) | |
| self.iou_mlp_out = nn.Sequential( | |
| nn.Linear(256, 256), | |
| nn.GELU(), | |
| nn.Linear(256, 256), | |
| nn.GELU(), | |
| nn.Linear(256, 3), | |
| ) | |
| self.iou_criterion = torch.nn.MSELoss() | |
| p3sam_model.build_P3SAM = build_P3SAM | |