73 lines
3.1 KiB
Python
73 lines
3.1 KiB
Python
import numpy as np
|
|
import torch
|
|
import comfy.samplers
|
|
import random
|
|
from ..components import utility
|
|
|
|
def prepare_noise(latent_image, seed, noise_inds=None, noise_device="cpu", variation_seed=None, variation_strength=None):
|
|
latent_size = latent_image.size()
|
|
latent_size_1batch = [1, latent_size[1], latent_size[2], latent_size[3]]
|
|
|
|
if variation_strength is not None and variation_strength > 0:
|
|
if noise_device == "cpu":
|
|
variation_generator = torch.manual_seed(variation_seed)
|
|
else:
|
|
torch.cuda.manual_seed(variation_seed)
|
|
variation_generator = None
|
|
|
|
variation_latent = torch.randn(latent_size_1batch, dtype=latent_image.dtype, layout=latent_image.layout, generator=variation_generator, device=noise_device)
|
|
else:
|
|
variation_latent = None
|
|
|
|
def apply_variation(input_latent):
|
|
if variation_latent is None:
|
|
return input_latent
|
|
else:
|
|
strength = variation_strength
|
|
|
|
variation_noise = variation_latent.expand(input_latent.size()[0], -1, -1, -1)
|
|
result = (1 - strength) * input_latent + strength * variation_noise
|
|
return result
|
|
|
|
if noise_device == "cpu":
|
|
generator = torch.manual_seed(seed)
|
|
else:
|
|
torch.cuda.manual_seed(seed)
|
|
generator = None
|
|
|
|
if noise_inds is None:
|
|
latents = torch.randn(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, generator=generator, device=noise_device)
|
|
latents = apply_variation(latents)
|
|
return latents
|
|
|
|
unique_inds, inverse = np.unique(noise_inds, return_inverse=True)
|
|
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=noise_device)
|
|
if i in unique_inds:
|
|
noises.append(noise)
|
|
noises = [noises[i] for i in inverse]
|
|
noises = torch.cat(noises, axis=0)
|
|
return noises
|
|
|
|
def noisy_samples(model, device, steps, cfg, sampler_name, scheduler_name, positive, negative, latent_image, denoise, seed, variation_extender):
|
|
torch_device = comfy.model_management.get_torch_device()
|
|
noise_device = "cpu" if device == "CPU" else torch_device
|
|
batch_inds = latent_image["batch_index"] if "batch_index" in latent_image else None
|
|
|
|
random.seed(seed)
|
|
variation_seed = random.randint(0, utility.MAX_SEED)
|
|
|
|
noise = prepare_noise(latent_image["samples"], seed, batch_inds, noise_device, variation_seed=variation_seed, variation_strength=variation_extender)
|
|
|
|
pbar = comfy.utils.ProgressBar(steps)
|
|
def callback(step, x0, x, total_steps):
|
|
preview_bytes = None
|
|
pbar.update_absolute(step + 1, total_steps, preview_bytes)
|
|
|
|
samples = comfy.sample.sample(model, noise, steps, cfg, sampler_name, scheduler_name, positive, negative, latent_image["samples"], denoise=denoise, disable_noise=False, start_step=None, last_step=None, force_full_denoise=False, noise_mask=None, callback=callback, disable_pbar=False, seed=seed)
|
|
out = latent_image.copy()
|
|
out["samples"] = samples
|
|
return (out,)
|
|
|