487 lines
18 KiB
Python
487 lines
18 KiB
Python
# (c) TenserTensor <tenser.tensor@proton.me> || Apache-2.0 (apache.org/licenses/LICENSE-2.0)
|
|
|
|
import fnmatch
|
|
import math
|
|
import os
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import folder_paths
|
|
import numpy
|
|
import torch
|
|
from PIL import Image, ImageOps, ImageSequence
|
|
from comfy import model_management, utils
|
|
from comfy_api.latest import ComfyExtension, io, ui
|
|
from node_helpers import conditioning_set_values, pillow
|
|
from spandrel import ModelLoader, ImageModelDescriptor
|
|
|
|
from .nodes_text_encoder import SingleCondCFGGuider
|
|
from .nodes_vae import vae_encode
|
|
|
|
CATEGORY = "TenserTensor/Image"
|
|
RESIZE_METHODS = ["nearest-exact", "bilinear", "area", "bicubic", "bislerp"]
|
|
FILENAME_FORMATS = ["name-###", "date-name-###", "name-datetime"]
|
|
FORMAT_EXT = {"PNG": ".png", "JPEG": ".jpg", "WEBP": ".webp", }
|
|
MEGAPIXELS = ["0.25 MP", "0.5 MP", "1 MP", "2 MP", "4 MP", "8 MP"]
|
|
|
|
|
|
def align_to_step(size, scale_factor, dimension_step):
|
|
return round(size * scale_factor / dimension_step) * dimension_step
|
|
|
|
|
|
# crop = "center" / "disabled"
|
|
def resize_image(timage, to_width, to_height, upscale_method, dimension_step=1, crop="disabled") -> torch.Tensor:
|
|
is_mask = timage.ndim == 3
|
|
if is_mask:
|
|
timage = timage.unsqueeze(-1)
|
|
|
|
samples = timage.movedim(-1, 1)
|
|
orig_height, orig_width = samples.shape[2:]
|
|
|
|
if crop == "disabled":
|
|
scale_factor = max(to_width / orig_width, to_height / orig_height)
|
|
final_width = align_to_step(orig_width, scale_factor, dimension_step)
|
|
final_height = align_to_step(orig_height, scale_factor, dimension_step)
|
|
else:
|
|
final_width = round(to_width / dimension_step) * dimension_step
|
|
final_height = round(to_height / dimension_step) * dimension_step
|
|
|
|
final_width, final_height = int(max(final_width, dimension_step)), int(max(final_height, dimension_step))
|
|
resized = utils.common_upscale(samples, final_width, final_height, upscale_method, crop).movedim(1, -1)
|
|
|
|
return resized.squeeze(-1) if is_mask else resized
|
|
|
|
|
|
def resize_image_to_megapixels(timage, resize_method, megapixels, dimension_step=1, crop="disabled") -> torch.Tensor:
|
|
is_mask = timage.ndim == 3
|
|
if is_mask:
|
|
timage = timage.unsqueeze(-1)
|
|
|
|
mp_value = float(megapixels.split()[0])
|
|
total_pixels = int(mp_value * 1_000_000)
|
|
samples = timage.movedim(-1, 1)
|
|
orig_height, orig_width = samples.shape[2:]
|
|
scale_factor = math.sqrt(total_pixels / (orig_width * orig_height))
|
|
|
|
if crop == "disabled":
|
|
final_width = align_to_step(orig_width, scale_factor, dimension_step)
|
|
final_height = align_to_step(orig_height, scale_factor, dimension_step)
|
|
else:
|
|
to_width, to_height = orig_width * scale_factor, orig_height * scale_factor
|
|
final_width = round(to_width / dimension_step) * dimension_step
|
|
final_height = round(to_height / dimension_step) * dimension_step
|
|
|
|
final_width, final_height = int(max(final_width, dimension_step)), int(max(final_height, dimension_step))
|
|
resized = utils.common_upscale(samples, final_width, final_height, resize_method, crop).movedim(1, -1)
|
|
|
|
return resized.squeeze(-1) if is_mask else resized
|
|
|
|
|
|
def rotate_image(pixels, turns):
|
|
return torch.rot90(pixels, k=turns, dims=[1, 2])
|
|
|
|
|
|
def flip_image(pixels, axis):
|
|
return torch.flip(pixels, dims=[2 if axis == "x" else 1])
|
|
|
|
|
|
def get_torch_device(device=None):
|
|
return (
|
|
model_management.get_torch_device()
|
|
if device is None or device == "default"
|
|
else torch.device(device)
|
|
)
|
|
|
|
|
|
def load_image(filename, device=None) -> tuple[torch.Tensor, torch.Tensor]:
|
|
tdevice = get_torch_device(device)
|
|
path = folder_paths.get_annotated_filepath(filename)
|
|
pil_image = pillow(Image.open, path)
|
|
|
|
images, masks = [], []
|
|
width, height = None, None
|
|
|
|
for frame in ImageSequence.Iterator(pil_image):
|
|
pil_frame = pillow(ImageOps.exif_transpose, frame)
|
|
|
|
if pil_frame.mode == "I":
|
|
pil_frame = pil_frame.point(lambda i: i * (1 / 255))
|
|
|
|
if len(images) == 0:
|
|
width, height = pil_frame.size
|
|
|
|
if pil_frame.size != (width, height,):
|
|
continue
|
|
|
|
alpha = numpy.zeros((height, width,), dtype=numpy.float32)
|
|
if 'A' in pil_frame.getbands():
|
|
alpha = numpy.array(pil_frame.getchannel('A'), dtype=numpy.float32) / 255.0
|
|
elif pil_frame.mode == 'P' and 'transparency' in pil_frame.info:
|
|
alpha = numpy.array(pil_frame.convert('RGBA').getchannel('A'), dtype=numpy.float32) / 255.0
|
|
tmask = torch.from_numpy(1.0 - alpha).to(tdevice)
|
|
masks.append(tmask)
|
|
|
|
rgb_frame = pil_frame.convert("RGB")
|
|
np_image = numpy.array(rgb_frame, dtype=numpy.float32) / 255.0
|
|
timage = torch.from_numpy(np_image)[None].to(tdevice)
|
|
|
|
images.append(timage)
|
|
|
|
if pil_image.format == "MPO":
|
|
break
|
|
|
|
if len(images) == 0:
|
|
raise ValueError(f"ERROR: No frames found in {path}")
|
|
|
|
o_image = torch.cat(images, dim=0) if len(images) > 1 else images[0]
|
|
|
|
match len(masks):
|
|
case 0:
|
|
o_mask = torch.zeros((1, height, width), dtype=torch.float32)
|
|
case 1:
|
|
o_mask = masks[0].unsqueeze(0)
|
|
case _:
|
|
o_mask = torch.cat(masks, dim=0)
|
|
|
|
return o_image, o_mask,
|
|
|
|
|
|
def get_image_files():
|
|
return sorted([
|
|
f.name for f in os.scandir(folder_paths.get_input_directory())
|
|
if f.is_file() and Path(f.name).suffix.lower() in (".png", ".jpg", ".webp", ".bmp")
|
|
])
|
|
|
|
|
|
class TT_ImageLoaderResizerNode(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="TT_ImageLoaderResizerNode",
|
|
display_name="TT Image Loader / Resizer",
|
|
category=CATEGORY,
|
|
description="",
|
|
inputs=[
|
|
io.Boolean.Input("resize_image", default=False, label_on="Scale", label_off="Skip"),
|
|
io.Combo.Input("resize_method", options=RESIZE_METHODS, advanced=True),
|
|
io.Int.Input("dimension_step", default=1, min=1, max=256, advanced=True),
|
|
io.Combo.Input("megapixels", options=MEGAPIXELS, advanced=True),
|
|
io.Combo.Input("image", options=get_image_files(), upload=io.UploadType.image)
|
|
],
|
|
outputs=[
|
|
io.Image.Output("IMAGE"),
|
|
io.Mask.Output("MASK"),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def fingerprint_inputs(cls, **kwargs) -> dict[str, Any]:
|
|
path = folder_paths.get_annotated_filepath(kwargs.get("image"))
|
|
resize_image = kwargs.get("resize_image")
|
|
result = {
|
|
"image": path + str(os.path.getmtime(path)),
|
|
"resize_image": resize_image,
|
|
}
|
|
|
|
if resize_image:
|
|
result["megapixels"] = kwargs.get("megapixels")
|
|
result["resize_method"] = kwargs.get("resize_method")
|
|
result["dimension_step"] = kwargs.get("dimension_step")
|
|
|
|
return result
|
|
|
|
@classmethod
|
|
def validate_inputs(cls, **kwargs) -> bool | str:
|
|
path = folder_paths.get_annotated_filepath(kwargs.get("image"))
|
|
if not os.path.isfile(path):
|
|
return f"Image file not found: {path}"
|
|
|
|
return True
|
|
|
|
@classmethod
|
|
def execute(cls, **kwargs) -> io.NodeOutput:
|
|
timage, tmask = load_image(kwargs.get("image"))
|
|
if kwargs.get("resize_image"):
|
|
resize_method, megapixels, dimension_step = (
|
|
kwargs.get("resize_method"),
|
|
kwargs.get("megapixels"),
|
|
kwargs.get("dimension_step")
|
|
)
|
|
timage = resize_image_to_megapixels(timage, resize_method, megapixels, dimension_step)
|
|
tmask = resize_image_to_megapixels(tmask, resize_method, megapixels, dimension_step)
|
|
|
|
return io.NodeOutput(timage, tmask, ui=ui.PreviewImage(timage, cls=cls))
|
|
|
|
|
|
class InvalidSavePathError(Exception):
|
|
pass
|
|
|
|
|
|
def build_filename(save_path, filename_mask) -> str:
|
|
count = sum(1 for f in os.scandir(save_path) if fnmatch.fnmatch(f.name, filename_mask)) + 1
|
|
|
|
return filename_mask.replace("*", f"{count:05}")
|
|
|
|
|
|
def build_save_path(**kwargs) -> tuple[str, str]:
|
|
output_path, subfolder, image_format, filename_prefix = (
|
|
kwargs.get("output_path"),
|
|
kwargs.get("subfolder"),
|
|
kwargs.get("image_format"),
|
|
kwargs.get("filename_prefix"),
|
|
)
|
|
save_path = os.path.join(output_path, subfolder)
|
|
|
|
if os.path.commonpath((output_path, os.path.abspath(save_path))) != output_path:
|
|
raise InvalidSavePathError("Saving image outside the output folder is not allowed.")
|
|
|
|
os.makedirs(save_path, exist_ok=True)
|
|
|
|
ext = FORMAT_EXT[image_format]
|
|
filename = ""
|
|
match kwargs.get("filename_format"):
|
|
case "name-###":
|
|
filename = build_filename(save_path, f"{filename_prefix}-*{ext}")
|
|
case "date-name-###":
|
|
date = datetime.now().strftime("%Y-%m-%d")
|
|
filename = build_filename(save_path, f"{date}-{filename_prefix}-*{ext}")
|
|
case "name-datetime":
|
|
timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
|
|
filename = f"{filename_prefix}-{timestamp}{ext}"
|
|
case _:
|
|
filename = f"{filename_prefix}{ext}"
|
|
|
|
return save_path, filename
|
|
|
|
|
|
def save_image(**kwargs) -> None:
|
|
kwargs["output_path"] = folder_paths.get_output_directory()
|
|
image = kwargs.get("image")
|
|
for frame in image:
|
|
np_frame = frame.cpu().numpy() * 255.0
|
|
tframe = Image.fromarray(numpy.clip(np_frame, 0, 255).astype(numpy.uint8))
|
|
save_path, filename = build_save_path(**kwargs)
|
|
image_format = kwargs.get("image_format")
|
|
params = (
|
|
{"compress_level": kwargs.get("compression_level")}
|
|
if kwargs.get("image_format") == "PNG"
|
|
else {"quality": kwargs.get("image_quality")}
|
|
)
|
|
tframe.save(os.path.join(save_path, filename), format=image_format, **params)
|
|
|
|
|
|
class TT_ImagePreviewSaveNode(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="TT_ImagePreviewSaveNode",
|
|
display_name="TT Image Preview / Save",
|
|
category=CATEGORY,
|
|
description="",
|
|
is_output_node=True,
|
|
inputs=[
|
|
io.Image.Input("image"),
|
|
io.Boolean.Input("save_image", default=True, label_on="Save image", label_off="Only preview"),
|
|
io.String.Input("filename_prefix", default="TT"),
|
|
io.Combo.Input("filename_format", options=FILENAME_FORMATS, default="name-datetime", advanced=True),
|
|
io.String.Input("subfolder", default="", advanced=True),
|
|
io.Combo.Input("image_format", options=list(FORMAT_EXT.keys()), default="PNG", advanced=True),
|
|
io.Int.Input("image_quality", default=75, min=1, max=100, advanced=True),
|
|
io.Int.Input("compression_level", default=6, min=0, max=9, advanced=True),
|
|
],
|
|
outputs=[]
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, **kwargs) -> io.NodeOutput:
|
|
timage = kwargs.get("image")
|
|
|
|
if kwargs.get("save_image"):
|
|
save_image(**kwargs)
|
|
|
|
return io.NodeOutput(ui=ui.PreviewImage(timage, cls=cls))
|
|
|
|
|
|
DEVICES = ["default", "cpu"]
|
|
|
|
|
|
def get_upscale_models() -> list[str]:
|
|
return folder_paths.get_filename_list("upscale_models")
|
|
|
|
|
|
def load_upscale_model(upscale_model) -> ImageModelDescriptor:
|
|
model_path = folder_paths.get_full_path_or_raise("upscale_models", upscale_model)
|
|
|
|
try:
|
|
sd = utils.load_torch_file(model_path, safe_load=True)
|
|
except Exception as e:
|
|
raise RuntimeError(f"Failed to load model file {model_path}: {e}")
|
|
|
|
sd = utils.state_dict_prefix_replace(sd, {"module.": ""})
|
|
upscale_model = ModelLoader().load_from_state_dict(sd).eval()
|
|
|
|
if not isinstance(upscale_model, ImageModelDescriptor):
|
|
raise TypeError(f"Expected ImageModelDescriptor, got {type(upscale_model).__name__}")
|
|
|
|
return upscale_model
|
|
|
|
|
|
BUFFER_FACTOR = 384.0
|
|
|
|
|
|
def upscale(**kwargs):
|
|
image, device_name, upscale_tile, upscale_overlap = (
|
|
kwargs.get("image"),
|
|
kwargs.get("upscaler_device"),
|
|
kwargs.get("upscale_tile"),
|
|
kwargs.get("upscale_overlap")
|
|
)
|
|
|
|
timage = image.clone()
|
|
device = get_torch_device(device_name)
|
|
model = load_upscale_model(kwargs.get("upscale_model"))
|
|
i_image = timage.movedim(-1, 1).to(device)
|
|
batch, _, height, width = i_image.shape
|
|
m_model = model_management.module_size(model.model)
|
|
m_upscale_tile = (upscale_tile * upscale_tile * 3) * i_image.element_size() * max(model.scale, 1) * BUFFER_FACTOR
|
|
m_total = m_model + m_upscale_tile + i_image.nelement() * i_image.element_size()
|
|
model_management.free_memory(m_total, device)
|
|
model.model.to(device)
|
|
|
|
progress_bar = None
|
|
try:
|
|
while True:
|
|
try:
|
|
steps = batch * utils.get_tiled_scale_steps(
|
|
width, height,
|
|
tile_x=upscale_tile,
|
|
tile_y=upscale_tile,
|
|
overlap=upscale_overlap
|
|
)
|
|
progress_bar = utils.ProgressBar(steps)
|
|
scaled = utils.tiled_scale(
|
|
i_image,
|
|
model,
|
|
tile_x=upscale_tile,
|
|
tile_y=upscale_tile,
|
|
overlap=upscale_overlap,
|
|
upscale_amount=model.scale,
|
|
pbar=progress_bar
|
|
)
|
|
break
|
|
except model_management.OOM_EXCEPTION as e:
|
|
progress_bar = None
|
|
upscale_tile //= 2
|
|
if upscale_tile < 128:
|
|
raise e
|
|
finally:
|
|
if device.type != "cpu":
|
|
model.model.to("cpu")
|
|
|
|
return torch.clamp(scaled.movedim(1, -1), min=0.0, max=1.0)
|
|
|
|
|
|
class TT_ImagePreviewUpscaleSaveNode(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="TT_ImagePreviewUpscaleSaveNode",
|
|
display_name="TT Image Preview / Upscale / Save",
|
|
category=CATEGORY,
|
|
description="",
|
|
is_output_node=True,
|
|
inputs=[
|
|
io.Image.Input("image"),
|
|
io.Boolean.Input("save_image", default=True, label_on="Save image", label_off="Only preview"),
|
|
io.Boolean.Input("upscale_image", default=True, label_on="Upscale image", label_off="Keep size"),
|
|
io.String.Input("filename_prefix", default="TT"),
|
|
io.Combo.Input("filename_format", options=FILENAME_FORMATS, default="name-datetime", advanced=True),
|
|
io.String.Input("subfolder", default="", advanced=True),
|
|
io.Combo.Input("image_format", options=list(FORMAT_EXT.keys()), default="PNG", advanced=True),
|
|
io.Int.Input("image_quality", default=75, min=1, max=100, advanced=True),
|
|
io.Int.Input("compression_level", default=6, min=0, max=9, advanced=True),
|
|
io.Combo.Input("upscaler_device", options=DEVICES, advanced=True),
|
|
io.Combo.Input("upscale_model", options=get_upscale_models(), advanced=True),
|
|
io.Int.Input("upscale_tile", default=512, min=128, max=4096, step=64, advanced=True),
|
|
io.Int.Input("upscale_overlap", default=64, min=8, max=256, step=8, advanced=True),
|
|
],
|
|
outputs=[]
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, **kwargs) -> io.NodeOutput:
|
|
image = kwargs.get("image")
|
|
timage = image.clone()
|
|
|
|
if kwargs.get("upscale_image"):
|
|
timage = upscale(**kwargs)
|
|
|
|
if kwargs.get("save_image"):
|
|
kwargs["image"] = timage
|
|
save_image(**kwargs)
|
|
|
|
return io.NodeOutput(ui=ui.PreviewImage(timage, cls=cls))
|
|
|
|
|
|
class TT_GuiderImageReferenceNode(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="TT_GuiderImageReferenceNode",
|
|
display_name="TT Guider Image Reference",
|
|
category=CATEGORY,
|
|
description="",
|
|
inputs=[
|
|
io.Vae.Input("vae"),
|
|
io.Guider.Input("guider"),
|
|
io.Combo.Input("megapixels", options=MEGAPIXELS),
|
|
io.Combo.Input("resize_method", options=RESIZE_METHODS, advanced=True),
|
|
io.Int.Input("dimension_step", default=1, min=1, max=256, advanced=True),
|
|
io.Combo.Input("image", options=get_image_files(), upload=io.UploadType.image)
|
|
],
|
|
outputs=[
|
|
io.Vae.Output("VAE"),
|
|
io.Guider.Output("GUIDER"),
|
|
]
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, **kwargs) -> io.NodeOutput:
|
|
timage, _ = load_image(kwargs.get("image"))
|
|
vae, guider, resize_method, megapixels, dimension_step = (
|
|
kwargs.get("vae"),
|
|
kwargs.get("guider"),
|
|
kwargs.get("resize_method"),
|
|
kwargs.get("megapixels"),
|
|
kwargs.get("dimension_step")
|
|
)
|
|
|
|
if not isinstance(guider, SingleCondCFGGuider):
|
|
guider = SingleCondCFGGuider.from_cfg_guider(guider)
|
|
|
|
resized = resize_image_to_megapixels(timage, resize_method, megapixels, dimension_step)
|
|
latent = vae_encode(resized, vae)
|
|
|
|
conditioning = guider.get_conds()
|
|
|
|
if conditioning is not None:
|
|
conditioning = conditioning_set_values(conditioning, {"reference_latents": [latent["samples"]]}, append=True)
|
|
guider.set_conds(conditioning)
|
|
else:
|
|
raise ValueError("ERROR: Guider has no conditioning")
|
|
|
|
return io.NodeOutput(vae, guider, ui=ui.PreviewImage(timage, cls=cls))
|
|
|
|
|
|
# ==============================================================================
|
|
# V3 entrypoint — registers context nodes with ComfyUI
|
|
# ==============================================================================
|
|
|
|
NODES = [
|
|
TT_ImageLoaderResizerNode,
|
|
TT_ImagePreviewSaveNode,
|
|
TT_ImagePreviewUpscaleSaveNode,
|
|
TT_GuiderImageReferenceNode,
|
|
]
|