Update gradio_tabs/img_edit.py

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):