121 lines
4.9 KiB
Python
121 lines
4.9 KiB
Python
import os
|
|
import torch
|
|
import folder_paths
|
|
from nodes import KSamplerAdvanced
|
|
import comfy.samplers
|
|
from safetensors.torch import save_file
|
|
|
|
from comfy_api.latest import io
|
|
try:
|
|
from comfy_extras.nodes_custom_sampler import SamplerCustomAdvanced
|
|
except ImportError:
|
|
import nodes
|
|
SamplerCustomAdvanced = nodes.NODE_CLASS_MAPPINGS["SamplerCustomAdvanced"]
|
|
|
|
class EvalSampler(KSamplerAdvanced):
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
input_types = KSamplerAdvanced.INPUT_TYPES()
|
|
new_inputs = {
|
|
"required": dict(input_types["required"])
|
|
}
|
|
if "optional" in input_types:
|
|
new_inputs["optional"] = dict(input_types["optional"])
|
|
|
|
new_inputs["required"]["output_folder"] = ("STRING", {"default": "eval_latents"})
|
|
return new_inputs
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
FUNCTION = "sample_eval"
|
|
CATEGORY = "sampling"
|
|
|
|
def sample_eval(self, model, add_noise, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, start_at_step, end_at_step, return_with_leftover_noise, output_folder, denoise=1.0):
|
|
model = model.clone()
|
|
|
|
out_dir = os.path.join(folder_paths.get_output_directory(), output_folder)
|
|
os.makedirs(out_dir, exist_ok=True)
|
|
|
|
step_counter = [0]
|
|
|
|
old_cfg_func = model.model_options.get("sampler_cfg_function", None)
|
|
|
|
def custom_cfg_hook(args):
|
|
cond_denoised = args["cond_denoised"]
|
|
uncond_denoised = args["uncond_denoised"]
|
|
|
|
step = step_counter[0]
|
|
|
|
batch_size = cond_denoised.shape[0]
|
|
for b_idx in range(batch_size):
|
|
b_seed = noise_seed + b_idx
|
|
|
|
cond_b = cond_denoised[b_idx:b_idx+1].contiguous()
|
|
save_file({"latent": cond_b}, os.path.join(out_dir, f"{b_seed}_S{step}.safetensors"))
|
|
|
|
uncond_b = uncond_denoised[b_idx:b_idx+1].contiguous()
|
|
save_file({"latent": uncond_b}, os.path.join(out_dir, f"{b_seed}_S{step}_uncond.safetensors"))
|
|
|
|
step_counter[0] += 1
|
|
|
|
if old_cfg_func is not None:
|
|
return old_cfg_func(args)
|
|
else:
|
|
cfg_result = args["uncond_denoised"] + (args["cond_denoised"] - args["uncond_denoised"]) * args["cond_scale"]
|
|
return args["input"] - cfg_result
|
|
|
|
model.set_model_sampler_cfg_function(custom_cfg_hook)
|
|
|
|
return super().sample(model, add_noise, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, start_at_step, end_at_step, return_with_leftover_noise, denoise=denoise)
|
|
|
|
|
|
class EvalSamplerAdvanced(SamplerCustomAdvanced):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
schema = super().define_schema()
|
|
schema.node_id = "EvalSamplerAdvanced"
|
|
schema.inputs.append(io.String.Input("output_folder", default="eval_latents"))
|
|
return schema
|
|
|
|
@classmethod
|
|
def execute(cls, noise, guider, sampler, sigmas, latent_image, output_folder) -> io.NodeOutput:
|
|
# Clone guider's model_patcher to avoid side effects on the original model
|
|
guider.model_patcher = guider.model_patcher.clone()
|
|
if hasattr(guider, "model_options"):
|
|
guider.model_options = guider.model_patcher.model_options
|
|
|
|
out_dir = os.path.join(folder_paths.get_output_directory(), output_folder)
|
|
os.makedirs(out_dir, exist_ok=True)
|
|
|
|
step_counter = [0]
|
|
noise_seed = getattr(noise, "seed", 0)
|
|
|
|
old_cfg_func = guider.model_patcher.model_options.get("sampler_cfg_function", None)
|
|
|
|
def custom_cfg_hook(args):
|
|
cond_denoised = args["cond_denoised"]
|
|
uncond_denoised = args["uncond_denoised"]
|
|
|
|
step = step_counter[0]
|
|
|
|
batch_size = cond_denoised.shape[0]
|
|
for b_idx in range(batch_size):
|
|
b_seed = noise_seed + b_idx
|
|
|
|
cond_b = cond_denoised[b_idx:b_idx+1].contiguous()
|
|
save_file({"latent": cond_b}, os.path.join(out_dir, f"{b_seed}_S{step}.safetensors"))
|
|
|
|
uncond_b = uncond_denoised[b_idx:b_idx+1].contiguous()
|
|
save_file({"latent": uncond_b}, os.path.join(out_dir, f"{b_seed}_S{step}_uncond.safetensors"))
|
|
|
|
step_counter[0] += 1
|
|
|
|
if old_cfg_func is not None:
|
|
return old_cfg_func(args)
|
|
else:
|
|
cfg_result = args["uncond_denoised"] + (args["cond_denoised"] - args["uncond_denoised"]) * args["cond_scale"]
|
|
return args["input"] - cfg_result
|
|
|
|
guider.model_patcher.set_model_sampler_cfg_function(custom_cfg_hook, disable_cfg1_optimization=True)
|
|
|
|
return super().execute(noise, guider, sampler, sigmas, latent_image)
|