44 lines
1.2 KiB
Python
44 lines
1.2 KiB
Python
import logging
|
|
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
|
log = logging.getLogger(__name__)
|
|
|
|
import torch
|
|
|
|
from ..utils.callback_utils import get_callback_fn
|
|
from ..utils.latent_utils import add_latent_noise
|
|
from ..utils.sampling_utils import get_model_fn, get_sample_args
|
|
|
|
|
|
def run_sampler(model, latents, positive, negative, sigmas, cfg, sampler_fn, add_noise=False, seed=0):
|
|
# seed
|
|
torch.manual_seed(seed)
|
|
torch.cuda.manual_seed(seed)
|
|
generator = torch.Generator(device=model.device)
|
|
generator.manual_seed(seed)
|
|
|
|
# prepare latents
|
|
latent_shape = latents.shape
|
|
|
|
if add_noise:
|
|
z = add_latent_noise(model, latent_shape, sigmas, latents, generator)
|
|
else:
|
|
z = latents.clone()
|
|
|
|
# prepare model and args
|
|
positive, negative = get_sample_args(model, positive, negative)
|
|
model_fn = get_model_fn(model)
|
|
|
|
# sampling
|
|
callback_fn = get_callback_fn(model, len(sigmas)-1)
|
|
extra_args = {
|
|
"positive": positive,
|
|
"negative": negative,
|
|
"cfg": cfg
|
|
}
|
|
z = sampler_fn(model_fn, z, sigmas, callback=callback_fn, extra_args=extra_args)
|
|
|
|
# cleanup
|
|
model.dit.to(model.offload_device)
|
|
|
|
return z
|