Update app.py

app.py

@@ -5,13 +5,11 @@ import imageio
 import math
 from math import ceil
 import matplotlib.pyplot as plt
-import matplotlib.animation as animation
 import numpy as np
 from PIL import Image
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from torch.autograd import Function
 
 
 class RelationModuleMultiScale(torch.nn.Module):
@@ -301,7 +299,7 @@ def MyPlot(frame_id, src_orig, tar_orig, src_recon, tar_recon, src_Zt, tar_Zt, s
     
 
 # == Load Model ==
-model = TransferVAE_Video(opt)
+model = TransferVAE_Video()
 model.load_state_dict(torch.load('TransferVAE.pth.tar', map_location=torch.device('cpu'))['state_dict'])
 model.eval()
     
@@ -393,24 +391,24 @@ def run(domain_source, action_source, hair_source, top_source, bottom_source, do
         tar_recon = tar_recon_sample[frame, :, :, :].detach().numpy().transpose((1, 2, 0))
     
         # Zt
-        f_expand_src = 0 * src_f_post.unsqueeze(1).expand(-1, 8, opt.f_dim)
+        f_expand_src = 0 * src_f_post.unsqueeze(1).expand(-1, 8, 512)
         zf_src = torch.cat((src_z_post, f_expand_src), dim=2)
         recon_x_src = model.decoder_frame(zf_src)
         src_Zt = recon_x_src.squeeze()[frame, :, :, :].detach().numpy().transpose((1, 2, 0))
     
-        f_expand_tar = 0 * tar_f_post.unsqueeze(1).expand(-1, 8, opt.f_dim)
-        zf_tar = torch.cat((tar_z_post, f_expand_tar), dim=2) # batch,frames,(z_dim+f_dim)
+        f_expand_tar = 0 * tar_f_post.unsqueeze(1).expand(-1, 8, 512)
+        zf_tar = torch.cat((tar_z_post, f_expand_tar), dim=2)
         recon_x_tar = model.decoder_frame(zf_tar)
         tar_Zt = recon_x_tar.squeeze()[frame, :, :, :].detach().numpy().transpose((1, 2, 0))
     
         # Zf_Zt
-        f_expand_src = src_f_post.unsqueeze(1).expand(-1, 8, opt.f_dim)
-        zf_srcZf_tarZt = torch.cat((tar_z_post, f_expand_src), dim=2) # batch,frames,(z_dim+f_dim)
+        f_expand_src = src_f_post.unsqueeze(1).expand(-1, 8, 512)
+        zf_srcZf_tarZt = torch.cat((tar_z_post, f_expand_src), dim=2)
         recon_x_srcZf_tarZt = model.decoder_frame(zf_srcZf_tarZt)
         src_Zf_tar_Zt = recon_x_srcZf_tarZt.squeeze()[frame, :, :, :].detach().numpy().transpose((1, 2, 0))
     
-        f_expand_tar = tar_f_post.unsqueeze(1).expand(-1, 8, opt.f_dim)
-        zf_tarZf_srcZt = torch.cat((src_z_post, f_expand_tar), dim=2) # batch,frames,(z_dim+f_dim)
+        f_expand_tar = tar_f_post.unsqueeze(1).expand(-1, 8, 512)
+        zf_tarZf_srcZt = torch.cat((src_z_post, f_expand_tar), dim=2)
         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))