File size: 8,049 Bytes
90ca518 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 | 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
|