| import torch |
| from torch import nn |
| from beartype import beartype |
| from modules.miche.encode import load_model |
| import os |
|
|
| _PKG_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) |
| |
|
|
| def exists(val): |
| return val is not None |
|
|
| def default(*values): |
| for value in values: |
| if exists(value): |
| return value |
| return None |
|
|
|
|
| |
| @beartype |
| class PointConditioner(torch.nn.Module): |
| def __init__( |
| self, |
| *, |
| dim_latent = None, |
| model_name = 'miche-256-feature', |
| cond_dim = 768, |
| feature_dim = 768, |
| freeze = True, |
| img_feat_dim: int = 0, |
| img_feat_extractor=None, |
| num_latents = None, |
| ): |
| super().__init__() |
|
|
| |
| if model_name == 'miche-256-feature': |
| ckpt_path = os.environ.get("SHAPEVAE_CKPT", |
| os.path.join(_PKG_ROOT, "ckpt", "shapevae-256.ckpt")) |
| if not os.path.exists(ckpt_path): |
| ckpt_path=None |
| print('[WARNING] Michelangelo ckpt not exist, please check if you are training') |
| config_path = os.path.join(_PKG_ROOT, 'modules', 'miche', 'shapevae-256.yaml') |
|
|
| self.feature_dim = feature_dim |
| self.cond_length = 257 |
| self.point_encoder = load_model(ckpt_path=ckpt_path, config_path=config_path) |
| if num_latents: |
| self._expand_latents(int(num_latents)) |
|
|
| |
| self.cond_head_proj = nn.Linear(cond_dim, self.feature_dim) |
| self.cond_proj = nn.Linear(cond_dim, self.feature_dim) |
|
|
| else: |
| raise NotImplementedError |
|
|
| |
| if freeze: |
| for parameter in self.point_encoder.parameters(): |
| parameter.requires_grad = False |
|
|
| self.freeze = freeze |
| self.model_name = model_name |
| self.dim_latent = self.feature_dim |
|
|
| |
| self.img_feat_dim = int(img_feat_dim or 0) |
| self.img_feat_extractor = img_feat_extractor |
| if self.img_feat_dim > 0: |
| self._widen_input_proj(self.img_feat_dim) |
|
|
| self.register_buffer('_device_param', torch.tensor(0.), persistent = False) |
|
|
|
|
| @property |
| def device(self): |
| return next(self.buffers()).device |
|
|
| def _expand_latents(self, num_latents: int): |
| """把 perceiver 可学习 query 扩到 [1头 + num_latents body]: 预训练 query 是 257=1+256(行0=头token), |
| 头保持 1 个, body 256 平铺复制 + 1% std 微噪打破对称(首份无噪)。attention/n_ctx 均动态读 shape, |
| 无其他结构依赖(SDPA 补丁同样形状无关)。训练热启动时上层 ckpt(若已含扩容 query)会原样覆盖。""" |
| enc = self.point_encoder.model.shape_model.encoder |
| old_q = enc.query.data |
| body = old_q[1:] |
| if body.shape[0] == num_latents: |
| self.cond_length = 1 + num_latents |
| return |
| assert num_latents % body.shape[0] == 0, f"num_latents={num_latents} 须为 body={body.shape[0]} 的整数倍" |
| reps = num_latents // body.shape[0] |
| new_body = body.repeat(reps, 1) |
| noise = torch.randn_like(new_body) * body.std() * 0.01 |
| noise[: body.shape[0]] = 0 |
| enc.query = nn.Parameter(torch.cat([old_q[:1], new_body + noise], dim=0)) |
| enc.num_latents = 1 + num_latents |
| self.cond_length = 1 + num_latents |
| print(f"[PointConditioner] query 扩容 1+{body.shape[0]} -> 1+{num_latents} (body tile+1%噪, 首份无噪)") |
|
|
| def _widen_input_proj(self, img_feat_dim: int): |
| """Replace Michelangelo's input_proj Linear with a wider one and zero-init |
| the new input columns so loading the pretrained ckpt doesn't change init |
| behavior (image features get multiplied by zero on construction).""" |
| enc = self.point_encoder.model.shape_model.encoder |
| if getattr(enc, "img_feat_dim", 0) == img_feat_dim and enc.input_proj.in_features == ( |
| enc.fourier_embedder.out_dim + enc.point_feats + img_feat_dim |
| ): |
| return |
| old = enc.input_proj |
| new_in = enc.fourier_embedder.out_dim + enc.point_feats + img_feat_dim |
| if new_in == old.in_features: |
| enc.img_feat_dim = img_feat_dim |
| return |
| new_proj = nn.Linear( |
| new_in, |
| old.out_features, |
| bias=old.bias is not None, |
| device=old.weight.device, |
| dtype=old.weight.dtype, |
| ) |
| with torch.no_grad(): |
| new_proj.weight.zero_() |
| new_proj.weight[:, : old.in_features].copy_(old.weight) |
| if old.bias is not None: |
| new_proj.bias.copy_(old.bias) |
| enc.input_proj = new_proj |
| enc.img_feat_dim = img_feat_dim |
| if self.freeze: |
| |
| |
| |
| |
| for p in new_proj.parameters(): |
| p.requires_grad = False |
|
|
|
|
| def embed_pc( |
| self, |
| pc_normal, |
| image=None, |
| transform_matrix=None, |
| camera_angle_x=None, |
| depth_map=None, |
| depth_valid=None, |
| img_feats=None, |
| ): |
| |
| if self.model_name == 'miche-256-feature': |
| if self.img_feat_dim > 0: |
| if img_feats is None: |
| if self.img_feat_extractor is None: |
| raise RuntimeError( |
| "img_feat_dim>0 but no img_feats provided and no extractor configured." |
| ) |
| if image is None or transform_matrix is None or camera_angle_x is None: |
| raise ValueError( |
| "image, transform_matrix, camera_angle_x are required when img_feat_extractor is set." |
| ) |
| img_feats = self.img_feat_extractor( |
| image, pc_normal[..., :3], transform_matrix, camera_angle_x, |
| depth_map=depth_map, depth_valid=depth_valid, |
| ) |
| img_feats = img_feats.to(dtype=pc_normal.dtype) |
| point_feature = self.point_encoder.encode_latents(pc_normal, img_feats=img_feats) |
| else: |
| point_feature = self.point_encoder.encode_latents(pc_normal) |
| pc_embed_head = self.cond_head_proj(point_feature[:, 0:1]) |
| pc_embed = self.cond_proj(point_feature[:, 1:]) |
| pc_embed = torch.cat([pc_embed_head, pc_embed], dim=1) |
|
|
| return pc_embed |
|
|
|
|
| def forward( |
| self, |
| pc = None, |
| pc_embeds = None, |
| image=None, |
| transform_matrix=None, |
| camera_angle_x=None, |
| depth_map=None, |
| depth_valid=None, |
| img_feats=None, |
| ): |
| if pc_embeds is None: |
| target_dtype = next(self.buffers()).dtype |
| pc_embeds = self.embed_pc( |
| pc.to(target_dtype), |
| image=image, |
| transform_matrix=transform_matrix, |
| camera_angle_x=camera_angle_x, |
| depth_map=depth_map, |
| depth_valid=depth_valid, |
| img_feats=img_feats, |
| ) |
|
|
| assert not torch.any(torch.isnan(pc_embeds)), 'NAN values in pc embedings' |
|
|
| return pc_embeds |
|
|
