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(