import torch
import torch.nn.functional as nnf
from torch import nn
import random
from transformers import AutoModelForCausalLM
from MeshAnything.miche.encode import load_model
from MeshAnything.models.shape_opt import ShapeOPTConfig
from einops import repeat, reduce, rearrange, pack, unpack
class MeshAnythingV2(nn.Module):
def __init__(self):
super().__init__()
self.point_encoder = load_model(ckpt_path=None)
self.n_discrete_size = 128
self.max_seq_ratio = 0.70
self.face_per_token = 9
self.cond_length = 257
self.cond_dim = 768
self.pad_id = -1
self.n_max_triangles = 1600
self.max_length = int(self.n_max_triangles * self.face_per_token * self.max_seq_ratio + 3 + self.cond_length) # add 1
self.coor_continuous_range = (-0.5, 0.5)
self.config = ShapeOPTConfig.from_pretrained(
"facebook/opt-350m",
n_positions=self.max_length,
max_position_embeddings=self.max_length,
vocab_size=self.n_discrete_size + 4,
_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._attn_implementation="flash_attention_2"
self.config.n_discrete_size = self.n_discrete_size
self.config.face_per_token = self.face_per_token
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.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 adjacent_detokenize(self, input_ids):
input_ids = input_ids.reshape(input_ids.shape[0], -1) # B x L
batch_size = input_ids.shape[0]
continuous_coors = torch.zeros((batch_size, self.n_max_triangles * 3 * 10, 3), device=input_ids.device)
continuous_coors[...] = float('nan')
for i in range(batch_size):
cur_ids = input_ids[i]
coor_loop_check = 0
vertice_count = 0
continuous_coors[i, :3, :] = torch.tensor([[-0.1, 0.0, 0.1], [-0.1, 0.1, 0.2], [-0.3, 0.3, 0.2]],
device=input_ids.device)
for id in cur_ids:
if id == self.pad_id:
break
elif id == self.n_discrete_size:
if coor_loop_check < 9:
break
if coor_loop_check % 3 !=0:
break
coor_loop_check = 0
else:
if coor_loop_check % 3 == 0 and coor_loop_check >= 9:
continuous_coors[i, vertice_count] = continuous_coors[i, vertice_count-2]
continuous_coors[i, vertice_count+1] = continuous_coors[i, vertice_count-1]
vertice_count += 2
continuous_coors[i, vertice_count, coor_loop_check % 3] = undiscretize(id, self.coor_continuous_range[0], self.coor_continuous_range[1], self.n_discrete_size)
if coor_loop_check % 3 == 2:
vertice_count += 1
coor_loop_check += 1
continuous_coors = rearrange(continuous_coors, 'b (nf nv) c -> b nf nv c', nv=3, c=3)
return continuous_coors # b, nf, 3, 3
def forward(self, data_dict: dict, is_eval: bool = False) -> dict:
if not is_eval:
return self.train_one_step(data_dict)
else:
return self.generate(data_dict)
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
# batch x ntokens
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.pad_id
outputs[outputs == self.eos_token_id] = self.pad_id
outputs[outputs == self.pad_token_id] = self.pad_id
outputs[outputs != self.pad_id] -= 3
gen_mesh = self.adjacent_detokenize(outputs)
return gen_mesh
def undiscretize(
t,
low,#-0.5
high,# 0.5
num_discrete
):
t = t.float() #[0, num_discrete-1]
t /= num_discrete # 0<=t<1
t = t * (high - low) + low # -0.5 <= t < 0.5
return t