custom shift

testing/test_custom_shift.py

@@ -0,0 +1,222 @@
+"""CPU checks: python -m unittest testing.test_custom_shift -v"""
+
+import ast
+from contextlib import nullcontext
+import math
+from pathlib import Path
+from types import SimpleNamespace
+import typing
+import unittest
+from unittest.mock import patch
+
+import torch
+
+from extensions.custom_shift import CustomShiftSampler
+from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler
+
+
+def load_train_config():
+    # Execute the real config classes without loading unrelated audio/model stacks.
+    path = Path(__file__).resolve().parents[1] / "toolkit" / "config_modules.py"
+    tree = ast.parse(path.read_text(encoding="utf-8"))
+    classes = [node for node in tree.body if isinstance(node, ast.ClassDef)
+               and node.name in {"TrainConfig", "EMAConfig", "ValidationConfig", "ValidationItem"}]
+    namespace = dict(vars(typing))
+    exec(compile(ast.Module(body=classes, type_ignores=[]), str(path), "exec"), namespace)
+    return namespace["TrainConfig"]
+
+
+TrainConfig = load_train_config()
+
+
+def load_prepare_timesteps():
+    # Exercise the actual training preparation block with a real scheduler and
+    # synthetic model context, without loading models or the training dependency stack.
+    path = Path(__file__).resolve().parents[1] / "jobs" / "process" / "BaseSDTrainProcess.py"
+    tree = ast.parse(path.read_text(encoding="utf-8"))
+
+    def timer_name(node):
+        if isinstance(node, ast.With):
+            call = node.items[0].context_expr
+            if isinstance(call, ast.Call) and isinstance(call.func, ast.Attribute) and call.func.attr == "timer":
+                return call.args[0].value
+        return None
+
+    for parent in ast.walk(tree):
+        for _, nodes in ast.iter_fields(parent):
+            if not isinstance(nodes, list):
+                continue
+            start = next((i for i, node in enumerate(nodes) if timer_name(node) == "prepare_scheduler"), None)
+            if start is None:
+                continue
+            stop = next(i for i in range(start, len(nodes)) if timer_name(nodes[i]) == "prepare_noise")
+            function = ast.parse("def prepare(self, batch, latents, is_reg): pass").body[0]
+            function.body = nodes[start:stop] + [ast.Return(value=ast.Name(id="timesteps", ctx=ast.Load()))]
+            module = ast.fix_missing_locations(ast.Module(body=[function], type_ignores=[]))
+            namespace = {"torch": torch}
+            exec(compile(module, str(path), "exec"), namespace)
+            return namespace["prepare"]
+    raise AssertionError("Training timestep preparation block not found")
+
+
+prepare_timesteps = load_prepare_timesteps()
+
+
+def sampler(**options):
+    return TrainConfig(noise_scheduler="flowmatch", timestep_type="custom_shift",
+                       timestep_sampler=options).custom_shift_sampler
+
+
+def cdf(t, shift=1.0, scale=1.0):
+    if t <= 0:
+        return 0.0
+    if t >= 1:
+        return 1.0
+    z = (math.log(t / (1 - t)) - math.log(shift)) / scale
+    return (1 + math.erf(z / math.sqrt(2))) / 2
+
+
+class CustomShiftTests(unittest.TestCase):
+    def setUp(self):
+        torch.manual_seed(42)
+
+    def test_original_types_stay_inactive_even_with_saved_options(self):
+        for mode in ("sigmoid", "linear", "shift", "weighted", "next_sample"):
+            config = TrainConfig(timestep_type=mode, timestep_sampler={"sigmoid_scale": -1})
+            self.assertIsNone(config.custom_shift_sampler)
+            self.assertEqual(config.scheduler_timestep_type, mode)
+
+    def test_custom_type_uses_linear_initialization(self):
+        config = TrainConfig(noise_scheduler="flowmatch", timestep_type="custom_shift")
+        self.assertEqual(config.timestep_type, "custom_shift")
+        self.assertEqual(config.scheduler_timestep_type, "linear")
+        self.assertIsInstance(config.custom_shift_sampler, CustomShiftSampler)
+
+    def test_matches_rational_shift_formula(self):
+        z = torch.tensor([-3., -1., 0., 1., 3.], dtype=torch.float64)
+        with patch("extensions.custom_shift.sampler.torch.randn", return_value=z):
+            actual = sampler(discrete_flow_shift=3.0, sigmoid_scale=1.7).sample(5, "cpu")
+        u = torch.sigmoid(z * 1.7)
+        expected = (3 * u / (1 + 2 * u) * 1000).float()
+        torch.testing.assert_close(actual, expected)
+        self.assertEqual(actual.dtype, torch.float32)
+
+    def test_scaled_distribution_matches_analytic_cdf(self):
+        values = sampler(discrete_flow_shift=3, sigmoid_scale=1.4,
+                         min_timestep=200, max_timestep=800).sample(50000, "cpu")
+        self.assertTrue(bool(((values >= 200) & (values <= 800)).all()))
+        for x in (300, 400, 600, 700):
+            self.assertAlmostEqual((values <= x).double().mean().item(),
+                                   cdf((x - 200) / 600, 3, 1.4), delta=0.01)
+
+    def test_preserved_distribution_matches_conditional_cdf(self):
+        values = sampler(discrete_flow_shift=3, sigmoid_scale=1.4, min_timestep=200,
+                         max_timestep=800, preserve_distribution_shape=True).sample(50000, "cpu")
+        self.assertTrue(bool(((values >= 200) & (values <= 800)).all()))
+        for x in (300, 400, 600, 700):
+            expected = (cdf(x / 1000, 3, 1.4) - cdf(.2, 3, 1.4)) / (cdf(.8, 3, 1.4) - cdf(.2, 3, 1.4))
+            self.assertAlmostEqual((values <= x).double().mean().item(), expected, delta=0.01)
+
+    def test_stage_conditions_scaled_distribution(self):
+        values = sampler(discrete_flow_shift=2, min_timestep=100, max_timestep=900).sample(
+            30000, "cpu", stage_range=(.6, 1))
+        self.assertTrue(bool(((values >= 600) & (values <= 900)).all()))
+        base_min = cdf((.6 - .1) / .8, 2)
+        expected = (cdf((.75 - .1) / .8, 2) - base_min) / (1 - base_min)
+        self.assertAlmostEqual((values <= 750).double().mean().item(), expected, delta=.015)
+
+    def test_no_stage_overlap_fails_before_drawing(self):
+        with patch("extensions.custom_shift.sampler.torch.randn") as draw:
+            with self.assertRaisesRegex(ValueError, "no overlap"):
+                sampler(min_timestep=0, max_timestep=500).sample(1, "cpu", stage_range=(.8, 1))
+            draw.assert_not_called()
+
+    def test_rejection_exhaustion_never_returns_out_of_range(self):
+        with patch("extensions.custom_shift.sampler.torch.randn", return_value=torch.zeros(1)):
+            with self.assertRaisesRegex(RuntimeError, "1000 attempts"):
+                sampler(min_timestep=900, max_timestep=1000,
+                        preserve_distribution_shape=True).sample(1, "cpu")
+
+    def test_invalid_options(self):
+        for options in ({"discrete_flow_shift": 0}, {"discrete_flow_shift": float("inf")},
+                        {"sigmoid_scale": float("nan")}, {"sigmoid_scale": -1},
+                        {"sigmoid_scale": True}, {"min_timestep": 1.5},
+                        {"min_timestep": 1000}, {"max_timestep": 1001},
+                        {"preserve_distribution_shape": "false"}, {"typo": 1}):
+            with self.subTest(options=options), self.assertRaises(ValueError):
+                sampler(**options)
+        with self.assertRaises(ValueError):
+            TrainConfig(noise_scheduler="flowmatch", timestep_type="custom_shift", timestep_sampler=[])
+
+    def test_incompatible_training_options(self):
+        for options in ({"noise_scheduler": "ddpm"}, {"num_train_timesteps": 500},
+                        {"linear_timesteps": True}, {"linear_timesteps2": True},
+                        {"first_timestep_chance": .1}, {"min_denoising_steps": 10},
+                        {"max_denoising_steps": 800}, {"loss_type": "stepped"},
+                        {"loss_type": "mean_flow"}, {"loss_target": "source"},
+                        {"diffusion_feature_extractor_path": "test"},
+                        {"train_turbo": True}, {"snr_gamma": 5}):
+            with self.subTest(options=options), self.assertRaises(ValueError):
+                TrainConfig(**{"noise_scheduler": "flowmatch", "timestep_type": "custom_shift", **options})
+
+    def test_model_stage_and_actual_scheduler_add_noise(self):
+        scheduler = CustomFlowMatchEulerDiscreteScheduler()
+        scheduler.set_train_timesteps(1000, "cpu", timestep_type="linear")
+        model = SimpleNamespace(is_flow_matching=True, is_multistage=True,
+                                multistage_boundaries=[.875, 0.0], noise_scheduler=scheduler)
+        for boundary_index, lower, upper in ((0, 875, 1000), (1, 0, 875)):
+            values = sampler().sample_for_model(model, 32, "cpu", boundary_index)
+            self.assertTrue(bool(((values >= lower) & (values <= upper)).all()))
+            clean = torch.ones(32, 4, 2, 2)
+            noise = torch.zeros_like(clean)
+            noisy = scheduler.add_noise(clean, noise, values.view(-1, 1, 1, 1))
+            expected = (1 - values / 1000).view(-1, 1, 1, 1).expand_as(clean)
+            torch.testing.assert_close(noisy, expected)
+
+    def test_custom_draw_does_not_change_inference_scheduler(self):
+        scheduler = CustomFlowMatchEulerDiscreteScheduler()
+        scheduler.set_timesteps(20)
+        times, sigmas = scheduler.timesteps.clone(), scheduler.sigmas.clone()
+        model = SimpleNamespace(is_flow_matching=True, is_multistage=False, noise_scheduler=scheduler)
+        sampler().sample_for_model(model, 4, "cpu")
+        torch.testing.assert_close(scheduler.timesteps, times)
+        torch.testing.assert_close(scheduler.sigmas, sigmas)
+
+    def test_scheduler_capability_checks(self):
+        for scheduler in (object(), CustomFlowMatchEulerDiscreteScheduler(invert_sigmas=True),
+                          CustomFlowMatchEulerDiscreteScheduler(num_train_timesteps=500)):
+            model = SimpleNamespace(is_flow_matching=True, is_multistage=False, noise_scheduler=scheduler)
+            with self.assertRaises(ValueError):
+                sampler().sample_for_model(model, 1, "cpu")
+
+    def test_training_entry_point_ignores_bias_and_respects_expert(self):
+        batch = SimpleNamespace(file_items=[None] * 16)
+        latents = torch.zeros(16, 4, 4, 4)
+        for is_reg in (False, True):
+            for stage, lo, hi in ((0, 875, 1000), (1, 200, 875)):
+                config = TrainConfig(noise_scheduler="flowmatch", timestep_type="custom_shift",
+                                     content_or_style="style", timestep_sampler={"min_timestep": 200})
+                model = SimpleNamespace(is_flow_matching=True, is_multistage=True, is_flux=True, arch="flux",
+                                        multistage_boundaries=[.875, 0],
+                                        noise_scheduler=CustomFlowMatchEulerDiscreteScheduler())
+                process = SimpleNamespace(train_config=config, sd=model, current_boundary_index=stage,
+                                          device_torch=torch.device("cpu"), timer=lambda _: nullcontext(),
+                                          model_config=SimpleNamespace(refiner_name_or_path=None))
+                values = prepare_timesteps(process, batch, latents, is_reg)
+                self.assertTrue(bool(((values >= lo) & (values <= hi)).all()))
+                self.assertEqual(values.shape, (16,))
+                self.assertEqual(model.noise_scheduler.timestep_type, "linear")
+
+    def test_original_training_entry_point_keeps_table_sampling(self):
+        config = TrainConfig(noise_scheduler="flowmatch", timestep_type="shift")
+        model = SimpleNamespace(is_flow_matching=True, is_multistage=False, is_flux=True, arch="flux",
+                                noise_scheduler=CustomFlowMatchEulerDiscreteScheduler(shift=3))
+        process = SimpleNamespace(train_config=config, sd=model, device_torch=torch.device("cpu"),
+                                  timer=lambda _: nullcontext(), model_config=SimpleNamespace(refiner_name_or_path=None))
+        values = prepare_timesteps(process, SimpleNamespace(file_items=[None] * 8), torch.zeros(8, 4, 4, 4), False)
+        self.assertEqual(model.noise_scheduler.timestep_type, "shift")
+        self.assertTrue(bool(torch.isin(values, model.noise_scheduler.timesteps).all()))
+
+
+if __name__ == "__main__":
+    unittest.main()