diff --git a/README.md b/README.md index 34b9754..c3b195e 100644 --- a/README.md +++ b/README.md @@ -55,7 +55,6 @@ Workflow JSON: [NetDistAdvancedV2.json](https://github.com/city96/ComfyUI_NetDis ![NetDistSaved](https://github.com/city96/ComfyUI_NetDist/assets/125218114/a39b5117-af1b-4f2c-a94e-5a330acc8ea4) - ### Remote images The `LoadImageUrl` ('Load Image (URL)') Node acts just like the normal 'Load Image' node. @@ -65,9 +64,25 @@ The `SaveImageUrl` ('Save Image (URL)') Node sends a POST request to the target - The filenames are **not** guaranteed to be unique across batches since they aren't saved locally. You should handle this server-side. - No data is written to disk on the server. +### Remote latents + +This node pack has a set of nodes which should (in theory) allow you to pass latents between the nodes seamlessly. A node to save the input latent as a `.npy` file is provided. This node also returns the filename of the saved latent, which can then be loaded by the other instance. + +To load a latent from the other instance, you can plug the filename into this URL: + +``` +# change the filename with a string replacement node. +http://127.0.0.1:8188/view?filename=ComfyUI_00001_.latent&type=output` +# To load them from the input folder instead, change type to 'input' +http://127.0.0.1:8188/view?filename=TestLatent.npy&type=input +``` + +The `LoadLatentNumpy` node can also load the default safetensor latents, the npy ones (simple numpy file containing just the latent in the standard torch format) as well as the sd_scripts npz cache files. + +![LatentSave](https://github.com/city96/ComfyUI_NetDist/assets/125218114/cd68d8dc-bd96-4018-82c9-400337fc5f80) ### Things you probably shouldn't do: -- Queue a workflow on the same client multiple times. +- Queue a workflow on the same remote worker multiple times from the same client. - ~~Expect this to work smoothly.~~ ## Roadmap diff --git a/__init__.py b/__init__.py index a7e64d5..98f0b8a 100644 --- a/__init__.py +++ b/__init__.py @@ -15,6 +15,9 @@ else: from .nodes.images import NODE_CLASS_MAPPINGS as ImgNodes NODE_CLASS_MAPPINGS.update(ImgNodes) + from .nodes.latents import NODE_CLASS_MAPPINGS as LatNodes + NODE_CLASS_MAPPINGS.update(LatNodes) + from .nodes.workflows import NODE_CLASS_MAPPINGS as WrkNodes NODE_CLASS_MAPPINGS.update(WrkNodes) diff --git a/nodes/latents.py b/nodes/latents.py new file mode 100644 index 0000000..8cad1bc --- /dev/null +++ b/nodes/latents.py @@ -0,0 +1,170 @@ +import os +import torch +import requests +import safetensors.torch +import numpy as np +from io import BytesIO + +import folder_paths + +class LoadLatentNumpy: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + exts = [".latent", ".safetensors", ".npy", ".npz"] + 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))] + files = [f for f in files if any([f.endswith(x) for x in exts])] + return { + "required": { + "latent": [sorted(files), ] + }, + } + + RETURN_TYPES = ("LATENT",) + FUNCTION = "load" + CATEGORY = "remote/latent" + TITLE = "Load Latent (Numpy)" + + def load_comfy(self, file): + # From default node - renamed safetensors file + if type(file) == str: + data = safetensors.torch.load_file(file) + else: + data = safetensors.torch.load(file) + + latent = data["latent_tensor"].to(torch.float32) + if "latent_format_version_0" not in data: + latent *= 1.0 / 0.18215 # XL? + return latent + + def load_numpy(self, file): + # plain npy file - saved as-is + return torch.from_numpy(np.load(file)) + + def load_koyha(self, file): + # generated by sd_scripts - npz + if "latents" in data.keys(): + latent = data["latents"] + else: + latent = [x for x in data.items() if x.shape > 3][0] + return torch.from_numpy(latent) + + def load(self, latent): + path = folder_paths.get_annotated_filepath(latent) + name, ext = os.path.splitext(latent) + + if ext in [".latent", ".safetensors"]: + latent = self.load_comfy(path) + elif ext == ".npy": + latent = self.load_numpy(path) + elif ext == ".npz": + latent = self.load_koyha(path) + else: + try: + latent = self.load_numpy(path) + except: + raise ValueError(f"Unknown latent extension '{ext}'") + + if len(latent.shape) == 3: + latent = latent.unsqueeze(0) + print("asdasd", latent.shape) + + return ({"samples": latent.to(torch.float32)},) + + @classmethod + def IS_CHANGED(s, latent): + 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): + if not folder_paths.exists_annotated_filepath(latent): + return f"Invalid latent file '{latent}'" + return True + +class LoadLatentUrl(LoadLatentNumpy): + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "url": ("STRING", { "multiline": False, }) + } + } + + RETURN_TYPES = ("LATENT",) + TITLE = "Load Latent (URL)" + + def load(self, url): + buffer = BytesIO() + with requests.get(url, stream=True, timeout=16) as r: + r.raise_for_status() + buffer.write(r.content) + buffer.seek(0) + + if ".latent" in url or ".safetensors" in url: + latent = self.load_comfy(buffer) + elif ".npy" in url: + latent = self.load_numpy(buffer) + elif ".npz" in url: + latent = self.load_koyha(buffer) + else: + try: + latent = self.load_comfy(buffer) + except: + raise ValueError(f"Unknown latent extension '{url}'") + + if len(latent.shape) == 3: + latent = latent.unsqueeze(0) + + del buffer + return ({"samples": latent.to(torch.float32)},) + + @classmethod + def IS_CHANGED(s, url): + return str(url) + + @classmethod + def VALIDATE_INPUTS(s, url): + return True + +class SaveLatentNumpy: + def __init__(self): + self.output_dir = folder_paths.get_output_directory() + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "samples": ("LATENT",), + "filename_prefix": ("STRING", {"default": "latents/ComfyUI"}) + } + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("filename",) + OUTPUT_NODE = True + FUNCTION = "save" + CATEGORY = "remote/latent" + TITLE = "Save Latent (Numpy)" + + def save(self, samples, filename_prefix="ComfyUI"): + full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir) + fname = f"{filename}_{counter:05}_.npy" + path = os.path.join(full_output_folder, fname) + np.save(path, samples["samples"].numpy()) + return (fname,) + +NODE_CLASS_MAPPINGS = { + "LoadLatentNumpy" : LoadLatentNumpy, + "LoadLatentUrl" : LoadLatentUrl, + "SaveLatentNumpy" : SaveLatentNumpy, +}