diff --git a/__init__.py b/__init__.py index c2e0edc..c0bb79e 100644 --- a/__init__.py +++ b/__init__.py @@ -137,6 +137,8 @@ NODE_CLASS_MAPPINGS = { "PreviewBridge": PreviewBridge, "ImageSender": ImageSender, "ImageReceiver": ImageReceiver, + "LatentSender": LatentSender, + "LatentReceiver": LatentReceiver, "ImageMaskSwitch": ImageMaskSwitch, "LatentSwitch": LatentSwitch, "SEGSSwitch": SEGSSwitch, diff --git a/js/impact-pack.js b/js/impact-pack.js index 8a64feb..b220919 100644 --- a/js/impact-pack.js +++ b/js/impact-pack.js @@ -104,7 +104,28 @@ function imgSendHandler(event) { for(let i in nodes) { if(nodes[i].type == 'ImageReceiver') { if(nodes[i].widgets[1].value == event.detail.link_id) { - nodes[i].widgets[0].value = filename; + nodes[i].widgets[0].value = `${data.subfolder}/${data.filename} [${data.type}]`; + let img = new Image(); + img.src = `/view?filename=${data.filename}&type=${data.type}&subfolder=${data.subfolder}`+app.getPreviewFormatParam(); + nodes[i].imgs = [img]; + nodes[i].size[1] = Math.max(200, nodes[i].size[1]); + } + } + } + } +} + + +function latentSendHandler(event) { + if(event.detail.images.length > 0){ + let data = event.detail.images[0]; + let filename = `${data.filename} [${data.type}]`; + + let nodes = app.graph._nodes; + for(let i in nodes) { + if(nodes[i].type == 'LatentReceiver') { + if(nodes[i].widgets[1].value == event.detail.link_id) { + nodes[i].widgets[0].value = `${data.subfolder}/${data.filename} [${data.type}]`; let img = new Image(); img.src = `/view?filename=${data.filename}&type=${data.type}&subfolder=${data.subfolder}`+app.getPreviewFormatParam(); nodes[i].imgs = [img]; @@ -117,6 +138,7 @@ function imgSendHandler(event) { var progressEventRegistered = false; var imgSendEventRegistered = false; +var latentSendEventRegistered = false; const impactProgressBadge = new ImpactProgressBadge(); app.registerExtension({ @@ -137,6 +159,13 @@ app.registerExtension({ imgSendEventRegistered = true; } } + + if (node.comfyClass == "LatentSender") { + if (!latentSendEventRegistered) { + api.addEventListener("latent-send", latentSendHandler); + latentSendEventRegistered = true; + } + } }, async beforeRegisterNodeDef(nodeType, nodeData, app) { diff --git a/latent.png b/latent.png new file mode 100644 index 0000000..19fed32 Binary files /dev/null and b/latent.png differ diff --git a/modules/impact/config.py b/modules/impact/config.py index 3871294..896b9ba 100644 --- a/modules/impact/config.py +++ b/modules/impact/config.py @@ -1,12 +1,14 @@ import configparser import os -version = "V2.14" +version = "V2.15" dependency_version = 1 my_path = os.path.dirname(__file__) config_path = os.path.join(my_path, "impact-pack.ini") +latent_letter_path = os.path.join(my_path, "..", "..", "latent.png") + MAX_RESOLUTION = 8192 def write_config(comfy_path): diff --git a/modules/impact/core.py b/modules/impact/core.py index 5070167..246fd5b 100644 --- a/modules/impact/core.py +++ b/modules/impact/core.py @@ -1157,4 +1157,33 @@ def update_node_status(node, text, progress=None): "node": node, "progress": progress, "text": text - }, PromptServer.instance.client_id) \ No newline at end of file + }, PromptServer.instance.client_id) + + +from comfy.cli_args import args, LatentPreviewMethod +import folder_paths +from latent_preview import TAESD, TAESDPreviewerImpl, Latent2RGBPreviewer +import comfy.latent_formats as latent_formats + +def get_previewer(device, latent_format=latent_formats.SD15(), force=False): + previewer = None + method = args.preview_method + if method != LatentPreviewMethod.NoPreviews or force: + # TODO previewer methods + taesd_decoder_path = folder_paths.get_full_path("vae_approx", latent_format.taesd_decoder_name) + + if method == LatentPreviewMethod.Auto: + method = LatentPreviewMethod.Latent2RGB + if taesd_decoder_path: + method = LatentPreviewMethod.TAESD + + if method == LatentPreviewMethod.TAESD: + if taesd_decoder_path: + taesd = TAESD(None, taesd_decoder_path).to(device) + previewer = TAESDPreviewerImpl(taesd) + else: + print("Warning: TAESD previews enabled, but could not find models/vae_approx/{}".format(latent_format.taesd_decoder_name)) + + if previewer is None: + previewer = Latent2RGBPreviewer(latent_format.latent_rgb_factors) + return previewer diff --git a/modules/impact/impact_pack.py b/modules/impact/impact_pack.py index 8719496..b465542 100644 --- a/modules/impact/impact_pack.py +++ b/modules/impact/impact_pack.py @@ -9,12 +9,14 @@ from segment_anything import sam_model_registry from impact.utils import * import impact.core as core from impact.core import SEG, NO_BBOX_DETECTOR, NO_SEGM_DETECTOR -from impact.config import MAX_RESOLUTION +from impact.config import MAX_RESOLUTION, latent_letter_path from PIL import Image import numpy as np import hashlib import json import safetensors.torch +from PIL.PngImagePlugin import PngInfo +import latent_preview warnings.filterwarnings('ignore', category=UserWarning, message='TypedStorage is deprecated') @@ -1495,6 +1497,198 @@ class ImageSender(nodes.PreviewImage): return result +from io import BytesIO +import piexif +import zipfile +from server import PromptServer + +class LatentReceiver: + def __init__(self): + self.input_dir = folder_paths.get_input_directory() + self.type = "input" + + @classmethod + def INPUT_TYPES(s): + def check_file_extension(x): + return x.endswith(".latent") or x.endswith(".latent.png") + + input_dir = folder_paths.get_input_directory() + files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f)) and check_file_extension(f)] + return {"required": { + "latent": (sorted(files), ), + "link_id": ("INT", {"default": 0, "min": 0, "max": sys.maxsize, "step": 1}), + }, + } + + FUNCTION = "doit" + + CATEGORY = "ImpactPack/Util" + + RETURN_TYPES = ("LATENT",) + + @staticmethod + def load_preview_latent(image_path): + image = Image.open(image_path) + exif_data = piexif.load(image.info["exif"]) + + if piexif.ExifIFD.UserComment in exif_data["Exif"]: + compressed_data = exif_data["Exif"][piexif.ExifIFD.UserComment] + compressed_data_io = BytesIO(compressed_data) + with zipfile.ZipFile(compressed_data_io, mode='r') as archive: + tensor_bytes = archive.read("latent") + tensor = safetensors.torch.load(tensor_bytes) + return {"samples": tensor['latent_tensor']} + return None + + def doit(self, latent, link_id): + latent_path = folder_paths.get_annotated_filepath(latent) + + if latent.endswith(".latent"): + latent = safetensors.torch.load_file(latent_path, device="cpu") + multiplier = 1.0 + if "latent_format_version_0" not in latent: + multiplier = 1.0 / 0.18215 + samples = {"samples": latent["latent_tensor"].float() * multiplier} + else: + samples = LatentReceiver.load_preview_latent(latent_path) + + preview = { + 'filename': latent_path, + 'subfolder': '', + 'type': self.type + } + + return { + 'ui': {"images": [preview]}, + 'result': (samples, ) + } + + @classmethod + def IS_CHANGED(s, latent, link_id): + image_path = folder_paths.get_annotated_filepath(latent) + m = hashlib.sha256() + with open(image_path, 'rb') as f: + m.update(f.read()) + return m.digest().hex() + + @classmethod + def VALIDATE_INPUTS(s, latent, link_id): + if not folder_paths.exists_annotated_filepath(latent): + return "Invalid latent file: {}".format(latent) + return True + + +class LatentSender(nodes.SaveLatent): + def __init__(self): + self.output_dir = folder_paths.get_temp_directory() + self.type = "temp" + + @classmethod + def INPUT_TYPES(s): + return {"required": { + "samples": ("LATENT", ), + "filename_prefix": ("STRING", {"default": "latents/LatentSender"}), + "link_id": ("INT", {"default": 0, "min": 0, "max": sys.maxsize, "step": 1}), }, + "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"}, + } + + OUTPUT_NODE = True + + RETURN_TYPES = () + + FUNCTION = "doit" + + CATEGORY = "ImpactPack/Util" + + @staticmethod + def save_to_file(tensor_bytes, prompt, extra_pnginfo, image, image_path): + compressed_data = BytesIO() + with zipfile.ZipFile(compressed_data, mode='w') as archive: + archive.writestr("latent", tensor_bytes) + image = image.copy() + exif_data = {"Exif": {piexif.ExifIFD.UserComment: compressed_data.getvalue()}} + + metadata = PngInfo() + if prompt is not None: + metadata.add_text("prompt", json.dumps(prompt)) + if extra_pnginfo is not None: + for x in extra_pnginfo: + metadata.add_text(x, json.dumps(extra_pnginfo[x])) + + exif_bytes = piexif.dump(exif_data) + image.save(image_path, format='png', exif=exif_bytes, pnginfo=metadata, optimize=True) + + @staticmethod + def prepare_preview(latent_tensor): + lower_bound = 128 + upper_bound = 256 + + previewer = core.get_previewer("cpu", force=True) + image = previewer.decode_latent_to_preview(latent_tensor) + min_size = min(image.size[0], image.size[1]) + max_size = max(image.size[0], image.size[1]) + + scale_factor = 1 + if max_size > upper_bound: + scale_factor = upper_bound/max_size + + # prevent too small preview + if min_size*scale_factor < lower_bound: + scale_factor = lower_bound/min_size + + w = int(image.size[0] * scale_factor) + h = int(image.size[1] * scale_factor) + + image = image.resize((w, h), resample=Image.NEAREST) + + return LatentSender.attach_format_text(image) + + @staticmethod + def attach_format_text(image): + width_a, height_a = image.size + + letter_image = Image.open(latent_letter_path) + width_b, height_b = letter_image.size + + new_width = max(width_a, width_b) + new_height = height_a + height_b + + new_image = Image.new('RGB', (new_width, new_height), (0, 0, 0)) + + offset_x = (new_width - width_b) // 2 + offset_y = (height_a + (new_height - height_a - height_b) // 2) + new_image.paste(letter_image, (offset_x, offset_y)) + + new_image.paste(image, (0, 0)) + + return new_image + + def doit(self, samples, filename_prefix="latents/LatentSender", link_id=0, prompt=None, extra_pnginfo=None): + full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir) + + # load preview + preview = LatentSender.prepare_preview(samples['samples']) + + # support save metadata for latent sharing + file = f"{filename}_{counter:05}_.latent.png" + fullpath = os.path.join(full_output_folder, file) + + output = {"latent_tensor": samples["samples"]} + + tensor_bytes = safetensors.torch.save(output) + LatentSender.save_to_file(tensor_bytes, prompt, extra_pnginfo, preview, fullpath) + + latent_path = { + 'filename': file, + 'subfolder': subfolder, + 'type': self.type + } + + PromptServer.instance.send_sync("latent-send", {"link_id": link_id, "images": [latent_path]}) + + return {'ui': {'images': [latent_path]}} + + class ImageMaskSwitch: @classmethod def INPUT_TYPES(s):