feat(training): dual-peak timestep sampling, cyclic LR schedules, que…

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.

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"]

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

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"]