|
import os
|
|
from PIL import ImageOps
|
|
from impact.utils import *
|
|
import latent_preview
|
|
|
|
|
|
|
|
|
|
from impact import core
|
|
|
|
import random
|
|
|
|
|
|
class PreviewBridge:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"images": ("IMAGE",),
|
|
"image": ("STRING", {"default": ""}),
|
|
},
|
|
"hidden": {"unique_id": "UNIQUE_ID"},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK", )
|
|
|
|
FUNCTION = "doit"
|
|
|
|
OUTPUT_NODE = True
|
|
|
|
CATEGORY = "ImpactPack/Util"
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.output_dir = folder_paths.get_temp_directory()
|
|
self.type = "temp"
|
|
self.prev_hash = None
|
|
|
|
@staticmethod
|
|
def load_image(pb_id):
|
|
is_fail = False
|
|
if pb_id not in core.preview_bridge_image_id_map:
|
|
is_fail = True
|
|
|
|
image_path, ui_item = core.preview_bridge_image_id_map[pb_id]
|
|
|
|
if not os.path.isfile(image_path):
|
|
is_fail = True
|
|
|
|
if not is_fail:
|
|
i = Image.open(image_path)
|
|
i = ImageOps.exif_transpose(i)
|
|
image = i.convert("RGB")
|
|
image = np.array(image).astype(np.float32) / 255.0
|
|
image = torch.from_numpy(image)[None,]
|
|
|
|
if 'A' in i.getbands():
|
|
mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0
|
|
mask = 1. - torch.from_numpy(mask)
|
|
else:
|
|
mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu")
|
|
else:
|
|
image = empty_pil_tensor()
|
|
mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu")
|
|
ui_item = {
|
|
"filename": 'empty.png',
|
|
"subfolder": '',
|
|
"type": 'temp'
|
|
}
|
|
|
|
return image, mask.unsqueeze(0), ui_item
|
|
|
|
def doit(self, images, image, unique_id):
|
|
need_refresh = False
|
|
|
|
if unique_id not in core.preview_bridge_cache:
|
|
need_refresh = True
|
|
|
|
elif core.preview_bridge_cache[unique_id][0] is not images:
|
|
need_refresh = True
|
|
|
|
if not need_refresh:
|
|
pixels, mask, path_item = PreviewBridge.load_image(image)
|
|
image = [path_item]
|
|
else:
|
|
res = nodes.PreviewImage().save_images(images, filename_prefix="PreviewBridge/PB-")
|
|
image2 = res['ui']['images']
|
|
pixels = images
|
|
mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu")
|
|
|
|
path = os.path.join(folder_paths.get_temp_directory(), 'PreviewBridge', image2[0]['filename'])
|
|
core.set_previewbridge_image(unique_id, path, image2[0])
|
|
core.preview_bridge_image_id_map[image] = (path, image2[0])
|
|
core.preview_bridge_image_name_map[unique_id, path] = (image, image2[0])
|
|
core.preview_bridge_cache[unique_id] = (images, image2)
|
|
|
|
image = image2
|
|
|
|
return {
|
|
"ui": {"images": image},
|
|
"result": (pixels, mask, ),
|
|
}
|
|
|
|
|
|
def decode_latent(latent, preview_method, vae_opt=None):
|
|
if vae_opt is not None:
|
|
image = nodes.VAEDecode().decode(vae_opt, latent)[0]
|
|
return image
|
|
|
|
from comfy.cli_args import LatentPreviewMethod
|
|
import comfy.latent_formats as latent_formats
|
|
|
|
if preview_method.startswith("TAE"):
|
|
decoder_name = None
|
|
|
|
if preview_method == "TAESD15":
|
|
decoder_name = "taesd"
|
|
elif preview_method == 'TAESDXL':
|
|
decoder_name = "taesdxl"
|
|
elif preview_method == 'TAESD3':
|
|
decoder_name = "taesd3"
|
|
|
|
if decoder_name:
|
|
vae = nodes.VAELoader().load_vae(decoder_name)[0]
|
|
image = nodes.VAEDecode().decode(vae, latent)[0]
|
|
return image
|
|
|
|
if preview_method == "Latent2RGB-SD15":
|
|
latent_format = latent_formats.SD15()
|
|
method = LatentPreviewMethod.Latent2RGB
|
|
elif preview_method == "Latent2RGB-SDXL":
|
|
latent_format = latent_formats.SDXL()
|
|
method = LatentPreviewMethod.Latent2RGB
|
|
elif preview_method == "Latent2RGB-SD3":
|
|
latent_format = latent_formats.SD3()
|
|
method = LatentPreviewMethod.Latent2RGB
|
|
elif preview_method == "Latent2RGB-SD-X4":
|
|
latent_format = latent_formats.SD_X4()
|
|
method = LatentPreviewMethod.Latent2RGB
|
|
elif preview_method == "Latent2RGB-Playground-2.5":
|
|
latent_format = latent_formats.SDXL_Playground_2_5()
|
|
method = LatentPreviewMethod.Latent2RGB
|
|
elif preview_method == "Latent2RGB-SC-Prior":
|
|
latent_format = latent_formats.SC_Prior()
|
|
method = LatentPreviewMethod.Latent2RGB
|
|
elif preview_method == "Latent2RGB-SC-B":
|
|
latent_format = latent_formats.SC_B()
|
|
method = LatentPreviewMethod.Latent2RGB
|
|
elif preview_method == "Latent2RGB-FLUX.1":
|
|
latent_format = latent_formats.Flux()
|
|
method = LatentPreviewMethod.Latent2RGB
|
|
else:
|
|
print(f"[Impact Pack] PreviewBridgeLatent: '{preview_method}' is unsupported preview method.")
|
|
latent_format = latent_formats.SD15()
|
|
method = LatentPreviewMethod.Latent2RGB
|
|
|
|
previewer = core.get_previewer("cpu", latent_format=latent_format, force=True, method=method)
|
|
samples = latent_format.process_in(latent['samples'])
|
|
|
|
pil_image = previewer.decode_latent_to_preview(samples)
|
|
pixels_size = pil_image.size[0]*8, pil_image.size[1]*8
|
|
resized_image = pil_image.resize(pixels_size, resample=LANCZOS)
|
|
|
|
return to_tensor(resized_image).unsqueeze(0)
|
|
|
|
|
|
class PreviewBridgeLatent:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"latent": ("LATENT",),
|
|
"image": ("STRING", {"default": ""}),
|
|
"preview_method": (["Latent2RGB-SD3", "Latent2RGB-SDXL", "Latent2RGB-SD15",
|
|
"Latent2RGB-SD-X4", "Latent2RGB-Playground-2.5",
|
|
"Latent2RGB-SC-Prior", "Latent2RGB-SC-B",
|
|
"Latent2RGB-FLUX.1",
|
|
"TAESD3", "TAESDXL", "TAESD15"],),
|
|
},
|
|
"optional": {
|
|
"vae_opt": ("VAE", )
|
|
},
|
|
"hidden": {"unique_id": "UNIQUE_ID"},
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT", "MASK", )
|
|
|
|
FUNCTION = "doit"
|
|
|
|
OUTPUT_NODE = True
|
|
|
|
CATEGORY = "ImpactPack/Util"
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.output_dir = folder_paths.get_temp_directory()
|
|
self.type = "temp"
|
|
self.prev_hash = None
|
|
self.prefix_append = "_temp_" + ''.join(random.choice("abcdefghijklmnopqrstupvxyz") for x in range(5))
|
|
|
|
@staticmethod
|
|
def load_image(pb_id):
|
|
is_fail = False
|
|
if pb_id not in core.preview_bridge_image_id_map:
|
|
is_fail = True
|
|
|
|
image_path, ui_item = core.preview_bridge_image_id_map[pb_id]
|
|
|
|
if not os.path.isfile(image_path):
|
|
is_fail = True
|
|
|
|
if not is_fail:
|
|
i = Image.open(image_path)
|
|
i = ImageOps.exif_transpose(i)
|
|
image = i.convert("RGB")
|
|
image = np.array(image).astype(np.float32) / 255.0
|
|
image = torch.from_numpy(image)[None,]
|
|
|
|
if 'A' in i.getbands():
|
|
mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0
|
|
mask = 1. - torch.from_numpy(mask)
|
|
else:
|
|
mask = None
|
|
else:
|
|
image = empty_pil_tensor()
|
|
mask = None
|
|
ui_item = {
|
|
"filename": 'empty.png',
|
|
"subfolder": '',
|
|
"type": 'temp'
|
|
}
|
|
|
|
return image, mask, ui_item
|
|
|
|
def doit(self, latent, image, preview_method, vae_opt=None, unique_id=None):
|
|
latent_channels = latent['samples'].shape[1]
|
|
preview_method_channels = 16 if 'SD3' in preview_method or 'SC-Prior' in preview_method or 'FLUX.1' in preview_method else 4
|
|
|
|
if vae_opt is None and latent_channels != preview_method_channels:
|
|
print(f"[PreviewBridgeLatent] The version of latent is not compatible with preview_method.\nSD3, SD1/SD2, SDXL, SC-Prior, SC-B and FLUX.1 are not compatible with each other.")
|
|
raise Exception("The version of latent is not compatible with preview_method.<BR>SD3, SD1/SD2, SDXL, SC-Prior, SC-B and FLUX.1 are not compatible with each other.")
|
|
|
|
need_refresh = False
|
|
|
|
if unique_id not in core.preview_bridge_cache:
|
|
need_refresh = True
|
|
|
|
elif (core.preview_bridge_cache[unique_id][0] is not latent
|
|
or (vae_opt is None and core.preview_bridge_cache[unique_id][2] is not None)
|
|
or (vae_opt is None and core.preview_bridge_cache[unique_id][1] != preview_method)
|
|
or (vae_opt is not None and core.preview_bridge_cache[unique_id][2] is not vae_opt)):
|
|
need_refresh = True
|
|
|
|
if not need_refresh:
|
|
pixels, mask, path_item = PreviewBridge.load_image(image)
|
|
|
|
if mask is None:
|
|
mask = torch.ones(latent['samples'].shape[2:], dtype=torch.float32, device="cpu").unsqueeze(0)
|
|
if 'noise_mask' in latent:
|
|
res_latent = latent.copy()
|
|
del res_latent['noise_mask']
|
|
else:
|
|
res_latent = latent
|
|
else:
|
|
res_latent = latent.copy()
|
|
res_latent['noise_mask'] = mask
|
|
|
|
res_image = [path_item]
|
|
else:
|
|
decoded_image = decode_latent(latent, preview_method, vae_opt)
|
|
|
|
if 'noise_mask' in latent:
|
|
mask = latent['noise_mask'].squeeze(0)
|
|
|
|
decoded_pil = to_pil(decoded_image)
|
|
|
|
inverted_mask = 1 - mask
|
|
resized_mask = resize_mask(inverted_mask, (decoded_image.shape[1], decoded_image.shape[2]))
|
|
result_pil = apply_mask_alpha_to_pil(decoded_pil, resized_mask)
|
|
|
|
full_output_folder, filename, counter, _, _ = folder_paths.get_save_image_path("PreviewBridge/PBL-"+self.prefix_append, folder_paths.get_temp_directory(), result_pil.size[0], result_pil.size[1])
|
|
file = f"{filename}_{counter}.png"
|
|
result_pil.save(os.path.join(full_output_folder, file), compress_level=4)
|
|
res_image = [{
|
|
'filename': file,
|
|
'subfolder': 'PreviewBridge',
|
|
'type': 'temp',
|
|
}]
|
|
else:
|
|
mask = torch.ones(latent['samples'].shape[2:], dtype=torch.float32, device="cpu").unsqueeze(0)
|
|
res = nodes.PreviewImage().save_images(decoded_image, filename_prefix="PreviewBridge/PBL-")
|
|
res_image = res['ui']['images']
|
|
|
|
path = os.path.join(folder_paths.get_temp_directory(), 'PreviewBridge', res_image[0]['filename'])
|
|
core.set_previewbridge_image(unique_id, path, res_image[0])
|
|
core.preview_bridge_image_id_map[image] = (path, res_image[0])
|
|
core.preview_bridge_image_name_map[unique_id, path] = (image, res_image[0])
|
|
core.preview_bridge_cache[unique_id] = (latent, preview_method, vae_opt, res_image)
|
|
|
|
res_latent = latent
|
|
|
|
return {
|
|
"ui": {"images": res_image},
|
|
"result": (res_latent, mask, ),
|
|
}
|
|
|
