From 93110ce5226ce1be49d849122f8c7bd1383203c9 Mon Sep 17 00:00:00 2001 From: auraithm Date: Fri, 13 Feb 2026 18:31:41 +0000 Subject: [PATCH 1/2] fix: convert boundary_ratio from timestep space to index space via inverse sigma_shift --- mova/diffusion/pipelines/mova_train.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/mova/diffusion/pipelines/mova_train.py b/mova/diffusion/pipelines/mova_train.py index 658913f..e11af66 100644 --- a/mova/diffusion/pipelines/mova_train.py +++ b/mova/diffusion/pipelines/mova_train.py @@ -1386,10 +1386,18 @@ def training_step( mode_scale=1.0, independent_timesteps=False, ) + # The sigma_shift formula in flow_match.py L57: + # sigma_shifted = shift * sigma_linear / (1 + (shift - 1) * sigma_linear) + # boundary_ratio is in shifted space; invert to get sigma_linear: + # sigma_linear = boundary_ratio / (shift - (shift - 1) * boundary_ratio) + # index_boundary = 1 - sigma_linear, which simplifies to: + # shift * (1 - boundary_ratio) / (1 + (shift - 1) * (1 - boundary_ratio)) + # e.g. shift=3, boundary_ratio=0.9 => sigma_linear=0.75 => index=0.25 + boundary = self.scheduler.shift * (1 - self.boundary_ratio) / (1 + (self.scheduler.shift - 1) * (1 - self.boundary_ratio)) if global_step % 2 == 0: - timestep_config.max_timestep_boundary = self.boundary_ratio + timestep_config.max_timestep_boundary = boundary else: - timestep_config.min_timestep_boundary = self.boundary_ratio + timestep_config.min_timestep_boundary = boundary timestep, audio_timestep = self.sample_timestep_pair(timestep_config) timestep = timestep.to(device=device) From 48e6c3bc33cb5a551f4fe6197a012c903f62370f Mon Sep 17 00:00:00 2001 From: auraithm Date: Sat, 14 Feb 2026 10:55:38 +0000 Subject: [PATCH 2/2] fix: use table lookup for boundary index instead of inverse sigma_shift formula Replace the analytical inverse sigma_shift formula with a direct table lookup on scheduler.timesteps. This avoids floating-point precision issues (e.g. 0.24999 vs 0.25) and the sigma_min != 0 approximation error, giving an exact boundary at the intended timestep threshold. --- mova/diffusion/pipelines/mova_train.py | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/mova/diffusion/pipelines/mova_train.py b/mova/diffusion/pipelines/mova_train.py index e11af66..4ddc983 100644 --- a/mova/diffusion/pipelines/mova_train.py +++ b/mova/diffusion/pipelines/mova_train.py @@ -1386,14 +1386,9 @@ def training_step( mode_scale=1.0, independent_timesteps=False, ) - # The sigma_shift formula in flow_match.py L57: - # sigma_shifted = shift * sigma_linear / (1 + (shift - 1) * sigma_linear) - # boundary_ratio is in shifted space; invert to get sigma_linear: - # sigma_linear = boundary_ratio / (shift - (shift - 1) * boundary_ratio) - # index_boundary = 1 - sigma_linear, which simplifies to: - # shift * (1 - boundary_ratio) / (1 + (shift - 1) * (1 - boundary_ratio)) - # e.g. shift=3, boundary_ratio=0.9 => sigma_linear=0.75 => index=0.25 - boundary = self.scheduler.shift * (1 - self.boundary_ratio) / (1 + (self.scheduler.shift - 1) * (1 - self.boundary_ratio)) + # Look up the index boundary directly from the scheduler's timestep table. + # boundary_ratio=0.9 means timestep 900; count how many timesteps >= 900 to get the index fraction. + boundary = (self.scheduler.timesteps >= self.boundary_ratio * self.scheduler.num_train_timesteps).sum().item() / self.scheduler.num_train_timesteps if global_step % 2 == 0: timestep_config.max_timestep_boundary = boundary else: