298977887/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-01-11 02:15:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

Some small changes to Load/SaveLatent.

Browse Source
This commit is contained in:
comfyanonymous 2023-05-17 23:04:40 -04:00
parent e7f2816c6f
commit a7375103b9
1 changed files with 8 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
nodes.py
View File

@ -11,6 +11,7 @@ import time
from PIL import Image from PIL import Image
from PIL.PngImagePlugin import PngInfo from PIL.PngImagePlugin import PngInfo
import numpy as np import numpy as np
import safetensors.torch
sys.path.insert(0, os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy")) sys.path.insert(0, os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy"))
@ -29,7 +30,6 @@ import importlib
import folder_paths import folder_paths
import safetensors.torch as sft
def before_node_execution(): def before_node_execution():
comfy.model_management.throw_exception_if_processing_interrupted() comfy.model_management.throw_exception_if_processing_interrupted()
@ -307,7 +307,10 @@ class SaveLatent:
file = f"{filename}_{counter:05}_.latent" file = f"{filename}_{counter:05}_.latent"
file = os.path.join(full_output_folder, file) file = os.path.join(full_output_folder, file)
sft.save_file(samples, file, metadata=metadata) output = {}
output["latent_tensor"] = samples["samples"]
safetensors.torch.save_file(output, file, metadata=metadata)
return {} return {}
@ -328,9 +331,10 @@ class LoadLatent:
def load(self, latent): def load(self, latent):
file = folder_paths.get_annotated_filepath(latent, self.input_dir) file = folder_paths.get_annotated_filepath(latent, self.input_dir)
latent = sft.load_file(file, device="cpu") latent = safetensors.torch.load_file(file, device="cpu")
samples = {"samples": latent["latent_tensor"]}
return (latent, ) return (samples, )
class CheckpointLoader: class CheckpointLoader: