skinning/skinning_models/models.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, einsum
import torch.nn.functional as F

from skinning_models.networks import PreNorm, FeedForward, Attention, PointEmbed
from third_partys.Michelangelo.encode import load_model
from third_partys.PartField.encode import partfield

class SkinningBlock(nn.Module):
    def __init__(self, dim, args, heads=8, dim_head=64, ff_mult=4):
        super().__init__()
        # -- Self-Attn (joint)
        self.self_attn_j   = PreNorm(dim, Attention(dim, dim, heads=heads, dim_head=dim_head), context_dim = dim)
        self.ff_j          = PreNorm(dim, FeedForward(dim, mult=ff_mult))

        # -- Cross-Attn (point -> shape)
        self.cross_attn_ps = PreNorm(dim, Attention(dim, dim, heads=heads, dim_head=dim_head), context_dim = dim)
        self.ff_ps         = PreNorm(dim, FeedForward(dim, mult=ff_mult))
        
        # -- Cross-Attn (joint -> shape)
        self.cross_attn_js = PreNorm(dim, Attention(dim, dim, heads=heads, dim_head=dim_head), context_dim = dim)
        self.ff_js         = PreNorm(dim, FeedForward(dim, mult=ff_mult))
        
        # -- Cross-Attn (joint -> point)
        self.cross_attn_jp = PreNorm(dim, Attention(dim, dim, heads=heads, dim_head=dim_head), context_dim = dim)
        self.ff_jp         = PreNorm(dim, FeedForward(dim, mult=ff_mult))

        # -- Cross-Attn (point -> joint)
        self.cross_attn_pj = PreNorm(dim, Attention(dim, dim, heads=heads, dim_head=dim_head), context_dim = dim)
        self.ff_pj         = PreNorm(dim, FeedForward(dim, mult=ff_mult))

        self.use_TAJA = args.use_TAJA
        if self.use_TAJA:
            self.rel_pos_embedding = nn.Embedding(10, dim // 4)  
            self.rel_pos_proj = nn.Linear(dim // 4, heads)
            self.rel_pos_scale = nn.Parameter(torch.ones(1) * 0.1)  # Initial value 0.1
 
    def forward(self, point, joint, shape, valid_mask=None, graph_dist=None):
        """
        point: (B, Np, dim)
        joint: (B, Nj, dim)
        shape: (B, Ns, dim)
        valid_mask: (B, Nj) or None
        return:
            updated_point, updated_joint
        """
        
        # 1) joint self-attention with TAJA
        if self.use_TAJA:
            batch_size, n_joints = joint.shape[0], joint.shape[1]
            
            dist_mask = valid_mask.unsqueeze(1) & valid_mask.unsqueeze(2) 
            safe_dist = torch.where(dist_mask, graph_dist, torch.zeros_like(graph_dist))
            distances_clamped = torch.clamp(safe_dist, 0, 9).long()  # (B, Nj, Nj)
     
            rel_pos_embeddings = self.rel_pos_embedding(distances_clamped)  # (B, Nj, Nj, dim//4)
            rel_pos_encoding = self.rel_pos_proj(rel_pos_embeddings) * self.rel_pos_scale  # (B, Nj, Nj, heads)
            
            rel_pos_encoding = torch.where(
                dist_mask.unsqueeze(-1).expand_as(rel_pos_encoding),
                rel_pos_encoding,
                torch.zeros_like(rel_pos_encoding)
            ) 
        else:
            rel_pos_encoding = None
        
        joint_enhance = self.self_attn_j(joint, context=joint, context_mask=valid_mask, rel_pos=rel_pos_encoding) + joint
        joint_enhance = self.ff_j(joint_enhance) + joint_enhance

        # 2) point->shape
        point_context = self.cross_attn_ps(point, context=shape) + point
        point_context = self.ff_ps(point_context) + point_context
        
        # 2) joint->shape
        joint_context = self.cross_attn_js(joint_enhance, context=shape, query_mask=valid_mask) + joint_enhance
        joint_context = self.ff_js(joint_context) + joint_context
        
        # 3) joint->point
        joint_refine = self.cross_attn_jp(joint_context, context=point_context, query_mask=valid_mask) + joint_context
        joint_refine = self.ff_jp(joint_refine) + joint_refine
        
        # 4) point->joint
        point_final = self.cross_attn_pj(point_context, context=joint_refine, context_mask=valid_mask) + point_context
        point_final = self.ff_pj(point_final) + point_final
        
        return point_final, joint_refine

class SkinningNetStacked(nn.Module):
    def __init__(self, args, dim=768, heads=8, dim_head=64, ff_mult=4, scale_init=1.):
        super().__init__()
        self.args = args
        self.max_joints = args.max_joints
        
        self.skeleton_condition = PointEmbed(dim=dim)
        self.scale = nn.Parameter(torch.tensor(scale_init), requires_grad=True)

        self.point_encoder = load_model()
        self.point_encoder.eval()
        for param in self.point_encoder.parameters():
            param.requires_grad = False
    
        self.point_embed_pe = PointEmbed(dim=dim)
        self.point_embed = partfield()
        self.proj = nn.Sequential(
                nn.Linear(448, dim),
                nn.SiLU(),
                nn.Linear(dim, dim),
            )
    
        # multiple blocks
        self.blocks = nn.ModuleList([
            SkinningBlock(dim, args, heads=heads, dim_head=dim_head, ff_mult=ff_mult)
            for _ in range(args.depth)
        ])

    def process_point_feature(self, point_feature):
        shape_latents = self.point_encoder.to_shape_latents(point_feature[:, 1:])
        point_feature_first_column = point_feature[:, 0:1]
        encode_feature = torch.cat([point_feature_first_column, shape_latents], dim=1)
        return encode_feature
    
    def forward(self,
                sample_points, skeleton, pc_w_norm, dist_graph, 
                valid_mask=None,  # (B, Nj)
                target=None):         # (B, Np, Nj)
        """
        cosine similarity + softmax + loss
        """
       
        point_out1 = self.point_embed(sample_points) 
        point_out1 = self.proj(point_out1)
        point_out2 = self.point_embed_pe(sample_points)
        point_out = point_out1 + point_out2 # (bs, 8192, 768)

        joint_out = self.skeleton_condition(skeleton) # (bs, args.max_joints, 768)

        point_feature = self.point_encoder.encode_latents(pc_w_norm)
        shape_feature = self.process_point_feature(point_feature=point_feature) # （bs, 257, 768）

        for block in self.blocks:
            point_out, joint_out = block(point_out, joint_out, shape_feature, valid_mask=valid_mask, graph_dist=dist_graph)

        point_norm = F.normalize(point_out, p=2, dim=-1)  # (B, Np, D)
        joint_norm = F.normalize(joint_out, p=2, dim=-1)  # (B, Nj, D)

        score_cos = einsum('b i d, b j d -> b i j', point_norm, joint_norm)
        score = (self.scale.abs()+1e-9) * score_cos
        
        skinning_weight = F.softmax(score, dim=-1) 
    
        if target is None:
            return skinning_weight