mirror of
https://github.com/vladmandic/automatic
synced 2026-09-06 13:00:44 +02:00
feat(samplers): add sigma-schedule parity to ER-SDE
ERSDEScheduler now accepts use_karras_sigmas, use_exponential_sigmas, use_beta_sigmas, and use_flow_sigmas, matching the other flow schedulers. The VP path derives alpha/sigma/lambda from the k-diffusion sigma so the karras/beta/exponential transforms can use fractional timesteps; the default schedule is numerically unchanged. use_flow_sigmas triggers flow mode and add_noise tolerates fractional timesteps. Wire the new keys into the ER-SDE presets so the sigma method selector drives them, and cover ER-SDE in the scheduler stability test.
This commit is contained in:
CalamitousFelicitousness
@@ -120,6 +120,9 @@ class ERSDEScheduler(SchedulerMixin, ConfigMixin):
|
||||
set_alpha_to_one: bool = True,
|
||||
lower_order_final: bool = True,
|
||||
use_karras_sigmas: bool = False,
|
||||
use_exponential_sigmas: bool = False,
|
||||
use_beta_sigmas: bool = False,
|
||||
use_flow_sigmas: bool = False,
|
||||
num_integration_steps: int = 100,
|
||||
shift: float = 1.0,
|
||||
use_dynamic_shifting: bool = False,
|
||||
@@ -130,6 +133,8 @@ class ERSDEScheduler(SchedulerMixin, ConfigMixin):
|
||||
raise ValueError(f"solver_order must be 1, 2, or 3, got {solver_order}")
|
||||
if func_type < 1 or func_type > 7:
|
||||
raise ValueError(f"func_type must be 1-7, got {func_type}")
|
||||
if sum([use_karras_sigmas, use_exponential_sigmas, use_beta_sigmas]) > 1:
|
||||
raise ValueError("Only one of use_karras_sigmas, use_exponential_sigmas, use_beta_sigmas can be enabled")
|
||||
|
||||
if trained_betas is not None:
|
||||
self.betas = torch.tensor(trained_betas, dtype=torch.float32)
|
||||
@@ -234,7 +239,7 @@ class ERSDEScheduler(SchedulerMixin, ConfigMixin):
|
||||
self.num_inference_steps = len(sigmas)
|
||||
sigmas = torch.from_numpy(sigmas).to(dtype=torch.float64, device=device)
|
||||
self._setup_flow(sigmas, device, mu)
|
||||
elif self.config.prediction_type == "flow_prediction":
|
||||
elif self.config.prediction_type == "flow_prediction" or self.config.use_flow_sigmas:
|
||||
# Flow-matching path: compute flow schedule internally
|
||||
self.num_inference_steps = num_inference_steps
|
||||
sigmas = np.linspace(1.0, 1.0 / self.config.num_train_timesteps, num_inference_steps)
|
||||
@@ -261,25 +266,40 @@ class ERSDEScheduler(SchedulerMixin, ConfigMixin):
|
||||
self.num_inference_steps = num_inference_steps
|
||||
|
||||
if self.config.timestep_spacing == "linspace":
|
||||
timesteps = np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps).round()[::-1].copy().astype(np.int64)
|
||||
timesteps = np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps).round()[::-1].copy().astype(np.float64)
|
||||
elif self.config.timestep_spacing == "leading":
|
||||
step_ratio = self.config.num_train_timesteps // num_inference_steps
|
||||
timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int64)
|
||||
timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.float64)
|
||||
timesteps += self.config.steps_offset
|
||||
elif self.config.timestep_spacing == "trailing":
|
||||
step_ratio = self.config.num_train_timesteps / num_inference_steps
|
||||
timesteps = np.round(np.arange(self.config.num_train_timesteps, 0, -step_ratio)).astype(np.int64)
|
||||
timesteps = np.round(np.arange(self.config.num_train_timesteps, 0, -step_ratio)).astype(np.float64)
|
||||
timesteps -= 1
|
||||
else:
|
||||
raise ValueError(f"{self.config.timestep_spacing} is not supported")
|
||||
|
||||
# k-diffusion sigmas across the full training schedule, sampled at the inference timesteps
|
||||
acp = self.alphas_cumprod.cpu().numpy()
|
||||
train_sigmas = ((1 - acp) / np.clip(acp, 1e-8, None)) ** 0.5
|
||||
log_sigmas = np.log(train_sigmas)
|
||||
sigmas = np.interp(timesteps, np.arange(len(train_sigmas)), train_sigmas)
|
||||
|
||||
# optional sigma-schedule transforms (VP models only); back-map to fractional timesteps
|
||||
if self.config.use_karras_sigmas:
|
||||
sigmas = self._convert_to_karras(sigmas, num_inference_steps)
|
||||
timesteps = np.array([self._sigma_to_t(s, log_sigmas) for s in sigmas])
|
||||
elif self.config.use_exponential_sigmas:
|
||||
sigmas = self._convert_to_exponential(sigmas, num_inference_steps)
|
||||
timesteps = np.array([self._sigma_to_t(s, log_sigmas) for s in sigmas])
|
||||
elif self.config.use_beta_sigmas:
|
||||
sigmas = self._convert_to_beta(sigmas, num_inference_steps)
|
||||
timesteps = np.array([self._sigma_to_t(s, log_sigmas) for s in sigmas])
|
||||
else:
|
||||
timesteps = timesteps.astype(np.int64) # integer timesteps for the default schedule
|
||||
|
||||
self.timesteps = torch.from_numpy(timesteps).to(device)
|
||||
sigmas_arr = []
|
||||
for t in timesteps:
|
||||
acp = self.alphas_cumprod[t].item()
|
||||
sigmas_arr.append(((1 - acp) / max(acp, 1e-8)) ** 0.5)
|
||||
sigmas_arr.append(0.0)
|
||||
self.sigmas = torch.tensor(sigmas_arr, dtype=torch.float64, device=device)
|
||||
sigmas = np.concatenate([sigmas, [0.0]]).astype(np.float64)
|
||||
self.sigmas = torch.from_numpy(sigmas).to(dtype=torch.float64, device=device)
|
||||
self._flow_alphas = None
|
||||
self._flow_sigmas = None
|
||||
self._flow_lambdas = None
|
||||
@@ -290,25 +310,72 @@ class ERSDEScheduler(SchedulerMixin, ConfigMixin):
|
||||
return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
|
||||
|
||||
def _get_alpha_sigma_lambda(self, step_idx):
|
||||
"""Get alpha, sigma, lambda values for a given step index."""
|
||||
"""Get alpha, sigma, lambda values for a given step index.
|
||||
|
||||
In VP mode the values are derived from the k-diffusion sigma stored in self.sigmas
|
||||
(sigma = sqrt((1-acp)/acp)) rather than from an integer alphas_cumprod lookup. This
|
||||
keeps the math independent of the timestep index so karras/beta/exponential schedules
|
||||
can use fractional timesteps: alpha = 1/sqrt(1+sigma^2), sigma_vp = sigma*alpha,
|
||||
lambda = sigma. For the default schedule this is algebraically identical to
|
||||
alpha = sqrt(acp), sigma_vp = sqrt(1-acp).
|
||||
"""
|
||||
if self._is_flow:
|
||||
alpha = self._flow_alphas[step_idx]
|
||||
sigma = self._flow_sigmas[step_idx]
|
||||
lam = self._flow_lambdas[step_idx]
|
||||
elif step_idx >= len(self.sigmas):
|
||||
# Past the last entry: fully denoised
|
||||
alpha = torch.tensor(1.0, dtype=torch.float64)
|
||||
sigma = torch.tensor(0.0, dtype=torch.float64)
|
||||
lam = torch.tensor(0.0, dtype=torch.float64)
|
||||
else:
|
||||
t = self.timesteps[step_idx] if step_idx < len(self.timesteps) else 0
|
||||
if step_idx >= len(self.timesteps):
|
||||
# Past the last timestep: fully denoised
|
||||
alpha = torch.tensor(1.0, dtype=torch.float64)
|
||||
sigma = torch.tensor(0.0, dtype=torch.float64)
|
||||
lam = torch.tensor(0.0, dtype=torch.float64)
|
||||
else:
|
||||
acp = self.alphas_cumprod[t].to(torch.float64)
|
||||
alpha = acp.sqrt()
|
||||
sigma = (1.0 - acp).sqrt()
|
||||
lam = sigma / alpha.clamp(min=1e-8)
|
||||
sig = self.sigmas[step_idx].to(torch.float64)
|
||||
alpha = 1.0 / torch.sqrt(1.0 + sig ** 2)
|
||||
sigma = sig * alpha
|
||||
lam = sig
|
||||
return alpha, sigma, lam
|
||||
|
||||
@staticmethod
|
||||
def _sigma_to_t(sigma, log_sigmas):
|
||||
"""Map a k-diffusion sigma back to a (fractional) training timestep via log-sigma interpolation."""
|
||||
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 = np.clip((low - log_sigma) / (low - high), 0, 1)
|
||||
t = (1 - w) * low_idx + w * high_idx
|
||||
return t.reshape(sigma.shape)
|
||||
|
||||
@staticmethod
|
||||
def _convert_to_karras(in_sigmas, num_inference_steps, rho=7.0):
|
||||
"""Karras et al. (2206.00364) noise schedule between sigma_min and sigma_max."""
|
||||
sigma_min = float(in_sigmas[-1])
|
||||
sigma_max = float(in_sigmas[0])
|
||||
ramp = np.linspace(0, 1, num_inference_steps)
|
||||
min_inv_rho = sigma_min ** (1 / rho)
|
||||
max_inv_rho = sigma_max ** (1 / rho)
|
||||
return (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
|
||||
|
||||
@staticmethod
|
||||
def _convert_to_exponential(in_sigmas, num_inference_steps):
|
||||
"""Exponential (log-linear) noise schedule between sigma_min and sigma_max."""
|
||||
sigma_min = float(in_sigmas[-1])
|
||||
sigma_max = float(in_sigmas[0])
|
||||
return np.exp(np.linspace(math.log(sigma_max), math.log(sigma_min), num_inference_steps))
|
||||
|
||||
@staticmethod
|
||||
def _convert_to_beta(in_sigmas, num_inference_steps, alpha=0.6, beta=0.6):
|
||||
"""Beta noise schedule (2407.12173) between sigma_min and sigma_max."""
|
||||
import scipy.stats
|
||||
sigma_min = float(in_sigmas[-1])
|
||||
sigma_max = float(in_sigmas[0])
|
||||
return np.array([
|
||||
sigma_min + ppf * (sigma_max - sigma_min)
|
||||
for ppf in [scipy.stats.beta.ppf(t, alpha, beta) for t in 1 - np.linspace(0, 1, num_inference_steps)]
|
||||
])
|
||||
|
||||
def _numerical_clip(self, x, eps=1e-6):
|
||||
"""Clip near-zero values to prevent negative variance from float errors."""
|
||||
if isinstance(x, torch.Tensor):
|
||||
@@ -541,6 +608,8 @@ class ERSDEScheduler(SchedulerMixin, ConfigMixin):
|
||||
def add_noise(self, original_samples: torch.Tensor, noise: torch.Tensor, timesteps: torch.Tensor) -> torch.Tensor:
|
||||
alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device, dtype=original_samples.dtype)
|
||||
timesteps = timesteps.to(original_samples.device)
|
||||
if torch.is_floating_point(timesteps): # karras/beta/exponential schedules use fractional timesteps
|
||||
timesteps = timesteps.round().long()
|
||||
sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5
|
||||
sqrt_alpha_prod = sqrt_alpha_prod.flatten()
|
||||
while len(sqrt_alpha_prod.shape) < len(original_samples.shape):
|
||||
|
||||
@@ -157,12 +157,12 @@ config.update({
|
||||
'UFOGen': { },
|
||||
'BDIA DDIM': { 'clip_sample': False, 'set_alpha_to_one': True, 'steps_offset': 0, 'clip_sample_range': 1.0, 'sample_max_value': 1.0, 'timestep_spacing': 'leading', 'rescale_betas_zero_snr': False, 'thresholding': False, 'gamma': 1.0 },
|
||||
|
||||
'ER-SDE': { 'solver_order': 1, 'func_type': 7, 'clip_sample': False, 'timestep_spacing': 'linspace', 'rescale_betas_zero_snr': False, 'thresholding': False, 'lower_order_final': True, 'shift': 1, 'use_dynamic_shifting': False, 'base_shift': 0.5, 'max_shift': 1.15 },
|
||||
'ER-SDE 2M': { 'solver_order': 2, 'func_type': 7, 'clip_sample': False, 'timestep_spacing': 'linspace', 'rescale_betas_zero_snr': False, 'thresholding': False, 'lower_order_final': True, 'shift': 1, 'use_dynamic_shifting': False, 'base_shift': 0.5, 'max_shift': 1.15 },
|
||||
'ER-SDE 3M': { 'solver_order': 3, 'func_type': 7, 'clip_sample': False, 'timestep_spacing': 'linspace', 'rescale_betas_zero_snr': False, 'thresholding': False, 'lower_order_final': True, 'shift': 1, 'use_dynamic_shifting': False, 'base_shift': 0.5, 'max_shift': 1.15 },
|
||||
'ER-SDE FlowMatch': { 'solver_order': 1, 'func_type': 7, 'clip_sample': False, 'thresholding': False, 'lower_order_final': True, 'shift': 1, 'use_dynamic_shifting': False, 'base_shift': 0.5, 'max_shift': 1.15 },
|
||||
'ER-SDE 2M FlowMatch': { 'solver_order': 2, 'func_type': 7, 'clip_sample': False, 'thresholding': False, 'lower_order_final': True, 'shift': 1, 'use_dynamic_shifting': False, 'base_shift': 0.5, 'max_shift': 1.15 },
|
||||
'ER-SDE 3M FlowMatch': { 'solver_order': 3, 'func_type': 7, 'clip_sample': False, 'thresholding': False, 'lower_order_final': True, 'shift': 1, 'use_dynamic_shifting': False, 'base_shift': 0.5, 'max_shift': 1.15 },
|
||||
'ER-SDE': { 'solver_order': 1, 'func_type': 7, 'clip_sample': False, 'timestep_spacing': 'linspace', 'rescale_betas_zero_snr': False, 'thresholding': False, 'lower_order_final': True, 'shift': 1, 'use_dynamic_shifting': False, 'base_shift': 0.5, 'max_shift': 1.15, 'use_karras_sigmas': False, 'use_exponential_sigmas': False, 'use_beta_sigmas': False, 'use_flow_sigmas': False },
|
||||
'ER-SDE 2M': { 'solver_order': 2, 'func_type': 7, 'clip_sample': False, 'timestep_spacing': 'linspace', 'rescale_betas_zero_snr': False, 'thresholding': False, 'lower_order_final': True, 'shift': 1, 'use_dynamic_shifting': False, 'base_shift': 0.5, 'max_shift': 1.15, 'use_karras_sigmas': False, 'use_exponential_sigmas': False, 'use_beta_sigmas': False, 'use_flow_sigmas': False },
|
||||
'ER-SDE 3M': { 'solver_order': 3, 'func_type': 7, 'clip_sample': False, 'timestep_spacing': 'linspace', 'rescale_betas_zero_snr': False, 'thresholding': False, 'lower_order_final': True, 'shift': 1, 'use_dynamic_shifting': False, 'base_shift': 0.5, 'max_shift': 1.15, 'use_karras_sigmas': False, 'use_exponential_sigmas': False, 'use_beta_sigmas': False, 'use_flow_sigmas': False },
|
||||
'ER-SDE FlowMatch': { 'solver_order': 1, 'func_type': 7, 'clip_sample': False, 'thresholding': False, 'lower_order_final': True, 'shift': 1, 'use_dynamic_shifting': False, 'base_shift': 0.5, 'max_shift': 1.15, 'use_flow_sigmas': False },
|
||||
'ER-SDE 2M FlowMatch': { 'solver_order': 2, 'func_type': 7, 'clip_sample': False, 'thresholding': False, 'lower_order_final': True, 'shift': 1, 'use_dynamic_shifting': False, 'base_shift': 0.5, 'max_shift': 1.15, 'use_flow_sigmas': False },
|
||||
'ER-SDE 3M FlowMatch': { 'solver_order': 3, 'func_type': 7, 'clip_sample': False, 'thresholding': False, 'lower_order_final': True, 'shift': 1, 'use_dynamic_shifting': False, 'base_shift': 0.5, 'max_shift': 1.15, 'use_flow_sigmas': False },
|
||||
|
||||
'PNDM': { 'skip_prk_steps': False, 'set_alpha_to_one': False, 'steps_offset': 0, 'timestep_spacing': 'linspace' },
|
||||
'IPNDM': { },
|
||||
|
||||
@@ -28,6 +28,7 @@ from modules.schedulers.scheduler_flashflow import FlashFlowMatchEulerDiscreteSc
|
||||
from modules.schedulers.scheduler_dpm_flowmatch import FlowMatchDPMSolverMultistepScheduler
|
||||
from modules.schedulers.scheduler_dc import DCSolverMultistepScheduler
|
||||
from modules.schedulers.scheduler_bdia import BDIA_DDIMScheduler
|
||||
from modules.schedulers.scheduler_ersde import ERSDEScheduler
|
||||
|
||||
def test_scheduler(name, scheduler_class, config):
|
||||
try:
|
||||
@@ -61,7 +62,7 @@ def test_scheduler(name, scheduler_class, config):
|
||||
# Re-introduce scaling calculation first
|
||||
scaled_sample = scheduler.scale_model_input(sample, t)
|
||||
|
||||
if config.get("prediction_type") == "flow_prediction" or name in ["UFOGenScheduler", "TDDScheduler", "TCDScheduler", "BDIA_DDIMScheduler", "DCSolverMultistepScheduler"]:
|
||||
if config.get("prediction_type") == "flow_prediction" or name in ["UFOGenScheduler", "TDDScheduler", "TCDScheduler", "BDIA_DDIMScheduler", "DCSolverMultistepScheduler", "ERSDEScheduler"]:
|
||||
# Some new schedulers don't use K-diffusion scaling
|
||||
expected_scale = 1.0
|
||||
else:
|
||||
@@ -239,6 +240,7 @@ def run_tests():
|
||||
RiemannianFlowScheduler,
|
||||
# sdnext schedulers
|
||||
FlowUniPCMultistepScheduler, FlashFlowMatchEulerDiscreteScheduler, FlowMatchDPMSolverMultistepScheduler,
|
||||
ERSDEScheduler,
|
||||
]
|
||||
for cls in flow_schedulers:
|
||||
test_scheduler(cls.__name__, cls, {"prediction_type": "flow_prediction", "use_flow_sigmas": True})
|
||||
@@ -250,7 +252,8 @@ def run_tests():
|
||||
TDDScheduler,
|
||||
TCDScheduler,
|
||||
DCSolverMultistepScheduler,
|
||||
BDIA_DDIMScheduler
|
||||
BDIA_DDIMScheduler,
|
||||
ERSDEScheduler
|
||||
]
|
||||
for prediction_type in ["epsilon", "v_prediction", "sample"]:
|
||||
for cls in extended_schedulers:
|
||||
|
||||
Reference in New Issue
Block a user