add

app.py

@@ -172,84 +172,7 @@ class BaseTrainer(object):
             self.args.vae_layer = 4
             self.vq_model_lower = getattr(vq_model_module, "VQVAEConvZero")(self.args).to(self.rank)
             other_tools.load_checkpoints(self.vq_model_lower, args.vqvae_lower_path, args.e_name)
-            
-        elif vq_type=="rvqvae":
-
-            args.num_quantizers = 6
-            args.shared_codebook =  False
-            args.quantize_dropout_prob = 0.2
-            args.mu = 0.99
-
-            args.nb_code = 512
-            args.code_dim = 512
-            args.code_dim = 512
-            args.down_t = 2
-            args.stride_t = 2
-            args.width = 512
-            args.depth = 3
-            args.dilation_growth_rate = 3
-            args.vq_act = "relu"
-            args.vq_norm = None
-                                
-            dim_pose = 78  
-            args.body_part = "upper"
-            self.vq_model_upper = RVQVAE(args,
-                                dim_pose,
-                                args.nb_code,
-                                args.code_dim,
-                                args.code_dim,
-                                args.down_t,
-                                args.stride_t,
-                                args.width,
-                                args.depth,
-                                args.dilation_growth_rate,
-                                args.vq_act,
-                                args.vq_norm)
-
-            dim_pose = 180
-            args.body_part = "hands"
-            self.vq_model_hands = RVQVAE(args,
-                                dim_pose,
-                                args.nb_code,
-                                args.code_dim,
-                                args.code_dim,
-                                args.down_t,
-                                args.stride_t,
-                                args.width,
-                                args.depth,
-                                args.dilation_growth_rate,
-                                args.vq_act,
-                                args.vq_norm)
-
-            dim_pose = 54
-            if args.use_trans:
-                dim_pose = 57
-                self.args.vqvae_lower_path = self.args.vqvae_lower_trans_path
-            args.body_part = "lower"
-            self.vq_model_lower = RVQVAE(args,
-                                dim_pose,
-                                args.nb_code,
-                                args.code_dim,
-                                args.code_dim,
-                                args.down_t,
-                                args.stride_t,
-                                args.width,
-                                args.depth,
-                                args.dilation_growth_rate,
-                                args.vq_act,
-                                args.vq_norm)
-                    
-            self.vq_model_upper.load_state_dict(torch.load(self.args.vqvae_upper_path)['net'])
-            self.vq_model_hands.load_state_dict(torch.load(self.args.vqvae_hands_path)['net'])
-            self.vq_model_lower.load_state_dict(torch.load(self.args.vqvae_lower_path)['net'])
-            
-            self.vqvae_latent_scale = self.args.vqvae_latent_scale 
 
-            self.vq_model_upper.eval().to(self.rank)
-            self.vq_model_hands.eval().to(self.rank)
-            self.vq_model_lower.eval().to(self.rank)
-            
-        
         
         
         
@@ -260,10 +183,7 @@ class BaseTrainer(object):
         self.args.vae_length = 240
 
 
-        # self.vq_model_face.eval()
-        self.vq_model_upper.eval()
-        self.vq_model_hands.eval()
-        self.vq_model_lower.eval()
+
 
         self.cls_loss = nn.NLLLoss().to(self.rank)
         self.reclatent_loss = nn.MSELoss().to(self.rank)
@@ -609,8 +529,91 @@ class BaseTrainer(object):
             'rec_exps': rec_exps,
         }
 
+
+    def _create_cuda_model(self):
+        args = self.args
+        args.num_quantizers = 6
+        args.shared_codebook =  False
+        args.quantize_dropout_prob = 0.2
+        args.mu = 0.99
+
+        args.nb_code = 512
+        args.code_dim = 512
+        args.code_dim = 512
+        args.down_t = 2
+        args.stride_t = 2
+        args.width = 512
+        args.depth = 3
+        args.dilation_growth_rate = 3
+        args.vq_act = "relu"
+        args.vq_norm = None
+                            
+        dim_pose = 78  
+        args.body_part = "upper"
+        self.vq_model_upper = RVQVAE(args,
+                            dim_pose,
+                            args.nb_code,
+                            args.code_dim,
+                            args.code_dim,
+                            args.down_t,
+                            args.stride_t,
+                            args.width,
+                            args.depth,
+                            args.dilation_growth_rate,
+                            args.vq_act,
+                            args.vq_norm)
+
+        dim_pose = 180
+        args.body_part = "hands"
+        self.vq_model_hands = RVQVAE(args,
+                            dim_pose,
+                            args.nb_code,
+                            args.code_dim,
+                            args.code_dim,
+                            args.down_t,
+                            args.stride_t,
+                            args.width,
+                            args.depth,
+                            args.dilation_growth_rate,
+                            args.vq_act,
+                            args.vq_norm)
+
+        dim_pose = 54
+        if args.use_trans:
+            dim_pose = 57
+            self.args.vqvae_lower_path = self.args.vqvae_lower_trans_path
+        args.body_part = "lower"
+        self.vq_model_lower = RVQVAE(args,
+                            dim_pose,
+                            args.nb_code,
+                            args.code_dim,
+                            args.code_dim,
+                            args.down_t,
+                            args.stride_t,
+                            args.width,
+                            args.depth,
+                            args.dilation_growth_rate,
+                            args.vq_act,
+                            args.vq_norm)
+                
+        self.vq_model_upper.load_state_dict(torch.load(self.args.vqvae_upper_path)['net'])
+        self.vq_model_hands.load_state_dict(torch.load(self.args.vqvae_hands_path)['net'])
+        self.vq_model_lower.load_state_dict(torch.load(self.args.vqvae_lower_path)['net'])
+        
+        self.vqvae_latent_scale = self.args.vqvae_latent_scale 
+
+        self.vq_model_upper.eval().to(self.rank)
+        self.vq_model_hands.eval().to(self.rank)
+        self.vq_model_lower.eval().to(self.rank)
+            
+        self.model = self.model.cuda()
+        self.model.eval()
+        
     @spaces.GPU(duration=149)
     def _warp(self, batch_data):
+        self._create_cuda_model()
+        
+        
         loaded_data = self._load_data(batch_data)    
         net_out = self._g_test(loaded_data)
         return net_out
@@ -634,8 +637,6 @@ class BaseTrainer(object):
         latent_ori = []
         l2_all = 0 
         lvel = 0
-        self.model = self.model.cuda()
-        self.model.eval()
         # self.eval_copy.eval()
         with torch.no_grad():
             for its, batch_data in enumerate(self.test_loader):

models/vq/quantizer.py

@@ -44,7 +44,7 @@ class QuantizeEMAReset(nn.Module):
         self.init = False
         self.code_sum = None
         self.code_count = None
-        self.register_buffer('codebook', torch.zeros(self.nb_code, self.code_dim, requires_grad=False))
+        self.register_buffer('codebook', torch.zeros(self.nb_code, self.code_dim, requires_grad=False).cuda())
 
     def _tile(self, x):
         nb_code_x, code_dim = x.shape