Merge branch 'main' of hf.co:spaces/jbilcke-hf/LIA-X-testing

gradio_tabs/img_edit.py

@@ -55,21 +55,31 @@ def img_preprocessing(img_path, size):
 	return imgs_norm, w, h
 
 
-def resize(img, size):
-	transform = torchvision.transforms.Compose([
-		torchvision.transforms.Resize((size,size), antialias=True),
-	])
+# Pre-compile resize transforms for better performance
+resize_transform_cache = {}
 
-	return transform(img)
+def get_resize_transform(size):
+    """Get cached resize transform - creates once, reuses many times"""
+    if size not in resize_transform_cache:
+        # Only create the transform if it doesn't exist in cache
+        resize_transform_cache[size] = torchvision.transforms.Resize(
+            size, 
+            interpolation=torchvision.transforms.InterpolationMode.BILINEAR,
+            antialias=True
+        )
+    return resize_transform_cache[size]
 
 
-def resize_back(img, w, h):
-	transform = torchvision.transforms.Compose([
-		torchvision.transforms.Resize((h, w), antialias=True),
-	])
+def resize(img, size):
+    """Use cached resize transform"""
+    transform = get_resize_transform((size, size))
+    return transform(img)
 
-	return transform(img)
 
+def resize_back(img, w, h):
+    """Use cached resize transform for back operation"""
+    transform = get_resize_transform((h, w))
+    return transform(img)
 
 def img_denorm(img):
 	img = img.clamp(-1, 1).cpu()
@@ -78,17 +88,23 @@ def img_denorm(img):
 	return img
 
 
-def img_postprocessing(image, w, h):
-
-	image = resize_back(image, w, h)
-	image = image.permute(0, 2, 3, 1)
-	edited_image = img_denorm(image)
-	img_output = (edited_image[0].numpy() * 255).astype(np.uint8)
-
-	with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_file:
-		imageio.imwrite(temp_file.name, img_output, quality=8)
-		return temp_file.name
-
+def img_postprocessing(img, w, h):
+    # Resize on GPU (using cached transform)
+    image = resize_back(image, w, h)
+    
+    # Denormalize ON GPU (avoid early CPU transfer)
+    image = image.clamp(-1, 1)  # Still on GPU
+    image = (image - image.min()) / (image.max() - image.min())  # Still on GPU
+    
+    # Single optimized CPU transfer
+    image = image.squeeze(0).permute(1, 2, 0).contiguous()  # contiguous() for fast transfer
+    img_output = (image.cpu().numpy() * 255).astype(np.uint8)  # Single CPU transfer
+    
+    # Use PIL directly (faster than imageio)
+    pil_image = Image.fromarray(img_output)
+
+    # return the PIL image directly
+    return pil_image
 
 def img_edit(gen, device):
 

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,9 +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')
+		# 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")
 
@@ -59,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, :, :, :]
@@ -77,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
@@ -114,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, :, :, :]
@@ -133,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
@@ -141,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
@@ -170,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
@@ -222,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