import torch from torch import nn, Tensor from transformers import AutoModelForCausalLM, AutoConfig, AutoModel from MeshAnything.miche.encode import load_model from MeshAnything.models.shape_opt import ShapeOPTConfig from einops.layers.torch import Rearrange from einops import rearrange, repeat, reduce, pack, unpack import torch.nn.functional as F class NoiseResistantDecoder(nn.Module): def __init__(self, args): super().__init__() self.args = args self.pad_id = -1 self.num_quantizers = 3 self.discrete_num = 128 self.codebook_size = args.codebook_size self.codebook_dim = args.codebook_dim config = AutoConfig.from_pretrained("bert-base-uncased") config.num_hidden_layers = 6 self.decoder= AutoModel.from_config(config=config).to_bettertransformer().encoder self.n_embd = self.decoder.config.hidden_size self.pos_embedding = nn.Embedding(18000, self.n_embd) self.layernorm = nn.LayerNorm(self.n_embd) self.point_layernorm = nn.LayerNorm(self.n_embd) self.cond_length = 257 self.cond_dim = 768 self.point_pe = nn.Embedding(self.cond_length, self.n_embd) self.cond_proj = nn.Linear(self.cond_dim, self.n_embd) self.cond_head_proj = nn.Linear(self.cond_dim, self.n_embd) self.project_down_codebook = nn.Linear(self.codebook_dim * 3, self.n_embd) self.to_coor_logits = nn.Sequential( nn.Linear(self.n_embd, self.discrete_num * 9), Rearrange('... (v c) -> ... v c', v = 9) ) def process_point_feature(self, encode_feature): point_feature = torch.zeros(encode_feature.shape[0], self.cond_length, self.n_embd, device=self.cond_head_proj.weight.device, dtype=self.cond_head_proj.weight.dtype) point_feature[:, 0] = self.cond_head_proj(encode_feature[:, 0]) point_feature[:, 1:] = self.cond_proj(encode_feature[:, 1:]) point_feature = self.point_layernorm(point_feature + self.point_pe.weight[None, :point_feature.shape[1]]) return point_feature def forward(self, input_ids, input_embeds, point_feature = None): input_ids = input_ids.reshape(input_ids.shape[0], -1) point_feature = self.process_point_feature(point_feature) face_embeds = rearrange(input_embeds, 'b (nf nv) d -> b nf (nv d)', nv = 3) face_embeds = self.project_down_codebook(face_embeds) face_mask = reduce(input_ids != self.pad_id, 'b (nf nv q) -> b nf', 'all', nv = 3, q = self.num_quantizers) face_embeds[~face_mask] = 0 face_embeds = self.layernorm(face_embeds + self.pos_embedding.weight[None, :face_embeds.shape[1]]) outputs = self.decoder( hidden_states=torch.concatenate([point_feature, face_embeds], dim=1), ) decoded = outputs.last_hidden_state[:, self.cond_length:] # batch x nfaces x dim decoded = decoded.masked_fill(~face_mask.unsqueeze(-1), 0.) # batch x nfaces x 9 -> batch x nfaces x 3 x 3 pred_face_logits = self.to_coor_logits(decoded) # batch x nfaces x 9 x ndiscrete pred_face_coords = rearrange(pred_face_logits.argmax(dim = -1), '... (v c) -> ... v c', v = 3) continuous_coors = undiscretize( pred_face_coords, num_discrete = self.discrete_num, low = -0.5, high = 0.5 ) continuous_coors = continuous_coors.masked_fill(~rearrange(face_mask, 'b nf -> b nf 1 1'), float('nan')) return continuous_coors class MeshAnything(nn.Module): def __init__(self, args): super().__init__() self.args = args self.point_encoder = load_model(ckpt_path=None) self.tokenizer = NoiseResistantDecoder(args) self.num_quantizers = 3 self.face_per_token = self.num_quantizers * 3 self.cond_length = 257 self.cond_dim = 768 self.max_length = args.n_max_triangles * self.face_per_token + 2 + self.cond_length self.config = ShapeOPTConfig.from_pretrained( args.llm, n_positions=18259, max_position_embeddings=18259, vocab_size=self.tokenizer.codebook_size + 3, _attn_implementation="flash_attention_2" ) self.bos_token_id = 0 self.eos_token_id = 1 self.pad_token_id = 2 self.config.bos_token_id = self.bos_token_id self.config.eos_token_id = self.eos_token_id self.config.pad_token_id = self.pad_token_id self.config.quantize_codebook_dim = self.tokenizer.codebook_dim self.config.face_per_token = self.face_per_token self.config._attn_implementation="flash_attention_2" self.config.cond_length = self.cond_length if self.config.word_embed_proj_dim != self.config.hidden_size: self.config.word_embed_proj_dim = self.config.hidden_size self.transformer = AutoModelForCausalLM.from_config( config=self.config, use_flash_attention_2 = True ) self.transformer.to_bettertransformer() self.transformer.model.decoder.quantize_codebooks = nn.Parameter(torch.zeros(1, self.tokenizer.codebook_size, self.tokenizer.codebook_dim)) self.cond_head_proj = nn.Linear(self.cond_dim, self.config.word_embed_proj_dim) self.cond_proj = nn.Linear(self.cond_dim * 2, self.config.word_embed_proj_dim) self.eval() def process_point_feature(self, point_feature): encode_feature = torch.zeros(point_feature.shape[0], self.cond_length, self.config.word_embed_proj_dim, device=self.cond_head_proj.weight.device, dtype=self.cond_head_proj.weight.dtype) encode_feature[:, 0] = self.cond_head_proj(point_feature[:, 0]) shape_latents = self.point_encoder.to_shape_latents(point_feature[:, 1:]) encode_feature[:, 1:] = self.cond_proj(torch.cat([point_feature[:, 1:], shape_latents], dim=-1)) return encode_feature @torch.no_grad() def forward(self, pc_normal, sampling=False) -> dict: batch_size = pc_normal.shape[0] point_feature = self.point_encoder.encode_latents(pc_normal) processed_point_feature = self.process_point_feature(point_feature) generate_length = self.max_length - self.cond_length net_device = next(self.parameters()).device outputs = torch.ones(batch_size, generate_length).long().to(net_device) * self.eos_token_id if not sampling: results = self.transformer.generate( inputs_embeds=processed_point_feature, max_new_tokens=generate_length, # all faces plus two num_beams=1, bos_token_id=self.bos_token_id, eos_token_id=self.eos_token_id, pad_token_id=self.pad_token_id, ) else: results = self.transformer.generate( inputs_embeds = processed_point_feature, max_new_tokens=generate_length, # all faces plus two do_sample=True, top_k=50, top_p=0.95, bos_token_id = self.bos_token_id, eos_token_id = self.eos_token_id, pad_token_id = self.pad_token_id, ) assert results.shape[1] <= generate_length # B x ID bos is not included since it's predicted outputs[:, :results.shape[1]] = results # batch x ntokens ====> batch x ntokens x D outputs = outputs[:, 1: -1] outputs[outputs == self.bos_token_id] = self.tokenizer.pad_id outputs[outputs == self.eos_token_id] = self.tokenizer.pad_id outputs[outputs == self.pad_token_id] = self.tokenizer.pad_id outputs[outputs != self.tokenizer.pad_id] -= 3 code_embed = self.get_codes(outputs) decoder_output = self.tokenizer(outputs, code_embed, point_feature=point_feature) return decoder_output def get_codes(self, indices): indices = indices.reshape(indices.shape[0], -1) indices = rearrange(indices, 'b (n q) -> b n q', q=self.num_quantizers) batch, quantize_dim = indices.shape[0], indices.shape[-1] # may also receive indices in the shape of 'b h w q' (accept_image_fmap) indices, ps = pack([indices], 'b * q') # because of quantize dropout, one can pass in indices that are coarse # and the network should be able to reconstruct if quantize_dim < self.num_quantizers: indices = F.pad(indices, (0, self.num_quantizers - quantize_dim), value = -1) # take care of quantizer dropout mask = indices == -1. indices = indices.masked_fill(mask, 0) # have it fetch a dummy code to be masked out later # dummy implementation for shared codebook all_codes = self.transformer.model.decoder.quantize_codebooks[0][indices] all_codes = all_codes.permute(2, 0, 1, 3) # mask out any codes that were dropout-ed all_codes = all_codes.masked_fill(rearrange(mask, 'b n q -> q b n 1'), 0.) # if (accept_image_fmap = True) then return shape (quantize, batch, height, width, dimension) codes, = unpack(all_codes, ps, 'q b * d') codes_summed = reduce(codes, 'q ... -> ...', 'sum') return codes_summed def undiscretize( t, low, high, num_discrete ) -> Tensor: t = t.float() t /= num_discrete return t * (high - low) + low