skeleton/skeleton_models/skeletongen.py

#  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_partys.Michelangelo.encode import load_model
from skeleton_models.skeleton_opt import SkeletonOPTConfig

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.joint_token = args.joint_token
    
        self.n_discrete_size = args.n_discrete_size
        
        if self.joint_token:
            self.bone_per_token = 4 # (x,y,z,parend_index)
            args.n_max_bones += 1 # add one for joints
        else:
            self.bone_per_token = 6  # (2 joints per bone, xyzxyz)
        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 = SkeletonOPTConfig.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.joint_token = self.joint_token

        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
        
        # target-aware indicator
        if self.args.seq_shuffle:
            self.feat_dim = self.config.word_embed_proj_dim
            self.target_aware_pos_embed = nn.Parameter(torch.zeros(1, args.n_max_bones, self.config.word_embed_proj_dim))
            nn.init.trunc_normal_(self.target_aware_pos_embed, 0., 0.02)
       
        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 detokenize_joint_token(self, input_ids):
        # input_ids: torch.Tensor of shape (batch_size, seq_length)
        batch_size = input_ids.size(0)

        bones_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 4 (xyz + parent index)
            if cur_ids.numel() % 4 != 0:
                return None

            num_joints = cur_ids.numel() // 4

            # Reshape into (num_joints, 4)
            joint_data = cur_ids.view(num_joints, 4)
            
            # Undiscretize the coordinates
            coords_discrete = joint_data[:, :3]   # shape: (num_joints, 3)
            coords_float = undiscretize(
                coords_discrete, 
                self.coor_continuous_range[0], 
                self.coor_continuous_range[1], 
                self.n_discrete_size
            )
            parents = joint_data[:, 3]
            
            ### recover bones
            bone_coords = []
            for child_idx in range(num_joints):
                p = parents[child_idx].item()
                if p > 0:
                    try:
                        parent_idx = p - 1
                        parent_coord = coords_float[parent_idx]
                        child_coord = coords_float[child_idx]
                        bone_coords.append([parent_coord, child_coord])
                    except:
                        return None
            try:
                bone_coords = torch.stack(
                    [torch.stack(pair, dim=0) for pair in bone_coords],
                    dim=0
                )  # shape: (num_bones, 2, 3)
            except:
                return None
            bones_coors_list.append(bone_coords)
            num_bones_list.append(bone_coords.size(0))
                
        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, :, :] = bones_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

        if self.args.seq_shuffle:
            num_joint_token = self.max_length - 2 - self.cond_length  # During inference, this is the total length to generate
            num_joints = num_joint_token // self.bone_per_token
            target_aware_pos_embed = self.target_aware_pos_embed.repeat(self.args.batchsize_per_gpu, 1, 1)  # [B, max_joint, embed_dim]
            
            cond_pos_embed = target_aware_pos_embed[:, 0:1, :]  # [B, 1, embed_dim]
            cond_pos_embed = cond_pos_embed.repeat(1, self.cond_length, 1)  # [B, cond_length, embed_dim]
            bone_pos_embed = target_aware_pos_embed[:, 1:num_joints, :]  # [B, num_joints-1, embed_dim]
            bone_pos_embed_expanded = bone_pos_embed.unsqueeze(2).repeat(1, 1, self.bone_per_token, 1)  # [B, num_joints-1, joint_per_token, embed_dim]
            bone_pos_embed_expanded = bone_pos_embed_expanded.view(self.args.batchsize_per_gpu, num_joint_token-self.bone_per_token, self.feat_dim)
            processed_point_feature += cond_pos_embed
        else:
            bone_pos_embed_expanded = None

        # 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,
                target_pos_embed=bone_pos_embed_expanded
            )
        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
        
        if self.joint_token:
            gen_joints = self.detokenize_joint_token(outputs)
        else:
            gen_joints = self.detokenize(outputs)

        return gen_joints