custom shift

.gitignore

@@ -179,6 +179,7 @@ cython_debug/
 !/output/.gitkeep
 /extensions/*
 !/extensions/example
+!/extensions/custom_shift
 /temp
 /wandb
 .vscode/settings.json
@@ -195,4 +196,4 @@ original_repo
 testing/.model_test_outputs
 testing/.engine_test_outputs
 testing/.*
-/issue_review
+/issue_review

extensions/custom_shift/README.md

@@ -0,0 +1,60 @@
+# Custom Shift
+
+Independent shifted logit-normal training timestep sampling. Musubi is an
+algorithm reference only: this extension does not import it or require its
+repository, packages, or installation.
+
+Select **Custom Shift** under **Timestep Type** in the training UI. Its five
+fields are saved with the job and retained when switching to another type.
+Alternatively:
+
+```yaml
+train:
+  noise_scheduler: flowmatch
+  timestep_type: custom_shift
+  timestep_sampler:
+    discrete_flow_shift: 1.0
+    sigmoid_scale: 1.0
+    min_timestep: 0
+    max_timestep: 1000
+    preserve_distribution_shape: false
+```
+
+All five settings are optional and default to the values above. Shift and scale
+must be finite and positive. Bounds must be integers satisfying
+`0 <= min_timestep < max_timestep <= 1000`.
+
+For `z ~ Normal(0, 1)`, draw `u = sigmoid(sigmoid_scale * z)`, then
+`t = shift * u / (1 + (shift - 1) * u)`. The implementation uses the numerically
+stable equivalent `sigmoid(sigmoid_scale * z + log(shift))`. Model timesteps are
+`1000 * t`, where 1000 is pure noise. The original **Shift** type is unchanged.
+
+- With shape preservation off, scale the distribution into the configured range.
+- With shape preservation on, reject values outside the range, yielding the
+  original distribution conditioned on that range. No boundary clamping is used.
+- For multistage models, further condition the final distribution on the current
+  expert's range. Empty intersections fail explicitly. Rejection sampling fails
+  after 1000 attempts rather than returning values outside the requested range.
+- Timestep Bias (`content_or_style`, including regularization data) is ignored.
+  Sampling applies to both ordinary and regularization batches.
+
+Supported path: the standard ai-toolkit flowmatch scheduler (including compatible
+subclasses), with 1000 training timesteps and non-inverted sigmas. Inference
+sampling settings are independent. Keep `min_denoising_steps/max_denoising_steps`
+at their defaults (0/999); use the new bounds instead.
+
+Discrete table weighting (`linear_timesteps`, `linear_timesteps2`),
+`first_timestep_chance`, SNR weighting, turbo/refiner training, source/unaugmented
+loss targets, diffusion feature extractor losses, and stepped/mean-flow losses are rejected because they require
+different timestep handling. These checks apply only to `custom_shift`.
+
+The integration is deliberately small: `TrainConfig` activates the extension,
+`BaseSDTrainProcess` requests final per-batch timesteps, and guidance initializes
+its scheduler with the resolved base type. The scheduler and model implementations
+are not patched or replaced. A `.gitignore` exception keeps this extension tracked.
+
+CPU validation (requires torch, diffusers and numpy):
+
+```shell
+python -m unittest testing.test_custom_shift -v
+```

extensions/custom_shift/__init__.py

@@ -0,0 +1,3 @@
+"""Independent custom-shift timestep sampling for ai-toolkit."""
+
+from .sampler import CustomShiftSampler

extensions/custom_shift/sampler.py

@@ -0,0 +1,109 @@
+"""Shifted logit-normal sampling. No dependency on musubi-tuner."""
+
+import math
+from collections.abc import Mapping
+
+import torch
+
+
+DEFAULTS = {
+    "discrete_flow_shift": 1.0,
+    "sigmoid_scale": 1.0,
+    "min_timestep": 0,
+    "max_timestep": 1000,
+    "preserve_distribution_shape": False,
+}
+
+
+class CustomShiftSampler:
+    def __init__(self, train_config):
+        raw = train_config.timestep_sampler
+        if raw is None:
+            raw = {}
+        if not isinstance(raw, Mapping):
+            raise ValueError("train.timestep_sampler must be a mapping for custom_shift")
+        unknown = set(raw) - set(DEFAULTS)
+        if unknown:
+            raise ValueError(f"Unknown custom_shift options: {sorted(unknown)}")
+        self.config = {**DEFAULTS, **raw}
+        for key in ("discrete_flow_shift", "sigmoid_scale", "min_timestep", "max_timestep"):
+            value = self.config[key]
+            if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value):
+                raise ValueError(f"custom_shift {key} must be a finite number")
+        for key in ("discrete_flow_shift", "sigmoid_scale"):
+            if self.config[key] <= 0:
+                raise ValueError(f"custom_shift {key} must be greater than zero")
+        lo, hi = self.config["min_timestep"], self.config["max_timestep"]
+        if not (0 <= lo < hi <= 1000) or int(lo) != lo or int(hi) != hi:
+            raise ValueError("custom_shift requires integer timesteps: 0 <= min_timestep < max_timestep <= 1000")
+        if not isinstance(self.config["preserve_distribution_shape"], bool):
+            raise ValueError("custom_shift preserve_distribution_shape must be a boolean")
+        if train_config.noise_scheduler != "flowmatch" or train_config.num_train_timesteps != 1000:
+            raise ValueError("custom_shift requires noise_scheduler: flowmatch and num_train_timesteps: 1000")
+        for key in ("linear_timesteps", "linear_timesteps2", "first_timestep_chance", "train_turbo",
+                    "learnable_snr_gos", "snr_gamma", "min_snr_gamma"):
+            if getattr(train_config, key, False):
+                raise ValueError(f"custom_shift is incompatible with {key}; disable it")
+        if train_config.min_denoising_steps != 0 or train_config.max_denoising_steps != 999:
+            raise ValueError("custom_shift uses min_timestep/max_timestep; leave min/max_denoising_steps at 0/999")
+        if train_config.loss_target in ("source", "unaugmented") or train_config.loss_type in ("stepped", "mean_flow"):
+            raise ValueError("custom_shift does not support losses requiring discrete scheduler lookup")
+        if getattr(train_config, "diffusion_feature_extractor_path", None) is not None:
+            raise ValueError("custom_shift does not support diffusion feature extractor losses")
+
+    def sample(self, count, device, *, stage_range=None):
+        """Sample final float32 timesteps; stage ranges condition, never rescale."""
+        lo = self.config["min_timestep"] / 1000.0
+        hi = self.config["max_timestep"] / 1000.0
+        lower, upper = lo, hi
+        if stage_range is not None:
+            stage_min, stage_max = stage_range
+            if not 0 <= stage_min < stage_max <= 1:
+                raise ValueError("custom_shift stage range must satisfy 0 <= min < max <= 1")
+            lower, upper = max(lo, stage_min), min(hi, stage_max)
+        if lower >= upper:
+            raise ValueError("custom_shift timestep range has no overlap with the current model stage")
+        if count < 1:
+            raise ValueError("custom_shift batch size must be positive")
+        scale = self.config["sigmoid_scale"]
+        log_shift = math.log(self.config["discrete_flow_shift"])
+        preserve = self.config["preserve_distribution_shape"]
+        draw_dtype = torch.float32 if torch.device(device).type == "mps" else torch.float64
+
+        def draw(n):
+            # Equivalent to s*u/(1+(s-1)*u), u=sigmoid(scale*z).
+            # Float64 keeps large valid shifts stable; output stays float32.
+            t = torch.sigmoid(scale * torch.randn(n, device=device, dtype=draw_dtype) + log_shift)
+            if not preserve:
+                t = lo + (hi - lo) * t
+            return t
+
+        accepted = []
+        remaining = count
+        for _ in range(1000):
+            candidates = draw(max(count, remaining))
+            valid = candidates[(candidates >= lower) & (candidates <= upper)]
+            taken = valid[:remaining]
+            accepted.append(taken)
+            remaining -= taken.numel()
+            if remaining == 0:
+                return (torch.cat(accepted) * 1000.0).to(torch.float32)
+        raise RuntimeError(
+            f"custom_shift could not sample {count} timesteps in [{lower * 1000}, {upper * 1000}] "
+            "after 1000 attempts; widen the range or adjust shift/scale"
+        )
+
+    def sample_for_model(self, model, count, device, boundary_index=0):
+        from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler
+
+        if not model.is_flow_matching or not isinstance(model.noise_scheduler, CustomFlowMatchEulerDiscreteScheduler):
+            raise ValueError("custom_shift requires the standard ai-toolkit flowmatch scheduler")
+        if model.noise_scheduler.config.num_train_timesteps != 1000:
+            raise ValueError("custom_shift requires a scheduler with 1000 training timesteps")
+        if getattr(model.noise_scheduler.config, "invert_sigmas", False):
+            raise ValueError("custom_shift does not support inverted-sigma schedulers")
+        stage_range = None
+        if model.is_multistage:
+            boundaries = [1] + model.multistage_boundaries
+            stage_range = (boundaries[boundary_index + 1], boundaries[boundary_index])
+        return self.sample(count, device, stage_range=stage_range)

extensions_built_in/diffusion_models/krea2/src/mmdit.py

@@ -31,7 +31,11 @@ from torch.utils.checkpoint import checkpoint
 
 
 def rope(pos: Tensor, dim: int, theta: float = 1e4, ntk: float = 1.0) -> Tensor:
-    scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim
+    # MPS has no float64
+    dt = torch.float32 if pos.device.type == "mps" else torch.float64
+    if dt is torch.float32:
+        pos = pos.to(dt)
+    scale = torch.arange(0, dim, 2, dtype=dt, device=pos.device) / dim
     omega = 1.0 / ((theta * ntk) ** scale)
     out = torch.einsum("...n,d->...nd", pos, omega)
     out = torch.stack(

extensions_built_in/sd_trainer/SDTrainer.py

@@ -170,8 +170,11 @@ class SDTrainer(BaseSDTrainProcess):
                 # needed so we can autoparse the prompt to handle flags
                 gen_img_config = GenerateImageConfig(
                     prompt=prompt, # it will autoparse the prompt
+                    width=sample_item.width,
+                    height=sample_item.height,
                     negative_prompt=sample_item.neg,
                     output_path=output_path,
+                    num_frames=sample_item.num_frames,
                     ctrl_img=sample_item.ctrl_img,
                     ctrl_img_1=sample_item.ctrl_img_1,
                     ctrl_img_2=sample_item.ctrl_img_2,

jobs/process/BaseSDTrainProcess.py

@@ -116,6 +116,8 @@ class BaseSDTrainProcess(BaseTrainProcess):
         # update modelconfig dtype to match train
         model_config['dtype'] = self.train_config.dtype
         self.model_config = ModelConfig(**model_config)
+        if self.train_config.custom_shift_sampler is not None and self.model_config.refiner_name_or_path is not None:
+            raise ValueError("custom_shift does not support refiner training")
 
         self.save_config = SaveConfig(**self.get_conf('save', {}))
         self.sample_config = SampleConfig(**self.get_conf('sample', {}))
@@ -1183,7 +1185,7 @@ class BaseSDTrainProcess(BaseTrainProcess):
                     
                     timestep_type = 'linear' if linear_timesteps else None
                     if timestep_type is None:
-                        timestep_type = self.train_config.timestep_type
+                        timestep_type = self.train_config.scheduler_timestep_type
                     
                     if self.train_config.timestep_type == 'next_sample':
                         # simulate a sample
@@ -1208,7 +1210,7 @@ class BaseSDTrainProcess(BaseTrainProcess):
                     self.sd.noise_scheduler.set_timesteps(
                         num_train_timesteps, device=self.device_torch
                     )
-            if self.sd.is_multistage:
+            if self.sd.is_multistage and self.train_config.custom_shift_sampler is None:
                 with self.timer('adjust_multistage_timesteps'):
                     # get our current sample range
                     boundaries = [1] + self.sd.multistage_boundaries
@@ -1232,7 +1234,10 @@ class BaseSDTrainProcess(BaseTrainProcess):
                 if is_reg:
                     content_or_style = self.train_config.content_or_style_reg
 
-                if self.train_config.timestep_type in ['two_step', 'four_step', 'eight_step']:
+                if self.train_config.custom_shift_sampler is not None:
+                    # Final continuous timesteps are drawn below, without index bias.
+                    timestep_indices = torch.zeros((batch_size,), device=self.device_torch, dtype=torch.long)
+                elif self.train_config.timestep_type in ['two_step', 'four_step', 'eight_step']:
                     if self.train_config.timestep_type == 'two_step':
                         indice_choices = [0, 499]
                     elif self.train_config.timestep_type == 'four_step':
@@ -1306,7 +1311,13 @@ class BaseSDTrainProcess(BaseTrainProcess):
                     timestep_indices = torch.where(force_first, torch.zeros_like(timestep_indices), timestep_indices)
             with self.timer('convert_timestep_indices_to_timesteps'):
                 # convert the timestep_indices to a timestep
-                timesteps = self.sd.noise_scheduler.timesteps[timestep_indices.long()]
+                if self.train_config.custom_shift_sampler is not None:
+                    timesteps = self.train_config.custom_shift_sampler.sample_for_model(
+                        self.sd, batch_size, self.device_torch,
+                        boundary_index=getattr(self, 'current_boundary_index', 0),
+                    )
+                else:
+                    timesteps = self.sd.noise_scheduler.timesteps[timestep_indices.long()]
                 
             with self.timer('prepare_noise'):
                 # get noise

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

toolkit/advanced_prompt_embeds.py

@@ -209,3 +209,4 @@ class AdvancedPromptEmbeds:
                 )
             for i in range(num_parts):
                 split_embeds[i]._store[key] = [values[i]]
+        return split_embeds

toolkit/config_modules.py

@@ -626,6 +626,14 @@ class TrainConfig:
         if validation is not None:
             self.validation_config: ValidationConfig = ValidationConfig(**validation)
 
+        self.timestep_sampler = kwargs.get('timestep_sampler', None)
+        self.custom_shift_sampler = None
+        self.scheduler_timestep_type = self.timestep_type
+        if self.timestep_type == 'custom_shift':
+            from extensions.custom_shift import CustomShiftSampler
+            self.custom_shift_sampler = CustomShiftSampler(self)
+            self.scheduler_timestep_type = 'linear'
+
 
 ModelArch = Literal['sd1', 'sd2', 'sd3', 'sdxl', 'pixart', 'pixart_sigma', 'auraflow', 'flux', 'flex1', 'flex2', 'lumina2', 'vega', 'ssd', 'wan21', 'anima']
 
@@ -1566,4 +1574,4 @@ def validate_configs(
             raise ValueError(f"Cannot cache unload text encoder with {model_config.arch} model. Control images are encoded with text embeddings. You can cache the text embeddings though")
     
     if train_config.diff_output_preservation and train_config.blank_prompt_preservation:
-        raise ValueError("Cannot use both differential output preservation and blank prompt preservation at the same time. Please set one of them to False.")
+        raise ValueError("Cannot use both differential output preservation and blank prompt preservation at the same time. Please set one of them to False.")

toolkit/dataloader_mixins.py

@@ -34,10 +34,13 @@ import albumentations as A
 from toolkit.print import print_acc
 from toolkit.accelerator import get_accelerator
 from toolkit.prompt_utils import PromptEmbeds
+from toolkit.advanced_prompt_embeds import AdvancedPromptEmbeds
 from torchvision.transforms import functional as TF
 
 from toolkit.train_tools import get_torch_dtype
 
+TEXT_EMBED_CACHE_BATCH = 16  # prompts per text-encoder forward when caching
+
 if TYPE_CHECKING:
     from toolkit.data_loader import AiToolkitDataset
     from toolkit.data_transfer_object.data_loader import FileItemDTO
@@ -2489,6 +2492,38 @@ class TextEmbeddingCachingMixin:
             
             did_move = False
 
+            # plain-text prompts are encoded TEXT_EMBED_CACHE_BATCH per forward: a 9B text
+            # encoder at batch 1 is kernel-launch bound (~15% GPU busy)
+            pending = []  # (path, caption), plain text only
+            pending_paths = set()
+            batch_ok = [None]  # None = untested; False = legacy PromptEmbeds, no per-item split
+
+            def flush_pending():
+                if not pending:
+                    return
+                # similar lengths together: a batch is right-padded to its longest prompt
+                pending.sort(key=lambda t: len(t[1]))
+                size = 1 if batch_ok[0] is False else TEXT_EMBED_CACHE_BATCH
+                start = 0
+                while start < len(pending):
+                    chunk = pending[start:start + size]
+                    embeds = self.sd.encode_prompt([c for _, c in chunk])
+                    if isinstance(embeds, AdvancedPromptEmbeds):
+                        batch_ok[0] = True
+                        parts = AdvancedPromptEmbeds.split_prompt_embeds(embeds, len(chunk))
+                    elif len(chunk) == 1:
+                        parts = [embeds]
+                    else:
+                        batch_ok[0] = False
+                        size = 1
+                        continue  # redo this chunk one prompt at a time
+                    for (path, _), part in zip(chunk, parts):
+                        part.save(path)
+                    del embeds, parts
+                    start += len(chunk)
+                pending.clear()
+                pending_paths.clear()
+
             # use tqdm to show progress
             i = 0
             for file_item in tqdm(self.file_list, desc='Caching text embeddings to disk'):
@@ -2616,9 +2651,11 @@ class TextEmbeddingCachingMixin:
                         file_item.tensor = None
                     else:
                         for path, caption in encode_targets:
-                            prompt_embeds: PromptEmbeds = self.sd.encode_prompt(caption)
-                            prompt_embeds.save(path)
-                            del prompt_embeds
+                            if path not in pending_paths:
+                                pending_paths.add(path)
+                                pending.append((path, caption))
+                        if len(pending) >= 4 * TEXT_EMBED_CACHE_BATCH:
+                            flush_pending()
                 if getattr(file_item, 'dopsd_self_ref', False):
                     # D-OPSD teacher embeds encode with the item's own media as the reference
                     control_video_paths = getattr(file_item, 'control_video_paths', None) or []
@@ -2665,6 +2702,7 @@ class TextEmbeddingCachingMixin:
                             del prompt_embeds
                 file_item.is_text_embedding_cached = True
                 i += 1
+            flush_pending()
             # restore device state
             # if did_move:
             #     self.sd.restore_device_state()

toolkit/guidance.py

@@ -437,7 +437,7 @@ def get_guided_loss_polarity(
             
             timestep_type = 'linear' if linear_timesteps else None
             if timestep_type is None:
-                timestep_type = train_config.timestep_type
+                timestep_type = train_config.scheduler_timestep_type
             
             sd.noise_scheduler.set_train_timesteps(
                 1000,

toolkit/memory_management/manager.py

@@ -457,15 +457,18 @@ class MemoryManager:
         gc.collect()
         torch.cuda.empty_cache()
 
-        # torch's pinned-host cache; private API, name varies by torch version
-        for fn_name in ("_accelerator_emptyHostCache", "_host_emptyCache"):
-            fn = getattr(torch._C, fn_name, None)
-            if fn is not None:
-                try:
-                    fn()
-                    break
-                except Exception:
-                    pass
+        # torch's pinned-host cache; private API, name varies by torch version.
+        # Skipped on MPS: it segfaults there (pytorch/pytorch#197561), not catchable
+        on_mps = getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_available()
+        if not on_mps:
+            for fn_name in ("_accelerator_emptyHostCache", "_host_emptyCache"):
+                fn = getattr(torch._C, fn_name, None)
+                if fn is not None:
+                    try:
+                        fn()
+                        break
+                    except Exception:
+                        pass
 
         # glibc keeps freed arenas mapped; CDLL(None) resolves malloc_trim in
         # the running process where glibc is present and fails cleanly on

toolkit/models/diffusion_feature_extraction.py

@@ -1221,7 +1221,7 @@ class DiffusionFeatureExtractor9(nn.Module):
         perceptual_loss = torch.nn.functional.mse_loss(
             pred.float(), target.float(), reduction="none"
         )
-        velocity_equiv_weight = (1.0 / torch.clamp(tv, min=0.1) ** 2)
+        velocity_equiv_weight = (1.0 / torch.clamp(tv, min=0.01) ** 2)
         loss_perceptual = (perceptual_loss * velocity_equiv_weight).mean()
         
         if self.do_partial_step:
@@ -1379,7 +1379,7 @@ class DiffusionFeatureExtractor10(nn.Module):
 
         pred_feats = self.get_lpips_features(pred_images.float())
 
-        velocity_equiv_weight = (1.0 / torch.clamp(tv, min=0.1) ** 2)
+        velocity_equiv_weight = (1.0 / torch.clamp(tv, min=0.01) ** 2)
 
         loss_perceptual = 0
         for idx, pred_feat in enumerate(pred_feats):

toolkit/models/v2/_mixin.py

@@ -903,6 +903,8 @@ class OstrisModelMixin:
                 value = value.to(dtype=dtype)
             state_dict[key] = value.detach().to("cpu").contiguous()
         state_dict.update(q_entries)
+        # block_compile wraps blocks in OptimizedModule: drop its "._orig_mod." key segment
+        state_dict = {k.replace("._orig_mod.", "."): v for k, v in state_dict.items()}
         state_dict = self.convert_state_dict_on_save(state_dict)
 
         parent = os.path.dirname(output_path)

toolkit/optimizer.py

@@ -103,6 +103,9 @@ def get_optimizer(
     elif lower_type == 'automagic3':
         from toolkit.optimizers.automagic3 import Automagic3
         optimizer = Automagic3(params, lr=float(learning_rate), **optimizer_params)
+    elif lower_type == 'automagic4':
+        from toolkit.optimizers.automagic4 import Automagic4
+        optimizer = Automagic4(params, lr=float(learning_rate), **optimizer_params)
     elif lower_type == 'automagicexperiment':
         from toolkit.optimizers.automagicEXPERIMENT import AutomagicEXPERIMENT
         optimizer = AutomagicEXPERIMENT(params, lr=float(learning_rate), **optimizer_params)

toolkit/optimizers/automagic4.py

@@ -0,0 +1,259 @@
+"""
+NOTE: This is experimental and under active development; expect breaking changes and bugs. Feedback welcome.
+"""
+
+import math
+import torch
+
+from toolkit.optimizers.automagic3 import Automagic3
+
+
+class Automagic4(Automagic3):
+    """
+    Automagic v4: Automagic3 with sensitivity-balanced updates.
+
+    v3 gives every element an Adam-normalized step of about lr, so how far one
+    step moves the OUTPUT scales with that weight's sensitivity. A single shared
+    lr cannot fit a model whose tensors differ in sensitivity by orders of
+    magnitude (a head reading a huge raw residual, a modulation shared by every
+    layer): at the lr the bulk needs, the sensitive ones overshoot.
+
+    v4 multiplies each element's (clipped) v3 update by
+
+        balance = min(1, reference / sensitivity)
+
+    sensitivity = sqrt(bias-corrected second moment). For a noise-dominated MSE
+    loss (diffusion) gradient power tracks the Gauss-Newton curvature, i.e. how
+    far the output moves per unit change of the weight. It comes from the same
+    row/col (or 1D) state v3 keeps, so it is free and per element: per column
+    for a head whose input dims differ in scale, per row for a modulation.
+    Above the reference the step becomes Newton-like (grad / v), which is
+    invariant to rescaling the weights -- including coupled Q/K-style
+    rescalings, so it cannot seed the fights the shared group lr prevents.
+
+    One-sided on purpose: low gradient power cannot be told apart from a weight
+    that is gated off (a LoRA down while its zero-initialized up is still ~0,
+    dead units), where Newton would take an unbounded step. Those keep v3's
+    normalized step, which is what gets zero-initialized LoRA pairs moving.
+
+    reference: the ``sensitivity_quantile`` (default 0.9) quantile of the
+    per-tensor geometric-mean sensitivity, weighted by element count, from the
+    previous step. High so a large near-dead mass (LoRA downs at init are about
+    half of a LoRA's elements) cannot drag it down and stall everything above
+    it; the outliers it targets are a few percent of elements.
+
+    The first step only measures (second moments and the reference): without a
+    reference the sensitive tensors would take a full v3 step.
+
+    lr controller: re-centered flip-rate vote. v3's neutral point is a 50% sign
+    flip rate, but a weight bouncing on gradient noise at its minimum has
+    lag-1 gradient correlation -effective_lr * curvature / 2 at any lr > 0, so
+    v3 only balances at lr = 0 and anneals once the signal runs out. v4 votes
+    (neutral - observed) flip rate per element, with
+
+        neutral = 0.5 + asin(step_target / 2) / pi
+
+    the flip rate of a stationary weight whose effective_lr * curvature equals
+    ``step_target`` (default 0.1; 2 is the stability edge). The lr then has a
+    fixed point set by step_target, with or without signal: trends flip less
+    and raise it, overshoot flips more and lowers it. Votes are weighted by the
+    balanced |update| (output impact). No gain: near the fixed point the flip
+    rate moves ~0.16 * step_target per e-fold of lr, so the loop gain over the
+    H-step window lag stays well below 1.
+    """
+
+    def __init__(
+        self,
+        params,
+        *args,
+        sensitivity_quantile: float = 0.9,
+        step_target: float = 0.1,
+        **kwargs,
+    ):
+        if not 0.0 < sensitivity_quantile <= 1.0:
+            raise ValueError(
+                f"sensitivity_quantile ({sensitivity_quantile}) must be in (0, 1]"
+            )
+        if not 0.0 < step_target < 2.0:
+            raise ValueError(f"step_target ({step_target}) must be in (0, 2)")
+        self.sensitivity_quantile = float(sensitivity_quantile)
+        self.step_target = float(step_target)
+        self._neutral_flip = 0.5 + math.asin(self.step_target / 2.0) / math.pi
+        super().__init__(params, *args, **kwargs)
+
+    # ------------------------------------------------------------------ utils
+
+    @staticmethod
+    def _ema_(
+        state_t: torch.Tensor, sample: torch.Tensor, beta2: float
+    ) -> torch.Tensor:
+        # fp32 beta2-EMA; low-precision state is written back
+        v = state_t if state_t.dtype == torch.float32 else state_t.to(torch.float32)
+        v.mul_(beta2).add_(sample, alpha=1.0 - beta2)
+        if v is not state_t:
+            state_t.copy_(v)
+        return v
+
+    def _rebuild_group_index(self) -> None:
+        super()._rebuild_group_index()
+        n = len(self.param_groups)
+        self._group_ref: list = [None] * n
+        self._group_sens_log: list = [[] for _ in range(n)]
+        self._group_sens_numel: list = [[] for _ in range(n)]
+
+    def _init_state(self, p: torch.Tensor, group: dict) -> None:
+        super()._init_state(p, group)
+        gi = self._param_group_index.get(p)
+        ref = self._group_ref[gi] if gi is not None else None
+        if ref is not None:
+            self.state[p]["sens_ref"] = ref.to(p.device, copy=True)
+
+    # -------------------------------------------------------------- per-param
+
+    @torch.no_grad()
+    def _update_param(self, p: torch.Tensor, group: dict) -> None:
+        if p.grad is None:
+            return
+        state = self.state[p]
+        if len(state) == 0:
+            self._init_state(p, group)
+
+        grad = p.grad
+        if grad.is_sparse:
+            raise RuntimeError("Automagic4 does not support sparse gradients.")
+        if grad.dtype != torch.float32:
+            grad = grad.to(torch.float32)
+        grad.nan_to_num_(nan=0.0, posinf=0.0, neginf=0.0)
+
+        beta2 = group["beta2"]
+        eps = group["eps"]
+        sq = grad * grad
+        if p.dim() >= 2:
+            row = self._ema_(state["exp_avg_sq_row"], sq.mean(dim=-1).add_(eps), beta2)
+            col = self._ema_(state["exp_avg_sq_col"], sq.mean(dim=-2).add_(eps), beta2)
+            del sq
+            # mean log of the factored v = row (x) col / mean(row)
+            log_v = (
+                row.log().mean(-1) - row.mean(-1).log() + col.log().mean(-1)
+            ).mean()
+            inv_sqrt_v = self._approx_sq_grad(row, col)
+        else:
+            v = self._ema_(state["exp_avg_sq"], sq, beta2).add(eps)
+            del sq
+            log_v = v.log().mean()
+            inv_sqrt_v = v.rsqrt_()
+        state["step"] += 1
+        bias_corr = 1.0 - beta2 ** state["step"]
+
+        gi = self._param_group_index.get(p)
+        if gi is not None:
+            self._group_sens_log[gi].append(0.5 * (log_v - math.log(bias_corr)))
+            self._group_sens_numel[gi].append(p.numel())
+
+        ref = state.get("sens_ref")
+        if ref is None:
+            # measurement-only first step (no reference yet)
+            p.grad = None
+            return
+
+        update = grad.mul_(inv_sqrt_v)
+        update.div_((self._rms(update) / group["clip_threshold"]).clamp_(min=1.0))
+        update.clamp_(-group["clip_threshold"], group["clip_threshold"])
+        # balance = min(1, ref / sensitivity), sensitivity = sqrt(v / bias_corr)
+        balance = inv_sqrt_v.mul_(ref * math.sqrt(bias_corr)).clamp_(max=1.0)
+        update.mul_(balance)
+        del balance, inv_sqrt_v
+
+        # flip-rate vote (see class docstring), pooled per group, applied in .step()
+        cur_bits = update.gt(0.0)
+        hist = state["sign_history"]
+        idx = state["hist_idx"]
+        H = hist.shape[0]
+        lr_t = state["lr"]
+        hist[idx].copy_(self._pack_bits(cur_bits))
+        state["hist_idx"] = (idx + 1) % H
+        fill = min(H, state["hist_fill"] + 1)
+        state["hist_fill"] = fill
+
+        if fill == H and gi is not None:
+            _, shifts = self._pack_consts(hist.device)
+            chron = torch.roll(hist, -state["hist_idx"], dims=0)
+            bits = (
+                (chron.unsqueeze(-1) >> shifts)
+                .bitwise_and_(1)
+                .view(H, -1)[:, : update.numel()]
+            )
+            flips = (bits[1:] ^ bits[:-1]).sum(0, dtype=torch.int16)
+            w = update.abs().view(-1)
+            den = w.sum()
+            num = den * self._neutral_flip - (w * flips).sum() / (H - 1)
+            if self._group_num[gi] is None:
+                self._group_num[gi] = num
+                self._group_den[gi] = den
+            else:
+                acc = self._group_num[gi]
+                if num.device != acc.device:
+                    num = num.to(acc.device)
+                    den = den.to(acc.device)
+                acc.add_(num)
+                self._group_den[gi].add_(den)
+
+        wd = group["weight_decay"]
+        if p.dtype == torch.float32:
+            if wd != 0.0:
+                update.add_(p, alpha=wd)
+            p.addcmul_(update, lr_t, value=-1.0)
+        else:
+            new_p_fp32 = p.to(torch.float32)
+            if wd != 0.0:
+                update.add_(new_p_fp32, alpha=wd)
+            new_p_fp32.addcmul_(update, lr_t, value=-1.0)
+            self._stochastic_copy_(p, new_p_fp32)
+
+        p.grad = None
+
+    # ----------------------------------------------------------- optimizer API
+
+    def _apply_group_votes(self) -> None:
+        self._update_sensitivity_refs()
+        super()._apply_group_votes()
+
+    def _update_sensitivity_refs(self) -> None:
+        # numel-weighted quantile of this step's per-tensor log sensitivities; no GPU sync
+        for gi, group in enumerate(self.param_groups):
+            logs = self._group_sens_log[gi]
+            if not logs:
+                continue
+            dev = logs[0].device
+            vals = torch.stack([t.to(dev) for t in logs])
+            weights = torch.tensor(
+                self._group_sens_numel[gi], dtype=torch.float64, device=dev
+            )
+            order = vals.argsort()
+            cum = weights[order].cumsum(0)
+            pos = torch.searchsorted(cum, cum[-1:] * self.sensitivity_quantile)
+            ref = vals[order][pos.clamp_(max=len(logs) - 1)].squeeze(0).exp_()
+            self._group_ref[gi] = ref
+            for p in group["params"]:
+                st = self.state.get(p)
+                if not st:
+                    continue
+                cur = st.get("sens_ref")
+                if cur is None:
+                    st["sens_ref"] = ref.to(p.device, copy=True)
+                else:
+                    cur.copy_(ref)
+            self._group_sens_log[gi] = []
+            self._group_sens_numel[gi] = []
+
+    def load_state_dict(self, state_dict):
+        super().load_state_dict(state_dict)
+        # parent cast the reference to param dtype; restore fp32 and the group mirror
+        for gi, group in enumerate(self.param_groups):
+            for p in group["params"]:
+                st = self.state.get(p)
+                ref = st.get("sens_ref") if st is not None else None
+                if isinstance(ref, torch.Tensor):
+                    st["sens_ref"] = ref.to(torch.float32)
+                    if self._group_ref[gi] is None:
+                        self._group_ref[gi] = st["sens_ref"].clone()

toolkit/util/unified_memory.py

@@ -1,13 +1,19 @@
-"""Unified memory (integrated GPU) support, e.g. NVIDIA DGX Spark (GB10).
+"""Unified memory (integrated GPU) support, e.g. NVIDIA DGX Spark (GB10) and
+the Windows RTX Spark (N1X).
 
 On these devices the CPU and GPU share the same physical memory, so moving a
 module to the CPU frees nothing. Worse, device -> host copies are extremely
-slow (~100MB/s on GB10), so every "offload to cpu" of a large transformer or
-text encoder stalls for minutes. When a unified memory device is detected we
-make module level moves to the CPU a no-op (dtype casts still apply). Tensor
-level .cpu()/.to('cpu') is left alone since saving, numpy, etc. need it.
+slow (~100MB/s on GB10 under Linux, ~0.5-2GB/s pageable on Windows/WDDM), so
+every "offload to cpu" of a large transformer or text encoder stalls. When a
+unified memory device is detected we make module level moves to the CPU a
+no-op (dtype casts still apply). Tensor level .cpu()/.to('cpu') is left alone
+since saving, numpy, etc. need it.
 
-Set AITK_DISABLE_UNIFIED_MEMORY=1 to turn this off.
+Note for Windows: WDDM gives CUDA a fixed budget (~80% of system RAM, e.g.
+~45GB of 54GB) that host allocations do not count against. nvidia-smi / Task
+Manager only show the small dedicated segment (8GB); the real pool is what
+torch.cuda.mem_get_info() reports. If a job needs to keep modules on the CPU
+to stay under that budget, set AITK_DISABLE_UNIFIED_MEMORY=1 to turn this off.
 """
 import os
 import sys
@@ -20,8 +26,8 @@ _patched = False
 def is_unified_memory() -> bool:
     if os.environ.get("AITK_DISABLE_UNIFIED_MEMORY", "0") == "1":
         return False
-    # linux only (DGX Spark). The Windows RTX Spark stack does not behave the same way
-    if not sys.platform.startswith("linux"):
+    # macOS has no CUDA; MPS unified memory is handled elsewhere
+    if sys.platform == "darwin":
         return False
     # NVIDIA only. ROCm APUs report integrated too, but may have a dedicated VRAM carveout
     if torch.version.hip is not None:

ui/src/app/api/files/info/route.ts

@@ -0,0 +1,111 @@
+/* eslint-disable */
+import { NextRequest, NextResponse } from 'next/server';
+import fs from 'fs';
+import path from 'path';
+import { getDatasetsRoot, getTrainingFolder } from '@/server/settings';
+
+// safetensors header: u64 LE byte length, then that many bytes of JSON.
+const MAX_HEADER_BYTES = 256 * 1024 * 1024;
+
+const DTYPE_BYTES: Record<string, number> = {
+  F64: 8, I64: 8, U64: 8,
+  F32: 4, I32: 4, U32: 4,
+  F16: 2, BF16: 2, I16: 2, U16: 2,
+  F8_E4M3: 1, F8_E5M2: 1, F8_E8M0: 1, I8: 1, U8: 1, BOOL: 1,
+  F4: 0.5, U4: 0.5, I4: 0.5,
+};
+
+export async function POST(request: NextRequest) {
+  try {
+    const body = await request.json();
+    const { filePath } = body;
+
+    if (!filePath || typeof filePath !== 'string') {
+      return new NextResponse('filePath is required', { status: 400 });
+    }
+
+    const decodedFilePath = decodeURIComponent(filePath);
+    const datasetRoot = await getDatasetsRoot();
+    const trainingRoot = await getTrainingFolder();
+    const allowedDirs = [datasetRoot, trainingRoot];
+
+    const resolvedFilePath = path.resolve(decodedFilePath);
+    const isAllowed = allowedDirs.some(
+      allowedDir => resolvedFilePath === allowedDir || resolvedFilePath.startsWith(allowedDir + path.sep),
+    );
+    if (!isAllowed) {
+      console.warn(`Access denied: ${resolvedFilePath} not in ${allowedDirs.join(', ')}`);
+      return new NextResponse('Access denied', { status: 403 });
+    }
+
+    let stat;
+    try {
+      stat = await fs.promises.stat(resolvedFilePath);
+    } catch {
+      return new NextResponse('File not found', { status: 404 });
+    }
+    if (!stat.isFile()) {
+      return new NextResponse('Not a file', { status: 400 });
+    }
+    if (stat.size < 8) {
+      return new NextResponse('Not a safetensors file', { status: 400 });
+    }
+
+    const handle = await fs.promises.open(resolvedFilePath, 'r');
+    let header: Record<string, any>;
+    let headerSize: number;
+    try {
+      const lenBuf = Buffer.alloc(8);
+      await handle.read(lenBuf, 0, 8, 0);
+      const headerLen = lenBuf.readBigUInt64LE(0);
+      if (headerLen > BigInt(MAX_HEADER_BYTES) || headerLen > BigInt(stat.size - 8)) {
+        return new NextResponse('Not a safetensors file (bad header length)', { status: 400 });
+      }
+      headerSize = Number(headerLen);
+      const headerBuf = Buffer.alloc(headerSize);
+      await handle.read(headerBuf, 0, headerSize, 8);
+      header = JSON.parse(headerBuf.toString('utf8'));
+    } catch (e) {
+      return new NextResponse('Not a safetensors file (unparseable header)', { status: 400 });
+    } finally {
+      await handle.close();
+    }
+
+    const metadata: Record<string, string> = header.__metadata__ ?? {};
+    const tensors: { name: string; dtype: string; shape: number[]; numel: number; bytes: number }[] = [];
+    const dtypeCounts: Record<string, { count: number; numel: number; bytes: number }> = {};
+    let totalParams = 0;
+    let totalBytes = 0;
+
+    for (const [name, info] of Object.entries(header)) {
+      if (name === '__metadata__') continue;
+      const t = info as { dtype: string; shape: number[]; data_offsets?: [number, number] };
+      const shape = Array.isArray(t.shape) ? t.shape : [];
+      const numel = shape.reduce((a, b) => a * b, 1);
+      const bytes = t.data_offsets ? t.data_offsets[1] - t.data_offsets[0] : numel * (DTYPE_BYTES[t.dtype] ?? 0);
+      tensors.push({ name, dtype: t.dtype, shape, numel, bytes });
+      totalParams += numel;
+      totalBytes += bytes;
+      const d = (dtypeCounts[t.dtype] ??= { count: 0, numel: 0, bytes: 0 });
+      d.count += 1;
+      d.numel += numel;
+      d.bytes += bytes;
+    }
+    tensors.sort((a, b) => a.name.localeCompare(b.name));
+
+    return NextResponse.json({
+      filePath: resolvedFilePath,
+      fileSize: stat.size,
+      headerSize,
+      metadata,
+      tensors,
+      tensorCount: tensors.length,
+      totalParams,
+      totalBytes,
+      dtypeCounts,
+    });
+  } catch (error) {
+    console.error('Error reading safetensors info:', error);
+    return new NextResponse('Internal Server Error', { status: 500 });
+  }
+}

