| import torch |
| import k_diffusion.sampling |
| from k_diffusion.sampling import to_d, default_noise_sampler, get_ancestral_step |
| from tqdm.auto import trange |
| from modules import sd_samplers, sd_samplers_common, sd_samplers_kdiffusion |
|
|
| _SAMPLER_REGISTERED = False |
|
|
|
|
| @torch.no_grad() |
| def sample_kohaku_lonyu_yog(model, x, sigmas, extra_args=None, callback=None, |
| disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), |
| s_noise=1., noise_sampler=None, eta=1.): |
| """ |
| Kohaku_LoNyu_Yog Sampler - Geometric Second-Order Method |
| """ |
| extra_args = {} if extra_args is None else extra_args |
| s_in = x.new_ones([x.shape[0]]) |
| noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler |
|
|
| steps_total = len(sigmas) - 1 |
| halfway_point = steps_total // 2 |
|
|
| for i in trange(steps_total, disable=disable, desc="Kohaku Sampling"): |
| gamma = min(s_churn / steps_total, 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0. |
| sigma_hat = sigmas[i] * (gamma + 1) |
|
|
| if gamma > 0: |
| eps = torch.randn_like(x) * s_noise |
| x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5 |
|
|
| denoised = model(x, sigma_hat * s_in, **extra_args) |
| d = to_d(x, sigma_hat, denoised) |
|
|
| sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta) |
| dt = sigma_down - sigmas[i] |
|
|
| if i <= halfway_point: |
| x_antipode = -x |
|
|
| denoised2 = model(x_antipode, sigma_hat * s_in, **extra_args) |
| d2 = to_d(x_antipode, sigma_hat, denoised2) |
|
|
| v_down = (d + d2) / 2 |
| x_closer = x + v_down * dt |
|
|
| denoised3 = model(x_closer, sigma_hat * s_in, **extra_args) |
| d3 = to_d(x_closer, sigma_hat, denoised3) |
|
|
| real_d = (d + d3) / 2 |
| x = x + real_d * dt |
|
|
| if sigma_up > 0: |
| x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up |
| else: |
| x = x + d * dt |
| if sigma_up > 0: |
| x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up |
|
|
| if callback is not None: |
| callback({ |
| 'x': x, |
| 'i': i, |
| 'sigma': sigmas[i], |
| 'sigma_hat': sigma_hat, |
| 'denoised': denoised |
| }) |
|
|
| return x |
|
|
|
|
| def register_kohaku_sampler(): |
| """Регистрирует Kohaku_LoNyu_Yog сэмплер в WebUI""" |
| global _SAMPLER_REGISTERED |
|
|
| if _SAMPLER_REGISTERED: |
| return |
|
|
| if any(s.name == 'Kohaku_LoNyu_Yog' for s in sd_samplers.all_samplers): |
| _SAMPLER_REGISTERED = True |
| return |
|
|
| if not hasattr(k_diffusion.sampling, 'sample_kohaku_lonyu_yog'): |
| setattr(k_diffusion.sampling, 'sample_kohaku_lonyu_yog', sample_kohaku_lonyu_yog) |
|
|
| sampler_data = sd_samplers_common.SamplerData( |
| name='Kohaku_LoNyu_Yog', |
| constructor=lambda model: sd_samplers_kdiffusion.KDiffusionSampler('sample_kohaku_lonyu_yog', model), |
| aliases=['kohaku', 'lonyu'], |
| options={'second_order': True} |
| ) |
|
|
| sd_samplers.all_samplers.append(sampler_data) |
| sd_samplers.all_samplers_map = {x.name: x for x in sd_samplers.all_samplers} |
|
|
| _SAMPLER_REGISTERED = True |
| print("✓ Kohaku_LoNyu_Yog sampler registered successfully!") |
|
|
