From 20f0ea25c3b0bba2bedaed2b9cb68e71983e4bc5 Mon Sep 17 00:00:00 2001 From: "Dr.Lt.Data" Date: Tue, 27 Jun 2023 00:37:40 +0900 Subject: [PATCH] add: LatentSender, LatentReceiver --- __init__.py | 2 + js/impact-pack.js | 31 +++++- latent.png | Bin 0 -> 2814 bytes modules/impact/config.py | 4 +- modules/impact/core.py | 31 +++++- modules/impact/impact_pack.py | 196 +++++++++++++++++++++++++++++++++- 6 files changed, 260 insertions(+), 4 deletions(-) create mode 100644 latent.png 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 0000000000000000000000000000000000000000..19fed324a25a7e1a2252400e7752ce5586742429 GIT binary patch literal 2814 zcmVEX>4Tx04R}tkv&MmKpe$iTT4Z%4y_>KkfAzR5EXIMDionYs1;guFuC*#nlvOS zE{=k0!NHHks)LKOt`4q(Aou~|>f)s6A|?JWDYS_3;J6>}?mh0_0scmXsb<$WplX(p zP9}tGZdC}qB8Uk3Fn}3}nfjb4rrU7TlmpZjz4DtVIuK7n|a>4rtTK|H-_ z>74h8qpTz;#OK6g23?T&k?XR{Z=8z``*~*6$fW0qqr^h7gXIopB|{~iB91AlM*04% z%L?Z$&T6H`TKD8H4Cl3#WvWO9|k z$gzMbR7j2={11M2YZfLa+@x><2)@|%#|RMG1)6o+{yw(t<_QpZ2ClTWzuEw1K1r{) zwb&8RzYSbmw>5bWxZD8-pLEHP94SE4Unl_YXY@@uAaV=zthv3l_Hp_EWT>mu4RCM> zj20<--Q(TeoxS~grq$mMkk@jv;{*#B00006VoOIv00000008+zyMF)x010qNS#tmY zE+YT{E+YYWr9XB6000McNliru=LHiE92X+v0}B8E02y>eSad^gZEa<4bO1wgWnpw> zWFU8GbZ8()Nlj2!fese{00`DeL_t(|+U;CxNL1e!U87c0ViN}Tp$BE<2P280qD(D{ zk{%*`C@K<%O6Z{!ge4~S@KsS!X3-0^f+)m-3I|cfQZtLPC>g>>mX1M=m}Ab%JI?xl zh`E2Jar|kfrrreyhJE%vXP-Okp0m&E$^Z}&#wjZmxAh>3~$XC91?j}IO{evH=ERtSOsl6d&=VZX~c zfBrm3qQ1V~txjZQBuK(&G&;$ZN+n3bVzCTbJH3YLf!@1!4_mixg{P;dlU$)tK%r1z z=FFK04i3i6n>YUn0$jg-9T5=`qrL|i5&~prXJh;J?Erw9ni@QP`V_^*#b|DB#^S|` z;qUK{;NW1)nKS2~c`%C?F9rZyy?PawE?q)RO^rl=Kf$V1tHeLQq@)C^SFd);|M20% zKlxx33I#~wz<~o&3@{MH#>PS}mqV}DBQi3w-}X5d)(?!3dqk}4yitX*~E_r6NnQPasB>+9>?wlgv^k|b#~8lB`yrP8VH!AOJN z1Ni#-ierC&ke;57(9lrW?RKcuYN%8yY}&L5xw*NpSS*+|YZfvyGhOWF&YcT`!GKez zP9Y>D1V@e>K}$;ue0+Rx<;oRwb#>wF*|P`>3qxF79CSJzfW$n0{|c|FK8;Fc|3P=O@Y|A|hzB+1S+7qqo<^e)gTB6?s;UZa-@e7<$&-51avV&}Juecx>DkM4*G0^8c!000320sXR@Fku1!;LV#iJ?8)X`O}Fg4*b5D z0sv;S83151nOyQ(T3P_WWqqPI+`W4j0I+V|I#IS_#R?dWMr34UIO#eTamI}s=eE45 zQ>Ov|Y&P5IBY;+`1pq8rvZPz{;lhOgfX>d&(MN!+tSnfqR`~h(^=m2h+qZAxwC#6X zzI+*zCQX9DU_fDE;jjcwg@uLi@bG|2rNV{{8(_EFk&~0-qId1uHJD5$czb)}!i5Vx zmbGTh8UR2;L&N9`0s$~5Cr6x4efso?=gysDcz8Gi0|Tj2shEH#e6XH*TcAzd!fx-OCp*UWo0Cjg1}h zAaF=2<>hQ?X`$6>rOjqzdwV+@8yl%ss~Hm$<6?9B_HF8P zI(Bt+(Ppzzuh+A%u#ifn(n)VP2(WG2Hj<=Xum8`s2LuE#KR=&3osJfZh27oV?Ck8M zR;y)dYHE)*qmBThE#16%lO&m+pFfh@A^nm4O+G$8Jbd`DIR1x+hf4$)E9v3Ghfd>n zX=&+*ZjTov{1Ge`3#zND5gQvjqT3;ZbTMR1Lpl^H5kP`O00|NSBuE4pG7R**dqF{g zcvf%k-o2tcJUm>K_cdL*bV>B*zfwX%f=e4n;>V94@b~xc(N~knB<}P7see)r@Eevc zT`KYz|4aO~KzVO(;Hqk|j$d0t^-o9Xf=G6DOjs zt`0>-Maa(1hRtTff&~k(Y}vA5^XuNddwBi&HN3sOk(`_?5n!;ea^*?@5a0Z=vNF`x z)?)nl@z}qA|A5P%Fkyn*Xxv~h^lE$V+&Kt>fHiB@AT%^oBEa9q=FOWiZ{9rEY&P7w zbxV|2R8#=?b)jQ8@#xVbR8&-8;>3wg$HgTA{3S$3N5jLz12r`@c=6(eD8F^<76d`S z?Afys6%{q0I;>W!coXRN)R%h$XU?2~)oMjpSQyr>T`LjbZ^ALgJKlH)Dl02dQ&R(( zOokmhb_}bJYPA|=Wo3}dGyquxHO6m+^VQf(0%MM8k!&v^1E_ zW`u-?Gyru QSO5S307*qoM6N<$g27oX=l}o! literal 0 HcmV?d00001 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):