added selector to ksamplers for choosing batch_behavior

This commit is contained in:
FizzleDorf 2023-12-04 01:28:00 -05:00
parent 9328cd4098
commit 093bb323d8
2 changed files with 21 additions and 17 deletions

View File

@ -6,10 +6,9 @@ import comfy.utils
import math
import numpy as np
def prepare_noise(latent_image, seeds, noise_inds=None):
def prepare_noise(latent_image, seeds, batch_behavior = "randomize"):
"""
Creates random noise given a latent image and a seed or a list of seeds.
Optional arg noise_inds can be used to select specific noise indices.
Creates noise based on the batch behavior, a latent image and a seed or a list of seeds.
"""
num_latents = latent_image.size(0)
@ -21,16 +20,20 @@ def prepare_noise(latent_image, seeds, noise_inds=None):
noises = []
for i in range(num_latents):
if i < len(seeds): # Use the provided seeds if available
for i in range(num_latents):
if i < len(seeds): # Use the provided seeds if available then follow behavior
seed = seeds[i]
else:
seed = seeds[-1] + i # Increment the last seed for additional latents
elif batch_behavior == "randomize" :
seed = torch.randint(0, 2 ** 32, (1,)).item()
elif batch_behavior == "fixed":
seed = seeds[-1]
elif batch_behavior == "increment":
seed = seeds[-1] + i
else :
seed = seeds[-1] - i
#else: maybe add this add a toggle or dropdown?
# seed = torch.randint(0, 2**32, (1,)).item() # Generate a random seed for additional latents
generator.manual_seed(seed)
print("seed:", seed) #get rid of this after testing
print("seed:", seed)
noise = torch.randn([1] + list(latent_image.size())[1:], dtype=latent_image.dtype, layout=latent_image.layout, device="cpu", generator=generator)
noises.append(noise)

View File

@ -1252,13 +1252,12 @@ class SetLatentNoiseMask:
s["noise_mask"] = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1]))
return (s,)
def common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent, denoise=1.0, disable_noise=False, start_step=None, last_step=None, force_full_denoise=False):
def common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent, denoise=1.0, batch_behavior = "randomize", disable_noise=False, start_step=None, last_step=None, force_full_denoise=False):
latent_image = latent["samples"]
if disable_noise:
noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, device="cpu")
else:
batch_inds = latent["batch_index"] if "batch_index" in latent else None
noise = comfy.sample.prepare_noise(latent_image, seed, batch_inds)
noise = comfy.sample.prepare_noise(latent_image, seed, batch_behavior)
noise_mask = None
if "noise_mask" in latent:
@ -1279,6 +1278,7 @@ class KSampler:
return {"required":
{"model": ("MODEL",),
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
"batch_behavior": (["randomize", "fixed", "increment", "decrement"],),
"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
"cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}),
"sampler_name": (comfy.samplers.KSampler.SAMPLERS, ),
@ -1295,8 +1295,8 @@ class KSampler:
CATEGORY = "sampling"
def sample(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=1.0):
return common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise)
def sample(self, model, seed, batch_behavior, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=1.0):
return common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise, batch_behavior=batch_behavior)
class KSamplerAdvanced:
@classmethod
@ -1305,6 +1305,7 @@ class KSamplerAdvanced:
{"model": ("MODEL",),
"add_noise": (["enable", "disable"], ),
"noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
"batch_behavior": (["randomize", "fixed", "increment", "decrement"],),
"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
"cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}),
"sampler_name": (comfy.samplers.KSampler.SAMPLERS, ),
@ -1323,14 +1324,14 @@ class KSamplerAdvanced:
CATEGORY = "sampling"
def sample(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, denoise=1.0):
def sample(self, model, add_noise, noise_seed, batch_behavior, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, start_at_step, end_at_step, return_with_leftover_noise, denoise=1.0):
force_full_denoise = True
if return_with_leftover_noise == "enable":
force_full_denoise = False
disable_noise = False
if add_noise == "disable":
disable_noise = True
return common_ksampler(model, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise, disable_noise=disable_noise, start_step=start_at_step, last_step=end_at_step, force_full_denoise=force_full_denoise)
return common_ksampler(model, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise, batch_behavior=batch_behavior, disable_noise=disable_noise, start_step=start_at_step, last_step=end_at_step, force_full_denoise=force_full_denoise)
class SaveImage:
def __init__(self):