|
|
|
@ -310,6 +310,24 @@ class SamplerDPMAdaptative:
|
|
|
|
|
"s_noise":s_noise }) |
|
|
|
|
return (sampler, ) |
|
|
|
|
|
|
|
|
|
class Noise_EmptyNoise: |
|
|
|
|
def __init__(self): |
|
|
|
|
self.seed = 0 |
|
|
|
|
|
|
|
|
|
def generate_noise(self, input_latent): |
|
|
|
|
latent_image = input_latent["samples"] |
|
|
|
|
return torch.zeros(shape, dtype=latent_image.dtype, layout=latent_image.layout, device="cpu") |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class Noise_RandomNoise: |
|
|
|
|
def __init__(self, seed): |
|
|
|
|
self.seed = seed |
|
|
|
|
|
|
|
|
|
def generate_noise(self, input_latent): |
|
|
|
|
latent_image = input_latent["samples"] |
|
|
|
|
batch_inds = input_latent["batch_index"] if "batch_index" in input_latent else None |
|
|
|
|
return comfy.sample.prepare_noise(latent_image, self.seed, batch_inds) |
|
|
|
|
|
|
|
|
|
class SamplerCustom: |
|
|
|
|
@classmethod |
|
|
|
|
def INPUT_TYPES(s): |
|
|
|
@ -337,10 +355,9 @@ class SamplerCustom:
|
|
|
|
|
latent = latent_image |
|
|
|
|
latent_image = latent["samples"] |
|
|
|
|
if not add_noise: |
|
|
|
|
noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, device="cpu") |
|
|
|
|
noise = Noise_EmptyNoise().generate_noise(latent) |
|
|
|
|
else: |
|
|
|
|
batch_inds = latent["batch_index"] if "batch_index" in latent else None |
|
|
|
|
noise = comfy.sample.prepare_noise(latent_image, noise_seed, batch_inds) |
|
|
|
|
noise = Noise_RandomNoise(noise_seed).generate_noise(latent) |
|
|
|
|
|
|
|
|
|
noise_mask = None |
|
|
|
|
if "noise_mask" in latent: |
|
|
|
@ -361,6 +378,100 @@ class SamplerCustom:
|
|
|
|
|
out_denoised = out |
|
|
|
|
return (out, out_denoised) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class CFGGuider: |
|
|
|
|
@classmethod |
|
|
|
|
def INPUT_TYPES(s): |
|
|
|
|
return {"required": |
|
|
|
|
{"model": ("MODEL",), |
|
|
|
|
"positive": ("CONDITIONING", ), |
|
|
|
|
"negative": ("CONDITIONING", ), |
|
|
|
|
"cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}), |
|
|
|
|
} |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
RETURN_TYPES = ("GUIDER",) |
|
|
|
|
|
|
|
|
|
FUNCTION = "get_guider" |
|
|
|
|
CATEGORY = "sampling/custom_sampling/guiders" |
|
|
|
|
|
|
|
|
|
def get_guider(self, model, positive, negative, cfg): |
|
|
|
|
guider = comfy.samplers.CFGGuider(model) |
|
|
|
|
guider.set_conds({"positive": positive, "negative": negative}) |
|
|
|
|
guider.set_cfg(cfg) |
|
|
|
|
return (guider,) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class DisableNoise: |
|
|
|
|
@classmethod |
|
|
|
|
def INPUT_TYPES(s): |
|
|
|
|
return {"required":{ |
|
|
|
|
} |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
RETURN_TYPES = ("NOISE",) |
|
|
|
|
FUNCTION = "get_noise" |
|
|
|
|
CATEGORY = "sampling/custom_sampling/noise" |
|
|
|
|
|
|
|
|
|
def get_noise(self, noise_seed): |
|
|
|
|
return (Noise_EmptyNoise(),) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class RandomNoise(DisableNoise): |
|
|
|
|
@classmethod |
|
|
|
|
def INPUT_TYPES(s): |
|
|
|
|
return {"required":{ |
|
|
|
|
"noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), |
|
|
|
|
} |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
def get_noise(self, noise_seed): |
|
|
|
|
return (Noise_RandomNoise(noise_seed),) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class SamplerCustomAdvanced: |
|
|
|
|
@classmethod |
|
|
|
|
def INPUT_TYPES(s): |
|
|
|
|
return {"required": |
|
|
|
|
{"noise": ("NOISE", ), |
|
|
|
|
"guider": ("GUIDER", ), |
|
|
|
|
"sampler": ("SAMPLER", ), |
|
|
|
|
"sigmas": ("SIGMAS", ), |
|
|
|
|
"latent_image": ("LATENT", ), |
|
|
|
|
} |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
RETURN_TYPES = ("LATENT","LATENT") |
|
|
|
|
RETURN_NAMES = ("output", "denoised_output") |
|
|
|
|
|
|
|
|
|
FUNCTION = "sample" |
|
|
|
|
|
|
|
|
|
CATEGORY = "sampling/custom_sampling" |
|
|
|
|
|
|
|
|
|
def sample(self, noise, guider, sampler, sigmas, latent_image): |
|
|
|
|
latent = latent_image |
|
|
|
|
latent_image = latent["samples"] |
|
|
|
|
|
|
|
|
|
noise_mask = None |
|
|
|
|
if "noise_mask" in latent: |
|
|
|
|
noise_mask = latent["noise_mask"] |
|
|
|
|
|
|
|
|
|
x0_output = {} |
|
|
|
|
callback = latent_preview.prepare_callback(guider.model_patcher, sigmas.shape[-1] - 1, x0_output) |
|
|
|
|
|
|
|
|
|
disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED |
|
|
|
|
samples = guider.sample(noise.generate_noise(latent), latent_image, sampler, sigmas, denoise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=noise.seed) |
|
|
|
|
samples = samples.to(comfy.model_management.intermediate_device()) |
|
|
|
|
|
|
|
|
|
out = latent.copy() |
|
|
|
|
out["samples"] = samples |
|
|
|
|
if "x0" in x0_output: |
|
|
|
|
out_denoised = latent.copy() |
|
|
|
|
out_denoised["samples"] = guider.model_patcher.model.process_latent_out(x0_output["x0"].cpu()) |
|
|
|
|
else: |
|
|
|
|
out_denoised = out |
|
|
|
|
return (out, out_denoised) |
|
|
|
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = { |
|
|
|
|
"SamplerCustom": SamplerCustom, |
|
|
|
|
"BasicScheduler": BasicScheduler, |
|
|
|
@ -378,4 +489,9 @@ NODE_CLASS_MAPPINGS = {
|
|
|
|
|
"SamplerDPMAdaptative": SamplerDPMAdaptative, |
|
|
|
|
"SplitSigmas": SplitSigmas, |
|
|
|
|
"FlipSigmas": FlipSigmas, |
|
|
|
|
|
|
|
|
|
"CFGGuider": CFGGuider, |
|
|
|
|
"RandomNoise": RandomNoise, |
|
|
|
|
"DisableNoise": DisableNoise, |
|
|
|
|
"SamplerCustomAdvanced": SamplerCustomAdvanced, |
|
|
|
|
} |
|
|
|
|