147 lines
4.5 KiB
Python
147 lines
4.5 KiB
Python
import io
|
|
from typing import List, Tuple
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
|
|
try:
|
|
import tinify
|
|
except Exception: # pragma: no cover
|
|
tinify = None
|
|
|
|
|
|
def _tensor_to_pil(image_tensor: torch.Tensor) -> Image.Image:
|
|
image_np = image_tensor.detach().cpu().numpy()
|
|
image_np = np.clip(image_np * 255.0, 0, 255).astype(np.uint8)
|
|
|
|
if image_np.ndim != 3:
|
|
raise ValueError(f"Expected HWC image tensor, got shape {image_np.shape}.")
|
|
|
|
channels = image_np.shape[2]
|
|
if channels == 1:
|
|
return Image.fromarray(image_np[:, :, 0], mode="L")
|
|
if channels == 3:
|
|
return Image.fromarray(image_np, mode="RGB")
|
|
if channels == 4:
|
|
return Image.fromarray(image_np, mode="RGBA")
|
|
|
|
raise ValueError(f"Unsupported channel count: {channels}")
|
|
|
|
|
|
def _pil_to_tensor(image: Image.Image) -> torch.Tensor:
|
|
# Keep alpha for true alpha modes and palette PNGs with transparency metadata.
|
|
has_alpha = "A" in image.getbands() or "transparency" in image.info
|
|
if has_alpha:
|
|
out = image.convert("RGBA")
|
|
else:
|
|
out = image.convert("RGB")
|
|
image_np = np.asarray(out).astype(np.float32) / 255.0
|
|
return torch.from_numpy(image_np)
|
|
|
|
|
|
def _tiny_resize(source, method: str, width: int, height: int):
|
|
if method == "scale":
|
|
options = {"method": method}
|
|
if width > 0:
|
|
options["width"] = width
|
|
elif height > 0:
|
|
options["height"] = height
|
|
else:
|
|
return source
|
|
return source.resize(**options)
|
|
|
|
if width <= 0 or height <= 0:
|
|
return source
|
|
|
|
return source.resize(method=method, width=width, height=height)
|
|
|
|
|
|
class TinyPNGCompressNode:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"api_key": ("STRING", {"default": "", "multiline": False}),
|
|
"operation": (["compress", "smart_resize"], {"default": "compress"}),
|
|
"resize_method": (["scale", "fit", "cover", "thumb"], {"default": "scale"}),
|
|
"width": ("INT", {"default": 0, "min": 0, "max": 16384, "step": 1}),
|
|
"height": ("INT", {"default": 0, "min": 0, "max": 16384, "step": 1}),
|
|
"preserve_copyright": ("BOOLEAN", {"default": False}),
|
|
"preserve_creation": ("BOOLEAN", {"default": False}),
|
|
"preserve_location": ("BOOLEAN", {"default": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "STRING")
|
|
RETURN_NAMES = ("image", "info")
|
|
FUNCTION = "run"
|
|
CATEGORY = "image/TinyPNG"
|
|
|
|
def run(
|
|
self,
|
|
image: torch.Tensor,
|
|
api_key: str,
|
|
operation: str,
|
|
resize_method: str,
|
|
width: int,
|
|
height: int,
|
|
preserve_copyright: bool,
|
|
preserve_creation: bool,
|
|
preserve_location: bool,
|
|
) -> Tuple[torch.Tensor, str]:
|
|
if tinify is None:
|
|
raise RuntimeError("Missing dependency: tinify. Please `pip install -r requirements.txt`.")
|
|
|
|
key = (api_key or "").strip()
|
|
if not key:
|
|
raise ValueError("TinyPNG API key is required.")
|
|
|
|
tinify.key = key
|
|
|
|
preserve_items: List[str] = []
|
|
if preserve_copyright:
|
|
preserve_items.append("copyright")
|
|
if preserve_creation:
|
|
preserve_items.append("creation")
|
|
if preserve_location:
|
|
preserve_items.append("location")
|
|
|
|
output_batch = []
|
|
total_in = 0
|
|
total_out = 0
|
|
|
|
for idx in range(image.shape[0]):
|
|
pil_image = _tensor_to_pil(image[idx])
|
|
raw_buf = io.BytesIO()
|
|
pil_image.save(raw_buf, format="PNG")
|
|
input_bytes = raw_buf.getvalue()
|
|
total_in += len(input_bytes)
|
|
|
|
source = tinify.from_buffer(input_bytes)
|
|
if preserve_items:
|
|
source = source.preserve(*preserve_items)
|
|
|
|
if operation == "smart_resize":
|
|
source = _tiny_resize(source, resize_method, width, height)
|
|
|
|
output_bytes = source.to_buffer()
|
|
total_out += len(output_bytes)
|
|
|
|
output_pil = Image.open(io.BytesIO(output_bytes))
|
|
output_batch.append(_pil_to_tensor(output_pil))
|
|
|
|
output_tensor = torch.stack(output_batch, dim=0)
|
|
|
|
ratio = 0.0
|
|
if total_in > 0:
|
|
ratio = 1.0 - (total_out / float(total_in))
|
|
|
|
info = (
|
|
f"batch={image.shape[0]}, input={total_in} bytes, output={total_out} bytes, "
|
|
f"saved={ratio * 100.0:.2f}%"
|
|
)
|
|
|
|
return (output_tensor, info)
|