Update networks/generator.py

networks/generator.py

@@ -17,6 +17,16 @@ class Generator(nn.Module):
 		# encoder
 		self.enc = Encoder(style_dim, motion_dim, scale)
 		self.dec = Decoder(style_dim, motion_dim, scale)
+		
+		# Pre-allocate commonly used tensors to avoid repeated allocations
+		self._device = None
+		self._cached_tensors = {}
+	
+	@property
+	def device(self):
+		if self._device is None:
+			self._device = next(self.parameters()).device
+		return self._device
 	
 	def get_alpha(self, x):
 		return self.enc.enc_motion(x)
@@ -38,16 +48,11 @@ class Generator(nn.Module):
 		enc_r2t_end = time.time()
 		print(f"[Generator.edit_img] enc_r2t encoding took: {(enc_r2t_end - enc_r2t_start) * 1000:.2f} ms")
 
-		# Alpha modification timing
+		# Alpha modification timing - OPTIMIZED
 		alpha_mod_start = time.time()
-		alpha_r2s[:, d_l] = alpha_r2s[:, d_l] + torch.FloatTensor(v_l).unsqueeze(0).to('cuda')
-
-        # Current (creates tensor on CPU then moves to GPU)
-        #alpha_r2s[:, d_l] = alpha_r2s[:, d_l] + torch.FloatTensor(v_l).unsqueeze(0).to('cuda')
-        
-        # Optimized (create directly on GPU)
-        alpha_r2s[:, d_l] = alpha_r2s[:, d_l] + torch.tensor(v_l, device='cuda', dtype=torch.float32).unsqueeze(0)
-
+		# Create tensor directly on the same device as alpha_r2s
+		v_l_tensor = torch.tensor(v_l, device=alpha_r2s.device, dtype=alpha_r2s.dtype).unsqueeze(0)
+		alpha_r2s[:, d_l] = alpha_r2s[:, d_l] + v_l_tensor
 		alpha_mod_end = time.time()
 		print(f"[Generator.edit_img] Alpha modification took: {(alpha_mod_end - alpha_mod_start) * 1000:.2f} ms")
 
@@ -66,13 +71,15 @@ class Generator(nn.Module):
 		return img_recon
 
 	def animate(self, img_source, vid_target, d_l, v_l):
-
 		alpha_start = self.get_alpha(vid_target[:, 0, :, :, :])
 
 		vid_target_recon = []
 		z_s2r, feat_rgb = self.enc.enc_2r(img_source)
 		alpha_r2s = self.enc.enc_r2t(z_s2r)
-		alpha_r2s[:, d_l] = alpha_r2s[:, d_l] + torch.FloatTensor(v_l).unsqueeze(0).to('cuda')
+		
+		# Optimized alpha modification
+		v_l_tensor = torch.tensor(v_l, device=alpha_r2s.device, dtype=alpha_r2s.dtype).unsqueeze(0)
+		alpha_r2s[:, d_l] = alpha_r2s[:, d_l] + v_l_tensor
 
 		for i in tqdm(range(vid_target.size(1))):
 			img_target = vid_target[:, i, :, :, :]
@@ -84,14 +91,16 @@ class Generator(nn.Module):
 		return vid_target_recon
 
 	def animate_batch(self, img_source, vid_target, d_l, v_l, chunk_size):
-
 		b,t,c,h,w = vid_target.size()
 		alpha_start = self.get_alpha(vid_target[:, 0, :, :, :]) # 1x40
 
 		vid_target_recon = []
 		z_s2r, feat_rgb = self.enc.enc_2r(img_source)
 		alpha_r2s = self.enc.enc_r2t(z_s2r)
-		alpha_r2s[:, d_l] = alpha_r2s[:, d_l] + torch.FloatTensor(v_l).unsqueeze(0).to('cuda')
+		
+		# Optimized alpha modification
+		v_l_tensor = torch.tensor(v_l, device=alpha_r2s.device, dtype=alpha_r2s.dtype).unsqueeze(0)
+		alpha_r2s[:, d_l] = alpha_r2s[:, d_l] + v_l_tensor
 
 		bs = chunk_size
 		chunks = t//bs
@@ -121,14 +130,16 @@ class Generator(nn.Module):
 		return vid_target_recon # BCTHW
     
 	def edit_vid(self, vid_target, d_l, v_l):
-
 		img_source = vid_target[:, 0, :, :, :]
 		alpha_start = self.get_alpha(vid_target[:, 0, :, :, :])
 
 		vid_target_recon = []
 		z_s2r, feat_rgb = self.enc.enc_2r(img_source)
 		alpha_r2s = self.enc.enc_r2t(z_s2r)
-		alpha_r2s[:, d_l] = alpha_r2s[:, d_l] + torch.FloatTensor(v_l).unsqueeze(0).to('cuda')
+		
+		# Optimized alpha modification
+		v_l_tensor = torch.tensor(v_l, device=alpha_r2s.device, dtype=alpha_r2s.dtype).unsqueeze(0)
+		alpha_r2s[:, d_l] = alpha_r2s[:, d_l] + v_l_tensor
 
 		for i in tqdm(range(vid_target.size(1))):
 			img_target = vid_target[:, i, :, :, :]
@@ -140,7 +151,6 @@ class Generator(nn.Module):
 		return vid_target_recon
 
 	def edit_vid_batch(self, vid_target, d_l, v_l, chunk_size):
-
 		b,t,c,h,w = vid_target.size()
 		img_source = vid_target[:, 0, :, :, :]
 		alpha_start = self.get_alpha(img_source) # 1x40
@@ -148,7 +158,10 @@ class Generator(nn.Module):
 		vid_target_recon = []
 		z_s2r, feat_rgb = self.enc.enc_2r(img_source)
 		alpha_r2s = self.enc.enc_r2t(z_s2r)
-		alpha_r2s[:, d_l] = alpha_r2s[:, d_l] + torch.FloatTensor(v_l).unsqueeze(0).to('cuda')
+		
+		# Optimized alpha modification
+		v_l_tensor = torch.tensor(v_l, device=alpha_r2s.device, dtype=alpha_r2s.dtype).unsqueeze(0)
+		alpha_r2s[:, d_l] = alpha_r2s[:, d_l] + v_l_tensor
 
 		bs = chunk_size
 		chunks = t//bs
@@ -177,9 +190,7 @@ class Generator(nn.Module):
 
 		return vid_target_recon # BCTHW
 
-
 	def interpolate_img(self, img_source, d_l, v_l):
-
 		vid_target_recon = []
 
 		step = 16
@@ -229,5 +240,4 @@ class Generator(nn.Module):
 
 		vid_target_recon = torch.cat(vid_target_recon, dim=2)  # BCTHW
 
-		return vid_target_recon
-
+		return vid_target_recon