Update nets/smplx_body_pixel.py

nets/smplx_body_pixel.py

@@ -1,21 +1,231 @@
-def __init__(self, args, config):
-    self.args = args
-    self.config = config
-    self.global_step = 0
-
-    # Force CPU device
-    self.device = torch.device('cpu')
-
-    self.convert_to_6d = self.config.Data.pose.convert_to_6d
-    self.expression = self.config.Data.pose.expression
-    self.epoch = 0
-    self.init_params()
-    self.num_classes = 4
-    self.audio = True
-    self.composition = self.config.Model.composition
-    self.bh_model = self.config.Model.bh_model
-
-    if self.audio:
-        self.audioencoder = AudioEncoder(in_dim=64, num_hiddens=256, num_residual_layers=2, num_residual_hiddens=256).to(self.device)
-    else:
-        self.audioencoder = None
+import os
+import sys
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torch.nn.functional as F
+from torch.optim.lr_scheduler import StepLR
+
+sys.path.append(os.getcwd())
+
+from nets.layers import *
+from nets.base import TrainWrapperBaseClass
+from nets.spg.gated_pixelcnn_v2 import GatedPixelCNN as pixelcnn
+from nets.spg.vqvae_1d import VQVAE as s2g_body, Wav2VecEncoder, AudioEncoder
+from nets.utils import parse_audio, denormalize
+from data_utils import get_mfcc, get_melspec, get_mfcc_old, get_mfcc_psf, get_mfcc_psf_min, get_mfcc_ta
+from data_utils.lower_body import c_index, c_index_3d, c_index_6d
+from data_utils.utils import smooth_geom, get_mfcc_sepa
+import numpy as np
+from sklearn.preprocessing import normalize
+
+
+class TrainWrapper(TrainWrapperBaseClass):
+    '''
+    a wrapper receiving a batch from data_utils and calculate loss
+    '''
+
+    def __init__(self, args, config):
+        self.args = args
+        self.config = config
+        self.global_step = 0
+
+        # Force CPU device
+        self.device = torch.device('cpu')
+
+        self.convert_to_6d = self.config.Data.pose.convert_to_6d
+        self.expression = self.config.Data.pose.expression
+        self.epoch = 0
+        self.init_params()
+        self.num_classes = 4
+        self.audio = True
+        self.composition = self.config.Model.composition
+        self.bh_model = self.config.Model.bh_model
+
+        if self.audio:
+            self.audioencoder = AudioEncoder(
+                in_dim=64,
+                num_hiddens=256,
+                num_residual_layers=2,
+                num_residual_hiddens=256
+            ).to(self.device)
+        else:
+            self.audioencoder = None
+
+        if self.convert_to_6d:
+            dim, layer = 512, 10
+        else:
+            dim, layer = 256, 15
+
+        self.generator = pixelcnn(2048, dim, layer, self.num_classes, self.audio, self.bh_model).to(self.device)
+        self.g_body = s2g_body(self.each_dim[1], embedding_dim=64, num_embeddings=config.Model.code_num, num_hiddens=1024,
+                               num_residual_layers=2, num_residual_hiddens=512).to(self.device)
+        self.g_hand = s2g_body(self.each_dim[2], embedding_dim=64, num_embeddings=config.Model.code_num, num_hiddens=1024,
+                               num_residual_layers=2, num_residual_hiddens=512).to(self.device)
+
+        model_path = self.config.Model.vq_path
+        model_ckpt = torch.load(model_path, map_location=torch.device('cpu'))
+        self.g_body.load_state_dict(model_ckpt['generator']['g_body'])
+        self.g_hand.load_state_dict(model_ckpt['generator']['g_hand'])
+
+        self.discriminator = None
+        if self.convert_to_6d:
+            self.c_index = c_index_6d
+        else:
+            self.c_index = c_index_3d
+
+        super().__init__(args, config)
+
+    def init_optimizer(self):
+        print('using Adam')
+        self.generator_optimizer = optim.Adam(
+            self.generator.parameters(),
+            lr=self.config.Train.learning_rate.generator_learning_rate,
+            betas=[0.9, 0.999]
+        )
+        if self.audioencoder is not None:
+            opt = self.config.Model.AudioOpt
+            if opt == 'Adam':
+                self.audioencoder_optimizer = optim.Adam(
+                    self.audioencoder.parameters(),
+                    lr=self.config.Train.learning_rate.generator_learning_rate,
+                    betas=[0.9, 0.999]
+                )
+            else:
+                print('using SGD')
+                self.audioencoder_optimizer = optim.SGD(
+                    filter(lambda p: p.requires_grad, self.audioencoder.parameters()),
+                    lr=self.config.Train.learning_rate.generator_learning_rate * 10,
+                    momentum=0.9,
+                    nesterov=False
+                )
+
+    def state_dict(self):
+        return {
+            'generator': self.generator.state_dict(),
+            'generator_optim': self.generator_optimizer.state_dict(),
+            'audioencoder': self.audioencoder.state_dict() if self.audio else None,
+            'audioencoder_optim': self.audioencoder_optimizer.state_dict() if self.audio else None,
+            'discriminator': self.discriminator.state_dict() if self.discriminator else None,
+            'discriminator_optim': self.discriminator_optimizer.state_dict() if self.discriminator else None
+        }
+
+    def load_state_dict(self, state_dict):
+        from collections import OrderedDict
+        new_state_dict = OrderedDict()
+        for k, v in state_dict.items():
+            sub_dict = OrderedDict()
+            if v is not None:
+                for k1, v1 in v.items():
+                    name = k1.replace('module.', '')
+                    sub_dict[name] = v1
+            new_state_dict[k] = sub_dict
+        state_dict = new_state_dict
+
+        if 'generator' in state_dict:
+            self.generator.load_state_dict(state_dict['generator'])
+        else:
+            self.generator.load_state_dict(state_dict)
+
+        if 'generator_optim' in state_dict and self.generator_optimizer is not None:
+            self.generator_optimizer.load_state_dict(state_dict['generator_optim'])
+
+        if self.discriminator is not None:
+            self.discriminator.load_state_dict(state_dict['discriminator'])
+            if 'discriminator_optim' in state_dict and self.discriminator_optimizer is not None:
+                self.discriminator_optimizer.load_state_dict(state_dict['discriminator_optim'])
+
+        if 'audioencoder' in state_dict and self.audioencoder is not None:
+            self.audioencoder.load_state_dict(state_dict['audioencoder'])
+
+    def init_params(self):
+        if self.config.Data.pose.convert_to_6d:
+            scale = 2
+        else:
+            scale = 1
+
+        global_orient = round(0 * scale)
+        leye_pose = reye_pose = round(0 * scale)
+        jaw_pose = round(0 * scale)
+        body_pose = round((63 - 24) * scale)
+        left_hand_pose = right_hand_pose = round(45 * scale)
+        if self.expression:
+            expression = 100
+        else:
+            expression = 0
+
+        b_j = 0
+        jaw_dim = jaw_pose
+        b_e = b_j + jaw_dim
+        eye_dim = leye_pose + reye_pose
+        b_b = b_e + eye_dim
+        body_dim = global_orient + body_pose
+        b_h = b_b + body_dim
+        hand_dim = left_hand_pose + right_hand_pose
+        b_f = b_h + hand_dim
+        face_dim = expression
+
+        self.dim_list = [b_j, b_e, b_b, b_h, b_f]
+        self.full_dim = jaw_dim + eye_dim + body_dim + hand_dim
+        self.pose = int(self.full_dim / round(3 * scale))
+        self.each_dim = [jaw_dim, eye_dim + body_dim, hand_dim, face_dim]
+
+    def __call__(self, bat):
+        self.global_step += 1
+        total_loss = None
+        loss_dict = {}
+
+        aud, poses = bat['aud_feat'].to(self.device).float(), bat['poses'].to(self.device).float()
+        id = bat['speaker'].to(self.device) - 20
+        poses = poses[:, self.c_index, :]
+        aud = aud.permute(0, 2, 1)
+        gt_poses = poses.permute(0, 2, 1)
+
+        with torch.no_grad():
+            self.g_body.eval()
+            self.g_hand.eval()
+            _, body_latents = self.g_body.encode(gt_poses=gt_poses[..., :self.each_dim[1]], id=id)
+            _, hand_latents = self.g_hand.encode(gt_poses=gt_poses[..., self.each_dim[1]:], id=id)
+            latents = torch.cat([body_latents.unsqueeze(-1), hand_latents.unsqueeze(-1)], dim=-1).detach()
+
+        if self.audio:
+            audio = self.audioencoder(aud.transpose(1, 2), frame_num=latents.shape[1]*4).unsqueeze(-1).repeat(1, 1, 1, 2)
+            logits = self.generator(latents, id, audio)
+        else:
+            logits = self.generator(latents, id)
+        logits = logits.permute(0, 2, 3, 1).contiguous()
+
+        self.generator_optimizer.zero_grad()
+        if self.audio:
+            self.audioencoder_optimizer.zero_grad()
+
+        loss = F.cross_entropy(logits.view(-1, logits.shape[-1]), latents.view(-1))
+        loss.backward()
+
+        grad = torch.nn.utils.clip_grad_norm(self.generator.parameters(), self.config.Train.max_gradient_norm)
+
+        loss_dict['grad'] = grad.item()
+        loss_dict['ce_loss'] = loss.item()
+        self.generator_optimizer.step()
+        if self.audio:
+            self.audioencoder_optimizer.step()
+
+        return total_loss, loss_dict
+
+# ----------------------------------------
+# 🚀 NEW SIMPLE WRAPPER CLASS for inference
+# ----------------------------------------
+
+class s2g_body_pixel(nn.Module):
+    def __init__(self, args, config):
+        super().__init__()
+        self.wrapper = TrainWrapper(args, config)
+
+    def infer_on_audio(self, *args, **kwargs):
+        return self.wrapper.infer_on_audio(*args, **kwargs)
+
+    def forward(self, *args, **kwargs):
+        return self.wrapper(*args, **kwargs)
+
+    def load_state_dict(self, *args, **kwargs):
+        return self.wrapper.load_state_dict(*args, **kwargs)