62 lines
2.5 KiB
Python
62 lines
2.5 KiB
Python
import os
|
|
import torch
|
|
import folder_paths
|
|
from nodes import KSamplerAdvanced
|
|
import comfy.samplers
|
|
from safetensors.torch import save_file
|
|
|
|
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)
|