From f7364728412a74234e4d236c1c126708906e6450 Mon Sep 17 00:00:00 2001 From: dualchimerra Date: Thu, 24 Sep 2026 15:55:44 +0200 Subject: [PATCH] feat(training): dual-peak timestep sampling, cyclic LR schedules, queue ETA Timestep sampling - New `dual_peak` mode backed by the standalone runtime/dual_peak_sampling package: a mixture of two log-SNR components with independent position, width and weight. Wired into the baseline and InfoNoise samplers and exposed in the schema as fields gated on timestep_sampling==dual_peak. LR schedulers - `cosine_cycles`: training is split into N equal cosine cycles, with a shared max/min LR or per-cycle lists. - `constant_then_cosine`: hold the LR, then cosine-decay over the final fraction of training (decay_start_ratio, decay_min_lr). - Warmup/restart fields are now shown only for the schedulers that use them. Web - FloatListField for per-cycle LR lists; ru/en/zh strings. - Projects and Queue show a training-finish forecast derived from queue order and live monitor progress (lib/queueEstimates). Projects can be sorted by queue position, finish time, updated, created or title. Tests cover the sampler, both schedulers, the schema bridge and the forecast. Co-Authored-By: Claude Opus 5.5 --- runtime/dual_peak_sampling/README.md | 122 +++++++++++++ runtime/dual_peak_sampling/__init__.py | 6 + runtime/dual_peak_sampling/core.py | 126 +++++++++++++ runtime/dual_peak_sampling/pyproject.toml | 21 +++ runtime/dual_peak_sampling/tests/test_core.py | 130 +++++++++++++ runtime/training/schedulers/__init__.py | 10 +- .../schedulers/constant_then_cosine.py | 60 ++++++ runtime/training/schedulers/cosine_cycles.py | 107 +++++++++++ .../training/timestep_samplers/baseline.py | 16 +- .../training/timestep_samplers/dual_peak.py | 27 +++ .../training/timestep_samplers/infonoise.py | 7 + runtime/training/timestep_sampling.py | 10 + studio/domain/training.py | 171 +++++++++++++++++- .../src/components/DualPeakSampling.test.tsx | 67 +++++++ studio/web/src/components/Field.test.tsx | 28 +++ studio/web/src/components/Field.tsx | 74 ++++++++ studio/web/src/i18n/locales/en.json | 40 +++- studio/web/src/i18n/locales/ru.json | 44 ++++- studio/web/src/i18n/locales/zh.json | 37 +++- studio/web/src/lib/queueEstimates.test.ts | 46 +++++ studio/web/src/lib/queueEstimates.ts | 69 +++++++ studio/web/src/lib/schema.test.ts | 4 +- studio/web/src/lib/schema.ts | 8 +- studio/web/src/pages/Projects.tsx | 106 ++++++++++- studio/web/src/pages/Queue.tsx | 58 ++++++ tests/test_argparse_bridge.py | 13 ++ tests/test_dual_peak_sampler.py | 109 +++++++++++ tests/test_plugin_registry.py | 63 ++++++- tests/test_studio_configs.py | 43 ++++- 29 files changed, 1596 insertions(+), 26 deletions(-) create mode 100644 runtime/dual_peak_sampling/README.md create mode 100644 runtime/dual_peak_sampling/__init__.py create mode 100644 runtime/dual_peak_sampling/core.py create mode 100644 runtime/dual_peak_sampling/pyproject.toml create mode 100644 runtime/dual_peak_sampling/tests/test_core.py create mode 100644 runtime/training/schedulers/constant_then_cosine.py create mode 100644 runtime/training/schedulers/cosine_cycles.py create mode 100644 runtime/training/timestep_samplers/dual_peak.py create mode 100644 studio/web/src/components/DualPeakSampling.test.tsx create mode 100644 studio/web/src/lib/queueEstimates.test.ts create mode 100644 studio/web/src/lib/queueEstimates.ts create mode 100644 tests/test_dual_peak_sampler.py diff --git a/runtime/dual_peak_sampling/README.md b/runtime/dual_peak_sampling/README.md new file mode 100644 index 0000000..beebbad --- /dev/null +++ b/runtime/dual_peak_sampling/README.md @@ -0,0 +1,122 @@ +# Dual Peak Sampling + +Experimental, configurable timestep sampling for flow-matching training. +The reusable package depends only on PyTorch, not on AnimaTrainHub, NumPy, +Pydantic, a model, or a GPU. It does **not** promise better style fidelity. + +## Standalone use + +This whole directory can become the root of a separate repository unchanged. +From that directory, run `pip install .` (or `pip install -e '.[test]'` for +development), then `python -m pytest`. No package installation is needed inside +AnimaTrainHub: its runtime directory already exposes this package. + +```python +import torch +from dual_peak_sampling import DualPeakConfig, sample_timesteps, density + +config = DualPeakConfig(peak1_position=0.525, peak2_position=0.85) +generator = torch.Generator(device="cpu").manual_seed(42) +t = sample_timesteps(4096, "cpu", config, generator=generator) +x = torch.linspace(0.001, 0.999, 10000, dtype=torch.float64) +p = density(x, config) # analytic curve for plotting, NOT loss weights +``` + +`sample_timesteps` returns float32 `[batch_size]`, on the requested device. +Without an explicit generator it uses the global PyTorch RNG, compatible with +the trainer's seed and RNG checkpoint restoration. CPU and CUDA are supported; +matching seeds do not imply identical results across devices/PyTorch versions. +There is no internal RNG, learned state, or sample history. + +## Distribution and agreed preset + +Convention: `x_t = (1-t)*data + t*noise`; `t=0` is clean, `t=1` is noise. +On a 1000-step display, `t=0.85` is 850. Do not confuse this flow noise +coefficient with noise-to-signal sigma `t/(1-t)`, or a width in log-SNR units. + +| Component | Location | Width in log-SNR | Relative weight | +|---|---|---|---| +| First peak | component mode t=0.525 | 0.55 | 0.15 | +| Second peak | component mode t=0.850 | 1.20 | 0.35 | +| Broad Style-Friendly background | log-SNR mean=-3.2 | 2.20 | 0.45 | +| Uniform coverage | entire interval | — | 0.05 | + +For each sample, choose **one** component with normalized weights, then draw +from that component. Do not average component samples: that produces a different +distribution. Increasing one weight reduces the normalized shares of the others; +weights represent probability mass, not independent absolute peak heights. +Weights can be zero, but cannot all be zero. They need not add to one. + +For a normal log-SNR component: + +``` +lambda ~ Normal(mean, width**2) +t = sigmoid(-lambda/2) +mean = -2 * (logit(position) - (width/2)**2 * (2*position - 1)) +``` + +The last equation places the **individual component's density mode** at +`position`, including the change-of-variable factor. Peak widths are restricted +to `(0, 2.8]` so that each peak remains unimodal; the broad background allows +widths up to 6 and is specified by mean, not mode. + +The sum has different maxima: with defaults, approximately **540 and 870**. +The broad background contributes over the whole interval, not only the tail. +Expected masses before optional trainer shifts (rounded): + +| t interval | Probability | +|---|---| +| 0–0.3 | 2.09% | +| 0.3–0.6 | 21.46% | +| 0.6–0.9 | 58.64% | +| 0.9–1 | 17.80% | + +The sampler clamps results to `[1e-4, 1-1e-4]` for numerical safety. `density` +describes the continuous mixture before that clamp, excluding its endpoint atoms. +Arbitrary settings can merge the two peaks; two visible maxima are not guaranteed. + +Inspired by the [Anima style experiment](https://note.com/kuon_noise/n/n82be977f167d) +and [Style-Friendly SNR](https://arxiv.org/abs/2411.14793). This is a custom +mixture, not an official implementation or a proven universal style optimum. + +## AnimaTrainHub integration + +Select **Dual Peak — два пика для стиля** in the timestep sampling section +(advanced settings). YAML/CLI mode: `dual_peak`. All package parameters are +prefixed with `dual_peak_` in the trainer: + +```yaml +timestep_sampling: dual_peak +dual_peak_peak1_position: 0.525 +dual_peak_peak1_width: 0.55 +dual_peak_peak1_weight: 0.15 +dual_peak_peak2_position: 0.85 +dual_peak_peak2_width: 1.20 +dual_peak_peak2_weight: 0.35 +dual_peak_background_mean: -3.2 +dual_peak_background_width: 2.20 +dual_peak_background_weight: 0.45 +dual_peak_uniform_weight: 0.05 + +# Preserve the agreed curve for a controlled comparison: +timestep_schedule_shift: 1.0 +timestep_shift_resolution_aware: false +infonoise_enabled: false +leap_enabled: false +loss_weighting: none +``` + +These are sampler settings, not a complete training configuration. Existing +configs/default sampling are not changed. `timestep_shift` is ignored in this +mode; schedule and resolution shifts still transform the final distribution. +InfoNoise uses this configuration only until its adaptive distribution is ready. +Leap uses it only for ordinary (non-Leap) training steps. Loss weighting changes +the relative training contribution after sampling, not the sampled histogram. + +CLI example: `--timestep-sampling dual_peak --dual-peak-peak2-position 0.85`. +Only the small adapter in `training/timestep_samplers/dual_peak.py` knows about +the trainer's prefixed arguments. To extract the package, keep that adapter in +the trainer and replace the bundled directory with the installed dependency. + +Before publishing a separate repository, choose a license and the final package +name/version. Nothing is published or uploaded by this integration. diff --git a/runtime/dual_peak_sampling/__init__.py b/runtime/dual_peak_sampling/__init__.py new file mode 100644 index 0000000..c393743 --- /dev/null +++ b/runtime/dual_peak_sampling/__init__.py @@ -0,0 +1,6 @@ +"""Portable, PyTorch-only dual-peak flow-matching timestep sampler.""" + +from .core import DualPeakConfig, density, sample_timesteps + +__version__ = "0.1.0" +__all__ = ["DualPeakConfig", "density", "sample_timesteps"] diff --git a/runtime/dual_peak_sampling/core.py b/runtime/dual_peak_sampling/core.py new file mode 100644 index 0000000..0354ff4 --- /dev/null +++ b/runtime/dual_peak_sampling/core.py @@ -0,0 +1,126 @@ +"""Independent of AnimaTrainHub: t=0 is data, t=1 is noise. + +Sample a mixture, NOT an average of samples (which would change its density). +The two peak locations describe individual components, not the final mixture. +""" + +from __future__ import annotations + +from dataclasses import dataclass, fields +import math + +import torch + + +@dataclass(frozen=True) +class DualPeakConfig: + """Widths are standard deviations in log-SNR; weights are relative masses. + + Peak widths <= 2.8 keep each peak component unimodal in t, making the + location parameter an actual mode, not a possible minimum between modes. + The broad background is an unconstrained log-SNR normal component. + """ + + peak1_position: float = 0.525 + peak1_width: float = 0.55 + peak1_weight: float = 0.15 + peak2_position: float = 0.850 + peak2_width: float = 1.20 + peak2_weight: float = 0.35 + background_mean: float = -3.2 + background_width: float = 2.20 + background_weight: float = 0.45 + uniform_weight: float = 0.05 + + def __post_init__(self) -> None: + for field in fields(self): + value = getattr(self, field.name) + if not math.isfinite(value): + raise ValueError(f"{field.name} must be finite") + for name in ("peak1_position", "peak2_position"): + if not 0.0 < getattr(self, name) < 1.0: + raise ValueError(f"{name} must be between 0 and 1 (exclusive)") + for name in ("peak1_width", "peak2_width"): + if not 0.0 < getattr(self, name) <= 2.8: + raise ValueError(f"{name} must be in (0, 2.8]") + if not 0.0 < self.background_width <= 6.0: + raise ValueError("background_width must be in (0, 6]") + if not -12.0 <= self.background_mean <= 6.0: + raise ValueError("background_mean must be in [-12, 6]") + for name in ("peak1_weight", "peak2_weight", "background_weight", "uniform_weight"): + if not 0.0 <= getattr(self, name) <= 1.0: + raise ValueError(f"{name} must be in [0, 1]") + if sum(self.weights) <= 0.0: + raise ValueError("At least one mixture weight must be positive") + + @property + def weights(self) -> tuple[float, float, float, float]: + return (self.peak1_weight, self.peak2_weight, + self.background_weight, self.uniform_weight) + + @property + def probabilities(self) -> tuple[float, ...]: + total = sum(self.weights) + return tuple(weight / total for weight in self.weights) + + @property + def components(self) -> tuple[tuple[float, float], ...]: + def logsnr_mean(position: float, width: float) -> float: + # d/dt log p(t) = 0 at position, including the transform Jacobian. + return -2.0 * (math.log(position) - math.log1p(-position) + - (width / 2.0) ** 2 * (2.0 * position - 1.0)) + + return ( + (logsnr_mean(self.peak1_position, self.peak1_width), self.peak1_width), + (logsnr_mean(self.peak2_position, self.peak2_width), self.peak2_width), + (self.background_mean, self.background_width), + ) + + +DEFAULT_CONFIG = DualPeakConfig() + + +@torch.no_grad() +def sample_timesteps( + batch_size: int, + device: torch.device | str = "cpu", + config: DualPeakConfig = DEFAULT_CONFIG, + *, + generator: torch.Generator | None = None, +) -> torch.Tensor: + """Return float32 [batch_size] on device, clamped to [1e-4, 1-1e-4]. + + Uses the caller's/global PyTorch RNG; never seeds it or uses Python/NumPy + randomness. For an explicit generator, its device must match `device`. + Float32 is intentional even inside mixed-precision training: bf16 would + round some high timesteps to exactly 1. No CPU/GPU synchronization needed. + """ + if isinstance(batch_size, bool) or not isinstance(batch_size, int) or batch_size < 0: + raise ValueError("batch_size must be a non-negative integer") + p1, p2, p_background, _ = config.probabilities + (m1, s1), (m2, s2), (mb, sb) = config.components + options = dict(device=device, dtype=torch.float32, generator=generator) + choose = torch.rand(batch_size, **options) + z = torch.randn(batch_size, **options) + mean = torch.where(choose < p1, m1, torch.where(choose < p1 + p2, m2, torch.full_like(choose, mb))) + width = torch.where(choose < p1, s1, torch.where(choose < p1 + p2, s2, torch.full_like(choose, sb))) + biased = torch.sigmoid(-0.5 * (mean + width * z)) + uniform = torch.rand(batch_size, **options) + t = torch.where(choose < p1 + p2 + p_background, biased, uniform) + return t.clamp(1e-4, 1.0 - 1e-4) + + +def density(t: torch.Tensor, config: DualPeakConfig = DEFAULT_CONFIG) -> torch.Tensor: + """Analytic density BEFORE safety clamping, for plots and distribution tests. + + Accepts floating tensors strictly inside (0, 1). Not a loss weight. + The safety clamp in sample_timesteps adds tiny endpoint point masses. + """ + if not t.is_floating_point() or not bool(((t > 0) & (t < 1)).all()): + raise ValueError("density expects floating t strictly inside (0, 1)") + lam = 2.0 * (torch.log1p(-t) - torch.log(t)) + result = torch.full_like(t, config.probabilities[3]) + for weight, (mean, width) in zip(config.probabilities[:3], config.components): + normal = torch.exp(-0.5 * ((lam - mean) / width).square()) + result = result + weight * normal / (width * math.sqrt(2.0 * math.pi)) * 2.0 / (t * (1.0 - t)) + return result diff --git a/runtime/dual_peak_sampling/pyproject.toml b/runtime/dual_peak_sampling/pyproject.toml new file mode 100644 index 0000000..3e8d37d --- /dev/null +++ b/runtime/dual_peak_sampling/pyproject.toml @@ -0,0 +1,21 @@ +[build-system] +requires = ["setuptools>=68"] +build-backend = "setuptools.build_meta" + +[project] +name = "dual-peak-sampling" +version = "0.1.0" +description = "Configurable dual-peak timestep sampling for flow-matching training" +readme = "README.md" +requires-python = ">=3.10" +dependencies = ["torch>=2.0"] + +[project.optional-dependencies] +test = ["pytest>=7"] + +[tool.setuptools] +packages = ["dual_peak_sampling"] +package-dir = {dual_peak_sampling = "."} + +[tool.pytest.ini_options] +testpaths = ["tests"] diff --git a/runtime/dual_peak_sampling/tests/test_core.py b/runtime/dual_peak_sampling/tests/test_core.py new file mode 100644 index 0000000..aa5cb9f --- /dev/null +++ b/runtime/dual_peak_sampling/tests/test_core.py @@ -0,0 +1,130 @@ +"""Portable tests: run after installing this directory as a package.""" + +from dataclasses import replace +import math + +import pytest +import torch + +from dual_peak_sampling import DualPeakConfig, density, sample_timesteps + + +def test_default_distribution_matches_agreed_curve(): + t = sample_timesteps(200_000, generator=torch.Generator().manual_seed(314)) + for lo, hi, expected in ( + (0, .3, .02093), (.3, .6, .21464), + (.6, .9, .58642), (.9, 1, .17800), + (.45, .6, .16832), (.8, .9, .28897), + ): + measured = float(((t >= lo) & (t < hi)).float().mean()) + assert abs(measured - expected) < .005 + + +def test_density_normalization_and_actual_mixture_peaks(): + t = torch.linspace(.00001, .99999, 100_000, dtype=torch.float64) + p = density(t) + assert abs(float(torch.trapezoid(p, t)) - 1.0) < 1e-4 + peaks = t[1:-1][(p[1:-1] > p[:-2]) & (p[1:-1] > p[2:])] + assert peaks.numel() == 2 + assert abs(float(peaks[0]) - .540) < .003 + assert abs(float(peaks[1]) - .870) < .003 + + +@pytest.mark.parametrize("position,width", [(.2, .55), (.525, .55), (.85, 1.2), (.7, 2.8)]) +def test_individual_peak_location_is_mode_not_median(position, width): + config = DualPeakConfig(peak1_position=position, peak1_width=width, + peak1_weight=1, peak2_weight=0, + background_weight=0, uniform_weight=0) + t = torch.linspace(.0001, .9999, 20000, dtype=torch.float64) + assert abs(float(t[density(t, config).argmax()]) - position) < .0001 + + +def test_independent_second_peak_and_width_controls(): + base = DualPeakConfig(peak1_weight=0, peak2_weight=1, + background_weight=0, uniform_weight=0) + for position, width in ((.65, .5), (.85, 1.2), (.9, 2)): + config = replace(base, peak2_position=position, peak2_width=width) + t = torch.linspace(.0001, .9999, 20000, dtype=torch.float64) + assert abs(float(t[density(t, config).argmax()]) - position) < .0001 + + +def test_zero_weights_and_normalization(): + config = DualPeakConfig(peak1_weight=0, peak2_weight=0, + background_weight=0, uniform_weight=.1) + t = sample_timesteps(30000, config=config, generator=torch.Generator().manual_seed(1)) + assert abs(float(t.mean()) - .5) < .01 + torch.testing.assert_close(density(torch.tensor([.1, .4, .9]), config), torch.ones(3)) + original = DualPeakConfig() + scaled = replace(original, peak1_weight=.3, peak2_weight=.7, + background_weight=.9, uniform_weight=.1) + torch.testing.assert_close( + sample_timesteps(256, config=original, generator=torch.Generator().manual_seed(7)), + sample_timesteps(256, config=scaled, generator=torch.Generator().manual_seed(7)), + rtol=0, atol=0, + ) + + +def test_pure_background_matches_logsnr_normal(): + config = DualPeakConfig(peak1_weight=0, peak2_weight=0, + background_weight=1, uniform_weight=0) + t = sample_timesteps(60000, config=config, generator=torch.Generator().manual_seed(8)) + lam = 2 * torch.log((1-t)/t) + assert abs(float(lam.mean()) - config.background_mean) < .04 + assert abs(float(lam.std()) - config.background_width) < .04 + + +@pytest.mark.parametrize("changes", [ + {"peak1_position": 0}, {"peak2_position": 1}, + {"peak1_width": 0}, {"peak2_width": 3}, {"background_width": -1}, + {"background_width": 7}, {"background_mean": -13}, + {"peak1_weight": -1}, {"uniform_weight": 2}, + {"peak2_position": math.nan}, {"background_mean": math.inf}, + {"uniform_weight": math.nan}, + {"peak1_weight": 0, "peak2_weight": 0, "background_weight": 0, "uniform_weight": 0}, +]) +def test_invalid_config_rejected(changes): + with pytest.raises(ValueError): + DualPeakConfig(**changes) + + +def test_reproducible_with_external_rng_state_restore(): + generator = torch.Generator().manual_seed(22) + sample_timesteps(13, generator=generator) + state = generator.get_state() + expected = sample_timesteps(512, generator=generator) + generator.set_state(state) + torch.testing.assert_close(expected, sample_timesteps(512, generator=generator), rtol=0, atol=0) + + +def test_empty_singleton_and_dtype_under_autocast(): + assert sample_timesteps(0).shape == (0,) + assert sample_timesteps(1).shape == (1,) + with torch.autocast("cpu", dtype=torch.bfloat16): + t = sample_timesteps(1000, config=DualPeakConfig(background_mean=-12, background_width=6)) + assert t.dtype == torch.float32 + assert not t.requires_grad + assert bool(torch.isfinite(t).all()) + assert bool(((t >= 1e-4) & (t <= 1-1e-4)).all()) + + +def test_float32_even_with_changed_global_default_dtype(): + previous = torch.get_default_dtype() + try: + torch.set_default_dtype(torch.float64) + assert sample_timesteps(16).dtype == torch.float32 + finally: + torch.set_default_dtype(previous) + + +@pytest.mark.parametrize("batch_size", [-1, 1.5, True]) +def test_invalid_batch_size(batch_size): + with pytest.raises(ValueError): + sample_timesteps(batch_size) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA unavailable") +def test_cuda(): + with torch.autocast("cuda", dtype=torch.bfloat16): + t = sample_timesteps(1024, "cuda", generator=torch.Generator(device="cuda").manual_seed(1)) + assert t.is_cuda and t.dtype == torch.float32 and t.shape == (1024,) + assert bool(((t > 0) & (t < 1)).all()) diff --git a/runtime/training/schedulers/__init__.py b/runtime/training/schedulers/__init__.py index 69ba7ea..aab289b 100644 --- a/runtime/training/schedulers/__init__.py +++ b/runtime/training/schedulers/__init__.py @@ -14,14 +14,22 @@ from typing import Callable, Optional -from training.schedulers import cosine, cosine_with_restart, cosine_with_warmup +from training.schedulers import ( + constant_then_cosine, + cosine, + cosine_cycles, + cosine_with_restart, + cosine_with_warmup, +) __all__ = ["BUILDERS", "build_scheduler", "validate_schema_consistency"] # "none" 不在 BUILDERS —— build_scheduler 显式判一下返回 None。 BUILDERS: dict[str, Callable] = { + "constant_then_cosine": constant_then_cosine.build, "cosine": cosine.build, + "cosine_cycles": cosine_cycles.build, "cosine_with_restart": cosine_with_restart.build, "cosine_with_warmup": cosine_with_warmup.build, } diff --git a/runtime/training/schedulers/constant_then_cosine.py b/runtime/training/schedulers/constant_then_cosine.py new file mode 100644 index 0000000..0e7d4a0 --- /dev/null +++ b/runtime/training/schedulers/constant_then_cosine.py @@ -0,0 +1,60 @@ +"""Constant learning rate followed by one final cosine decay.""" + +from __future__ import annotations + +import logging +import math +from typing import Optional + + +logger = logging.getLogger(__name__) + + +def build(args, optimizer, total_steps: Optional[int]): + if total_steps is None or int(total_steps) <= 0: + logger.warning("constant_then_cosine requires total_steps; falling back to none") + return None + + import torch + + total_steps = int(total_steps) + decay_start = float( + getattr(args, "lr_scheduler_decay_start_ratio", 2.0 / 3.0) + ) + min_lr = float(getattr(args, "lr_scheduler_decay_min_lr", 0.0) or 0.0) + if not 0.0 <= decay_start < 1.0: + raise ValueError( + "constant_then_cosine: lr_scheduler_decay_start_ratio must be in [0, 1)" + ) + if min_lr < 0.0: + raise ValueError("constant_then_cosine: lr_scheduler_decay_min_lr must be >= 0") + + lambdas = [] + for group in optimizer.param_groups: + base_lr = float(group["lr"]) + if min_lr > base_lr: + raise ValueError( + "constant_then_cosine: lr_scheduler_decay_min_lr " + f"({min_lr}) exceeds parameter-group LR ({base_lr})" + ) + min_factor = min_lr / base_lr if base_lr > 0.0 else 0.0 + + def lr_lambda(step: int, *, _min_factor=min_factor) -> float: + progress = min(1.0, max(0.0, float(step) / float(total_steps))) + if progress <= decay_start: + return 1.0 + decay_progress = (progress - decay_start) / (1.0 - decay_start) + cosine = 0.5 * (1.0 + math.cos(math.pi * decay_progress)) + return _min_factor + (1.0 - _min_factor) * cosine + + lambdas.append(lr_lambda) + + scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambdas) + logger.info( + "LR scheduler: constant_then_cosine " + "(total_steps=%s, decay_start_ratio=%s, min_lr=%s)", + total_steps, + decay_start, + min_lr, + ) + return scheduler diff --git a/runtime/training/schedulers/cosine_cycles.py b/runtime/training/schedulers/cosine_cycles.py new file mode 100644 index 0000000..e0267fb --- /dev/null +++ b/runtime/training/schedulers/cosine_cycles.py @@ -0,0 +1,107 @@ +"""Equal-length cosine cycles with configurable LR bounds per cycle.""" + +from __future__ import annotations + +import logging +import math +from typing import Optional + + +logger = logging.getLogger(__name__) + + +def _float_list(args, name: str) -> list[float]: + value = getattr(args, name, None) + if value is None: + return [] + return [float(item) for item in value] + + +def _cycle_values( + *, + count: int, + common: float, + overrides: list[float], + name: str, +) -> list[float]: + if not overrides: + return [common] * count + if len(overrides) != count: + raise ValueError( + f"cosine_cycles: {name} must contain exactly {count} values; " + f"got {len(overrides)}" + ) + return overrides + + +def build(args, optimizer, total_steps: Optional[int]): + """Build hard-restart cosine cycles of equal length. + + ``cycle_max_lrs`` / ``cycle_min_lrs`` are absolute learning rates. Empty + lists fall back to the common values. A missing common maximum preserves + each optimizer parameter group's own base LR. + """ + if total_steps is None or int(total_steps) <= 0: + logger.warning("cosine_cycles requires total_steps; falling back to none") + return None + + import torch + + total_steps = int(total_steps) + cycle_count = int(getattr(args, "lr_scheduler_cycle_count", 3) or 3) + if cycle_count < 1: + raise ValueError("cosine_cycles: lr_scheduler_cycle_count must be >= 1") + + common_max = getattr(args, "lr_scheduler_cycle_max_lr", None) + common_max = None if common_max is None else float(common_max) + common_min = float(getattr(args, "lr_scheduler_cycle_min_lr", 0.0) or 0.0) + max_overrides = _float_list(args, "lr_scheduler_cycle_max_lrs") + min_overrides = _float_list(args, "lr_scheduler_cycle_min_lrs") + + lambdas = [] + resolved_for_log: list[tuple[list[float], list[float]]] = [] + for group in optimizer.param_groups: + base_lr = float(group["lr"]) + max_lrs = _cycle_values( + count=cycle_count, + common=base_lr if common_max is None else common_max, + overrides=max_overrides, + name="lr_scheduler_cycle_max_lrs", + ) + min_lrs = _cycle_values( + count=cycle_count, + common=common_min, + overrides=min_overrides, + name="lr_scheduler_cycle_min_lrs", + ) + for index, (min_lr, max_lr) in enumerate(zip(min_lrs, max_lrs), start=1): + if max_lr <= 0.0: + raise ValueError(f"cosine_cycles: cycle {index} max LR must be > 0") + if min_lr < 0.0: + raise ValueError(f"cosine_cycles: cycle {index} min LR must be >= 0") + if min_lr > max_lr: + raise ValueError( + f"cosine_cycles: cycle {index} min LR ({min_lr}) exceeds " + f"max LR ({max_lr})" + ) + + def lr_lambda(step: int, *, _base=base_lr, _max=max_lrs, _min=min_lrs) -> float: + progress = min(1.0, max(0.0, float(step) / float(total_steps))) + scaled = progress * cycle_count + cycle_index = min(int(scaled), cycle_count - 1) + phase = min(1.0, max(0.0, scaled - cycle_index)) + cosine = 0.5 * (1.0 + math.cos(math.pi * phase)) + lr = _min[cycle_index] + (_max[cycle_index] - _min[cycle_index]) * cosine + return lr / _base if _base > 0.0 else 0.0 + + lambdas.append(lr_lambda) + resolved_for_log.append((max_lrs, min_lrs)) + + scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambdas) + logger.info( + "LR scheduler: cosine_cycles (total_steps=%s, cycles=%s, group_bounds=%s)", + total_steps, + cycle_count, + resolved_for_log, + ) + return scheduler diff --git a/runtime/training/timestep_samplers/baseline.py b/runtime/training/timestep_samplers/baseline.py index 52fc24f..6651981 100644 --- a/runtime/training/timestep_samplers/baseline.py +++ b/runtime/training/timestep_samplers/baseline.py @@ -1,14 +1,18 @@ """Baseline timestep 采样器：包装 training.timestep_sampling.sample_t。 非自适应；record / maybe_refresh 是 no-op。 -覆盖 7 种 mode：logit_normal / uniform / logit_normal_low / mode / -mixed_uniform_low / mixed_uniform_logit / style_friendly。 +覆盖 8 种 mode：logit_normal / uniform / logit_normal_low / mode / +mixed_uniform_low / mixed_uniform_logit / style_friendly / dual_peak。 """ from __future__ import annotations +from dataclasses import asdict + import torch +from dual_peak_sampling import DualPeakConfig +from training.timestep_samplers.dual_peak import config_from_args from training.timestep_sampling import sample_t @@ -23,6 +27,7 @@ def __init__( timestep_schedule_shift: float = 1.0, style_snr_mean: float = -6.0, style_snr_sigma: float = 2.0, + dual_peak_config: DualPeakConfig | None = None, ): self.mode = mode self.shift = shift @@ -30,6 +35,10 @@ def __init__( self.timestep_schedule_shift = timestep_schedule_shift self.style_snr_mean = style_snr_mean self.style_snr_sigma = style_snr_sigma + self.dual_peak_config = ( + (dual_peak_config or DualPeakConfig()) + if (mode or "logit_normal").lower() == "dual_peak" else None + ) def sample(self, bs: int, device, *, token_counts=None) -> torch.Tensor: return sample_t( @@ -41,6 +50,7 @@ def sample(self, bs: int, device, *, token_counts=None) -> torch.Tensor: timestep_schedule_shift=self.timestep_schedule_shift, style_snr_mean=self.style_snr_mean, style_snr_sigma=self.style_snr_sigma, + dual_peak_config=self.dual_peak_config, ) def record(self, t: torch.Tensor, raw_mse: torch.Tensor) -> None: @@ -58,6 +68,7 @@ def status(self) -> dict: "timestep_schedule_shift": self.timestep_schedule_shift, "style_snr_mean": self.style_snr_mean, "style_snr_sigma": self.style_snr_sigma, + **({"dual_peak": asdict(self.dual_peak_config)} if self.dual_peak_config is not None else {}), } # ─── pause/resume（ADR 0006 Addendum 1）：无状态采样器是真的无状态，no-op ─── @@ -79,6 +90,7 @@ def build(args, total_steps) -> BaselineTimestepSampler: # 显式 None 检查而非 `or`：mean 的合法值里有 0.0，sigma 没有但保持同一写法 style_snr_mean=_float_or(getattr(args, "style_snr_mean", None), -6.0), style_snr_sigma=_float_or(getattr(args, "style_snr_sigma", None), 2.0), + dual_peak_config=config_from_args(args), ) diff --git a/runtime/training/timestep_samplers/dual_peak.py b/runtime/training/timestep_samplers/dual_peak.py new file mode 100644 index 0000000..db17699 --- /dev/null +++ b/runtime/training/timestep_samplers/dual_peak.py @@ -0,0 +1,27 @@ +"""AnimaTrainHub config adapter; the reusable package knows nothing about Studio.""" + +from dataclasses import fields +import logging + +from dual_peak_sampling import DualPeakConfig + +logger = logging.getLogger(__name__) + + +def config_from_args(args) -> DualPeakConfig | None: + if str(getattr(args, "timestep_sampling", "")).lower() != "dual_peak": + return None + values = {} + for field in fields(DualPeakConfig): + value = getattr(args, f"dual_peak_{field.name}", None) + if value is not None: + values[field.name] = float(value) + config = DualPeakConfig(**values) + if getattr(args, "infonoise_enabled", False): + logger.warning("Dual Peak: InfoNoise uses this curve only until its adaptive distribution is ready.") + if getattr(args, "leap_enabled", False): + logger.warning("Dual Peak: the curve applies only to non-Leap steps; Leap samples its own timesteps.") + if (getattr(args, "timestep_schedule_shift", 1.0) not in (None, 1.0) + or getattr(args, "timestep_shift_resolution_aware", False)): + logger.warning("Dual Peak: schedule/resolution shifts are enabled and will move the configured peaks.") + return config diff --git a/runtime/training/timestep_samplers/infonoise.py b/runtime/training/timestep_samplers/infonoise.py index d3807e5..61c0599 100644 --- a/runtime/training/timestep_samplers/infonoise.py +++ b/runtime/training/timestep_samplers/infonoise.py @@ -38,6 +38,9 @@ import numpy as np import torch +from dual_peak_sampling import DualPeakConfig +from training.timestep_samplers.dual_peak import config_from_args + logger = logging.getLogger(__name__) @@ -63,6 +66,7 @@ def __init__( baseline_timestep_schedule_shift: float = 1.0, baseline_style_snr_mean: float = -6.0, baseline_style_snr_sigma: float = 2.0, + baseline_dual_peak_config: DualPeakConfig | None = None, ): if not 0.0 < p_onset < 1.0: raise ValueError(f"p_onset 必须 ∈ (0,1)，得到 {p_onset}") @@ -81,6 +85,7 @@ def __init__( self.baseline_timestep_schedule_shift = baseline_timestep_schedule_shift self.baseline_style_snr_mean = baseline_style_snr_mean self.baseline_style_snr_sigma = baseline_style_snr_sigma + self.baseline_dual_peak_config = baseline_dual_peak_config self._internal_step = 0 sigma_min = t_min / (1.0 - t_min) @@ -125,6 +130,7 @@ def _sample_baseline(self, bs: int, device) -> torch.Tensor: timestep_schedule_shift=self.baseline_timestep_schedule_shift, style_snr_mean=self.baseline_style_snr_mean, style_snr_sigma=self.baseline_style_snr_sigma, + dual_peak_config=self.baseline_dual_peak_config, ) def record(self, t: torch.Tensor, raw_mse: torch.Tensor): @@ -312,6 +318,7 @@ def build(args, total_steps: Optional[int]) -> InfoNoiseScheduler: gate_pivot_c=float(getattr(args, "infonoise_gate_pivot_c", 0.15) or 0.0), baseline_shift=float(getattr(args, "timestep_shift", 3.0) or 3.0), baseline_mode=str(getattr(args, "timestep_sampling", "logit_normal") or "logit_normal"), + baseline_dual_peak_config=config_from_args(args), baseline_mix_low_prob=float(getattr(args, "timestep_mix_low_prob", 0.0) or 0.0), baseline_timestep_schedule_shift=float(getattr(args, "timestep_schedule_shift", 1.0) or 1.0), # style_friendly 作为热身 baseline 时也要拿到用户设的窗口（0.0 是合法 mean） diff --git a/runtime/training/timestep_sampling.py b/runtime/training/timestep_sampling.py index 4494474..580c708 100644 --- a/runtime/training/timestep_sampling.py +++ b/runtime/training/timestep_sampling.py @@ -11,6 +11,8 @@ import torch +from dual_peak_sampling import DualPeakConfig, sample_timesteps + from training.families.anima import ANIMA_SPEC as _ANIMA_SPEC @@ -23,10 +25,13 @@ def sample_t( timestep_schedule_shift: float = 1.0, style_snr_mean: float = -6.0, style_snr_sigma: float = 2.0, + dual_peak_config: DualPeakConfig | None = None, ) -> torch.Tensor: """采样 Flow Matching 时间步 t ∈ (0, 1)。 mode: + dual_peak — Portable custom style mixture; independent peak positions, + widths and weights. Ignores shift; see dual_peak_config. logit_normal — SD3/Anima 默认，偏向中间 t；shift>1 推向高噪声端 uniform — 均匀采样，对细节端和结构端覆盖更均衡 logit_normal_low — logit-normal 反向 shift，偏向低噪声/细节端 @@ -47,6 +52,11 @@ def sample_t( """ mode = (mode or "logit_normal").lower() + if mode == "dual_peak": + # Like style_friendly, positions are explicit: ignore timestep_shift. + t = sample_timesteps(bs, device, dual_peak_config or DualPeakConfig()) + return _apply_timestep_schedule_shift(t, timestep_schedule_shift) + if mode == "style_friendly": t = sample_t_style_friendly(bs, device, mean=style_snr_mean, sigma=style_snr_sigma) return _apply_timestep_schedule_shift(t, timestep_schedule_shift) diff --git a/studio/domain/training.py b/studio/domain/training.py index 75b26c4..6ca32ad 100644 --- a/studio/domain/training.py +++ b/studio/domain/training.py @@ -418,7 +418,14 @@ class TrainingConfig(BaseModel): disable_hint="此优化器自动管理学习率", ), ) - lr_scheduler: Literal["none", "cosine", "cosine_with_restart", "cosine_with_warmup"] = Field( + lr_scheduler: Literal[ + "none", + "cosine", + "cosine_with_restart", + "cosine_with_warmup", + "cosine_cycles", + "constant_then_cosine", + ] = Field( "none", description="学习率调度（none = 常数；Prodigy / PPSF / Automagic / SOAP-SF 固定为 none）", json_schema_extra=_meta( @@ -441,13 +448,52 @@ class TrainingConfig(BaseModel): lr_scheduler_eta_min: float = Field( 1e-6, ge=0.0, description="学习率衰减下限：cosine 调度到此值后不再下降；通常远小于初始 lr（如初始 1e-4 配 1e-6）", - json_schema_extra=_meta("training", show_when="lr_scheduler!=none", advanced=True), + json_schema_extra=_meta( + "training", + show_when="lr_scheduler==cosine||lr_scheduler==cosine_with_restart||lr_scheduler==cosine_with_warmup", + advanced=True, + ), ) lr_scheduler_warmup_steps: int = Field( 100, ge=0, description="cosine_with_warmup 预热步数", json_schema_extra=_meta("training", show_when="lr_scheduler==cosine_with_warmup", advanced=True), ) + lr_scheduler_cycle_count: int = Field( + 3, ge=1, le=100, + description="cosine_cycles 的峰数；总训练步数会被等分为这么多个周期", + json_schema_extra=_meta("training", show_when="lr_scheduler==cosine_cycles"), + ) + lr_scheduler_cycle_max_lr: Optional[float] = Field( + None, gt=0.0, + description="所有周期共用的最高 LR；留空时使用 learning_rate", + json_schema_extra=_meta("training", show_when="lr_scheduler==cosine_cycles"), + ) + lr_scheduler_cycle_min_lr: float = Field( + 0.0, ge=0.0, + description="所有周期共用的最低 LR；逐周期列表中的值会覆盖它", + json_schema_extra=_meta("training", show_when="lr_scheduler==cosine_cycles"), + ) + lr_scheduler_cycle_max_lrs: list[float] = Field( + default_factory=list, + description="Каждому пику свой максимум LR через запятую; пусто = общий максимум. Число значений должно совпадать с числом пиков", + json_schema_extra=_meta("training", show_when="lr_scheduler==cosine_cycles", advanced=True), + ) + lr_scheduler_cycle_min_lrs: list[float] = Field( + default_factory=list, + description="Каждому пику свой минимум LR через запятую; пусто = общий минимум. Число значений должно совпадать с числом пиков", + json_schema_extra=_meta("training", show_when="lr_scheduler==cosine_cycles", advanced=True), + ) + lr_scheduler_decay_start_ratio: float = Field( + 2.0 / 3.0, ge=0.0, lt=1.0, + description="constant_then_cosine: доля обучения, после которой начинается плавное снижение; 0.6667 = последняя треть", + json_schema_extra=_meta("training", show_when="lr_scheduler==constant_then_cosine"), + ) + lr_scheduler_decay_min_lr: float = Field( + 0.0, ge=0.0, + description="constant_then_cosine: LR в самом конце обучения", + json_schema_extra=_meta("training", show_when="lr_scheduler==constant_then_cosine"), + ) optimizer_type: Literal["adamw", "adamw8bit", "automagic", "came", "lion", "prodigy", "prodigy_plus_schedulefree", "soap", "soap_sf"] = Field( "adamw", description="优化器。adamw 标准基线；adamw8bit 同 AdamW 但状态量化到 int8（state 显存约 AdamW 的 1/4，超参照搬不用换算，需装 bitsandbytes）；automagic 自适应每参数 lr（推荐 lr=1e-6）；came 置信度引导 + 分解二阶矩（state 显存低于 AdamW，lr 同 AdamW 量级）；lion 显存约 AdamW 一半（推荐 lr=AdamW lr / 3）；prodigy / prodigy_plus_schedulefree 自适应估 lr（lr 填 1.0）；soap Adam-in-Shampoo-eigenbasis 二阶预条件（拟合更快，lr 同 AdamW 量级）；soap_sf SOAP + Schedule-Free（lr_scheduler 固定 none）", @@ -713,6 +759,7 @@ class TrainingConfig(BaseModel): "mixed_uniform_logit", "krea2_shift", "style_friendly", + "dual_peak", ] = Field( "logit_normal", description="采样分布。logit_normal 偏中段（SD3/Anima 默认）；krea2_shift 按每图 token 数动态 shift；uniform 等概率；mode 单峰偏移；mixed_* 混合 uniform 与偏置端（比例由 timestep_mix_low_prob 控制）；style_friendly 直接在 log-SNR 轴正态采样、把训练集中到风格成形的高噪声窗口（arXiv 2411.14793，风格 LoRA 专用，忽略 timestep_shift）", @@ -732,7 +779,7 @@ class TrainingConfig(BaseModel): json_schema_extra=_meta( "timestep_sampling", # style_friendly 的偏移完全由 style_snr_mean 给定，两者叠加会双重偏移 - show_when="timestep_sampling!=uniform&×tep_sampling!=style_friendly", + show_when="timestep_sampling!=uniform&×tep_sampling!=style_friendly&×tep_sampling!=dual_peak", alt_description="InfoNoise 开启时作为热身阶段的 baseline shift，正式阶段由自适应 CDF 接管；Leap 启用时 leap 路径恒用 U(0,1)，本字段仅作用于 (1-leap_ratio) 比例的标准 step", alt_description_when="infonoise_enabled==true||leap_enabled==true", advanced=True, @@ -743,7 +790,7 @@ class TrainingConfig(BaseModel): description="mixed_* 模式下走偏置端的样本比例：0 = 全 uniform；典型 0.15-0.30", json_schema_extra=_meta( "timestep_sampling", - show_when="timestep_sampling!=uniform", + show_when="timestep_sampling!=uniform&×tep_sampling!=dual_peak", alt_description="InfoNoise 开启 + mixed_* baseline 时，热身阶段混合比例，正式阶段由自适应 CDF 接管；Leap 启用时 leap 路径恒用 U(0,1)，本字段仅作用于 (1-leap_ratio) 比例的标准 step", alt_description_when="infonoise_enabled==true||leap_enabled==true", advanced=True, @@ -782,6 +829,76 @@ class TrainingConfig(BaseModel): advanced=True, ), ) + dual_peak_peak1_position: float = Field( + 0.525, gt=0.0, lt=1.0, allow_inf_nan=False, + description="First component peak in flow t: 0.525 = timestep 525/1000. The mixture maximum can move.", + json_schema_extra=_meta( + "timestep_sampling", show_when="timestep_sampling==dual_peak", advanced=True, + ), + ) + dual_peak_peak1_width: float = Field( + 0.55, gt=0.0, le=2.8, allow_inf_nan=False, + description="First peak width (log-SNR standard deviation). Larger = wider, lower peak at the same weight.", + json_schema_extra=_meta( + "timestep_sampling", show_when="timestep_sampling==dual_peak", advanced=True, + ), + ) + dual_peak_peak1_weight: float = Field( + 0.15, ge=0.0, le=1.0, allow_inf_nan=False, + description="First peak relative weight. All four weights are normalized together; zero disables this component.", + json_schema_extra=_meta( + "timestep_sampling", show_when="timestep_sampling==dual_peak", advanced=True, + ), + ) + dual_peak_peak2_position: float = Field( + 0.85, gt=0.0, lt=1.0, allow_inf_nan=False, + description="Second component peak in flow t: 0.85 = timestep 850/1000. Defaults put the mixture maximum near 870.", + json_schema_extra=_meta( + "timestep_sampling", show_when="timestep_sampling==dual_peak", advanced=True, + ), + ) + dual_peak_peak2_width: float = Field( + 1.2, gt=0.0, le=2.8, allow_inf_nan=False, + description="Second peak width (log-SNR standard deviation). Independent of the first peak.", + json_schema_extra=_meta( + "timestep_sampling", show_when="timestep_sampling==dual_peak", advanced=True, + ), + ) + dual_peak_peak2_weight: float = Field( + 0.35, ge=0.0, le=1.0, allow_inf_nan=False, + description="Second peak relative weight. Normalized with the other three weights.", + json_schema_extra=_meta( + "timestep_sampling", show_when="timestep_sampling==dual_peak", advanced=True, + ), + ) + dual_peak_background_mean: float = Field( + -3.2, ge=-12.0, le=6.0, allow_inf_nan=False, + description="Broad Style-Friendly component log-SNR mean. Lower moves it toward high noise. This covers the whole interval, not only the right tail.", + json_schema_extra=_meta( + "timestep_sampling", show_when="timestep_sampling==dual_peak", advanced=True, + ), + ) + dual_peak_background_width: float = Field( + 2.2, gt=0.0, le=6.0, allow_inf_nan=False, + description="Broad Style-Friendly component log-SNR standard deviation.", + json_schema_extra=_meta( + "timestep_sampling", show_when="timestep_sampling==dual_peak", advanced=True, + ), + ) + dual_peak_background_weight: float = Field( + 0.45, ge=0.0, le=1.0, allow_inf_nan=False, + description="Broad background relative weight. Zero removes the broad component.", + json_schema_extra=_meta( + "timestep_sampling", show_when="timestep_sampling==dual_peak", advanced=True, + ), + ) + dual_peak_uniform_weight: float = Field( + 0.05, ge=0.0, le=1.0, allow_inf_nan=False, + description="Uniform coverage relative weight. Default 0.05, normalized with other weights.", + json_schema_extra=_meta( + "timestep_sampling", show_when="timestep_sampling==dual_peak", advanced=True, + ), + ) timestep_shift_resolution_aware: bool = Field( False, description="按每图 token 数对采样后的 t 做分辨率修正（SD3 式 s=sqrt(该图 token 数/基准档 token 数)，" @@ -1209,6 +1326,52 @@ def _enforce_disable_rules(self) -> "TrainingConfig": raise ValueError("A cross-field constraint is not satisfied — " + "; ".join(lines)) return self + @model_validator(mode="after") + def _validate_dual_peak_weights(self) -> "TrainingConfig": + if self.timestep_sampling == "dual_peak" and ( + self.dual_peak_peak1_weight + self.dual_peak_peak2_weight + + self.dual_peak_background_weight + self.dual_peak_uniform_weight + ) <= 0.0: + raise ValueError("dual_peak: at least one mixture weight must be positive") + return self + + @model_validator(mode="after") + def _validate_custom_lr_schedulers(self) -> "TrainingConfig": + if self.lr_scheduler == "cosine_cycles": + count = self.lr_scheduler_cycle_count + for name, values in ( + ("lr_scheduler_cycle_max_lrs", self.lr_scheduler_cycle_max_lrs), + ("lr_scheduler_cycle_min_lrs", self.lr_scheduler_cycle_min_lrs), + ): + if values and len(values) != count: + raise ValueError( + f"{name} must be empty or contain exactly {count} values" + ) + + common_max = self.lr_scheduler_cycle_max_lr or self.learning_rate + max_lrs = self.lr_scheduler_cycle_max_lrs or [common_max] * count + min_lrs = self.lr_scheduler_cycle_min_lrs or [self.lr_scheduler_cycle_min_lr] * count + for index, (min_lr, max_lr) in enumerate(zip(min_lrs, max_lrs), start=1): + if min_lr < 0.0 or max_lr <= 0.0: + raise ValueError( + f"cosine_cycles cycle {index}: min LR must be >= 0 and max LR > 0" + ) + if min_lr > max_lr: + raise ValueError( + f"cosine_cycles cycle {index}: min LR ({min_lr}) " + f"cannot exceed max LR ({max_lr})" + ) + + if ( + self.lr_scheduler == "constant_then_cosine" + and self.lr_scheduler_decay_min_lr > self.learning_rate + ): + raise ValueError( + "lr_scheduler_decay_min_lr cannot exceed learning_rate for " + "constant_then_cosine" + ) + return self + @model_validator(mode="after") def _validate_detail_inv_t_range(self) -> "TrainingConfig": """detail_inv_t 加权曲线的 min 必须 <= max；fail-fast 取代历史的静默 swap。""" diff --git a/studio/web/src/components/DualPeakSampling.test.tsx b/studio/web/src/components/DualPeakSampling.test.tsx new file mode 100644 index 0000000..2cc9a1c --- /dev/null +++ b/studio/web/src/components/DualPeakSampling.test.tsx @@ -0,0 +1,67 @@ +import { fireEvent, render, screen } from '@testing-library/react' +import { afterEach, describe, expect, it, vi } from 'vitest' +import type { SchemaResponse } from '../api/client' +import i18n from '../i18n' +import { schemaEnumLabel } from '../lib/schema' +import SchemaForm from './SchemaForm' + +const parameters = { + peak1_position: 0.525, peak1_width: 0.55, peak1_weight: 0.15, + peak2_position: 0.85, peak2_width: 1.2, peak2_weight: 0.35, + background_mean: -3.2, background_width: 2.2, background_weight: 0.45, + uniform_weight: 0.05, +} + +const schema: SchemaResponse = { + groups: [{ key: 'timestep_sampling', label: 'Timestep sampling' }], + schema: { properties: { + timestep_sampling: { + type: 'string', enum: ['logit_normal', 'style_friendly', 'dual_peak'], + default: 'logit_normal', group: 'timestep_sampling', advanced: true, + }, + timestep_shift: { + type: 'number', default: 3, group: 'timestep_sampling', advanced: true, + show_when: 'timestep_sampling!=uniform&×tep_sampling!=style_friendly&×tep_sampling!=dual_peak', + }, + ...Object.fromEntries(Object.entries(parameters).map(([key, value]) => [ + `dual_peak_${key}`, { + type: 'number', default: value, group: 'timestep_sampling', advanced: true, + show_when: 'timestep_sampling==dual_peak', + }, + ])), + } }, +} + +afterEach(async () => { await i18n.changeLanguage('zh') }) + +describe('Dual Peak sampling controls', () => { + it.each(['ru', 'en', 'zh'])('has a translated option and parameter help in %s', async (lang) => { + await i18n.changeLanguage(lang) + expect(schemaEnumLabel('timestep_sampling', 'dual_peak', i18n.t)).toContain('Dual Peak') + for (const key of Object.keys(parameters)) { + expect(i18n.exists(`schema.descriptions.dual_peak_${key}`, { lng: lang, fallbackLng: false })).toBe(true) + } + }) + + it('shows custom controls only for dual_peak, preserves values and edits the second peak', async () => { + await i18n.changeLanguage('ru') + const onChange = vi.fn() + const values = { + timestep_sampling: 'dual_peak', + ...Object.fromEntries(Object.entries(parameters).map(([k, v]) => [`dual_peak_${k}`, v])), + } + const { rerender } = render() + expect(screen.getByRole('option', { name: 'Dual Peak — два пика для стиля' })).toBeInTheDocument() + expect(screen.queryByText('Timestep Shift')).not.toBeInTheDocument() + expect(screen.getAllByRole('textbox')).toHaveLength(10) + const secondPeak = screen.getByDisplayValue('0.85') + fireEvent.change(secondPeak, { target: { value: '0.82' } }) + fireEvent.blur(secondPeak) + expect(onChange).toHaveBeenCalledWith({ ...values, dual_peak_peak2_position: 0.82 }) + onChange.mockClear() + rerender() + expect(screen.queryByText('Dual Peak Peak1 Position')).not.toBeInTheDocument() + expect(screen.getByText('Timestep Shift')).toBeInTheDocument() + expect(onChange).not.toHaveBeenCalled() + }) +}) diff --git a/studio/web/src/components/Field.test.tsx b/studio/web/src/components/Field.test.tsx index 043b8c9..b079c3d 100644 --- a/studio/web/src/components/Field.test.tsx +++ b/studio/web/src/components/Field.test.tsx @@ -298,6 +298,34 @@ describe('Field string-list input (sample_prompts 换行)', () => { }) }) +const floatListProp: SchemaProperty = { + type: 'array', + items: { type: 'number' }, + default: [], + group: 'training', + description: '', +} + +describe('Field float-list input', () => { + it('parses decimal and scientific-notation learning rates', async () => { + const onChange = vi.fn() + const user = userEvent.setup() + render( + , + ) + + const input = screen.getByRole('textbox') as HTMLInputElement + await user.type(input, '1e-4, 8e-5, 0.00006') + fireEvent.blur(input) + expect(onChange).toHaveBeenLastCalledWith([1e-4, 8e-5, 6e-5]) + }) +}) + describe('Field code input', () => { it('renders object values as formatted JSON and commits parsed objects on blur', () => { const onChange = vi.fn() diff --git a/studio/web/src/components/Field.tsx b/studio/web/src/components/Field.tsx index 451f01f..a85c805 100644 --- a/studio/web/src/components/Field.tsx +++ b/studio/web/src/components/Field.tsx @@ -208,6 +208,21 @@ export default function Field({ ) } + // float-list (e.g. per-cycle learning rates) ------------------------- + if (kind === 'float-list') { + return ( + + ) + } + // code ---------------------------------------------------------------- if (kind === 'code') { return ( @@ -474,6 +489,65 @@ function IntListField({ ) } +/** Decimal-number list. Scientific notation is accepted, which is important + * for LR values such as 1e-4 and 5e-5. */ +function FloatListField({ + label, helpNode, value, defaultValue, onChange, disabled = false, hintNode, +}: IntListFieldProps) { + const fmt = (v: unknown) => + Array.isArray(v) ? (v as number[]).join(', ') : v === null || v === undefined ? '' : String(v) + const [raw, setRaw] = useState(() => fmt(value)) + const inputRef = useRef(null) + const placeholder = fmt(defaultValue) + + useEffect(() => { + if (document.activeElement !== inputRef.current) setRaw(fmt(value)) + }, [value]) + + const commit = () => { + const nums = raw + .split(/[,\s]+/) + .map((s) => s.trim()) + .filter((s) => s.length > 0) + .map((s) => Number(s)) + .filter((n) => Number.isFinite(n)) + if (nums.length === 0) { + const fallback = Array.isArray(defaultValue) ? defaultValue : [] + onChange(fallback) + setRaw(fmt(fallback)) + return + } + onChange(nums) + setRaw(nums.join(', ')) + } + + return ( +
+
+ {label}{hintNode} +
+ setRaw(e.target.value)} + onBlur={commit} + onKeyDown={(e) => { + if (e.key === 'Enter') { + e.preventDefault() + commit() + } + }} + disabled={disabled} + className="input input-mono" style={fieldStyle(disabled)} + placeholder={placeholder} + /> + {helpNode} +
+ ) +} + interface NumberFieldProps { label: string kind: 'int' | 'float' diff --git a/studio/web/src/i18n/locales/en.json b/studio/web/src/i18n/locales/en.json index f5597c9..32e474e 100644 --- a/studio/web/src/i18n/locales/en.json +++ b/studio/web/src/i18n/locales/en.json @@ -310,7 +310,9 @@ "none": "Constant", "cosine": "Cosine", "cosineWithRestart": "Cosine with restart", - "cosineWithWarmup": "Cosine with warmup" + "cosineWithWarmup": "Cosine with warmup", + "cosineCycles": "Custom cosine cycles", + "constantThenCosine": "Constant, then cosine decay" }, "optimizerType": { "adamw": "AdamW", @@ -326,7 +328,8 @@ "uniform": "Uniform", "logitNormalLow": "Logit Normal Low", "mode": "Mode", - "styleFriendly": "Style-Friendly SNR" + "styleFriendly": "Style-Friendly SNR", + "dualPeak": "Dual Peak — custom style" }, "lossWeighting": { "none": "None", @@ -382,6 +385,16 @@ "infonoise_enabled": "Reset noise_enhancement / loss_weighting / loss_type / schedule_shift to defaults to enable (schema mutex)" }, "descriptions": { + "dual_peak_peak1_position": "First component peak in flow t: 0.525 = timestep 525/1000. The mixture maximum can move.", + "dual_peak_peak1_width": "First peak width (log-SNR standard deviation). Larger = wider, lower peak at the same weight.", + "dual_peak_peak1_weight": "First peak relative weight. All four weights are normalized together; zero disables this component.", + "dual_peak_peak2_position": "Second component peak in flow t: 0.85 = timestep 850/1000. Defaults put the mixture maximum near 870.", + "dual_peak_peak2_width": "Second peak width (log-SNR standard deviation). Independent of the first peak.", + "dual_peak_peak2_weight": "Second peak relative weight. Normalized with the other three weights.", + "dual_peak_background_mean": "Broad Style-Friendly component log-SNR mean. Lower moves it toward high noise. This covers the whole interval, not only the right tail.", + "dual_peak_background_width": "Broad Style-Friendly component log-SNR standard deviation.", + "dual_peak_background_weight": "Broad background relative weight. Zero removes the broad component.", + "dual_peak_uniform_weight": "Uniform coverage relative weight. Default 0.05, normalized with other weights.", "model_family": "Model family. Decides which model implementation training uses; other fields show/hide by family capability. Switching recalculates model paths and resets family-specific defaults", "transformer_path": "Main diffusion model weights (.safetensors)", "vae_path": "VAE weights (.safetensors)", @@ -424,6 +437,13 @@ "lr_scheduler_t_mult": "cosine_with_restart cycle-length multiplier relative to the previous cycle (>1 = lengthening cycles)", "lr_scheduler_eta_min": "LR decay floor for cosine schedules; should be much smaller than the initial lr (e.g. 1e-6 for an initial lr of 1e-4)", "lr_scheduler_warmup_steps": "cosine_with_warmup warmup steps", + "lr_scheduler_cycle_count": "Number of equal-length cosine peaks. Default: 3.", + "lr_scheduler_cycle_max_lr": "Maximum LR shared by all peaks. Leave empty to use learning_rate.", + "lr_scheduler_cycle_min_lr": "Minimum LR shared by all peaks. Per-peak values override it.", + "lr_scheduler_cycle_max_lrs": "Optional maximum LR for each peak, comma-separated. The value count must match the peak count.", + "lr_scheduler_cycle_min_lrs": "Optional minimum LR for each peak, comma-separated. The value count must match the peak count.", + "lr_scheduler_decay_start_ratio": "Training fraction after which the final cosine decay starts. 0.6667 means the last third.", + "lr_scheduler_decay_min_lr": "Learning rate at the end of the final cosine decay.", "optimizer_type": "Optimizer. adamw is the baseline; automagic adapts per-parameter lr (recommend lr=1e-6); came uses confidence-guided factored second moments (less state VRAM than AdamW, lr at AdamW scale); lion uses ~half the state of AdamW (recommend lr = AdamW lr / 3); prodigy / prodigy_plus_schedulefree self-estimate lr (fill 1.0)", "prodigy_d_coef": "Overall scaling coefficient on Prodigy's estimated d; larger = larger effective lr. Raise (2.0+) when underfitting; lower (0.5) when overfitting or on small datasets", "prodigy_safeguard_warmup": "Prevent d from being inflated by early high gradients during Prodigy warmup; default on for stability", @@ -458,7 +478,7 @@ "noise_offset": "DC bias strength (0-0.2, 0=off). Shifts noise mean away from 0, giving the model exposure to extreme brightness scenes (pure black / pure white / high contrast). Typical range 0.05-0.1; below 0.05 the noise is nearly identical to baseline, above 0.1 the initial loss rises substantially", "pyramid_noise_iters": "Number of pyramid noise levels (0-6, 0=off). Each level injects at spatial // 2^(k+1). Effective strength is controlled by pyramid_noise_discount — iters alone determines which frequency bands are covered; layer count makes little difference when discount is low", "pyramid_noise_discount": "Per-level decay coefficient (0.1-0.9). Key knob for low-frequency strength: this implementation normalizes overall noise std to 1, so discount sets the low-frequency share. 0.1-0.4 leaves noise close to standard Gaussian after normalization (equivalent to off); 0.5-0.7 significantly alters low-frequency structure", - "timestep_sampling": "Sampling distribution. logit_normal biases the middle (SD3/Anima default); uniform is equal probability; mode is a single-peak shift; mixed_* blends uniform with a biased end (ratio controlled by timestep_mix_low_prob)", + "timestep_sampling": "Sampling distribution. dual_peak has two configurable peaks and a broad style background; style_friendly emphasizes a high-noise window; logit_normal is the Anima default; uniform covers all t equally. dual_peak ignores timestep_shift, but schedule and resolution shifts still move the curve.", "timestep_shift": "Internal distribution shift for logit-normal / mode: >1 biases the high-noise end (coarse structure), <1 biases the low-noise end (details)", "timestep_mix_low_prob": "Ratio of samples sent to the biased end in mixed_* modes: 0 = all uniform; typical 0.15-0.30", "infonoise_enabled": "[InfoNoise] Enable adaptive timestep sampling: dynamically adjusts the t distribution during training based on information content, focusing on more useful training intervals", @@ -596,7 +616,7 @@ "epochs": "aim for 2000–3000 total steps", "mixed_precision": "bf16", "num_workers": "0 on Windows — required", - "timestep_sampling": "style_friendly for a style LoRA", + "timestep_sampling": "For style: style_friendly or experimental dual_peak with two peaks", "style_snr_mean": "-6 is the paper recipe; try -4 in an A/B if line work matters", "style_snr_sigma": "2.0", "timestep_shift": "3.0; larger pushes toward light and composition, smaller toward fine texture", @@ -712,6 +732,11 @@ "sort_updated": "Recently active", "sort_created": "Newest created", "sort_title": "By name", + "sort_queue": "Queued first", + "sort_finish": "By training finish", + "trainingNow": "Training now", + "queuePosition": "Queued · #{{position}}", + "finishUnknown": "estimating", "archivedToggle": "Archived ({{n}})", "archivedToggleHint": "Toggle: show only archived projects", "archiveProjectTitle": "Archive project (recoverable)", @@ -729,6 +754,13 @@ "evaluatingHint": "Training finished; metrics are being computed against the validation set for each checkpoint (sampling + CLIP / DINO). This disappears when done. See the Eval tab in task detail for details." }, "queue": { + "summaryTitle": "Training queue now", + "summaryCount_one": "{{count}} task", + "summaryCount_other": "{{count}} tasks", + "summaryBreakdown": "Running: {{running}} · waiting: {{waiting}}", + "summaryFinish": "Queue finishes", + "summaryTaskFinish": "Estimated finish", + "summaryUnknown": "Estimating", "title": "GPU task queue", "description": "Train / reg AI / test render · one at a time, next auto-starts on completion", "titleJobs": "Data task queue", diff --git a/studio/web/src/i18n/locales/ru.json b/studio/web/src/i18n/locales/ru.json index bff467c..628a860 100644 --- a/studio/web/src/i18n/locales/ru.json +++ b/studio/web/src/i18n/locales/ru.json @@ -503,6 +503,11 @@ "sort_updated": "Недавно активные", "sort_created": "Новые", "sort_title": "По названию", + "sort_queue": "Сначала в очереди", + "sort_finish": "По времени окончания", + "trainingNow": "Тренируется сейчас", + "queuePosition": "В очереди · №{{position}}", + "finishUnknown": "время уточняется", "archivedToggle": "В архиве ({{n}})", "archivedToggleHint": "Показать только архивные проекты", "archiveProjectTitle": "Архивировать проект (обратимо)", @@ -659,7 +664,26 @@ "noRuns": "Прогонов пока нет." }, "schema": { + "enums": { + "lrScheduler": { + "cosineCycles": "Настраиваемые косинусные циклы", + "constantThenCosine": "Постоянная LR → плавное снижение" + }, + "timestepSampling": { + "dualPeak": "Dual Peak — два пика для стиля" + } + }, "descriptions": { + "dual_peak_peak1_position": "Положение первого пика: 0.525 = 525 из 1000. Это пик отдельной составляющей; максимум общей кривой может немного сместиться.", + "dual_peak_peak1_width": "Ширина первого пика. Больше — шире и ниже при том же весе. Это ширина в log-SNR, а не положение по оси t.", + "dual_peak_peak1_weight": "Вес первого пика. Четыре веса автоматически приводятся к сумме 1. Ноль отключает эту составляющую.", + "dual_peak_peak2_position": "Положение второго пика: 0.85 = 850 из 1000. Со стандартным широким фоном максимум общей кривой находится примерно на 870.", + "dual_peak_peak2_width": "Ширина второго пика. Настраивается независимо от первого. Больше — шире и ниже при том же весе.", + "dual_peak_peak2_weight": "Вес второго пика. Больше — выше его доля после пересчёта всех четырёх весов.", + "dual_peak_background_mean": "Центр широкого фона в log-SNR. Меньше — ближе к шумному краю. Фон заполняет промежуток между пиками и правый склон, но действует по всей оси.", + "dual_peak_background_width": "Ширина широкого фона. По умолчанию 2.2 — как у style-friendly в нашем сравнении.", + "dual_peak_background_weight": "Вес широкого фона. Ноль убирает его; высоты и положения общих максимумов при этом изменятся.", + "dual_peak_uniform_weight": "Вес равномерного покрытия всей оси. По умолчанию 0.05; доля зависит от суммы всех четырёх весов.", "text_encoder_cache": "Кеширует выход текстового энкодера: одинаковые капшены кодируются один раз, а не каждый шаг.", "bucket_min_reso": "Наименьшая сторона корзины. Пусто — выводится из resolution и aspect_ratio_limit.", "bucket_max_reso": "Наибольшая сторона корзины. Пусто — выводится из resolution и aspect_ratio_limit.", @@ -756,6 +780,13 @@ "lr_scheduler_t_mult": "Во сколько раз каждый следующий цикл cosine_with_restart длиннее предыдущего (>1 — циклы растут)", "lr_scheduler_eta_min": "Нижний предел затухания для косинусных расписаний; должен быть сильно меньше стартовой lr (например 1e-6 при старте 1e-4)", "lr_scheduler_warmup_steps": "Длина разогрева для cosine_with_warmup, в шагах", + "lr_scheduler_cycle_count": "Число равных косинусных циклов (пиков). По умолчанию 3.", + "lr_scheduler_cycle_max_lr": "Общий максимум LR для всех пиков. Оставьте пустым, чтобы использовать learning_rate.", + "lr_scheduler_cycle_min_lr": "Общий минимум LR для всех пиков. Индивидуальные значения имеют приоритет.", + "lr_scheduler_cycle_max_lrs": "Максимумы LR для каждого пика через запятую. Число значений должно совпадать с числом пиков.", + "lr_scheduler_cycle_min_lrs": "Минимумы LR для каждого пика через запятую. Число значений должно совпадать с числом пиков.", + "lr_scheduler_decay_start_ratio": "Доля обучения, после которой начинается финальное плавное снижение. 0.6667 означает последнюю треть.", + "lr_scheduler_decay_min_lr": "Скорость обучения в конце плавного снижения.", "optimizer_type": "Оптимизатор. adamw — база; automagic подбирает скорость на каждый параметр (рекомендуется lr=1e-6); came использует факторизованные вторые моменты с оценкой уверенности (памяти меньше, чем у AdamW, lr того же порядка); lion расходует примерно вдвое меньше состояния (рекомендуется lr AdamW делённая на 3); prodigy и prodigy_plus_schedulefree сами оценивают lr (ставьте 1.0)", "prodigy_d_coef": "Общий множитель к оценке d у Prodigy; больше — выше фактическая скорость. Поднимать (2.0+) при недообучении, снижать (0.5) при переобучении и на маленьких датасетах", "prodigy_safeguard_warmup": "Не давать большим градиентам в начале раздуть оценку d у Prodigy; по умолчанию включено ради устойчивости", @@ -790,7 +821,7 @@ "noise_offset": "Сила постоянного смещения (0–0.2, 0 — выключено). Сдвигает среднее шума от нуля, показывая модели крайние по яркости сцены (чёрное, белое, высокий контраст). Обычно 0.05–0.1; ниже 0.05 почти неотличимо от базы, выше 0.1 заметно поднимает стартовый лосс", "pyramid_noise_iters": "Число уровней пирамидального шума (0–6, 0 — выключено). Каждый уровень добавляется на масштабе в 2^(k+1) раз меньше. Фактическую силу задаёт pyramid_noise_discount — число уровней лишь определяет охват частот и при малом discount почти ничего не меняет", "pyramid_noise_discount": "Коэффициент затухания по уровням (0.1–0.9). Главная ручка силы низких частот: реализация нормирует общий разброс шума к 1, поэтому discount задаёт долю низких частот. 0.1–0.4 после нормировки почти неотличимо от обычного гауссова шума; 0.5–0.7 заметно меняет низкочастотную структуру", - "timestep_sampling": "Распределение выборки. logit_normal смещено к середине (по умолчанию для SD3/Anima); uniform — равновероятно; mode — одногорбое смещение; mixed_* смешивает равномерное со смещённым краем (доля задаётся timestep_mix_low_prob)", + "timestep_sampling": "Распределение выборки. dual_peak — два настраиваемых пика с широким фоном для стиля; style_friendly — широкое окно на шумном краю; logit_normal — базовый вариант Anima; uniform — равномерно; mixed_* — смешанные варианты. В dual_peak обычный timestep_shift не действует, но schedule shift и поправка на разрешение сдвигают кривую.", "timestep_shift": "Сдвиг распределения внутри logit-normal / mode: больше 1 — к сильному шуму (крупная структура), меньше 1 — к слабому шуму (детали)", "timestep_mix_low_prob": "Доля сэмплов, уходящих в смещённый край в режимах mixed_*: 0 — всё равномерно; обычно 0.15–0.30", "infonoise_enabled": "InfoNoise: адаптивная выборка шагов времени — по ходу тренировки распределение t подстраивается по информативности и концентрируется на полезных участках", @@ -868,7 +899,7 @@ "epochs": "целиться в 2000–3000 шагов суммарно", "mixed_precision": "bf16", "num_workers": "0 на Windows — обязательно", - "timestep_sampling": "style_friendly для стилевой лоры", + "timestep_sampling": "Для стиля: style_friendly или экспериментальный dual_peak с двумя пиками", "style_snr_mean": "-6 — рецепт статьи; сравните с -4, если важна линия", "style_snr_sigma": "2.0", "timestep_shift": "3.0; больше — в сторону света и композиции, меньше — в мелкую фактуру", @@ -1051,6 +1082,15 @@ "saveFailed": "Не удалось сохранить триггер-слово: {{error}}" }, "queue": { + "summaryTitle": "Очередь тренировок сейчас", + "summaryCount_one": "{{count}} задача", + "summaryCount_few": "{{count}} задачи", + "summaryCount_many": "{{count}} задач", + "summaryCount_other": "{{count}} задач", + "summaryBreakdown": "Выполняется: {{running}} · ожидает: {{waiting}}", + "summaryFinish": "Очередь завершится", + "summaryTaskFinish": "Примерное окончание", + "summaryUnknown": "Время уточняется", "note": "Заметка", "noteAdd": "Добавить заметку", "noteEdit": "Изменить заметку", diff --git a/studio/web/src/i18n/locales/zh.json b/studio/web/src/i18n/locales/zh.json index b114d32..34ee9f7 100644 --- a/studio/web/src/i18n/locales/zh.json +++ b/studio/web/src/i18n/locales/zh.json @@ -310,7 +310,9 @@ "none": "常数", "cosine": "Cosine", "cosineWithRestart": "Cosine 重启", - "cosineWithWarmup": "Cosine 预热" + "cosineWithWarmup": "Cosine 预热", + "cosineCycles": "自定义 Cosine 周期", + "constantThenCosine": "常数后 Cosine 衰减" }, "optimizerType": { "adamw": "AdamW", @@ -326,7 +328,8 @@ "uniform": "Uniform", "logitNormalLow": "Logit Normal Low", "mode": "Mode", - "styleFriendly": "Style-Friendly SNR（风格档）" + "styleFriendly": "Style-Friendly SNR（风格档）", + "dualPeak": "Dual Peak — 自定义双峰" }, "lossWeighting": { "none": "无", @@ -382,6 +385,16 @@ "infonoise_enabled": "需先把 noise_enhancement / loss_weighting / loss_type / schedule_shift 全部置默认才能启用（schema 互斥）" }, "descriptions": { + "dual_peak_peak1_position": "第一个分量峰的位置：0.525 = 525/1000。混合后的峰可能略有偏移。", + "dual_peak_peak1_width": "第一个峰的宽度（log-SNR 标准差）。权重不变时越大越宽、越低。", + "dual_peak_peak1_weight": "第一个峰的相对权重。四个权重自动归一化；0 关闭此分量。", + "dual_peak_peak2_position": "第二个分量峰的位置：0.85 = 850/1000。默认混合曲线峰值约为 870。", + "dual_peak_peak2_width": "第二个峰的宽度（log-SNR 标准差），独立于第一个峰。", + "dual_peak_peak2_weight": "第二个峰的相对权重，与其他三个权重一起归一化。", + "dual_peak_background_mean": "宽背景分量的 log-SNR 均值。越小越偏高噪声，覆盖整个区间而不只是右侧尾部。", + "dual_peak_background_width": "宽背景分量的 log-SNR 标准差，默认 2.2。", + "dual_peak_background_weight": "宽背景的相对权重。0 关闭；混合峰的位置和高度也会改变。", + "dual_peak_uniform_weight": "全区间均匀覆盖的相对权重，默认 0.05，与其他权重归一化。", "model_family": "模型族。决定训练走哪套模型实现；其余字段按族能力自动显隐。切换将重算模型路径并重置族相关默认值", "transformer_path": "主扩散模型权重（.safetensors）", "vae_path": "VAE 权重（.safetensors）", @@ -424,6 +437,13 @@ "lr_scheduler_t_mult": "cosine_with_restart 每次重启后周期相对上轮的倍数（>1 周期递增）", "lr_scheduler_eta_min": "学习率衰减下限：cosine 调度到此值后不再下降；通常远小于初始 lr（如初始 1e-4 配 1e-6）", "lr_scheduler_warmup_steps": "cosine_with_warmup 预热步数", + "lr_scheduler_cycle_count": "等长 Cosine 峰的数量，默认 3。", + "lr_scheduler_cycle_max_lr": "所有峰共用的最高 LR；留空时使用 learning_rate。", + "lr_scheduler_cycle_min_lr": "所有峰共用的最低 LR；逐峰列表会覆盖此值。", + "lr_scheduler_cycle_max_lrs": "每个峰的最高 LR，逗号分隔；数量必须与峰数一致。", + "lr_scheduler_cycle_min_lrs": "每个峰的最低 LR，逗号分隔；数量必须与峰数一致。", + "lr_scheduler_decay_start_ratio": "最终 Cosine 衰减开始时的训练比例；0.6667 表示最后三分之一。", + "lr_scheduler_decay_min_lr": "最终 Cosine 衰减结束时的 LR。", "optimizer_type": "优化器。adamw 标准基线；automagic 自适应每参数 lr（推荐 lr=1e-6）；came 置信度引导 + 分解二阶矩（state 显存低于 AdamW，lr 同 AdamW 量级）；lion 显存约 AdamW 一半（推荐 lr=AdamW lr / 3）；prodigy / prodigy_plus_schedulefree 自适应估 lr（lr 填 1.0）", "prodigy_d_coef": "Prodigy 估出的 d 整体缩放系数；越大有效 lr 越大。欠拟合时调高（2.0+），过拟合 / 小数据集时调低（0.5）", "prodigy_safeguard_warmup": "Prodigy warmup 期间防止 d 被初期高梯度推高；默认开启更稳", @@ -458,7 +478,7 @@ "noise_offset": "DC 偏置强度（0-0.2，0=关闭）。让噪声 mean 偏离 0，让模型有机会学习生成极端亮度场景（pure black / pure white / 强对比）。典型范围 0.05-0.1；0.05 以下噪声场跟 baseline 几乎一样，超过 0.1 起点 loss 会显著偏高", "pyramid_noise_iters": "金字塔噪声层数（0-6，0=关闭）。每层在 spatial // 2^(k+1) 尺度注入。实际效果强度由 pyramid_noise_discount 决定 —— iters 单独决定覆盖的频段范围，discount 低时层数多少差异很小", "pyramid_noise_discount": "每层相对衰减系数（0.1-0.9）。控制低频强度的核心参数：本实现把整体噪声 std 归一化到 1，所以 discount 决定低频占比。0.1-0.4 归一化后噪声接近标准高斯，等价于关闭；0.5-0.7 显著改变低频结构", - "timestep_sampling": "采样分布。logit_normal 偏中段（SD3/Anima 默认）；uniform 等概率；mode 单峰偏移；mixed_* 混合 uniform 与偏置端（比例由 timestep_mix_low_prob 控制）", + "timestep_sampling": "采样分布。dual_peak 为可调双峰加宽背景；style_friendly 偏高噪声风格窗口；logit_normal 是 Anima 默认；uniform 等概率。dual_peak 忽略 timestep_shift，但 schedule shift 和分辨率修正仍会移动曲线。", "timestep_shift": "logit-normal / mode 内部的分布偏移：>1 偏向高噪声端（粗结构），<1 偏向低噪声端（细节）", "timestep_mix_low_prob": "mixed_* 模式下走偏置端的样本比例：0 = 全 uniform；典型 0.15-0.30", "schedule_shift": "采样后对 t 做的额外 σ schedule 偏移：1.0 = 无偏移；越大整体偏向高噪声端。与 timestep_shift 区别：作用于最终 t 而非 logit-normal 内部", @@ -593,6 +613,11 @@ "sort_updated": "最近活跃", "sort_created": "最新创建", "sort_title": "按名称", + "sort_queue": "队列中的优先", + "sort_finish": "按训练结束时间", + "trainingNow": "正在训练", + "queuePosition": "队列第 {{position}} 位", + "finishUnknown": "正在估算", "archivedToggle": "已归档（{{n}}）", "archivedToggleHint": "切换：只显示已归档项目", "archiveProjectTitle": "归档项目（可恢复）", @@ -610,6 +635,12 @@ "evaluatingHint": "训练已结束，正在用验证集对各 checkpoint 算指标（出图 + CLIP / DINO），完成后这里消失。明细看任务详情的「评估」tab。" }, "queue": { + "summaryTitle": "当前训练队列", + "summaryCount": "{{count}} 个任务", + "summaryBreakdown": "运行中：{{running}} · 等待：{{waiting}}", + "summaryFinish": "队列预计结束", + "summaryTaskFinish": "预计结束", + "summaryUnknown": "正在估算", "title": "GPU 任务队列", "description": "训练 / 正则 AI / 测试出图 · 同一时刻仅运行一个，完成后自动启动下一个", "titleJobs": "数据任务队列", diff --git a/studio/web/src/lib/queueEstimates.test.ts b/studio/web/src/lib/queueEstimates.test.ts new file mode 100644 index 0000000..6f7b2f7 --- /dev/null +++ b/studio/web/src/lib/queueEstimates.test.ts @@ -0,0 +1,46 @@ +import { describe, expect, it } from 'vitest' +import type { Task } from '../api/client' +import { buildTrainingForecast, typicalTrainingDuration } from './queueEstimates' + +function task(over: Partial): Task { + return { + id: 1, + name: 'training', + config_name: 'train', + task_type: 'train', + status: 'pending', + priority: 0, + created_at: 1, + started_at: null, + finished_at: null, + pid: null, + exit_code: null, + output_dir: null, + error_msg: null, + ...over, + } +} + +describe('queue training forecast', () => { + it('uses the median completed duration for waiting tasks', () => { + const tasks = [ + task({ id: 1, status: 'done', started_at: 10, finished_at: 110 }), + task({ id: 2, status: 'done', started_at: 10, finished_at: 310 }), + task({ id: 3, status: 'done', started_at: 10, finished_at: 210 }), + ] + expect(typicalTrainingDuration(tasks)).toBe(200) + }) + + it('orders by scheduler priority and adds estimates cumulatively', () => { + const tasks = [ + task({ id: 1, status: 'done', started_at: 0.1, finished_at: 100.1 }), + task({ id: 2, status: 'running', started_at: 900, created_at: 1 }), + task({ id: 3, priority: 1, created_at: 3 }), + task({ id: 4, priority: 0, created_at: 2 }), + ] + const result = buildTrainingForecast(tasks, { step: 50, total_steps: 100, speed: 1 }, 1000) + + expect(result.map(({ task }) => task.id)).toEqual([2, 3, 4]) + expect(result.map(({ finishesAt }) => finishesAt)).toEqual([1050, 1150, 1250]) + }) +}) diff --git a/studio/web/src/lib/queueEstimates.ts b/studio/web/src/lib/queueEstimates.ts new file mode 100644 index 0000000..09c742c --- /dev/null +++ b/studio/web/src/lib/queueEstimates.ts @@ -0,0 +1,69 @@ +import type { MonitorState, Task } from '../api/client' + +export interface TaskForecast { + task: Task + finishesAt: number | null +} + +const isTrainingTask = (task: Task) => (task.task_type ?? 'train') === 'train' + +export function typicalTrainingDuration(tasks: Task[]): number | null { + const durations = tasks + .filter((task) => isTrainingTask(task) && task.status === 'done' && task.started_at && task.finished_at) + .map((task) => task.finished_at! - task.started_at!) + .filter((seconds) => seconds > 0) + .sort((a, b) => a - b) + + if (!durations.length) return null + const middle = Math.floor(durations.length / 2) + return durations.length % 2 + ? durations[middle] + : (durations[middle - 1] + durations[middle]) / 2 +} + +export function buildTrainingForecast( + tasks: Task[], + monitor: MonitorState | null, + now = Date.now() / 1000, +): TaskForecast[] { + const live = tasks + .filter((task) => isTrainingTask(task) && (task.status === 'running' || task.status === 'pending')) + .sort((a, b) => { + if (a.status !== b.status) return a.status === 'running' ? -1 : 1 + return b.priority - a.priority || a.created_at - b.created_at + }) + const typical = typicalTrainingDuration(tasks) + let cursor = now + let canEstimate = true + + return live.map((task) => { + let remaining: number | null = typical + if (task.status === 'running') { + if ( + monitor?.speed && monitor.speed > 0 && + monitor.step != null && monitor.total_steps != null && + monitor.total_steps >= monitor.step + ) { + remaining = (monitor.total_steps - monitor.step) / monitor.speed + } else if (typical != null && task.started_at) { + remaining = Math.max(0, typical - (now - task.started_at)) + } + } + + if (!canEstimate || remaining == null) { + canEstimate = false + return { task, finishesAt: null } + } + cursor += remaining + return { task, finishesAt: cursor } + }) +} + +export function latestProjectFinish(tasks: Task[], projectId: number): number | null { + let latest: number | null = null + for (const task of tasks) { + if (!isTrainingTask(task) || task.project_id !== projectId || !task.finished_at) continue + latest = latest == null ? task.finished_at : Math.max(latest, task.finished_at) + } + return latest +} diff --git a/studio/web/src/lib/schema.test.ts b/studio/web/src/lib/schema.test.ts index a87fe3e..30ee779 100644 --- a/studio/web/src/lib/schema.test.ts +++ b/studio/web/src/lib/schema.test.ts @@ -33,13 +33,13 @@ describe('controlKind', () => { expect(controlKind({ type: 'string' } as SchemaProperty)).toBe('string') }) - it('maps integer arrays (e.g. resolution) to int-list, string arrays to string-list', () => { + it('maps integer, float, and string arrays to matching list controls', () => { expect( controlKind({ type: 'array', items: { type: 'integer' } } as SchemaProperty) ).toBe('int-list') expect( controlKind({ type: 'array', items: { type: 'number' } } as SchemaProperty) - ).toBe('int-list') + ).toBe('float-list') expect( controlKind({ type: 'array', items: { type: 'string' } } as SchemaProperty) ).toBe('string-list') diff --git a/studio/web/src/lib/schema.ts b/studio/web/src/lib/schema.ts index cd72678..7204d3a 100644 --- a/studio/web/src/lib/schema.ts +++ b/studio/web/src/lib/schema.ts @@ -14,6 +14,7 @@ export type ControlKind = | 'code' | 'string-list' | 'int-list' + | 'float-list' /** * 推断字段的控件类型。优先用 schema 里 control 自定义元字段，否则按 @@ -47,7 +48,9 @@ export function controlKind(prop: SchemaProperty): ControlKind { if (type === 'number') return 'float' if (type === 'array') { const itemType = prop.items?.type - return itemType === 'integer' || itemType === 'number' ? 'int-list' : 'string-list' + if (itemType === 'integer') return 'int-list' + if (itemType === 'number') return 'float-list' + return 'string-list' } return 'string' } @@ -125,6 +128,8 @@ export const SCHEMA_ENUM_LABEL_KEYS: Record> = { cosine: 'schema.enums.lrScheduler.cosine', cosine_with_restart: 'schema.enums.lrScheduler.cosineWithRestart', cosine_with_warmup: 'schema.enums.lrScheduler.cosineWithWarmup', + cosine_cycles: 'schema.enums.lrScheduler.cosineCycles', + constant_then_cosine: 'schema.enums.lrScheduler.constantThenCosine', }, optimizer_type: { adamw: 'schema.enums.optimizerType.adamw', @@ -141,6 +146,7 @@ export const SCHEMA_ENUM_LABEL_KEYS: Record> = { logit_normal_low: 'schema.enums.timestepSampling.logitNormalLow', mode: 'schema.enums.timestepSampling.mode', style_friendly: 'schema.enums.timestepSampling.styleFriendly', + dual_peak: 'schema.enums.timestepSampling.dualPeak', }, loss_weighting: { none: 'schema.enums.lossWeighting.none', diff --git a/studio/web/src/pages/Projects.tsx b/studio/web/src/pages/Projects.tsx index 4d36098..f7af7ce 100644 --- a/studio/web/src/pages/Projects.tsx +++ b/studio/web/src/pages/Projects.tsx @@ -1,7 +1,7 @@ -import { useEffect, useState } from 'react' +import { useEffect, useMemo, useState } from 'react' import { useTranslation } from 'react-i18next' import { useNavigate } from 'react-router-dom' -import { api, type BundleImportResult, type ProjectSummary } from '../api/client' +import { api, type BundleImportResult, type ProjectSummary, type Task } from '../api/client' import PageHeader from '../components/PageHeader' import PathPicker from '../components/PathPicker' import UploadProgressBar from '../components/UploadProgressBar' @@ -10,10 +10,22 @@ import { useDialog } from '../components/Dialog' import { useToast } from '../components/Toast' import { useEventStream } from '../lib/useEventStream' import { useUploadProgress } from '../lib/useUploadProgress' +import { useMonitorProgress } from '../lib/useMonitorProgress' +import { buildTrainingForecast, latestProjectFinish } from '../lib/queueEstimates' + +type ProjectSort = 'queue' | 'finish' | 'updated' | 'created' | 'title' + +function formatProjectTime(ts: number): string { + return new Intl.DateTimeFormat(undefined, { + day: '2-digit', month: 'short', hour: '2-digit', minute: '2-digit', + }).format(new Date(ts * 1000)) +} export default function ProjectsPage() { const { t } = useTranslation() const [items, setItems] = useState([]) + const [tasks, setTasks] = useState([]) + const [sort, setSort] = useState('queue') const [loading, setLoading] = useState(true) const [error, setError] = useState(null) const [creating, setCreating] = useState(false) @@ -28,8 +40,12 @@ export default function ProjectsPage() { const refresh = async () => { try { - const list = await api.listProjects() + const [list, queue] = await Promise.all([ + api.listProjects(), + api.listQueue().catch(() => [] as Task[]), + ]) setItems(list) + setTasks(queue) setError(null) } catch (e) { setError(String(e)) @@ -41,9 +57,51 @@ export default function ProjectsPage() { useEffect(() => { void refresh() }, []) useEventStream((evt) => { - if (evt.type === 'project_state_changed') void refresh() + if (evt.type === 'project_state_changed' || evt.type === 'task_state_changed') void refresh() }) + const runningTaskId = tasks.find((task) => task.status === 'running')?.id ?? null + const { state: monitor } = useMonitorProgress(runningTaskId) + const forecast = useMemo(() => buildTrainingForecast(tasks, monitor), [tasks, monitor]) + const forecastByProject = useMemo(() => { + const map = new Map() + forecast.forEach(({ task, finishesAt }, order) => { + if (task.project_id != null && !map.has(task.project_id)) { + map.set(task.project_id, { order, finishesAt, status: task.status }) + } + }) + return map + }, [forecast]) + const sortedItems = useMemo(() => [...items].sort((a, b) => { + const aq = forecastByProject.get(a.id) + const bq = forecastByProject.get(b.id) + if (sort === 'queue') { + if (aq && bq) return aq.order - bq.order + if (aq) return -1 + if (bq) return 1 + return b.updated_at - a.updated_at + } + if (sort === 'finish') { + if (aq || bq) { + if (!aq) return 1 + if (!bq) return -1 + if (aq.finishesAt != null && bq.finishesAt != null) return aq.finishesAt - bq.finishesAt + if (aq.finishesAt != null) return -1 + if (bq.finishesAt != null) return 1 + return aq.order - bq.order + } + const af = latestProjectFinish(tasks, a.id) + const bf = latestProjectFinish(tasks, b.id) + if (af != null && bf != null) return bf - af + if (af != null) return -1 + if (bf != null) return 1 + return b.updated_at - a.updated_at + } + if (sort === 'created') return b.created_at - a.created_at + if (sort === 'title') return a.title.localeCompare(b.title) + return b.updated_at - a.updated_at + }), [items, tasks, forecastByProject, sort]) + const handleCreate = async (form: NewProjectForm) => { setBusy(true) try { @@ -162,6 +220,26 @@ export default function ProjectsPage() {
{error}
)} + {!loading && items.length > 0 && ( +
+ + +
+ )} + {loading ? (
{[1, 2, 3].map(i => ( @@ -178,12 +256,13 @@ export default function ProjectsPage() {
) : (
- {items.map((p) => ( + {sortedItems.map((p) => ( openProject(p)} onDelete={(e) => handleDelete(p, e)} + queueInfo={forecastByProject.get(p.id)} /> ))}
@@ -298,10 +377,12 @@ function ProjectCard({ project: p, onClick, onDelete, + queueInfo, }: { project: ProjectSummary onClick: () => void onDelete: (e: React.MouseEvent) => void + queueInfo?: { order: number; finishesAt: number | null; status: Task['status'] } }) { const { t } = useTranslation() @@ -329,6 +410,21 @@ function ProjectCard({

)} + {queueInfo && ( +
+ + {queueInfo.status === 'running' + ? t('projects.trainingNow') + : t('projects.queuePosition', { position: queueInfo.order + 1 })} + + + {queueInfo.finishesAt + ? `≈ ${formatProjectTime(queueInfo.finishesAt)}` + : t('projects.finishUnknown')} + +
+ )} +
{/* active version 名（直接版本，无前缀文本） */} {p.active_version_label ? ( diff --git a/studio/web/src/pages/Queue.tsx b/studio/web/src/pages/Queue.tsx index 65b47f9..57b70f8 100644 --- a/studio/web/src/pages/Queue.tsx +++ b/studio/web/src/pages/Queue.tsx @@ -11,6 +11,7 @@ import { useDialog } from '../components/Dialog' import { useToast } from '../components/Toast' import { useEventStream } from '../lib/useEventStream' import { useMonitorProgress } from '../lib/useMonitorProgress' +import { buildTrainingForecast } from '../lib/queueEstimates' /** 备注输入上限 —— 与后端 _MAX_NOTE_LEN 对齐（超了后端截断，这里先拦住）。 */ const MAX_NOTE_LEN = 500 @@ -62,6 +63,12 @@ function fmtDurationShort(ms: number): string { return `${(ms / 3600e3).toFixed(1)}h` } +function fmtFinishTime(ts: number): string { + return new Intl.DateTimeFormat(undefined, { + day: '2-digit', month: 'short', hour: '2-digit', minute: '2-digit', + }).format(new Date(ts * 1000)) +} + export default function QueuePage() { const { t } = useTranslation() const [tasks, setTasks] = useState([]) @@ -174,6 +181,15 @@ export default function QueuePage() { // 切换时 hook 自动清状态 + 重拉 /api/state 冷启动；不需要本组件再写清理逻辑。 const { state: monitor } = useMonitorProgress(runningTaskId) + const queueForecast = useMemo( + () => buildTrainingForecast(tasks, monitor), + [tasks, monitor], + ) + const waitingCount = queueForecast.filter(({ task }) => task.status === 'pending').length + const queueFinish = queueForecast.length + ? queueForecast[queueForecast.length - 1].finishesAt + : null + const sorted = useMemo(() => [...tasks].sort((a, b) => b.id - a.id), [tasks]) const prevCount = useCallback((taskId: number): number => { @@ -398,6 +414,48 @@ export default function QueuePage() { } >
+ {loaded && queueForecast.length > 0 && ( +
+
+
+
+ {t('queue.summaryTitle')} +
+
+ {t('queue.summaryCount', { count: queueForecast.length })} +
+
+ {t('queue.summaryBreakdown', { + running: queueForecast.length - waitingCount, + waiting: waitingCount, + })} +
+
+
+
{t('queue.summaryFinish')}
+
+ {queueFinish ? `≈ ${fmtFinishTime(queueFinish)}` : t('queue.summaryUnknown')} +
+
+
+
+ {queueForecast.map(({ task, finishesAt }) => ( + + {task.name} + {task.status === 'running' && {t('status.running')}} + + ))} +
+
+ )} {/* ADR §4.1 队列挂起 banner — 仅 held=true 时显示，sticky 顶部。 */} {holdState?.held && (
None: assert ns.lr_scheduler_eta_min == 1e-7 +def test_training_config_cli_cosine_cycles_float_lists() -> None: + parser = bridge.build_parser(TrainingConfig) + ns = parser.parse_args([ + "--lr-scheduler", "cosine_cycles", + "--lr-scheduler-cycle-count", "3", + "--lr-scheduler-cycle-max-lrs", "1e-4", "8e-5", "6e-5", + "--lr-scheduler-cycle-min-lrs", "5e-5", "4e-5", "0", + ]) + assert ns.lr_scheduler == "cosine_cycles" + assert ns.lr_scheduler_cycle_max_lrs == [1e-4, 8e-5, 6e-5] + assert ns.lr_scheduler_cycle_min_lrs == [5e-5, 4e-5, 0.0] + + def test_training_config_yaml_round_trip() -> None: """走完 CLI → YAML 合并这一条路径，确认 yaml_dict 字段都能被读进 args。""" args = bridge.namespace_from_config( diff --git a/tests/test_dual_peak_sampler.py b/tests/test_dual_peak_sampler.py new file mode 100644 index 0000000..67631af --- /dev/null +++ b/tests/test_dual_peak_sampler.py @@ -0,0 +1,109 @@ +"""Schema -> CLI/YAML -> registry -> real sample_t integration.""" + +from dataclasses import asdict, fields +from types import SimpleNamespace + +import pytest +import torch +import yaml +from pydantic import ValidationError + +from dual_peak_sampling import DualPeakConfig, sample_timesteps +from studio.schema import TrainingConfig +from studio.infrastructure.argparse_bridge import build_parser +from training.timestep_sampling import sample_t, apply_resolution_shift +from training.timestep_samplers import build_timestep_sampler +from training.timestep_samplers.protocol import TimestepSamplerProtocol + + +def test_schema_defaults_match_portable_package_and_roundtrip(): + cfg = TrainingConfig(timestep_sampling="dual_peak") + for field in fields(DualPeakConfig): + key = f"dual_peak_{field.name}" + assert getattr(cfg, key) == getattr(DualPeakConfig(), field.name) + assert TrainingConfig.model_fields[key].json_schema_extra["show_when"] == "timestep_sampling==dual_peak" + restored = TrainingConfig.model_validate(yaml.safe_load(yaml.safe_dump(cfg.model_dump()))) + assert restored == cfg + assert "timestep_sampling!=dual_peak" in TrainingConfig.model_fields["timestep_shift"].json_schema_extra["show_when"] + assert TrainingConfig().timestep_sampling == "logit_normal" + + +def test_cli_to_registry_preserves_all_custom_parameters(): + expected = DualPeakConfig(peak1_position=.48, peak1_width=.7, peak1_weight=.2, + peak2_position=.82, peak2_width=1.4, peak2_weight=.4, + background_mean=0, background_width=2.4, + background_weight=.3, uniform_weight=0) + argv = ["--timestep-sampling", "dual_peak"] + for name, value in asdict(expected).items(): + argv.extend(["--dual-peak-" + name.replace("_", "-"), str(value)]) + args = build_parser(TrainingConfig, add_config_arg=False).parse_args(argv) + sampler = build_timestep_sampler(args, 100) + assert isinstance(sampler, TimestepSamplerProtocol) + assert sampler.status()["dual_peak"] == asdict(expected) + torch.manual_seed(31) + actual = sampler.sample(1024, "cpu") + torch.manual_seed(31) + torch.testing.assert_close(actual, sample_timesteps(1024, config=expected), rtol=0, atol=0) + + +def test_shift_is_ignored_but_schedule_and_resolution_shifts_compose(): + torch.manual_seed(7) + plain = sample_t(512, "cpu", mode="dual_peak", shift=.5) + torch.manual_seed(7) + other_shift = sample_t(512, "cpu", mode="dual_peak", shift=9) + torch.testing.assert_close(plain, other_shift, rtol=0, atol=0) + torch.manual_seed(7) + shifted = sample_t(512, "cpu", mode="dual_peak", timestep_schedule_shift=2) + expected = (2 * plain / (1 + plain)).clamp(1e-4, 1-1e-4) + torch.testing.assert_close(shifted, expected) + torch.testing.assert_close(apply_resolution_shift(plain, [1024]*512, 1024), plain) + + +def test_infonoise_warmup_uses_custom_config(): + args = SimpleNamespace(timestep_sampling="dual_peak", infonoise_enabled=True, + dual_peak_peak1_weight=0, dual_peak_peak2_weight=0, + dual_peak_background_weight=0, dual_peak_uniform_weight=1) + sampler = build_timestep_sampler(args, 100) + assert sampler.baseline_dual_peak_config.uniform_weight == 1 + torch.manual_seed(4) + actual = sampler.sample(512, "cpu") + torch.manual_seed(4) + expected = sample_timesteps(512, config=sampler.baseline_dual_peak_config) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_stateless_resume_uses_trainer_global_rng(): + args = SimpleNamespace(timestep_sampling="dual_peak") + first = build_timestep_sampler(args, 100) + first.sample(5, "cpu") + rng = torch.get_rng_state() + state = first.state_dict() + expected = first.sample(64, "cpu") + resumed = build_timestep_sampler(args, 100) + resumed.load_state_dict(state) + torch.set_rng_state(rng) + torch.testing.assert_close(resumed.sample(64, "cpu"), expected, rtol=0, atol=0) + + +def test_overrides_warn_without_changing_user_settings(caplog): + args = SimpleNamespace(timestep_sampling="dual_peak", leap_enabled=True, + timestep_schedule_shift=2, timestep_shift_resolution_aware=True) + sampler = build_timestep_sampler(args, 100) + assert "non-Leap" in caplog.text + assert "move the configured peaks" in caplog.text + assert sampler.timestep_schedule_shift == 2 + assert args.leap_enabled and args.timestep_shift_resolution_aware + + +@pytest.mark.parametrize("changes", [ + {"dual_peak_peak1_position": 0}, {"dual_peak_peak2_position": 1}, + {"dual_peak_peak1_width": 0}, {"dual_peak_peak2_width": 3}, + {"dual_peak_background_mean": float("nan")}, {"dual_peak_uniform_weight": -1}, + {"dual_peak_peak1_weight": 0, "dual_peak_peak2_weight": 0, + "dual_peak_background_weight": 0, "dual_peak_uniform_weight": 0}, +]) +def test_schema_and_runtime_reject_invalid_custom_config(changes): + with pytest.raises(ValidationError): + TrainingConfig(timestep_sampling="dual_peak", **changes) + with pytest.raises(ValueError): + build_timestep_sampler(SimpleNamespace(timestep_sampling="dual_peak", **changes), 100) diff --git a/tests/test_plugin_registry.py b/tests/test_plugin_registry.py index c9b8141..7d95ffe 100644 --- a/tests/test_plugin_registry.py +++ b/tests/test_plugin_registry.py @@ -103,7 +103,13 @@ def test_adamw8bit_validate_rejects_missing_bitsandbytes(monkeypatch) -> None: def test_scheduler_builders_dict_excludes_none() -> None: from training.schedulers import BUILDERS, SCHEMA_ONLY_OPTIONS - assert set(BUILDERS) == {"cosine", "cosine_with_restart", "cosine_with_warmup"} + assert set(BUILDERS) == { + "constant_then_cosine", + "cosine", + "cosine_cycles", + "cosine_with_restart", + "cosine_with_warmup", + } assert SCHEMA_ONLY_OPTIONS == {"none"} @@ -255,6 +261,61 @@ def test_cosine_with_warmup_clamps_negative_eta_min() -> None: scheduler.step() +def test_cosine_cycles_supports_per_peak_min_and_max_lr() -> None: + torch = pytest.importorskip("torch") + from training.schedulers import build_scheduler + + param = torch.nn.Parameter(torch.tensor([1.0])) + optimizer = torch.optim.SGD([param], lr=1e-3) + args = argparse.Namespace( + lr_scheduler="cosine_cycles", + lr_scheduler_cycle_count=3, + lr_scheduler_cycle_max_lr=None, + lr_scheduler_cycle_min_lr=0.0, + lr_scheduler_cycle_max_lrs=[1e-3, 8e-4, 6e-4], + lr_scheduler_cycle_min_lrs=[5e-4, 4e-4, 1e-5], + ) + scheduler = build_scheduler(args, optimizer, total_steps=12) + + lrs = [optimizer.param_groups[0]["lr"]] + for _ in range(12): + optimizer.step() + scheduler.step() + lrs.append(optimizer.param_groups[0]["lr"]) + + assert lrs[0] == pytest.approx(1e-3) + assert lrs[4] == pytest.approx(8e-4) + assert lrs[8] == pytest.approx(6e-4) + assert lrs[12] == pytest.approx(1e-5) + assert min(lrs[0:4]) >= 5e-4 + assert min(lrs[4:8]) >= 4e-4 + + +def test_constant_then_cosine_stays_flat_then_decays() -> None: + torch = pytest.importorskip("torch") + from training.schedulers import build_scheduler + + param = torch.nn.Parameter(torch.tensor([1.0])) + optimizer = torch.optim.SGD([param], lr=1e-3) + args = argparse.Namespace( + lr_scheduler="constant_then_cosine", + lr_scheduler_decay_start_ratio=2.0 / 3.0, + lr_scheduler_decay_min_lr=1e-5, + ) + scheduler = build_scheduler(args, optimizer, total_steps=12) + + lrs = [optimizer.param_groups[0]["lr"]] + for _ in range(12): + optimizer.step() + scheduler.step() + lrs.append(optimizer.param_groups[0]["lr"]) + + assert lrs[:9] == pytest.approx([1e-3] * 9) + assert lrs[9] < 1e-3 + assert lrs[9] > lrs[10] > lrs[11] > lrs[12] + assert lrs[12] == pytest.approx(1e-5) + + def test_ppsf_zero_prodigy_steps_disables_freeze(monkeypatch) -> None: """ppsf_prodigy_steps=0 keeps the upstream PPSF meaning: never freeze d.""" from training.optimizers import prodigy_plus_schedulefree as ppsf diff --git a/tests/test_studio_configs.py b/tests/test_studio_configs.py index c357308..6bbb8b6 100644 --- a/tests/test_studio_configs.py +++ b/tests/test_studio_configs.py @@ -39,7 +39,9 @@ def test_schema_is_complete() -> None: "transformer_path", "data_dir", "lora_type", "lora_rank", "epochs", "tlora_min_rank", "tlora_alpha_rank_scale", "tlora_use_ortho", "optimizer_type", "prodigy_d_coef", "prodigy_safeguard_warmup", - "lr_scheduler_warmup_steps", + "lr_scheduler_warmup_steps", "lr_scheduler_cycle_count", + "lr_scheduler_cycle_max_lrs", "lr_scheduler_cycle_min_lrs", + "lr_scheduler_decay_start_ratio", "lr_scheduler_decay_min_lr", "lion_beta1", "lion_beta2", "came_beta1", "came_beta2", "came_beta3", "came_eps1", "came_eps2", "came_clip_threshold", @@ -64,6 +66,8 @@ def test_schema_is_complete() -> None: scheduler_annotation = fields["lr_scheduler"].annotation scheduler_options = getattr(scheduler_annotation, "__args__", ()) assert "cosine_with_warmup" in scheduler_options + assert "cosine_cycles" in scheduler_options + assert "constant_then_cosine" in scheduler_options # optimizer_type Literal 包含 Lion / PPSF optimizer_annotation = fields["optimizer_type"].annotation # Literal 的 __args__ 包含所有合法值 @@ -110,6 +114,8 @@ def test_schema_carries_ui_metadata(client: TestClient) -> None: assert "automagic" not in props["learning_rate"]["disable_when"] assert props["lr_scheduler"]["disable_when"] == "optimizer_type==automagic||optimizer_type==prodigy||optimizer_type==prodigy_plus_schedulefree||optimizer_type==soap_sf" assert props["lr_scheduler_warmup_steps"]["show_when"] == "lr_scheduler==cosine_with_warmup" + assert props["lr_scheduler_cycle_count"]["show_when"] == "lr_scheduler==cosine_cycles" + assert props["lr_scheduler_decay_start_ratio"]["show_when"] == "lr_scheduler==constant_then_cosine" assert props["automagic_min_lr"]["show_when"] == "optimizer_type==automagic" assert props["automagic_max_lr"]["show_when"] == "optimizer_type==automagic" assert props["automagic_variant"]["show_when"] == "optimizer_type==automagic" @@ -129,6 +135,41 @@ def test_extra_fields_are_silently_ignored() -> None: assert not hasattr(cfg, "learning_ratee") +def test_cosine_cycles_validates_per_peak_lr_lists() -> None: + cfg = TrainingConfig.model_validate({ + "lr_scheduler": "cosine_cycles", + "learning_rate": 1e-4, + "lr_scheduler_cycle_count": 3, + "lr_scheduler_cycle_max_lrs": [1e-4, 8e-5, 6e-5], + "lr_scheduler_cycle_min_lrs": [5e-5, 4e-5, 0.0], + }) + assert cfg.lr_scheduler_cycle_count == 3 + + with pytest.raises(Exception, match="exactly 3 values"): + TrainingConfig.model_validate({ + "lr_scheduler": "cosine_cycles", + "lr_scheduler_cycle_count": 3, + "lr_scheduler_cycle_max_lrs": [1e-4, 8e-5], + }) + + with pytest.raises(Exception, match="cannot exceed max LR"): + TrainingConfig.model_validate({ + "lr_scheduler": "cosine_cycles", + "lr_scheduler_cycle_count": 1, + "lr_scheduler_cycle_max_lrs": [1e-5], + "lr_scheduler_cycle_min_lrs": [2e-5], + }) + + +def test_constant_then_cosine_min_cannot_exceed_base_lr() -> None: + with pytest.raises(Exception, match="cannot exceed learning_rate"): + TrainingConfig.model_validate({ + "lr_scheduler": "constant_then_cosine", + "learning_rate": 1e-4, + "lr_scheduler_decay_min_lr": 2e-4, + }) + + def test_ppsf_rejects_non_none_scheduler() -> None: """PPSF + lr_scheduler != none 应该在 pydantic 层就被拒。""" payload = TrainingConfig().model_dump(mode="python")