ui/src/app/api/jobs/[jobID]/loss/route.ts

@@ -5,6 +5,8 @@ import fs from 'fs';
 import { getTrainingFolder } from '@/server/settings';
 
 import sqlite3 from 'sqlite3';
+import zlib from 'zlib';
+import { promisify } from 'util';
 
 export const runtime = 'nodejs';
 
@@ -29,6 +31,31 @@ function closeDb(db: sqlite3.Database) {
   });
 }
 
+// Serialize a metric series to JSON arrays inside SQLite. Building the arrays
+// with json_group_array avoids marshaling one JS object per row, which is the
+// dominant cost on multi-million-point logs. Values are truncated to 7
+// significant digits (float32-level precision) to shrink the wire payload;
+// value_text rows are never numeric (see _coerce_value in logging_aitk.py)
+// so they serialize as null.
+const SERIES_SQL = `
+  SELECT json_group_array(step) AS steps, json_group_array(v) AS vals
+  FROM (
+    SELECT step,
+           CASE WHEN value_real IS NULL THEN NULL
+                ELSE CAST(printf('%.7g', value_real) AS REAL) END AS v
+    FROM metrics
+    WHERE key = ? AND step >= ?
+    ORDER BY step ASC
+  )
+`;
+
+const gzip = promisify(zlib.gzip);
+
+// GET /api/jobs/:id/loss?since_step=N
+// Returns every metric key in one response: {keys, since_step, series: {key:
+// {steps: number[], values: (number|null)[]}}}. Points with step >= since_step
+// are included (default 0 = everything). The body is gzipped when the client
+// accepts it; a 150k-step run with a dozen keys is tens of MB uncompressed.
 export async function GET(request: NextRequest, { params }: { params: { jobID: string } }) {
   // this must be awaited to avoid TS error
   const { jobID } = await params;
@@ -43,58 +70,46 @@ export async function GET(request: NextRequest, { params }: { params: { jobID: s
   try {
     await fs.promises.access(logPath);
   } catch {
-    return NextResponse.json({ keys: [], key: 'loss', points: [] });
+    return NextResponse.json({ keys: [], since_step: 0, series: {} });
   }
 
   const url = new URL(request.url);
-  const key = url.searchParams.get('key') ?? 'loss';
-  const limit = Math.min(Number(url.searchParams.get('limit') ?? 2000), 20000);
-  const sinceStepParam = url.searchParams.get('since_step');
-  const sinceStep = sinceStepParam != null ? Number(sinceStepParam) : null;
-  const stride = Math.max(1, Number(url.searchParams.get('stride') ?? 1));
+  const sinceStepParam = Number(url.searchParams.get('since_step') ?? 0);
+  const sinceStep = Number.isFinite(sinceStepParam) ? sinceStepParam : 0;
 
   const db = openDb(logPath);
 
+  let body: string;
   try {
     const keysRows = await all<{ key: string }>(db, `SELECT key FROM metric_keys ORDER BY key ASC`);
     const keys = keysRows.map((r) => r.key);
 
-    const points = await all<{
-      step: number;
-      wall_time: number;
-      value: number | null;
-      value_text: string | null;
-    }>(
-      db,
-      `
-      SELECT
-        m.step AS step,
-        s.wall_time AS wall_time,
-        m.value_real AS value,
-        m.value_text AS value_text
-      FROM metrics m
-      JOIN steps s ON s.step = m.step
-      WHERE m.key = ?
-        AND (? IS NULL OR m.step > ?)
-        AND (m.step % ?) = 0
-      ORDER BY m.step ASC
-      LIMIT ?
-      `,
-      [key, sinceStep, sinceStep, stride, limit]
-    );
-
-    return NextResponse.json({
-      key,
-      keys,
-      points: points.map((p) => ({
-        step: p.step,
-        wall_time: p.wall_time,
-        value: p.value ?? (p.value_text ? Number(p.value_text) : null),
-      })),
-    });
+    const parts: string[] = [];
+    for (const key of keys) {
+      const rows = await all<{ steps: string; vals: string }>(db, SERIES_SQL, [key, sinceStep]);
+      const steps = rows[0]?.steps ?? '[]';
+      const vals = rows[0]?.vals ?? '[]';
+      parts.push(`${JSON.stringify(key)}:{"steps":${steps},"values":${vals}}`);
+    }
+    body = `{"keys":${JSON.stringify(keys)},"since_step":${sinceStep},"series":{${parts.join(',')}}}`;
   } finally {
     await closeDb(db);
   }
+
+  const headers: Record<string, string> = {
+    'Content-Type': 'application/json',
+    'Cache-Control': 'no-store',
+    Vary: 'Accept-Encoding',
+  };
+  const acceptEncoding = request.headers.get('accept-encoding') ?? '';
+  if (/\bgzip\b/.test(acceptEncoding)) {
+    // Level 1: the payload is mostly digits, so higher levels buy ~5% for 3x
+    // the CPU time.
+    const gz = await gzip(Buffer.from(body), { level: 1 });
+    headers['Content-Encoding'] = 'gzip';
+    return new NextResponse(gz, { headers });
+  }
+  return new NextResponse(body, { headers });
 }
 
 function run(db: sqlite3.Database, sql: string, params: any[] = []) {

ui/src/app/jobs/new/SimpleJob.tsx

@@ -22,6 +22,7 @@ import {
   CreatableSelectInput,
 } from '@/components/formInputs';
 import Card from '@/components/Card';
+import CustomShiftFields, { customShiftDefaults } from '@/components/CustomShiftFields';
 import { X, Copy, Wand2, SquareDashed, Info } from 'lucide-react';
 import { openDoc } from '@/components/DocModal';
 import { openUpsamplePromptsModal, toAspectRatio } from '@/components/UpsamplePromptsModal';
@@ -745,6 +746,7 @@ export default function SimpleJob({
                     { value: 'automagic', label: 'Automagic' },
                     { value: 'automagic2', label: 'Automagic v2' },
                     { value: 'automagic3', label: 'Automagic v3' },
+                    { value: 'automagic4', label: 'Automagic v4' },
                     { value: 'prodigyopt', label: 'Prodigy' },
                     { value: 'prodigy8bit', label: 'Prodigy8Bit' },
                   ]}
@@ -775,17 +777,36 @@ export default function SimpleJob({
                       label="Timestep Type"
                       value={jobConfig.config.process[0].train.timestep_type}
                       disabled={disableSections.includes('train.timestep_type') || false}
-                      onChange={value => setJobConfig(value, 'config.process[0].train.timestep_type')}
+                      onChange={value => {
+                        if (value === 'custom_shift') {
+                          setJobConfig(
+                            { ...customShiftDefaults, ...jobConfig.config.process[0].train.timestep_sampler },
+                            'config.process[0].train.timestep_sampler',
+                          );
+                        }
+                        setJobConfig(value, 'config.process[0].train.timestep_type');
+                      }}
                       options={[
                         { value: 'sigmoid', label: 'Sigmoid' },
                         { value: 'linear', label: 'Linear' },
                         { value: 'shift', label: 'Shift' },
+                        ...(jobConfig.config.process[0].train.noise_scheduler === 'flowmatch' ||
+                        jobConfig.config.process[0].train.timestep_type === 'custom_shift'
+                          ? [{ value: 'custom_shift', label: 'Custom Shift' }] : []),
                         { value: 'weighted', label: 'Weighted' },
                       ]}
                     />
                   )}
+                  {!disableSections.includes('train.timestep_type') &&
+                    jobConfig.config.process[0].train.timestep_type === 'custom_shift' && (
+                      <CustomShiftFields
+                        value={jobConfig.config.process[0].train.timestep_sampler}
+                        onChange={value => setJobConfig(value, 'config.process[0].train.timestep_sampler')}
+                      />
+                    )}
                   <SelectInput
                     label="Timestep Bias"
+                    disabled={jobConfig.config.process[0].train.timestep_type === 'custom_shift'}
                     className="pt-2"
                     value={jobConfig.config.process[0].train.content_or_style}
                     onChange={value => setJobConfig(value, 'config.process[0].train.content_or_style')}

ui/src/app/layout.tsx

@@ -10,6 +10,7 @@ import DocModal from '@/components/DocModal';
 import os from 'os';
 import { CaptionDatasetModal } from '@/components/CaptionDatasetModal';
 import MergeLoRAsModal from '@/components/MergeLoRAsModal';
+import SafetensorsInfoModal from '@/components/SafetensorsInfoModal';
 import UpsamplePromptsModal from '@/components/UpsamplePromptsModal';
 import PromptBoxEditorModal from '@/components/PromptBoxEditorModal';
 
@@ -65,6 +66,7 @@ export default function RootLayout({ children }: { children: React.ReactNode })
         <DocModal />
         <CaptionDatasetModal />
         <MergeLoRAsModal />
+        <SafetensorsInfoModal />
         <UpsamplePromptsModal />
         <PromptBoxEditorModal />
       </body>

ui/src/components/CustomShiftFields.tsx

@@ -0,0 +1,84 @@
+'use client';
+
+import { useEffect, useRef, useState } from 'react';
+import { Checkbox } from '@/components/formInputs';
+import { CustomShiftConfig } from '@/types';
+
+export const customShiftDefaults: Required<CustomShiftConfig> = {
+  discrete_flow_shift: 1.0,
+  sigmoid_scale: 1.0,
+  min_timestep: 0,
+  max_timestep: 1000,
+  preserve_distribution_shape: false,
+};
+
+function NumericField({
+  label, value, onChange, min, max, integer = false, error = '',
+}: {
+  label: string;
+  value: number;
+  onChange: (value: number) => void;
+  min: number;
+  max?: number;
+  integer?: boolean;
+  error?: string;
+}) {
+  const ref = useRef<HTMLInputElement>(null);
+  const [text, setText] = useState(String(value));
+  useEffect(() => setText(String(value)), [value]);
+  useEffect(() => { ref.current?.setCustomValidity(error); }, [error]);
+
+  return (
+    <label className="block pt-2 text-xs text-gray-300">
+      {label}
+      <input
+        ref={ref}
+        type="number"
+        className="mt-1 w-full text-sm px-3 py-1 bg-gray-950 dark:bg-gray-800 border border-gray-700 rounded-sm text-gray-100"
+        value={text}
+        onChange={event => {
+          const next = event.target.value;
+          setText(next);
+          if (next !== '' && Number.isFinite(Number(next))) onChange(Number(next));
+        }}
+        min={min}
+        max={max}
+        step={integer ? 1 : 'any'}
+        required
+      />
+    </label>
+  );
+}
+
+export default function CustomShiftFields({
+  value, onChange,
+}: {
+  value?: CustomShiftConfig;
+  onChange: (value: CustomShiftConfig) => void;
+}) {
+  const config = { ...customShiftDefaults, ...value };
+  const update = (key: keyof CustomShiftConfig, next: number | boolean) => onChange({ ...config, [key]: next });
+  const rangeError = config.min_timestep >= config.max_timestep ? 'Min Timestep must be less than Max Timestep.' : '';
+
+  return (
+    <div className="pt-2">
+      <NumericField label="Discrete Flow Shift" value={config.discrete_flow_shift} min={0}
+        error={config.discrete_flow_shift <= 0 ? 'Discrete Flow Shift must be greater than zero.' : ''}
+        onChange={next => update('discrete_flow_shift', next)} />
+      <NumericField label="Sigmoid Scale" value={config.sigmoid_scale} min={0}
+        error={config.sigmoid_scale <= 0 ? 'Sigmoid Scale must be greater than zero.' : ''}
+        onChange={next => update('sigmoid_scale', next)} />
+      <NumericField label="Min Timestep" value={config.min_timestep} min={0} max={999} integer
+        error={rangeError} onChange={next => update('min_timestep', next)} />
+      <NumericField label="Max Timestep" value={config.max_timestep} min={1} max={1000} integer
+        error={rangeError} onChange={next => update('max_timestep', next)} />
+      <Checkbox label="Preserve Distribution Shape" className="pt-3"
+        checked={config.preserve_distribution_shape}
+        onChange={next => update('preserve_distribution_shape', next)} />
+      <p className="pt-2 text-xs text-gray-400">
+        When enabled, samples outside the range are rejected. Otherwise, the distribution is scaled into the range.
+        Model stage boundaries still apply. Timestep Bias is ignored.
+      </p>
+    </div>
+  );
+}

ui/src/components/FilesWidget.tsx

@@ -1,7 +1,8 @@
 import React from 'react';
 import useFilesList from '@/hooks/useFilesList';
-import { Loader2, AlertCircle, Download, Box, Brain, Trash2, SlidersHorizontal } from 'lucide-react';
+import { Loader2, AlertCircle, Download, Box, Brain, Trash2, SlidersHorizontal, Info } from 'lucide-react';
 import { openMergeLoRAsModal } from './MergeLoRAsModal';
+import { openSafetensorsInfoModal } from './SafetensorsInfoModal';
 import { getFilename, getFoldername, encodeFilePathForUrl } from '@/utils/basic';
 import { openConfirm } from './ConfirmModal';
 import { apiClient } from '@/utils/api';
@@ -114,6 +115,16 @@ export default function FilesWidget({ jobID, jobName }: { jobID: string; jobName
                   </a>
                   <div className="flex items-center space-x-3 flex-shrink-0">
                     <span className="text-xs text-gray-400">{cleanSize(file.size)}</span>
+                    {fileName.endsWith('.safetensors') && (
+                      <button
+                        type="button"
+                        onClick={() => openSafetensorsInfoModal(file.path)}
+                        className="bg-sky-500 bg-opacity-0 group-hover:bg-opacity-10 hover:!bg-opacity-30 rounded-full p-1 transition-all"
+                        title="Safetensors info"
+                      >
+                        <Info className="w-3 h-3 text-sky-400" />
+                      </button>
+                    )}
                     <a
                       target="_blank"
                       href={`/api/files/${encodeFilePathForUrl(file.path)}`}

ui/src/components/JobLossGraph.tsx

@@ -1,7 +1,7 @@
 'use client';
 
 import { Job } from '@prisma/client';
-import useJobLossLog, { LossPoint } from '@/hooks/useJobLossLog';
+import useJobLossLog, { LossSeries } from '@/hooks/useJobLossLog';
 import { useMemo, useState, useEffect, useRef, useCallback } from 'react';
 import uPlot from 'uplot';
 import 'uplot/dist/uPlot.min.css';
@@ -233,7 +233,16 @@ export default function JobLossGraph({ job }: Props) {
     });
   }, [lossKeys]);
 
-  const activeKeys = useMemo(() => lossKeys.filter(k => enabled[k] !== false), [lossKeys, enabled]);
+  // Keys not yet in `enabled` (data and keys now arrive in the same render,
+  // before the sync effect above has run) take their default here rather than
+  // rendering as enabled and getting pruned a frame later.
+  const activeKeys = useMemo(
+    () =>
+      lossKeys.filter(
+        k => enabled[k] ?? persistedEnabledRef.current?.[k] ?? (k === 'loss/loss' || k === 'val/loss'),
+      ),
+    [lossKeys, enabled],
+  );
 
   // Assign palette slots by position in the (sorted) lossKeys list rather than
   // by hashing the key name — hashing let different keys collide onto the same
@@ -257,11 +266,11 @@ export default function JobLossGraph({ job }: Props) {
     // Union of all steps across active series.
     const stepSet = new Set<number>();
     for (const key of activeKeys) {
-      const pts: LossPoint[] = series[key] ?? [];
-      for (const p of pts) {
-        if (p.value === null || !Number.isFinite(p.value as number)) continue;
-        if (useLogScale && (p.value as number) <= 0) continue;
-        stepSet.add(p.step);
+      const s: LossSeries | undefined = series[key];
+      if (!s) continue;
+      for (let i = 0; i < s.steps.length; i++) {
+        if (useLogScale && s.values[i] <= 0) continue;
+        stepSet.add(s.steps[i]);
       }
     }
     let xs = Array.from(stepSet).sort((a, b) => a - b);
@@ -291,13 +300,13 @@ export default function JobLossGraph({ job }: Props) {
     for (let ki = 0; ki < activeKeys.length; ki++) {
       const key = activeKeys[ki];
       const scaleKey = `y::${key}`;
-      const pts: LossPoint[] = series[key] ?? [];
+      const s: LossSeries = series[key] ?? { steps: [], values: [] };
       const map = new Map<number, number>();
-      for (const p of pts) {
-        if (p.value === null || !Number.isFinite(p.value as number)) continue;
-        if (useLogScale && (p.value as number) <= 0) continue;
-        if (!xsSet.has(p.step)) continue;
-        map.set(p.step, p.value as number);
+      for (let i = 0; i < s.steps.length; i++) {
+        const v = s.values[i];
+        if (useLogScale && v <= 0) continue;
+        if (!xsSet.has(s.steps[i])) continue;
+        map.set(s.steps[i], v);
       }
       const raw: (number | null)[] = xs.map(s => (map.has(s) ? (map.get(s) as number) : null));
       // Metrics logged less often than every step are null at most x positions.

ui/src/components/SafetensorsInfoModal.tsx

@@ -0,0 +1,281 @@
+'use client';
+import React, { useEffect, useMemo, useState } from 'react';
+import { createGlobalState } from 'react-global-hooks';
+import { Loader2, AlertCircle, Search } from 'lucide-react';
+import { Modal } from './Modal';
+import { getFilename } from '@/utils/basic';
+import { apiClient } from '@/utils/api';
+
+export interface SafetensorsInfoModalState {
+  filePath: string;
+}
+
+interface TensorInfo {
+  name: string;
+  dtype: string;
+  shape: number[];
+  numel: number;
+  bytes: number;
+}
+
+interface SafetensorsInfo {
+  filePath: string;
+  fileSize: number;
+  headerSize: number;
+  metadata: Record<string, string>;
+  tensors: TensorInfo[];
+  tensorCount: number;
+  totalParams: number;
+  totalBytes: number;
+  dtypeCounts: Record<string, { count: number; numel: number; bytes: number }>;
+}
+
+export const safetensorsInfoModalState = createGlobalState<SafetensorsInfoModalState | null>(null);
+
+export const openSafetensorsInfoModal = (filePath: string) => {
+  safetensorsInfoModalState.set({ filePath });
+};
+
+const fmtBytes = (size: number) => {
+  if (size < 1024) return `${size} B`;
+  if (size < 1024 * 1024) return `${(size / 1024).toFixed(1)} KB`;
+  if (size < 1024 * 1024 * 1024) return `${(size / (1024 * 1024)).toFixed(1)} MB`;
+  return `${(size / (1024 * 1024 * 1024)).toFixed(2)} GB`;
+};
+
+const fmtCount = (n: number) => {
+  if (n < 1000) return `${n}`;
+  if (n < 1e6) return `${(n / 1e3).toFixed(1)}K`;
+  if (n < 1e9) return `${(n / 1e6).toFixed(1)}M`;
+  return `${(n / 1e9).toFixed(2)}B`;
+};
+
+// metadata values are strings; many are JSON blobs worth pretty-printing
+const fmtMetaValue = (value: string) => {
+  const trimmed = value.trim();
+  if ((trimmed.startsWith('{') && trimmed.endsWith('}')) || (trimmed.startsWith('[') && trimmed.endsWith(']'))) {
+    try {
+      return JSON.stringify(JSON.parse(trimmed), null, 2);
+    } catch {
+      /* not json */
+    }
+  }
+  return value;
+};
+
+const MAX_ROWS = 2000;
+
+const SafetensorsInfoModal: React.FC = () => {
+  const [modalInfo, setModalInfo] = safetensorsInfoModalState.use();
+  const isOpen = modalInfo !== null;
+  const [info, setInfo] = useState<SafetensorsInfo | null>(null);
+  const [error, setError] = useState<string | null>(null);
+  const [isLoading, setIsLoading] = useState(false);
+  const [filter, setFilter] = useState('');
+  const [tab, setTab] = useState<'tensors' | 'metadata'>('tensors');
+
+  useEffect(() => {
+    if (!modalInfo) {
+      setInfo(null);
+      setError(null);
+      setIsLoading(false);
+      setFilter('');
+      setTab('tensors');
+      return;
+    }
+    let cancelled = false;
+    setIsLoading(true);
+    setInfo(null);
+    setError(null);
+    apiClient
+      .post('/api/files/info', { filePath: modalInfo.filePath })
+      .then(res => {
+        if (cancelled) return;
+        setInfo(res.data as SafetensorsInfo);
+      })
+      .catch(err => {
+        if (cancelled) return;
+        setError(err?.response?.data || err?.message || 'Failed to read file');
+      })
+      .finally(() => {
+        if (!cancelled) setIsLoading(false);
+      });
+    return () => {
+      cancelled = true;
+    };
+  }, [modalInfo]);
+
+  const onClose = () => setModalInfo(null);
+
+  const filteredTensors = useMemo(() => {
+    if (!info) return [];
+    const q = filter.trim().toLowerCase();
+    if (!q) return info.tensors;
+    return info.tensors.filter(t => t.name.toLowerCase().includes(q) || t.dtype.toLowerCase().includes(q));
+  }, [info, filter]);
+
+  const metadataEntries = info ? Object.entries(info.metadata) : [];
+  const title = modalInfo ? getFilename(modalInfo.filePath) : 'Safetensors Info';
+
+  return (
+    <Modal isOpen={isOpen} onClose={onClose} title={title} size="xl">
+      {isLoading && (
+        <div className="flex flex-col items-center justify-center py-16 space-y-3 text-gray-300">
+          <Loader2 className="w-8 h-8 text-purple-400 animate-spin" />
+          <span className="text-sm">Gathering info...</span>
+        </div>
+      )}
+
+      {!isLoading && error && (
+        <div className="flex items-center justify-center py-10 text-rose-400 space-x-2">
+          <AlertCircle className="w-5 h-5" />
+          <span className="text-sm">{String(error)}</span>
+        </div>
+      )}
+
+      {!isLoading && info && (
+        <div className="space-y-3">
+          <div className="grid grid-cols-2 sm:grid-cols-4 gap-2 text-xs">
+            <div className="bg-gray-800 rounded-lg px-3 py-2">
+              <div className="text-gray-500">File size</div>
+              <div className="text-gray-100 font-mono">{fmtBytes(info.fileSize)}</div>
+            </div>
+            <div className="bg-gray-800 rounded-lg px-3 py-2">
+              <div className="text-gray-500">Tensors</div>
+              <div className="text-gray-100 font-mono">{info.tensorCount.toLocaleString()}</div>
+            </div>
+            <div className="bg-gray-800 rounded-lg px-3 py-2">
+              <div className="text-gray-500">Parameters</div>
+              <div className="text-gray-100 font-mono" title={info.totalParams.toLocaleString()}>
+                {fmtCount(info.totalParams)}
+              </div>
+            </div>
+            <div className="bg-gray-800 rounded-lg px-3 py-2">
+              <div className="text-gray-500">Header</div>
+              <div className="text-gray-100 font-mono">{fmtBytes(info.headerSize)}</div>
+            </div>
+          </div>
+
+          <div className="flex flex-wrap gap-1.5 text-xs">
+            {Object.entries(info.dtypeCounts)
+              .sort((a, b) => b[1].bytes - a[1].bytes)
+              .map(([dtype, d]) => (
+                <span
+                  key={dtype}
+                  className="px-2 py-0.5 rounded-full bg-purple-500/10 text-purple-300 font-mono"
+                  title={`${d.count} tensors · ${d.numel.toLocaleString()} params`}
+                >
+                  {dtype} · {d.count} · {fmtBytes(d.bytes)}
+                </span>
+              ))}
+          </div>
+
+          <div className="flex items-center justify-between border-b border-gray-800">
+            <div className="flex space-x-1">
+              <button
+                type="button"
+                onClick={() => setTab('tensors')}
+                className={`px-3 py-1.5 text-sm border-b-2 ${
+                  tab === 'tensors'
+                    ? 'border-purple-500 text-gray-100'
+                    : 'border-transparent text-gray-400 hover:text-gray-200'
+                }`}
+              >
+                Tensors
+              </button>
+              <button
+                type="button"
+                onClick={() => setTab('metadata')}
+                className={`px-3 py-1.5 text-sm border-b-2 ${
+                  tab === 'metadata'
+                    ? 'border-purple-500 text-gray-100'
+                    : 'border-transparent text-gray-400 hover:text-gray-200'
+                }`}
+              >
+                Metadata
+                <span className="ml-1.5 px-1.5 py-0.5 bg-gray-700 rounded-full text-xs text-gray-300">
+                  {metadataEntries.length}
+                </span>
+              </button>
+            </div>
+            {tab === 'tensors' && (
+              <div className="relative">
+                <Search className="w-3.5 h-3.5 text-gray-500 absolute left-2 top-1/2 -translate-y-1/2" />
+                <input
+                  type="text"
+                  value={filter}
+                  onChange={e => setFilter(e.target.value)}
+                  placeholder="Filter tensors"
+                  className="text-xs pl-7 pr-2 py-1 w-56 bg-gray-950 dark:bg-gray-800 border border-gray-700 rounded-md text-gray-100 focus:ring-1 focus:ring-gray-600 focus:outline-none"
+                />
+              </div>
+            )}
+          </div>
+
+          {tab === 'tensors' && (
+            <div className="overflow-auto max-h-[60vh] rounded-md border border-gray-800">
+              <table className="w-full text-xs font-mono">
+                <thead className="sticky top-0 bg-gray-800 text-gray-400">
+                  <tr>
+                    <th className="text-left px-2 py-1.5 font-normal">Name</th>
+                    <th className="text-left px-2 py-1.5 font-normal">Shape</th>
+                    <th className="text-left px-2 py-1.5 font-normal">Dtype</th>
+                    <th className="text-right px-2 py-1.5 font-normal">Params</th>
+                    <th className="text-right px-2 py-1.5 font-normal">Size</th>
+                  </tr>
+                </thead>
+                <tbody>
+                  {filteredTensors.slice(0, MAX_ROWS).map(t => (
+                    <tr key={t.name} className="border-t border-gray-800/60 hover:bg-gray-800/50">
+                      <td className="px-2 py-1 text-gray-200 break-all">{t.name}</td>
+                      <td className="px-2 py-1 text-gray-300 whitespace-nowrap">[{t.shape.join(', ')}]</td>
+                      <td className="px-2 py-1 text-purple-300 whitespace-nowrap">{t.dtype}</td>
+                      <td className="px-2 py-1 text-gray-400 text-right whitespace-nowrap">
+                        {t.numel.toLocaleString()}
+                      </td>
+                      <td className="px-2 py-1 text-gray-400 text-right whitespace-nowrap">{fmtBytes(t.bytes)}</td>
+                    </tr>
+                  ))}
+                  {filteredTensors.length === 0 && (
+                    <tr>
+                      <td colSpan={5} className="px-2 py-4 text-center text-gray-500">
+                        No tensors match
+                      </td>
+                    </tr>
+                  )}
+                </tbody>
+              </table>
+              {filteredTensors.length > MAX_ROWS && (
+                <div className="px-2 py-1.5 text-xs text-gray-500 bg-gray-800 border-t border-gray-700">
+                  Showing {MAX_ROWS.toLocaleString()} of {filteredTensors.length.toLocaleString()} — narrow with the
+                  filter
+                </div>
+              )}
+            </div>
+          )}
+
+          {tab === 'metadata' && (
+            <div className="overflow-auto max-h-[60vh] rounded-md border border-gray-800">
+              {metadataEntries.length === 0 ? (
+                <div className="px-2 py-4 text-center text-gray-500 text-xs">No metadata</div>
+              ) : (
+                <div className="divide-y divide-gray-800/60">
+                  {metadataEntries.map(([key, value]) => (
+                    <div key={key} className="px-2 py-1.5">
+                      <div className="text-xs text-purple-300 font-mono">{key}</div>
+                      <pre className="text-xs text-gray-200 font-mono whitespace-pre-wrap break-all mt-0.5">
+                        {fmtMetaValue(String(value))}
+                      </pre>
+                    </div>
+                  ))}
+                </div>
+              )}
+            </div>
+          )}
+        </div>
+      )}
+    </Modal>
+  );
+};
+
+export default SafetensorsInfoModal;

ui/src/hooks/useJobLossLog.tsx

@@ -4,13 +4,20 @@ import { useEffect, useState, useRef, useCallback, useMemo } from 'react';
 import { apiClient } from '@/utils/api';
 import usePollLoop from '@/hooks/usePollLoop';
 
-export interface LossPoint {
-  step: number;
-  wall_time?: number;
-  value: number | null;
+// Columnar series: steps[i] pairs with values[i]. Null / non-finite values are
+// dropped at ingest, so every value here is a finite number.
+export interface LossSeries {
+  steps: number[];
+  values: number[];
 }
 
-type SeriesMap = Record<string, LossPoint[]>;
+type SeriesMap = Record<string, LossSeries>;
+
+interface LossResponse {
+  keys?: string[];
+  since_step?: number;
+  series?: Record<string, { steps: number[]; values: (number | null)[] }>;
+}
 
 export default function useJobLossLog(jobID: string, reloadInterval: null | number = null) {
   const [series, setSeries] = useState<SeriesMap>({});
@@ -20,8 +27,10 @@ export default function useJobLossLog(jobID: string, reloadInterval: null | numb
   const didInitialLoadRef = useRef(false);
   const inFlightRef = useRef(false);
 
-  // track last step per key so polling is incremental per series
-  const lastStepByKeyRef = useRef<Record<string, number | null>>({});
+  // Highest step seen across every key. Polls ask for step >= this (not >) so a
+  // key logged at that step after the previous poll is still picked up; the
+  // per-key merge below drops the one re-sent point.
+  const lastStepRef = useRef<number | null>(null);
 
   const lossKeys = useMemo(() => {
     const base = keys ?? [];
@@ -40,69 +49,60 @@ export default function useJobLossLog(jobID: string, reloadInterval: null | numb
     setStatus(loadStatus);
 
     try {
-      // Step 1: get key list (we can do this by calling endpoint once; it returns keys)
-      // Keep it cheap: limit=1.
-      const first = await apiClient
-        .get(`/api/jobs/${jobID}/loss`, { params: { key: 'loss', limit: 1 } })
-        .then(res => res.data as { keys?: string[] });
+      // One request returns every key. Incremental after the initial load.
+      const params: Record<string, number> = {};
+      if (reloadInterval && didInitialLoadRef.current && lastStepRef.current != null) {
+        params.since_step = lastStepRef.current;
+      }
+      const res = await apiClient.get(`/api/jobs/${jobID}/loss`, { params });
+      const data = res.data as LossResponse;
 
-      const newKeys = first.keys ?? [];
+      const newKeys = data.keys ?? [];
       setKeys(newKeys);
-
       const wantedLossKeys = (newKeys.length ? [...newKeys] : ['loss']).sort();
-
-      // Step 2: fetch each loss key incrementally (since_step per key if polling)
-      const requests = wantedLossKeys.map(k => {
-        const params: Record<string, any> = { key: k };
-
-        if (reloadInterval && lastStepByKeyRef.current[k] != null) {
-          params.since_step = lastStepByKeyRef.current[k];
-        }
-
-        params.limit = 1000000;
-
-        return apiClient
-          .get(`/api/jobs/${jobID}/loss`, { params })
-          .then(res => res.data as { key: string; points?: LossPoint[] });
-      });
-
-      const results = await Promise.all(requests);
+      const incoming = data.series ?? {};
+      const initial = !didInitialLoadRef.current;
 
       setSeries(prev => {
-        const next: SeriesMap = { ...prev };
+        const next: SeriesMap = initial ? {} : { ...prev };
+        let maxStep = lastStepRef.current ?? null;
 
-        for (const r of results) {
-          const k = r.key;
-          const newPoints = (r.points ?? []).filter(p => p.value !== null);
+        for (const k of wantedLossKeys) {
+          const inc = incoming[k];
+          const existing = initial ? undefined : next[k];
+          const prevLast = existing && existing.steps.length ? existing.steps[existing.steps.length - 1] : null;
 
-          if (!didInitialLoadRef.current) {
-            // initial: replace
-            next[k] = newPoints;
-          } else if (newPoints.length) {
-            const existing = next[k] ?? [];
-            const prevLast = existing.length ? existing[existing.length - 1].step : null;
-            const filtered = prevLast == null ? newPoints : newPoints.filter(p => p.step > prevLast);
-            next[k] = filtered.length ? [...existing, ...filtered] : existing;
-          } else {
-            // no new points: keep existing
-            next[k] = next[k] ?? [];
+          const steps: number[] = existing ? existing.steps.slice() : [];
+          const values: number[] = existing ? existing.values.slice() : [];
+          let appended = false;
+          if (inc) {
+            const n = Math.min(inc.steps.length, inc.values.length);
+            for (let i = 0; i < n; i++) {
+              const v = inc.values[i];
+              if (v === null || !Number.isFinite(v)) continue;
+              const s = inc.steps[i];
+              if (prevLast != null && s <= prevLast) continue;
+              steps.push(s);
+              values.push(v as number);
+              appended = true;
+            }
           }
+          if (initial || appended) next[k] = { steps, values };
+          else next[k] = existing ?? { steps: [], values: [] };
 
-          // update last step per key
-          const finalArr = next[k] ?? [];
-          lastStepByKeyRef.current[k] = finalArr.length
-            ? finalArr[finalArr.length - 1].step
-            : (lastStepByKeyRef.current[k] ?? null);
+          const arr = next[k].steps;
+          if (arr.length) {
+            const last = arr[arr.length - 1];
+            if (maxStep == null || last > maxStep) maxStep = last;
+          }
         }
 
         // remove stale keys that no longer exist (rare, but keeps UI clean)
         for (const existingKey of Object.keys(next)) {
-          if (!wantedLossKeys.includes(existingKey)) {
-            delete next[existingKey];
-            delete lastStepByKeyRef.current[existingKey];
-          }
+          if (!wantedLossKeys.includes(existingKey)) delete next[existingKey];
         }
 
+        lastStepRef.current = maxStep;
         return next;
       });
 
@@ -124,13 +124,24 @@ export default function useJobLossLog(jobID: string, reloadInterval: null | numb
       await apiClient.delete(`/api/jobs/${jobID}/loss`, { data: { min_step: minStep, max_step: maxStep } });
       setSeries(prev => {
         const next: SeriesMap = {};
+        let last: number | null = null;
         for (const k of Object.keys(prev)) {
-          const kept = prev[k].filter(p => p.step < minStep || p.step > maxStep);
-          next[k] = kept;
-          // Re-anchor incremental polling to the new tail so a deleted tail
-          // isn't treated as already-fetched.
-          lastStepByKeyRef.current[k] = kept.length ? kept[kept.length - 1].step : null;
+          const steps: number[] = [];
+          const values: number[] = [];
+          const src = prev[k];
+          for (let i = 0; i < src.steps.length; i++) {
+            const s = src.steps[i];
+            if (s < minStep || s > maxStep) {
+              steps.push(s);
+              values.push(src.values[i]);
+            }
+          }
+          next[k] = { steps, values };
+          if (steps.length && (last == null || steps[steps.length - 1] > last)) last = steps[steps.length - 1];
         }
+        // Re-anchor incremental polling to the new tail so a deleted tail
+        // isn't treated as already-fetched.
+        lastStepRef.current = last;
         return next;
       });
     },
@@ -141,7 +152,7 @@ export default function useJobLossLog(jobID: string, reloadInterval: null | numb
   // before the first fetch when jobID changes.
   useEffect(() => {
     didInitialLoadRef.current = false;
-    lastStepByKeyRef.current = {};
+    lastStepRef.current = null;
     setSeries({});
     setKeys([]);
     setStatus('idle');

ui/src/server/monitor.ts

@@ -36,11 +36,20 @@ const NV_BATCH_FLUSH_MS = 100;
 const CPU_TEMP_REFRESH_MS = 5000;
 
 // Power limits for GPUs whose driver reports none (nvidia-smi gives [N/A]).
-// GB10 (DGX Spark) is a unified SoC with no settable GPU limit; 140 W is
-// NVIDIA's rated TDP for the GB10 superchip.
-const KNOWN_POWER_LIMITS_W: Record<string, number> = {
-  'NVIDIA GB10': 140,
-};
+// These unified SoCs have no settable GPU limit, so show the rated TDP of the
+// whole chip. GB10 (DGX Spark) is 140 W. RTX Spark N1X: the 5120-core part is
+// laptop-only, rated 45-80 W (OEM configurable); the 6144-core part is 140 W
+// on desktop and 45-80 W in laptops but shares its name, so it falls through
+// to the desktop figure.
+const KNOWN_POWER_LIMITS_W: [RegExp, number][] = [
+  [/\bGB10\b/, 140],
+  [/RTX Spark.*5120-core/, 80],
+  [/RTX Spark/, 140],
+];
+
+function knownPowerLimit(name: string): number {
+  return KNOWN_POWER_LIMITS_W.find(([re]) => re.test(name))?.[1] ?? NaN;
+}
 
 function parseGpuLine(line: string): GpuInfo | null {
   const [
@@ -76,7 +85,7 @@ function parseGpuLine(line: string): GpuInfo | null {
     },
     power: {
       draw: parseFloat(powerDraw),
-      limit: isNaN(parseFloat(powerLimit)) ? (KNOWN_POWER_LIMITS_W[name] ?? NaN) : parseFloat(powerLimit),
+      limit: isNaN(parseFloat(powerLimit)) ? knownPowerLimit(name) : parseFloat(powerLimit),
     },
     clocks: {
       graphics: parseInt(clockGraphics),
@@ -88,19 +97,28 @@ function parseGpuLine(line: string): GpuInfo | null {
   };
 }
 
+// Integrated NVIDIA GPUs that share system RAM with the CPU.
+const UNIFIED_MEMORY_GPU = /\b(GB10|RTX Spark)\b/i;
+
 /**
- * Unified-memory GPUs (GB10 / DGX Spark) have no dedicated VRAM, so
- * nvidia-smi reports memory as [N/A]. Report the system RAM pool they share
- * instead, like the Mac path does for Apple silicon.
+ * Unified-memory GPUs (GB10 / DGX Spark, RTX Spark) have no dedicated VRAM.
+ * On Linux nvidia-smi reports their memory as [N/A]; on Windows (WDDM) it
+ * reports only the small dedicated segment (8 GB) even though CUDA can use
+ * most of system RAM. Report the system RAM pool they share instead, like
+ * the Mac path does for Apple silicon.
  */
+function isSharedMemoryGpu(g: GpuInfo): boolean {
+  return isNaN(g.memory.total) || UNIFIED_MEMORY_GPU.test(g.name);
+}
+
 function withSharedMemory(gpu: GPUApiResponse, cpu: CpuInfo | null): GPUApiResponse {
-  if (!cpu || !gpu.gpus.some(g => isNaN(g.memory.total))) return gpu;
+  if (!cpu || !gpu.gpus.some(isSharedMemoryGpu)) return gpu;
   const total = Math.round(cpu.totalMemory);
   const free = Math.round(cpu.availableMemory);
   return {
     ...gpu,
     gpus: gpu.gpus.map(g =>
-      isNaN(g.memory.total) ? { ...g, memory: { total, free, used: total - free, shared: true } } : g,
+      isSharedMemoryGpu(g) ? { ...g, memory: { total, free, used: total - free, shared: true } } : g,
     ),
   };
 }

ui/src/types.ts

@@ -155,6 +155,14 @@ export interface ValidationConfig {
   validation_sigmas?: number[];
 }
 
+export interface CustomShiftConfig {
+  discrete_flow_shift?: number;
+  sigmoid_scale?: number;
+  min_timestep?: number;
+  max_timestep?: number;
+  preserve_distribution_shape?: boolean;
+}
+
 export interface TrainConfig {
   batch_size: number;
   bypass_guidance_embedding?: boolean;
@@ -165,6 +173,7 @@ export interface TrainConfig {
   gradient_checkpointing: boolean;
   noise_scheduler: string;
   timestep_type: string;
+  timestep_sampler?: CustomShiftConfig;
   content_or_style: string;
   optimizer: string;
   lr: number;

version.py

@@ -1 +1 @@
-VERSION = "0.13.37"
+VERSION = "0.13.40"