Update gradio_tabs/img_edit.py

gradio_tabs/img_edit.py

@@ -42,6 +42,7 @@ def load_image(img, size):
 	w, h = img.size
 	img = img.resize((size, size))
 	img = np.asarray(img)
+	img = np.copy(img)
 	img = np.transpose(img, (2, 0, 1))	# 3 x 256 x 256
 
 	return img / 255.0, w, h
@@ -55,40 +56,76 @@ 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),
-	])
+# def resize(img, size):
+# 	transform = torchvision.transforms.Compose([
+# 		torchvision.transforms.Resize((size,size), antialias=True),
+# 	])
 
-	return transform(img)
+# 	return transform(img)
 
 
-def resize_back(img, w, h):
-	transform = torchvision.transforms.Compose([
-		torchvision.transforms.Resize((h, w), antialias=True),
-	])
+# def resize_back(img, w, h):
+# 	transform = torchvision.transforms.Compose([
+# 		torchvision.transforms.Resize((h, w), antialias=True),
+# 	])
+
+# 	return transform(img)
+
+# Pre-compile resize transforms for better performance
+resize_transform_cache = {}
+
+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(img, size):
+	"""Use cached resize transform"""
+	transform = get_resize_transform((size, size))
 	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()
+	img = img.clamp(-1, 1)
 	img = (img - img.min()) / (img.max() - img.min())
 
 	return img
 
 
-def img_postprocessing(image, w, h):
+# 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
+
 
-	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)
+def img_postprocessing(img, w, h):
 
-	with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_file:
-		imageio.imwrite(temp_file.name, img_output, quality=8)
-		return temp_file.name
+	img = resize_back(img, w, h)
+	img = img_denorm(img)
+	img = img.squeeze(0).permute(1, 2, 0).contiguous()	# contiguous() for fast transfer
+	img_output = (img.cpu().numpy() * 255).astype(np.uint8)
 
+	return img_output
+    
 
 def img_edit(gen, device):