Update app.py

app.py

@@ -13,10 +13,6 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 
-def log_gpu_memory():
-    print(subprocess.check_output('nvidia-smi').decode('utf-8'))
-    
-
 class RelationModuleMultiScale(torch.nn.Module):
 
     def __init__(self, img_feature_dim, num_bottleneck, num_frames):
@@ -129,42 +125,19 @@ class TransferVAE_Video(nn.Module):
     
           
     def encode_and_sample_post(self, x):
-        if isinstance(x, list):
-            conv_x = self.encoder_frame(x[0])
-        else:
-            conv_x = self.encoder_frame(x)
-
+        conv_x = self.encoder_frame(x)
         lstm_out, _ = self.z_lstm(conv_x)
-
         backward = lstm_out[:, 0, self.hidden_dim:2 * self.hidden_dim]
         frontal = lstm_out[:, self.frames - 1, 0:self.hidden_dim]
         lstm_out_f = torch.cat((frontal, backward), dim=1)
         f_mean = self.f_mean(lstm_out_f)
         f_logvar = self.f_logvar(lstm_out_f)
         f_post = self.reparameterize(f_mean, f_logvar, random_sampling=False)
-
         features, _ = self.z_rnn(lstm_out)
         z_mean = self.z_mean(features)
         z_logvar = self.z_logvar(features)
         z_post = self.reparameterize(z_mean, z_logvar, random_sampling=False)
-
-        if isinstance(x, list):
-            f_mean_list = [f_mean]
-            f_post_list = [f_post]
-            for t in range(1,3,1):
-                conv_x = self.encoder_frame(x[t])
-                lstm_out, _ = self.z_lstm(conv_x)
-                backward = lstm_out[:, 0, self.hidden_dim:2 * self.hidden_dim]
-                frontal = lstm_out[:, self.frames - 1, 0:self.hidden_dim]
-                lstm_out_f = torch.cat((frontal, backward), dim=1)
-                f_mean = self.f_mean(lstm_out_f)
-                f_logvar = self.f_logvar(lstm_out_f)
-                f_post = self.reparameterize(f_mean, f_logvar, random_sampling=False)
-                f_mean_list.append(f_mean)
-                f_post_list.append(f_post)
-            f_mean = f_mean_list
-            f_post = f_post_list
-        return f_mean, f_logvar, f_post, z_mean, z_logvar, z_post
+        return f_post, z_post
 
     
     def decoder_frame(self,zf):
@@ -190,7 +163,7 @@ class TransferVAE_Video(nn.Module):
 
 
     def forward(self, x, beta):
-        _, _, f_post, _, _, z_post = self.encode_and_sample_post(x)
+        f_post, z_post = self.encode_and_sample_post(x)
         if isinstance(f_post, list):
             f_expand = f_post[0].unsqueeze(1).expand(-1, self.frames, self.f_dim)
         else:
@@ -269,15 +242,11 @@ def MyPlot(frame_id, src_orig, tar_orig, src_recon, tar_recon, src_Zt, tar_Zt, s
     plt.savefig(save_name, dpi=200, format='png', bbox_inches='tight', pad_inches=0.0)
     
 
-log_gpu_memory()
-
 # == Load Model ==
 model = TransferVAE_Video()
 model.load_state_dict(torch.load('TransferVAE.pth.tar', map_location=torch.device('cpu'))['state_dict'])
 model.eval()
 
-log_gpu_memory()
-
   
 def run(source, action_source, hair_source, top_source, bottom_source, target, action_target, hair_target, top_target, bottom_target):
 
@@ -338,12 +307,11 @@ def run(source, action_source, hair_source, top_source, bottom_source, target, a
     images_target = name2seq(file_name_target)
     x = torch.cat((images_source, images_target), dim=0)
     
-    
-    log_gpu_memory()
+ 
     # == Forward ==
     with torch.no_grad():
         f_post, z_post, recon_x = model(x, [0]*3)
-    log_gpu_memory()
+    
     
     src_orig_sample = x[0, :, :, :, :]
     src_recon_sample = recon_x[0, :, :, :, :]
@@ -389,9 +357,7 @@ def run(source, action_source, hair_source, top_source, bottom_source, target, a
         recon_x_tarZf_srcZt = model.decoder_frame(zf_tarZf_srcZt)
         tar_Zf_src_Zt = recon_x_tarZf_srcZt.squeeze()[frame, :, :, :].detach().numpy().transpose((1, 2, 0))
 
-        log_gpu_memory()
         MyPlot(frame, src_orig, tar_orig, src_recon, tar_recon, src_Zt, tar_Zt, src_Zf_tar_Zt, tar_Zf_src_Zt)
-        log_gpu_memory()
         
     a = concat('MyPlot_')