Spaces:
Running
Running
File size: 3,550 Bytes
c37b2dd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 |
# Due to the current lack of maintenance for the `ComfyUI_Noise` extension,
# I have copied the code from the applied PR.
# https://github.com/BlenderNeko/ComfyUI_Noise/pull/13/files
import comfy
import torch
from comfy import sampler_helpers
class Unsampler:
@classmethod
def INPUT_TYPES(s):
return {"required":
{"model": ("MODEL",),
"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
"end_at_step": ("INT", {"default": 0, "min": 0, "max": 10000}),
"cfg": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0}),
"sampler_name": (comfy.samplers.KSampler.SAMPLERS,),
"scheduler": (comfy.samplers.KSampler.SCHEDULERS,),
"normalize": (["disable", "enable"],),
"positive": ("CONDITIONING",),
"negative": ("CONDITIONING",),
"latent_image": ("LATENT",),
}}
RETURN_TYPES = ("LATENT",)
FUNCTION = "unsampler"
CATEGORY = "sampling"
def unsampler(self, model, cfg, sampler_name, steps, end_at_step, scheduler, normalize, positive, negative,
latent_image):
normalize = normalize == "enable"
device = comfy.model_management.get_torch_device()
latent = latent_image
latent_image = latent["samples"]
end_at_step = min(end_at_step, steps - 1)
end_at_step = steps - end_at_step
noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, device="cpu")
noise_mask = None
if "noise_mask" in latent:
noise_mask = comfy.sampler_helpers.prepare_mask(latent["noise_mask"], noise.shape, device)
noise = noise.to(device)
latent_image = latent_image.to(device)
conds0 = \
{"positive": comfy.sampler_helpers.convert_cond(positive),
"negative": comfy.sampler_helpers.convert_cond(negative)}
conds = {}
for k in conds0:
conds[k] = list(map(lambda a: a.copy(), conds0[k]))
models, inference_memory = comfy.sampler_helpers.get_additional_models(conds, model.model_dtype())
comfy.model_management.load_models_gpu([model] + models, model.memory_required(noise.shape) + inference_memory)
sampler = comfy.samplers.KSampler(model, steps=steps, device=device, sampler=sampler_name,
scheduler=scheduler, denoise=1.0, model_options=model.model_options)
sigmas = sampler.sigmas.flip(0) + 0.0001
pbar = comfy.utils.ProgressBar(steps)
def callback(step, x0, x, total_steps):
pbar.update_absolute(step + 1, total_steps)
samples = sampler.sample(noise, positive, negative, cfg=cfg, latent_image=latent_image,
force_full_denoise=False, denoise_mask=noise_mask, sigmas=sigmas, start_step=0,
last_step=end_at_step, callback=callback)
if normalize:
# technically doesn't normalize because unsampling is not guaranteed to end at a std given by the schedule
samples -= samples.mean()
samples /= samples.std()
samples = samples.cpu()
comfy.sampler_helpers.cleanup_additional_models(models)
out = latent.copy()
out["samples"] = samples
return (out,)
|