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__))) # helper functions def exists(val): return val is not None def default(*values): for value in values: if exists(value): return value return None # point-cloud encoder from Michelangelo @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, # None=预训练默认(256); >256 时把可学习 query 平铺+微噪扩容(预训练权重原样热启动) ): super().__init__() # open-source version of miche 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 # embedding dimension self.cond_length = 257 # length of embedding self.point_encoder = load_model(ckpt_path=ckpt_path, config_path=config_path) if num_latents: self._expand_latents(int(num_latents)) # additional layers to connect miche and GPT self.cond_head_proj = nn.Linear(cond_dim, self.feature_dim) self.cond_proj = nn.Linear(cond_dim, self.feature_dim) else: raise NotImplementedError # whether to finetuen point-cloud encoder 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 # ---------- image-conditioning extension ---------- 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 # [1+B, w] 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 # 首份 body 保持原 query 精确不动 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 # already widened (e.g. resume from a widened ckpt) 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: # input_proj is part of the shape encoder, which we keep frozen by # default; only the new image-feature columns ever produce nonzero # gradients on top of frozen weights. If callers want to fine-tune # the whole input_proj they can flip freeze=False. 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, ): # encode point cloud to embeddings 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