"""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()