xmarre
ef8a641be0
* Update inner_sample signature * Allow optional context arguments in NAG cross attention * Stop passing kwargs through outer sample executor * Set appended query count from origin batch * Make NAG attention mask slices contiguous
34 lines
1.9 KiB
Python
34 lines
1.9 KiB
Python
import comfy
|
|
from .samplers import KSamplerWithNAG
|
|
from .samplers import sample_with_nag as samplers_sample_with_nag
|
|
|
|
|
|
def sample_with_nag(
|
|
model, noise, steps, cfg, nag_scale, nag_tau, nag_alpha, nag_sigma_end, sampler_name, scheduler, positive, negative, nag_negative, latent_image, denoise=1.0, disable_noise=False, start_step=None, last_step=None, force_full_denoise=False, noise_mask=None, sigmas=None, callback=None, disable_pbar=False, seed=None, latent_shapes=None, **kwargs):
|
|
sampler = KSamplerWithNAG(model, steps=steps, device=model.load_device, sampler=sampler_name, scheduler=scheduler, denoise=denoise, model_options=model.model_options)
|
|
|
|
samples = sampler.sample(
|
|
noise, positive, negative, nag_negative,
|
|
cfg=cfg, nag_scale=nag_scale, nag_tau=nag_tau, nag_alpha=nag_alpha, nag_sigma_end=nag_sigma_end,
|
|
latent_image=latent_image,
|
|
start_step=start_step, last_step=last_step, force_full_denoise=force_full_denoise,
|
|
denoise_mask=noise_mask, sigmas=sigmas, callback=callback, disable_pbar=disable_pbar, seed=seed,
|
|
latent_shapes=latent_shapes, **kwargs,
|
|
)
|
|
samples = samples.to(comfy.model_management.intermediate_device())
|
|
return samples
|
|
|
|
|
|
def sample_custom_with_nag(
|
|
model, noise, cfg, nag_scale, nag_tau, nag_alpha, nag_sigma_end, sampler, sigmas, positive, negative, nag_negative, latent_image, noise_mask=None, callback=None, disable_pbar=False, seed=None, latent_shapes=None, **kwargs):
|
|
samples = samplers_sample_with_nag(
|
|
model, noise, positive, negative, nag_negative,
|
|
cfg, nag_scale, nag_tau, nag_alpha, nag_sigma_end,
|
|
model.load_device, sampler, sigmas,
|
|
model_options=model.model_options, latent_image=latent_image, denoise_mask=noise_mask,
|
|
callback=callback, disable_pbar=disable_pbar, seed=seed,
|
|
latent_shapes=latent_shapes, **kwargs,
|
|
)
|
|
samples = samples.to(comfy.model_management.intermediate_device())
|
|
return samples
|