# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import torch from torch import nn from transformers import AutoModelForCausalLM from third_party.Michelangelo.encode import load_model from skeleton_models.shape_opt import ShapeOPTConfig def undiscretize(t, low, high, num_discrete): assert (t >= 0).all() and (t <= num_discrete-1).all() assert high > low t = t.float() t /= num_discrete t = t * (high - low) + low assert (t < high).all() and (t >= low).all() return t class SkeletonGPT(nn.Module): def __init__(self, args): super().__init__() self.args = args self.point_encoder = load_model() self.cond_length = 257 self.cond_dim = 768 self.n_discrete_size = args.n_discrete_size self.bone_per_token = 6 # (2 joints per bone) self.max_length = int(args.n_max_bones * self.bone_per_token + 2 + self.cond_length) self.pad_id = -1 self.coor_continuous_range = (-0.5, 0.5) vocab_size = self.n_discrete_size + 3 # 3 for bos, eos, pad self.config = ShapeOPTConfig.from_pretrained( args.llm, n_positions=self.max_length, max_position_embeddings=self.max_length, vocab_size = vocab_size, _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.bone_per_token = self.bone_per_token self.config.cond_length = self.cond_length self.config.word_embed_proj_dim = self.config.hidden_size # 1024 self.transformer = AutoModelForCausalLM.from_config( config=self.config, attn_implementation="flash_attention_2") self.cond_head_proj = nn.Linear(self.cond_dim, self.config.word_embed_proj_dim) self.cond_proj = nn.Linear(self.cond_dim, self.config.word_embed_proj_dim) self.eval() def detokenize(self, input_ids): # input_ids: torch.Tensor of shape (batch_size, seq_length) batch_size = input_ids.size(0) continuous_coors_list = [] num_bones_list = [] for i in range(batch_size): cur_ids = input_ids[i] # Shape: (seq_length,) # Remove padding tokens cur_ids = cur_ids[cur_ids != self.pad_id] # Shape: (effective_seq_length,) # Check if length is a multiple of 6 (2 joints * 3 coordinates) if cur_ids.numel() % 6 != 0: return None # raise ValueError(f"Invalid length of input_ids in sample {i}. It should be a multiple of 6.") num_bones = cur_ids.numel() // 6 num_bones_list.append(num_bones) # Reshape into (num_bones, 6) bone_coords = cur_ids.view(num_bones, 6) # Shape: (num_bones, 6) # Undiscretize the coordinates # Initialize tensor to hold bone coordinates bones_coors = torch.zeros((num_bones, 2, 3), dtype=torch.float16, device=cur_ids.device) for j in range(num_bones): bone_coord = bone_coords[j] # Shape: (6,) # Split into two joints joint1_ids = bone_coord[:3] joint2_ids = bone_coord[3:] # Undiscretize joint coordinates joint1_coords = undiscretize(joint1_ids, self.coor_continuous_range[0], self.coor_continuous_range[1], self.n_discrete_size) joint2_coords = undiscretize(joint2_ids, self.coor_continuous_range[0], self.coor_continuous_range[1], self.n_discrete_size) # Assign to bones_coors bones_coors[j, 0, :] = joint1_coords bones_coors[j, 1, :] = joint2_coords continuous_coors_list.append(bones_coors) max_num_bones = max(num_bones_list) # Initialize the continuous_coors tensor with NaNs continuous_coors = torch.full( (batch_size, max_num_bones, 2, 3), float('nan'), dtype=torch.float16, device=input_ids.device ) # Place the bones_coors into continuous_coors for i in range(batch_size): num_bones = num_bones_list[i] continuous_coors[i, :num_bones, :, :] = continuous_coors_list[i] return continuous_coors # Shape: (batch_size, max_num_bones, 2, 3) # def forward(self, data_dict: dict, is_eval: bool = False) -> dict: # return self.generate(data_dict) def process_point_feature(self, point_feature): encode_feature = torch.zeros(self.args.batchsize_per_gpu, 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(shape_latents) return encode_feature @torch.no_grad() def generate(self, data_dict) -> dict: point_feature = self.point_encoder.encode_latents(data_dict["pc_normal"]) processed_point_feature = self.process_point_feature(point_feature=point_feature) generate_length = self.max_length - self.cond_length net_device = next(self.parameters()).device outputs = torch.ones(self.args.batchsize_per_gpu, generate_length).long().to(net_device) * self.eos_token_id # batch x ntokens if self.args.num_beams is not None and "pc_normal" in data_dict: results = self.transformer.generate( inputs_embeds=processed_point_feature, max_new_tokens=generate_length, # all faces plus two num_beams=self.args.num_beams, 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] # eos and bos removed 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_joints = self.detokenize(outputs) return gen_joints