seed scheduling added

This commit is contained in:
FizzleDorf 2023-12-01 09:15:05 -05:00
parent ec7a00aa96
commit 860894d83f

View File

@ -6,22 +6,32 @@ import comfy.utils
import math import math
import numpy as np import numpy as np
def prepare_noise(latent_image, seed, noise_inds=None): def prepare_noise(latent_image, seeds, noise_inds=None):
""" """
creates random noise given a latent image and a seed. Creates random noise given a latent image and a seed or a list of seeds.
optional arg skip can be used to skip and discard x number of noise generations for a given seed Optional arg noise_inds can be used to select specific noise indices.
""" """
generator = torch.manual_seed(seed) num_latents = latent_image.size(0)
if noise_inds is None:
return torch.randn(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, generator=generator, device="cpu") if not isinstance(seeds, list):
seeds = [seeds]
generator = torch.Generator()
generator.manual_seed(seeds[0]) # Use the first seed as the default generator seed
unique_inds, inverse = np.unique(noise_inds, return_inverse=True)
noises = [] noises = []
for i in range(unique_inds[-1]+1):
noise = torch.randn([1] + list(latent_image.size())[1:], dtype=latent_image.dtype, layout=latent_image.layout, generator=generator, device="cpu") for i in range(num_latents):
if i in unique_inds: if i < len(seeds): # Use the provided seeds if available
noises.append(noise) seed = seeds[i]
noises = [noises[i] for i in inverse] else:
seed = torch.randint(0, 2**32, (1,)).item() # Generate a random seed for additional latents
generator.manual_seed(seed)
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)
noises = [noises[i] for i in range(num_latents)]
noises = torch.cat(noises, axis=0) noises = torch.cat(noises, axis=0)
return noises return noises