mirror of
https://github.com/vladmandic/automatic
synced 2026-08-28 08:00:59 +02:00
Vladimir Mandic
181 lines
7.0 KiB
Python
181 lines
7.0 KiB
Python
import math
|
|
from typing import Literal
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
try:
|
|
import scipy.stats
|
|
_scipy_available = True
|
|
except ImportError:
|
|
_scipy_available = False
|
|
|
|
def betas_for_alpha_bar(
|
|
num_diffusion_timesteps: int,
|
|
max_beta: float = 0.999,
|
|
alpha_transform_type: Literal["cosine", "exp", "laplace"] = "cosine",
|
|
dtype: torch.dtype = torch.float32,
|
|
) -> torch.Tensor:
|
|
if alpha_transform_type == "cosine":
|
|
def alpha_bar_fn(t):
|
|
return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2
|
|
elif alpha_transform_type == "laplace":
|
|
def alpha_bar_fn(t):
|
|
lmb = -0.5 * math.copysign(1, 0.5 - t) * math.log(1 - 2 * math.fabs(0.5 - t) + 1e-6)
|
|
snr = math.exp(lmb)
|
|
return math.sqrt(snr / (1 + snr))
|
|
elif alpha_transform_type == "exp":
|
|
def alpha_bar_fn(t):
|
|
return math.exp(t * -12.0)
|
|
else:
|
|
raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}")
|
|
|
|
betas = []
|
|
for i in range(num_diffusion_timesteps):
|
|
t1 = i / num_diffusion_timesteps
|
|
t2 = (i + 1) / num_diffusion_timesteps
|
|
betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta))
|
|
return torch.tensor(betas, dtype=dtype)
|
|
|
|
def rescale_zero_terminal_snr(betas: torch.Tensor) -> torch.Tensor:
|
|
alphas = 1.0 - betas
|
|
alphas_cumprod = torch.cumprod(alphas, dim=0)
|
|
alphas_bar_sqrt = alphas_cumprod.sqrt()
|
|
alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone()
|
|
alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone()
|
|
alphas_bar_sqrt -= alphas_bar_sqrt_T
|
|
alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T)
|
|
alphas_bar = alphas_bar_sqrt**2
|
|
alphas = alphas_bar[1:] / alphas_bar[:-1]
|
|
alphas = torch.cat([alphas_bar[0:1], alphas])
|
|
betas = 1 - alphas
|
|
return betas
|
|
|
|
def get_sigmas_karras(n, sigma_min, sigma_max, rho=7.0, device="cpu", dtype: torch.dtype = torch.float32):
|
|
ramp = np.linspace(0, 1, n)
|
|
min_inv_rho = sigma_min ** (1 / rho)
|
|
max_inv_rho = sigma_max ** (1 / rho)
|
|
sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
|
|
return torch.from_numpy(sigmas).to(dtype=dtype, device=device)
|
|
|
|
def get_sigmas_exponential(n, sigma_min, sigma_max, device="cpu", dtype: torch.dtype = torch.float32):
|
|
sigmas = np.exp(np.linspace(math.log(sigma_max), math.log(sigma_min), n))
|
|
return torch.from_numpy(sigmas).to(dtype=dtype, device=device)
|
|
|
|
def get_sigmas_beta(n, sigma_min, sigma_max, alpha=0.6, beta=0.6, device="cpu", dtype: torch.dtype = torch.float32):
|
|
if not _scipy_available:
|
|
raise ImportError("scipy is required for beta sigmas")
|
|
sigmas = np.array(
|
|
[
|
|
sigma_min + (ppf * (sigma_max - sigma_min))
|
|
for ppf in [
|
|
scipy.stats.beta.ppf(timestep, alpha, beta)
|
|
for timestep in 1 - np.linspace(0, 1, n)
|
|
]
|
|
]
|
|
)
|
|
return torch.from_numpy(sigmas).to(dtype=dtype, device=device)
|
|
|
|
def get_sigmas_flow(n, sigma_min, sigma_max, device="cpu", dtype: torch.dtype = torch.float32):
|
|
# Linear flow sigmas
|
|
sigmas = np.linspace(sigma_max, sigma_min, n)
|
|
return torch.from_numpy(sigmas).to(dtype=dtype, device=device)
|
|
|
|
def apply_shift(sigmas, shift):
|
|
return shift * sigmas / (1 + (shift - 1) * sigmas)
|
|
|
|
|
|
def get_base_sigmas(alphas_cumprod: torch.Tensor) -> np.ndarray:
|
|
"""Return the standard diffusion sigmas derived from alphas_cumprod."""
|
|
return np.array(((1 - alphas_cumprod.cpu().numpy()) / alphas_cumprod.cpu().numpy()) ** 0.5, dtype=np.float32)
|
|
|
|
|
|
def sigma_to_t(sigma: np.ndarray, log_sigmas: np.ndarray) -> np.ndarray:
|
|
"""Convert sigma values to corresponding discrete timesteps via interpolation."""
|
|
sigma = np.array(sigma, dtype=np.float32)
|
|
log_sigma = np.log(np.maximum(sigma, 1e-10))
|
|
dists = log_sigma - log_sigmas[:, np.newaxis]
|
|
low_idx = np.cumsum((dists >= 0), axis=0).argmax(axis=0).clip(max=log_sigmas.shape[0] - 2)
|
|
high_idx = low_idx + 1
|
|
low = log_sigmas[low_idx]
|
|
high = log_sigmas[high_idx]
|
|
w = (low - log_sigma) / (low - high)
|
|
w = np.clip(w, 0, 1)
|
|
t = (1 - w) * low_idx + w * high_idx
|
|
return t.reshape(sigma.shape).astype(np.float32)
|
|
|
|
|
|
def get_dynamic_shift(mu, base_shift, max_shift, base_seq_len, max_seq_len):
|
|
m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
|
|
b = base_shift - m * base_seq_len
|
|
return m * mu + b
|
|
|
|
def index_for_timestep(timestep, timesteps):
|
|
# Normalize inputs to numpy arrays for a robust, device-agnostic argmin
|
|
if isinstance(timestep, torch.Tensor):
|
|
timestep_np = timestep.detach().cpu().numpy()
|
|
else:
|
|
timestep_np = np.array(timestep)
|
|
|
|
if isinstance(timesteps, torch.Tensor):
|
|
timesteps_np = timesteps.detach().cpu().numpy()
|
|
else:
|
|
timesteps_np = np.array(timesteps)
|
|
|
|
# Use numpy argmin on absolute difference for stability
|
|
idx = np.abs(timesteps_np - timestep_np).argmin()
|
|
return int(idx)
|
|
|
|
def add_noise_to_sample(
|
|
original_samples: torch.Tensor,
|
|
noise: torch.Tensor,
|
|
sigmas: torch.Tensor,
|
|
timestep: torch.Tensor,
|
|
timesteps: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
step_index = index_for_timestep(timestep, timesteps)
|
|
sigma = sigmas[step_index].to(original_samples.dtype)
|
|
|
|
noisy_samples = original_samples + sigma * noise
|
|
return noisy_samples
|
|
|
|
|
|
def validate_custom_schedule_args(timesteps, sigmas):
|
|
if timesteps is not None and sigmas is not None:
|
|
raise ValueError("Only one of `timesteps` or `sigmas` should be set.")
|
|
if timesteps is None and sigmas is None:
|
|
raise ValueError("Must pass exactly one of `timesteps` or `sigmas`.")
|
|
|
|
|
|
def prepare_res4lyf_timesteps_and_sigmas(
|
|
config,
|
|
alphas_cumprod: torch.Tensor,
|
|
num_inference_steps: int | None = None,
|
|
timesteps=None,
|
|
sigmas=None,
|
|
device: str | torch.device = None,
|
|
dtype: torch.dtype = torch.float32,
|
|
):
|
|
validate_custom_schedule_args(timesteps, sigmas)
|
|
|
|
base_sigmas = get_base_sigmas(alphas_cumprod)
|
|
|
|
if timesteps is not None:
|
|
if getattr(config, "use_karras_sigmas", False) or getattr(config, "use_exponential_sigmas", False) or getattr(config, "use_beta_sigmas", False):
|
|
raise ValueError("Cannot set `timesteps` when karras/exponential/beta sigmas are enabled.")
|
|
if getattr(config, "use_flow_sigmas", False):
|
|
raise ValueError("Cannot set `timesteps` when `use_flow_sigmas` is enabled.")
|
|
|
|
timesteps_array = np.array(timesteps, dtype=np.float32)
|
|
sigmas_array = np.interp(timesteps_array, np.arange(len(base_sigmas), dtype=np.float32), base_sigmas)
|
|
num_inference_steps = len(timesteps_array)
|
|
sigmas_array = np.concatenate([sigmas_array, [0.0]])
|
|
return num_inference_steps, timesteps_array, sigmas_array
|
|
|
|
sigmas_array = np.array(sigmas, dtype=np.float32)
|
|
if num_inference_steps is None:
|
|
num_inference_steps = len(sigmas_array) - 1
|
|
log_sigmas = np.log(base_sigmas)
|
|
timesteps_array = np.array([sigma_to_t(sigma, log_sigmas) for sigma in sigmas_array[:-1]], dtype=np.float32)
|
|
return num_inference_steps, timesteps_array, sigmas_array
|