melMass
64b2c72cf4
Supports group but create trimming mask for each group for some reason. I also want a better logic to group batches.
135 lines
3.5 KiB
Python
135 lines
3.5 KiB
Python
from PIL import Image
|
|
import numpy as np
|
|
import torch
|
|
from pathlib import Path
|
|
import sys
|
|
|
|
from typing import Union, List
|
|
from pytoshop.user import nested_layers
|
|
from pytoshop import enums
|
|
from .log import log
|
|
|
|
|
|
def add_path(path, prepend=False):
|
|
if isinstance(path, list):
|
|
for p in path:
|
|
add_path(p, prepend)
|
|
return
|
|
|
|
if isinstance(path, Path):
|
|
path = path.resolve().as_posix()
|
|
|
|
if path not in sys.path:
|
|
if prepend:
|
|
sys.path.insert(0, path)
|
|
else:
|
|
sys.path.append(path)
|
|
|
|
|
|
# Get the absolute path of the parent directory of the current script
|
|
here = Path(__file__).parent.resolve()
|
|
|
|
# Construct the absolute path to the ComfyUI directory
|
|
comfy_dir = here.parent.parent
|
|
|
|
# Construct the path to the font file
|
|
font_path = here / "font.ttf"
|
|
|
|
# Add exteextern folder to path
|
|
extern_root = here / "extern"
|
|
add_path(extern_root)
|
|
for pth in extern_root.iterdir():
|
|
if pth.is_dir():
|
|
add_path(pth)
|
|
|
|
|
|
# Add the ComfyUI directory and custom nodes path to the sys.path list
|
|
add_path(comfy_dir)
|
|
add_path((comfy_dir / "custom_nodes"))
|
|
|
|
|
|
def tensor2pil(image: torch.Tensor) -> List[Image.Image]:
|
|
batch_count = 1
|
|
if len(image.shape) > 3:
|
|
batch_count = image.size(0)
|
|
|
|
if batch_count > 1:
|
|
out = []
|
|
out.extend([tensor2pil(image[i]) for i in range(batch_count)])
|
|
return out
|
|
|
|
return [
|
|
Image.fromarray(
|
|
np.clip(255.0 * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)
|
|
)
|
|
]
|
|
|
|
|
|
def pil2tensor(image: Image.Image | List[Image.Image]) -> torch.Tensor:
|
|
if isinstance(image, list):
|
|
return torch.cat([pil2tensor(img) for img in image], dim=0)
|
|
|
|
return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
|
|
|
|
|
|
def np2tensor(img_np: np.ndarray | List[np.ndarray]) -> torch.Tensor:
|
|
if isinstance(img_np, list):
|
|
return torch.cat([np2tensor(img) for img in img_np], dim=0)
|
|
|
|
return torch.from_numpy(img_np.astype(np.float32) / 255.0).unsqueeze(0)
|
|
|
|
|
|
def tensor2np(tensor: torch.Tensor) -> Union[np.ndarray, List[np.ndarray]]:
|
|
batch_count = 1
|
|
if len(tensor.shape) > 3:
|
|
batch_count = tensor.size(0)
|
|
if batch_count > 1:
|
|
out = []
|
|
out.extend([tensor2np(tensor[i]) for i in range(batch_count)])
|
|
return out
|
|
|
|
return [np.clip(255.0 * tensor.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)]
|
|
|
|
|
|
def tensor2pytolayer(
|
|
tensor: torch.Tensor,
|
|
name: str,
|
|
visible: bool = True,
|
|
opacity: int = 255,
|
|
group_id: int = 0,
|
|
blend_mode=enums.BlendMode.normal,
|
|
x: int = 0,
|
|
y: int = 0,
|
|
# channels: int = 3,
|
|
metadata: dict = {},
|
|
layer_color=0,
|
|
color_mode=None,
|
|
) -> nested_layers.Image:
|
|
batch_count = 1
|
|
if len(tensor.shape) > 3:
|
|
batch_count = tensor.size(0)
|
|
|
|
if batch_count > 1:
|
|
raise Exception(
|
|
f"Only one image is supported (batch size is currently {batch_count})"
|
|
)
|
|
out_channels = tensor2pil(tensor)
|
|
arr = np.array(out_channels)
|
|
|
|
# the array is currently H, W, C but we want C, H, W
|
|
# out_channels = np.transpose(out_channels, (2, 0, 1))
|
|
channels = [arr[:, :, 0], arr[:, :, 1], arr[:, :, 2]]
|
|
return nested_layers.Image(
|
|
name=name,
|
|
visible=visible,
|
|
opacity=opacity,
|
|
group_id=group_id,
|
|
blend_mode=blend_mode,
|
|
top=y,
|
|
left=x,
|
|
channels=channels,
|
|
metadata=metadata,
|
|
layer_color=layer_color,
|
|
color_mode=color_mode,
|
|
)
|