From 32556949ce2ffcb7e74e79fc64943ea363304f5a Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 4 Jul 2026 03:29:45 +0000 Subject: [PATCH 1/7] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/routes/models.py | 4 +--- studio/backend/tests/test_diffusion_dataset_api.py | 6 ++++-- studio/backend/tests/test_diffusion_dit_trainer.py | 5 ++++- studio/backend/tests/test_diffusion_training.py | 6 ++++-- studio/backend/tests/test_local_model_format.py | 13 +++++++++---- 5 files changed, 22 insertions(+), 12 deletions(-) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 54bc490052..d4501517c8 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -911,9 +911,7 @@ async def list_local_models( models = collect_local_models(models_root) # Tag each model with its task so the Images picker can filter to diffusion # (GGUF by architecture; local diffusers checkpoints by pipeline / family). - models = [ - m.model_copy(update = {"task": _local_model_task(m)}) for m in models - ] + models = [m.model_copy(update = {"task": _local_model_task(m)}) for m in models] return LocalModelListResponse( models_dir = str(models_root), diff --git a/studio/backend/tests/test_diffusion_dataset_api.py b/studio/backend/tests/test_diffusion_dataset_api.py index e1c2c0a186..b10dd04f4c 100644 --- a/studio/backend/tests/test_diffusion_dataset_api.py +++ b/studio/backend/tests/test_diffusion_dataset_api.py @@ -64,8 +64,10 @@ def test_list_images_caption_precedence(client, ds_root): # a.png -> sidecar (an explicit edit beats the metadata row), b.png -> metadata-only, # c.png -> none. (folder / "metadata.jsonl").write_text( - json.dumps({"file_name": "a.png", "text": "from metadata"}) + "\n" - + json.dumps({"file_name": "b.png", "text": "from metadata"}) + "\n", + json.dumps({"file_name": "a.png", "text": "from metadata"}) + + "\n" + + json.dumps({"file_name": "b.png", "text": "from metadata"}) + + "\n", encoding = "utf-8", ) (folder / "a.txt").write_text("edited sidecar", encoding = "utf-8") diff --git a/studio/backend/tests/test_diffusion_dit_trainer.py b/studio/backend/tests/test_diffusion_dit_trainer.py index e1aa7868a3..63cc01922c 100644 --- a/studio/backend/tests/test_diffusion_dit_trainer.py +++ b/studio/backend/tests/test_diffusion_dit_trainer.py @@ -58,7 +58,10 @@ def test_select_lora_targets_explicit_override_wins(): base_model = "black-forest-labs/FLUX.1-dev", data_dir = "d", output_dir = "o" ).normalized() assert cfg.lora_target_modules == DEFAULT_LORA_TARGETS - assert _select_lora_targets(cfg.lora_target_modules, _SPECS["flux.1"].lora_targets) == _FLUX_TARGETS + assert ( + _select_lora_targets(cfg.lora_target_modules, _SPECS["flux.1"].lora_targets) + == _FLUX_TARGETS + ) @pytest.mark.parametrize( diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index 803d225fa8..099d0f0b97 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -471,8 +471,10 @@ def test_diffusion_info_counts_metadata_captions(client, dataset_roots): (folder / "c.png").write_bytes(b"x") # a.png + b.png via metadata; a.png also has a sidecar (must count once); c.png none. (folder / "metadata.jsonl").write_text( - json.dumps({"file_name": "a.png", "text": "cap a"}) + "\n" - + json.dumps({"file_name": "b.png", "text": "cap b"}) + "\n", + json.dumps({"file_name": "a.png", "text": "cap a"}) + + "\n" + + json.dumps({"file_name": "b.png", "text": "cap b"}) + + "\n", encoding = "utf-8", ) (folder / "a.txt").write_text("edited a", encoding = "utf-8") diff --git a/studio/backend/tests/test_local_model_format.py b/studio/backend/tests/test_local_model_format.py index dd888eb50d..f9184774d6 100644 --- a/studio/backend/tests/test_local_model_format.py +++ b/studio/backend/tests/test_local_model_format.py @@ -121,7 +121,14 @@ def test_scan_models_dir_classifies_root_gguf_with_config(tmp_path): from models.models import LocalModelInfo # noqa: E402 -def _local(path, *, model_format = None, model_id = None, display_name = "m", id = "m"): +def _local( + path, + *, + model_format = None, + model_id = None, + display_name = "m", + id = "m", +): return LocalModelInfo( id = id, display_name = display_name, @@ -147,9 +154,7 @@ def test_local_task_tags_diffusers_by_family_id(tmp_path): d = tmp_path / "flux-checkpoint" _touch(d / "flux1-dev.safetensors") assert ( - models_route._local_model_task( - _local(d, model_id = "black-forest-labs/FLUX.1-dev") - ) + models_route._local_model_task(_local(d, model_id = "black-forest-labs/FLUX.1-dev")) == "text-to-image" ) From f2af2874dfcdf0f0f9491efefd1b22cc4b73174a Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 4 Jul 2026 03:30:16 +0000 Subject: [PATCH 2/7] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/routes/models.py | 4 +--- studio/backend/tests/test_diffusion_dataset_api.py | 6 ++++-- studio/backend/tests/test_diffusion_dit_trainer.py | 5 ++++- studio/backend/tests/test_diffusion_training.py | 6 ++++-- studio/backend/tests/test_local_model_format.py | 13 +++++++++---- 5 files changed, 22 insertions(+), 12 deletions(-) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 54bc490052..d4501517c8 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -911,9 +911,7 @@ async def list_local_models( models = collect_local_models(models_root) # Tag each model with its task so the Images picker can filter to diffusion # (GGUF by architecture; local diffusers checkpoints by pipeline / family). - models = [ - m.model_copy(update = {"task": _local_model_task(m)}) for m in models - ] + models = [m.model_copy(update = {"task": _local_model_task(m)}) for m in models] return LocalModelListResponse( models_dir = str(models_root), diff --git a/studio/backend/tests/test_diffusion_dataset_api.py b/studio/backend/tests/test_diffusion_dataset_api.py index e1c2c0a186..b10dd04f4c 100644 --- a/studio/backend/tests/test_diffusion_dataset_api.py +++ b/studio/backend/tests/test_diffusion_dataset_api.py @@ -64,8 +64,10 @@ def test_list_images_caption_precedence(client, ds_root): # a.png -> sidecar (an explicit edit beats the metadata row), b.png -> metadata-only, # c.png -> none. (folder / "metadata.jsonl").write_text( - json.dumps({"file_name": "a.png", "text": "from metadata"}) + "\n" - + json.dumps({"file_name": "b.png", "text": "from metadata"}) + "\n", + json.dumps({"file_name": "a.png", "text": "from metadata"}) + + "\n" + + json.dumps({"file_name": "b.png", "text": "from metadata"}) + + "\n", encoding = "utf-8", ) (folder / "a.txt").write_text("edited sidecar", encoding = "utf-8") diff --git a/studio/backend/tests/test_diffusion_dit_trainer.py b/studio/backend/tests/test_diffusion_dit_trainer.py index e1aa7868a3..63cc01922c 100644 --- a/studio/backend/tests/test_diffusion_dit_trainer.py +++ b/studio/backend/tests/test_diffusion_dit_trainer.py @@ -58,7 +58,10 @@ def test_select_lora_targets_explicit_override_wins(): base_model = "black-forest-labs/FLUX.1-dev", data_dir = "d", output_dir = "o" ).normalized() assert cfg.lora_target_modules == DEFAULT_LORA_TARGETS - assert _select_lora_targets(cfg.lora_target_modules, _SPECS["flux.1"].lora_targets) == _FLUX_TARGETS + assert ( + _select_lora_targets(cfg.lora_target_modules, _SPECS["flux.1"].lora_targets) + == _FLUX_TARGETS + ) @pytest.mark.parametrize( diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index 803d225fa8..099d0f0b97 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -471,8 +471,10 @@ def test_diffusion_info_counts_metadata_captions(client, dataset_roots): (folder / "c.png").write_bytes(b"x") # a.png + b.png via metadata; a.png also has a sidecar (must count once); c.png none. (folder / "metadata.jsonl").write_text( - json.dumps({"file_name": "a.png", "text": "cap a"}) + "\n" - + json.dumps({"file_name": "b.png", "text": "cap b"}) + "\n", + json.dumps({"file_name": "a.png", "text": "cap a"}) + + "\n" + + json.dumps({"file_name": "b.png", "text": "cap b"}) + + "\n", encoding = "utf-8", ) (folder / "a.txt").write_text("edited a", encoding = "utf-8") diff --git a/studio/backend/tests/test_local_model_format.py b/studio/backend/tests/test_local_model_format.py index dd888eb50d..f9184774d6 100644 --- a/studio/backend/tests/test_local_model_format.py +++ b/studio/backend/tests/test_local_model_format.py @@ -121,7 +121,14 @@ def test_scan_models_dir_classifies_root_gguf_with_config(tmp_path): from models.models import LocalModelInfo # noqa: E402 -def _local(path, *, model_format = None, model_id = None, display_name = "m", id = "m"): +def _local( + path, + *, + model_format = None, + model_id = None, + display_name = "m", + id = "m", +): return LocalModelInfo( id = id, display_name = display_name, @@ -147,9 +154,7 @@ def test_local_task_tags_diffusers_by_family_id(tmp_path): d = tmp_path / "flux-checkpoint" _touch(d / "flux1-dev.safetensors") assert ( - models_route._local_model_task( - _local(d, model_id = "black-forest-labs/FLUX.1-dev") - ) + models_route._local_model_task(_local(d, model_id = "black-forest-labs/FLUX.1-dev")) == "text-to-image" ) From d8495d058f5039759a066d5982489a06403ffa97 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 4 Jul 2026 03:30:49 +0000 Subject: [PATCH 3/7] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/routes/models.py | 4 +--- studio/backend/tests/test_diffusion_dataset_api.py | 6 ++++-- studio/backend/tests/test_diffusion_dit_trainer.py | 5 ++++- studio/backend/tests/test_diffusion_training.py | 6 ++++-- studio/backend/tests/test_local_model_format.py | 13 +++++++++---- 5 files changed, 22 insertions(+), 12 deletions(-) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 54bc490052..d4501517c8 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -911,9 +911,7 @@ async def list_local_models( models = collect_local_models(models_root) # Tag each model with its task so the Images picker can filter to diffusion # (GGUF by architecture; local diffusers checkpoints by pipeline / family). - models = [ - m.model_copy(update = {"task": _local_model_task(m)}) for m in models - ] + models = [m.model_copy(update = {"task": _local_model_task(m)}) for m in models] return LocalModelListResponse( models_dir = str(models_root), diff --git a/studio/backend/tests/test_diffusion_dataset_api.py b/studio/backend/tests/test_diffusion_dataset_api.py index e1c2c0a186..b10dd04f4c 100644 --- a/studio/backend/tests/test_diffusion_dataset_api.py +++ b/studio/backend/tests/test_diffusion_dataset_api.py @@ -64,8 +64,10 @@ def test_list_images_caption_precedence(client, ds_root): # a.png -> sidecar (an explicit edit beats the metadata row), b.png -> metadata-only, # c.png -> none. (folder / "metadata.jsonl").write_text( - json.dumps({"file_name": "a.png", "text": "from metadata"}) + "\n" - + json.dumps({"file_name": "b.png", "text": "from metadata"}) + "\n", + json.dumps({"file_name": "a.png", "text": "from metadata"}) + + "\n" + + json.dumps({"file_name": "b.png", "text": "from metadata"}) + + "\n", encoding = "utf-8", ) (folder / "a.txt").write_text("edited sidecar", encoding = "utf-8") diff --git a/studio/backend/tests/test_diffusion_dit_trainer.py b/studio/backend/tests/test_diffusion_dit_trainer.py index e1aa7868a3..63cc01922c 100644 --- a/studio/backend/tests/test_diffusion_dit_trainer.py +++ b/studio/backend/tests/test_diffusion_dit_trainer.py @@ -58,7 +58,10 @@ def test_select_lora_targets_explicit_override_wins(): base_model = "black-forest-labs/FLUX.1-dev", data_dir = "d", output_dir = "o" ).normalized() assert cfg.lora_target_modules == DEFAULT_LORA_TARGETS - assert _select_lora_targets(cfg.lora_target_modules, _SPECS["flux.1"].lora_targets) == _FLUX_TARGETS + assert ( + _select_lora_targets(cfg.lora_target_modules, _SPECS["flux.1"].lora_targets) + == _FLUX_TARGETS + ) @pytest.mark.parametrize( diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index 1efd19f18f..517186a196 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -547,8 +547,10 @@ def test_diffusion_info_counts_metadata_captions(client, dataset_roots): (folder / "c.png").write_bytes(b"x") # a.png + b.png via metadata; a.png also has a sidecar (must count once); c.png none. (folder / "metadata.jsonl").write_text( - json.dumps({"file_name": "a.png", "text": "cap a"}) + "\n" - + json.dumps({"file_name": "b.png", "text": "cap b"}) + "\n", + json.dumps({"file_name": "a.png", "text": "cap a"}) + + "\n" + + json.dumps({"file_name": "b.png", "text": "cap b"}) + + "\n", encoding = "utf-8", ) (folder / "a.txt").write_text("edited a", encoding = "utf-8") diff --git a/studio/backend/tests/test_local_model_format.py b/studio/backend/tests/test_local_model_format.py index dd888eb50d..f9184774d6 100644 --- a/studio/backend/tests/test_local_model_format.py +++ b/studio/backend/tests/test_local_model_format.py @@ -121,7 +121,14 @@ def test_scan_models_dir_classifies_root_gguf_with_config(tmp_path): from models.models import LocalModelInfo # noqa: E402 -def _local(path, *, model_format = None, model_id = None, display_name = "m", id = "m"): +def _local( + path, + *, + model_format = None, + model_id = None, + display_name = "m", + id = "m", +): return LocalModelInfo( id = id, display_name = display_name, @@ -147,9 +154,7 @@ def test_local_task_tags_diffusers_by_family_id(tmp_path): d = tmp_path / "flux-checkpoint" _touch(d / "flux1-dev.safetensors") assert ( - models_route._local_model_task( - _local(d, model_id = "black-forest-labs/FLUX.1-dev") - ) + models_route._local_model_task(_local(d, model_id = "black-forest-labs/FLUX.1-dev")) == "text-to-image" ) From 48e1cbc242893d2be5dc21b9d9e62e3f00c0d7a9 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 4 Jul 2026 03:31:20 +0000 Subject: [PATCH 4/7] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/routes/models.py | 4 +--- studio/backend/tests/test_diffusion_dataset_api.py | 6 ++++-- studio/backend/tests/test_diffusion_dit_trainer.py | 5 ++++- studio/backend/tests/test_diffusion_training.py | 6 ++++-- studio/backend/tests/test_local_model_format.py | 13 +++++++++---- 5 files changed, 22 insertions(+), 12 deletions(-) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 54bc490052..d4501517c8 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -911,9 +911,7 @@ async def list_local_models( models = collect_local_models(models_root) # Tag each model with its task so the Images picker can filter to diffusion # (GGUF by architecture; local diffusers checkpoints by pipeline / family). - models = [ - m.model_copy(update = {"task": _local_model_task(m)}) for m in models - ] + models = [m.model_copy(update = {"task": _local_model_task(m)}) for m in models] return LocalModelListResponse( models_dir = str(models_root), diff --git a/studio/backend/tests/test_diffusion_dataset_api.py b/studio/backend/tests/test_diffusion_dataset_api.py index e1c2c0a186..b10dd04f4c 100644 --- a/studio/backend/tests/test_diffusion_dataset_api.py +++ b/studio/backend/tests/test_diffusion_dataset_api.py @@ -64,8 +64,10 @@ def test_list_images_caption_precedence(client, ds_root): # a.png -> sidecar (an explicit edit beats the metadata row), b.png -> metadata-only, # c.png -> none. (folder / "metadata.jsonl").write_text( - json.dumps({"file_name": "a.png", "text": "from metadata"}) + "\n" - + json.dumps({"file_name": "b.png", "text": "from metadata"}) + "\n", + json.dumps({"file_name": "a.png", "text": "from metadata"}) + + "\n" + + json.dumps({"file_name": "b.png", "text": "from metadata"}) + + "\n", encoding = "utf-8", ) (folder / "a.txt").write_text("edited sidecar", encoding = "utf-8") diff --git a/studio/backend/tests/test_diffusion_dit_trainer.py b/studio/backend/tests/test_diffusion_dit_trainer.py index a727f0f13b..c40a0575a3 100644 --- a/studio/backend/tests/test_diffusion_dit_trainer.py +++ b/studio/backend/tests/test_diffusion_dit_trainer.py @@ -61,7 +61,10 @@ def test_select_lora_targets_explicit_override_wins(): base_model = "black-forest-labs/FLUX.1-dev", data_dir = "d", output_dir = "o" ).normalized() assert cfg.lora_target_modules == DEFAULT_LORA_TARGETS - assert _select_lora_targets(cfg.lora_target_modules, _SPECS["flux.1"].lora_targets) == _FLUX_TARGETS + assert ( + _select_lora_targets(cfg.lora_target_modules, _SPECS["flux.1"].lora_targets) + == _FLUX_TARGETS + ) @pytest.mark.parametrize( diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index 1efd19f18f..517186a196 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -547,8 +547,10 @@ def test_diffusion_info_counts_metadata_captions(client, dataset_roots): (folder / "c.png").write_bytes(b"x") # a.png + b.png via metadata; a.png also has a sidecar (must count once); c.png none. (folder / "metadata.jsonl").write_text( - json.dumps({"file_name": "a.png", "text": "cap a"}) + "\n" - + json.dumps({"file_name": "b.png", "text": "cap b"}) + "\n", + json.dumps({"file_name": "a.png", "text": "cap a"}) + + "\n" + + json.dumps({"file_name": "b.png", "text": "cap b"}) + + "\n", encoding = "utf-8", ) (folder / "a.txt").write_text("edited a", encoding = "utf-8") diff --git a/studio/backend/tests/test_local_model_format.py b/studio/backend/tests/test_local_model_format.py index dd888eb50d..f9184774d6 100644 --- a/studio/backend/tests/test_local_model_format.py +++ b/studio/backend/tests/test_local_model_format.py @@ -121,7 +121,14 @@ def test_scan_models_dir_classifies_root_gguf_with_config(tmp_path): from models.models import LocalModelInfo # noqa: E402 -def _local(path, *, model_format = None, model_id = None, display_name = "m", id = "m"): +def _local( + path, + *, + model_format = None, + model_id = None, + display_name = "m", + id = "m", +): return LocalModelInfo( id = id, display_name = display_name, @@ -147,9 +154,7 @@ def test_local_task_tags_diffusers_by_family_id(tmp_path): d = tmp_path / "flux-checkpoint" _touch(d / "flux1-dev.safetensors") assert ( - models_route._local_model_task( - _local(d, model_id = "black-forest-labs/FLUX.1-dev") - ) + models_route._local_model_task(_local(d, model_id = "black-forest-labs/FLUX.1-dev")) == "text-to-image" ) From 8ad8a587429acb8f8dda5911d32a484e1d5a2178 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 4 Jul 2026 04:31:04 +0000 Subject: [PATCH 5/7] Add grad norm chart, clearer completion state, Windows caption keys, GGUF compute copy - Trainers emit the pre-clip gradient norm; the service keeps a bounded grad_norm history and the Train tab renders a Grad Norm chart next to Loss and LR - Completed runs show 'Training complete' with a celebratory marker in the success color instead of a plain status word - metadata.jsonl caption keys now match on Windows (as_posix relative paths) in both the trainer discovery and the dataset image records - RMSNorm eager patch skips installation on torch builds without F.rms_norm instead of failing at forward time - GGUF compute description no longer says the GGUF is dequantised: the INT8/FP8/FP4 modes load the base model's bf16 transformer and quantise that directly; label no longer wraps in the Advanced panel --- .../core/inference/diffusion_eager_patches.py | 5 +++ .../core/training/diffusion_dit_trainer.py | 6 ++- .../core/training/diffusion_lora_trainer.py | 5 ++- .../core/training/diffusion_train_common.py | 5 ++- .../training/diffusion_training_service.py | 45 ++++++++++++++----- studio/backend/models/training.py | 9 +++- studio/backend/routes/training.py | 9 ++++ studio/frontend/src/features/images/api.ts | 1 + .../src/features/images/images-page.tsx | 4 +- .../images/train/diffusion-charts.tsx | 36 +++++++++++++-- .../images/train/diffusion-train-panel.tsx | 25 ++++++++++- 11 files changed, 126 insertions(+), 24 deletions(-) diff --git a/studio/backend/core/inference/diffusion_eager_patches.py b/studio/backend/core/inference/diffusion_eager_patches.py index 593f545984..fff77eb139 100644 --- a/studio/backend/core/inference/diffusion_eager_patches.py +++ b/studio/backend/core/inference/diffusion_eager_patches.py @@ -181,6 +181,11 @@ def install_compile_safe_patches() -> int: for cls, new_fn in _specs(): if cls is None: continue + # torch < 2.4 has no F.rms_norm: leave diffusers' original RMSNorm.forward in + # place rather than installing a patch whose fast path would AttributeError. + if cls is _RMSNorm and not hasattr(F, "rms_norm"): + logger.info("eager-patch: skipping RMSNorm (this torch has no F.rms_norm)") + continue # Capture the live original BEFORE patching so the RMSNorm fast path can fall back # to it for the uncommon (NPU / bias / fp32-weight / tuple-dim) cases. if cls is _RMSNorm: diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index d4ddf9d6d4..4f112b0a27 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -640,8 +640,11 @@ def run_dit_lora_training( (loss / cfg.gradient_accumulation_steps).backward() step_loss += float(loss.detach()) / cfg.gradient_accumulation_steps + grad_norm: Optional[float] = None if cfg.max_grad_norm and cfg.max_grad_norm > 0: - torch.nn.utils.clip_grad_norm_(lora_params, cfg.max_grad_norm) + # clip_grad_norm_ returns the PRE-clip total norm: the signal the grad-norm + # chart wants (spikes stay visible even when clipping flattens the update). + grad_norm = float(torch.nn.utils.clip_grad_norm_(lora_params, cfg.max_grad_norm)) optimizer.step() running_loss += step_loss @@ -662,6 +665,7 @@ def run_dit_lora_training( loss = round(step_loss, 5), avg_loss = round(running_loss / done, 5), learning_rate = cfg.learning_rate, + grad_norm = round(grad_norm, 5) if grad_norm is not None else None, samples_per_second = sps, peak_memory_gb = peak_gb or None, ) diff --git a/studio/backend/core/training/diffusion_lora_trainer.py b/studio/backend/core/training/diffusion_lora_trainer.py index deef4d03ce..db2eca1011 100644 --- a/studio/backend/core/training/diffusion_lora_trainer.py +++ b/studio/backend/core/training/diffusion_lora_trainer.py @@ -338,8 +338,10 @@ def run_diffusion_lora_training( # max_grad_norm <= 0 means "disable clipping" (the Studio payload sends 0.0 for that); # passing 0.0 to clip_grad_norm_ would scale every gradient to zero (no learning). + grad_norm: Optional[float] = None if cfg.max_grad_norm and cfg.max_grad_norm > 0: - torch.nn.utils.clip_grad_norm_(lora_params, cfg.max_grad_norm) + # clip_grad_norm_ returns the PRE-clip total norm (the grad-norm chart signal). + grad_norm = float(torch.nn.utils.clip_grad_norm_(lora_params, cfg.max_grad_norm)) optimizer.step() lr_sched.step() @@ -364,6 +366,7 @@ def run_diffusion_lora_training( loss = round(step_loss, 5), avg_loss = round(running_loss / done, 5), learning_rate = lr_sched.get_last_lr()[0], + grad_norm = round(grad_norm, 5) if grad_norm is not None else None, samples_per_second = samples_per_second, peak_memory_gb = peak_gb or None, ) diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index fe077020ac..ef63735185 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -338,9 +338,10 @@ def discover_image_caption_pairs( if sidecar.is_file(): caption = sidecar.read_text(encoding = "utf-8").strip() break - # 2. metadata row keyed by file name (basename or the name as written). + # 2. metadata row keyed by file name (basename or the relative path; as_posix so a + # Windows backslash path still matches the jsonl's forward-slash keys). if caption is None: - caption = meta_caption.get(img.name) or meta_caption.get(str(img.relative_to(root))) + caption = meta_caption.get(img.name) or meta_caption.get(img.relative_to(root).as_posix()) # 3. dreambooth instance prompt. if caption is None and instance_prompt: caption = instance_prompt diff --git a/studio/backend/core/training/diffusion_training_service.py b/studio/backend/core/training/diffusion_training_service.py index 2ad30e8347..f53fd4623e 100644 --- a/studio/backend/core/training/diffusion_training_service.py +++ b/studio/backend/core/training/diffusion_training_service.py @@ -67,6 +67,7 @@ def _idle_state() -> dict[str, Any]: "loss": None, "avg_loss": None, "learning_rate": None, + "grad_norm": None, "num_images": None, "in_model_load": False, "output_dir": None, @@ -82,16 +83,19 @@ def _idle_state() -> dict[str, Any]: "metric_steps": [], "metric_loss": [], "metric_lr": [], + "metric_grad_norm": [], } -def _append_metric(state: dict[str, Any], step: Any, loss: Any, lr: Any) -> None: - """Append one (step, loss, lr) point to the bounded history arrays on ``state``. +def _append_metric(state: dict[str, Any], step: Any, loss: Any, lr: Any, grad_norm: Any = None) -> None: + """Append one (step, loss, lr, grad_norm) point to the bounded history arrays on + ``state``. Only records finite, positive-step points (mirrors the LLM trainer, which logs history only for step > 0 with a real loss). When the arrays hit ``_METRIC_CAP`` they are decimated in place (keep every other point) so appends stay bounded without losing the - curve's shape. lr may be None (kept as None so the LR series can be sparse).""" + curve's shape. lr / grad_norm may be None (kept as None so those series can be + sparse while staying index-aligned with ``steps``).""" try: istep = int(step) except (TypeError, ValueError): @@ -104,22 +108,34 @@ def _append_metric(state: dict[str, Any], step: Any, loss: Any, lr: Any) -> None return if floss != floss: # NaN guard return - flr: Optional[float] - try: - flr = float(lr) if lr is not None else None - except (TypeError, ValueError): - flr = None + + def _opt_float(v: Any) -> Optional[float]: + try: + return float(v) if v is not None else None + except (TypeError, ValueError): + return None + + flr = _opt_float(lr) + fgn = _opt_float(grad_norm) steps = state["metric_steps"] losses = state["metric_loss"] lrs = state["metric_lr"] + gns = state["metric_grad_norm"] if len(steps) >= _METRIC_CAP: state["metric_steps"] = steps[::2] state["metric_loss"] = losses[::2] state["metric_lr"] = lrs[::2] - steps, losses, lrs = state["metric_steps"], state["metric_loss"], state["metric_lr"] + state["metric_grad_norm"] = gns[::2] + steps, losses, lrs, gns = ( + state["metric_steps"], + state["metric_loss"], + state["metric_lr"], + state["metric_grad_norm"], + ) steps.append(istep) losses.append(floss) lrs.append(flr) + gns.append(fgn) class DiffusionTrainingService: @@ -288,6 +304,7 @@ class DiffusionTrainingService: loss = ev.get("loss", s["loss"]), avg_loss = ev.get("avg_loss", s["avg_loss"]), learning_rate = ev.get("learning_rate", s["learning_rate"]), + grad_norm = ev.get("grad_norm", s["grad_norm"]), message = "Training...", ) # Fold optional perf fields (emitted by the trainers) so the UI can show @@ -296,8 +313,14 @@ class DiffusionTrainingService: s["samples_per_second"] = ev.get("samples_per_second") if ev.get("peak_memory_gb") is not None: s["peak_memory_gb"] = ev.get("peak_memory_gb") - # Retain a bounded (step, loss, lr) history for the live loss chart. - _append_metric(s, ev.get("step"), ev.get("loss"), ev.get("learning_rate")) + # Retain a bounded (step, loss, lr, grad_norm) history for the live charts. + _append_metric( + s, + ev.get("step"), + ev.get("loss"), + ev.get("learning_rate"), + ev.get("grad_norm"), + ) elif etype == "complete": # Reset in_model_load: a stop during model load emits complete without a # preceding model_load_completed, which would otherwise leave a stale diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 2479571357..965a854042 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -725,12 +725,14 @@ class DiffusionTrainingStartResponse(BaseModel): class DiffusionMetricHistory(BaseModel): - """Paired step-indexed history arrays for the live training charts. ``lr`` entries may - be null so a sparse learning-rate series still aligns with ``steps`` by index.""" + """Paired step-indexed history arrays for the live training charts. ``lr`` and + ``grad_norm`` entries may be null so those sparse series still align with ``steps`` + by index.""" steps: List[int] = Field(default_factory = list) loss: List[float] = Field(default_factory = list) lr: List[Optional[float]] = Field(default_factory = list) + grad_norm: List[Optional[float]] = Field(default_factory = list) class DiffusionTrainingStatusResponse(BaseModel): @@ -745,6 +747,9 @@ class DiffusionTrainingStatusResponse(BaseModel): loss: Optional[float] = None avg_loss: Optional[float] = None learning_rate: Optional[float] = None + # Pre-clip gradient norm from the trainer's progress events (None when clipping is + # disabled), feeding the grad-norm chart. + grad_norm: Optional[float] = None num_images: Optional[int] = None in_model_load: bool = False output_dir: Optional[str] = None diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index e61042f1f6..54324dd544 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -1267,6 +1267,7 @@ async def diffusion_training_status(current_subject: str = Depends(get_current_s steps = snap.pop("metric_steps", []), loss = snap.pop("metric_loss", []), lr = snap.pop("metric_lr", []), + grad_norm = snap.pop("metric_grad_norm", []), ) return DiffusionTrainingStatusResponse(**snap, metric_history = metric_history) @@ -1503,7 +1504,15 @@ def _image_record( caption = None break if caption is None: + # Basename first, then the relative path as written in the jsonl (as_posix so a + # Windows backslash path still matches forward-slash keys) -- the same lookup + # order discover_image_caption_pairs uses. meta = meta_captions.get(image_path.name) + if meta is None: + try: + meta = meta_captions.get(image_path.relative_to(folder).as_posix()) + except ValueError: + meta = None if meta is not None: caption = meta source = "metadata" diff --git a/studio/frontend/src/features/images/api.ts b/studio/frontend/src/features/images/api.ts index 911a2ff426..6ff2c72b77 100644 --- a/studio/frontend/src/features/images/api.ts +++ b/studio/frontend/src/features/images/api.ts @@ -291,6 +291,7 @@ export interface DiffusionMetricHistory { steps: number[]; loss: number[]; lr: Array; + grad_norm: Array; } // A snapshot of the current diffusion training job (GET /api/train/diffusion/status). diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index e712f24c21..1b1d2c366d 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -481,7 +481,7 @@ function AdvancedSelect({ return (
- + {label} {hint && {hint}} @@ -1872,7 +1872,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) { {!status?.loaded || status.model_kind === "gguf" ? ( setTransformerQuant(v as typeof transformerQuant)} diff --git a/studio/frontend/src/features/images/train/diffusion-charts.tsx b/studio/frontend/src/features/images/train/diffusion-charts.tsx index c4378156cf..a59dc3f7fb 100644 --- a/studio/frontend/src/features/images/train/diffusion-charts.tsx +++ b/studio/frontend/src/features/images/train/diffusion-charts.tsx @@ -9,6 +9,8 @@ import type { TrainingSeriesPoint } from "@/features/training"; // which are meaningless for diffusion LoRA training and showed as an empty card and an // "Evaluation not configured" placeholder. This is a diffusion-only two-card layout. // eslint-disable-next-line no-restricted-imports +import { GradNormChartCard } from "@/features/studio/sections/charts/grad-norm-chart-card"; +// eslint-disable-next-line no-restricted-imports import { LearningRateChartCard } from "@/features/studio/sections/charts/learning-rate-chart-card"; // eslint-disable-next-line no-restricted-imports import { TrainingLossChartCard } from "@/features/studio/sections/charts/training-loss-chart-card"; @@ -43,14 +45,17 @@ function fullStepDomain(steps: number[]): [number, number] { return [min, max]; } -// A diffusion-only metrics view: just Training Loss and Learning Rate, side by side, with a -// note under the loss card explaining why per-step loss looks noisy. +// A diffusion-only metrics view: Training Loss and Learning Rate side by side, plus Grad +// Norm (the pre-clip total gradient norm; spikes flag instability that raw loss noise +// hides), with a note under the loss card explaining why per-step loss looks noisy. export function DiffusionCharts({ lossHistory, lrHistory, + gradNormHistory = [], }: { lossHistory: TrainingSeriesPoint[]; lrHistory: TrainingSeriesPoint[]; + gradNormHistory?: TrainingSeriesPoint[]; }): ReactElement | null { const lossItems = useMemo(() => toLossItems(lossHistory), [lossHistory]); const smoothed = useMemo( @@ -82,12 +87,24 @@ export function DiffusionCharts({ [lrHistory], ); + const gradNormData = useMemo( + () => + compressSeries( + gradNormHistory + .filter((p) => Number.isFinite(p.value)) + .map((p) => ({ step: p.step, gradNorm: p.value, displayGradNorm: p.value })), + MAX_RENDER_POINTS, + ), + [gradNormHistory], + ); + const steps = useMemo(() => { const set = new Set(); for (const p of lossData) set.add(p.step); for (const p of lrData) set.add(p.step); + for (const p of gradNormData) set.add(p.step); return Array.from(set).sort((a, b) => a - b); - }, [lossData, lrData]); + }, [lossData, lrData, gradNormData]); const stepDomain = useMemo(() => fullStepDomain(steps), [steps]); const xAxisTicks = useMemo( @@ -103,6 +120,10 @@ export function DiffusionCharts({ () => buildYDomain(lrData.map((p) => p.displayLr)), [lrData], ); + const gradNormDomain = useMemo( + () => buildYDomain(gradNormData.map((p) => p.displayGradNorm)), + [gradNormData], + ); const avgRaw = lossItems.length > 0 @@ -138,6 +159,15 @@ export function DiffusionCharts({ xAxisTicks={xAxisTicks} scale="linear" /> + {gradNormData.length > 0 && ( + + )}
); } diff --git a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx index 45df1d50a6..0f67a3aebe 100644 --- a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx +++ b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx @@ -359,6 +359,13 @@ export function DiffusionTrainPanel({ .map((step, i) => ({ step, value: h.lr[i] })) .filter((p): p is TrainingSeriesPoint => p.value != null); }, [status?.metric_history]); + const gradNormHistory: TrainingSeriesPoint[] = useMemo(() => { + const h = status?.metric_history; + if (!h) return []; + return h.steps + .map((step, i) => ({ step, value: h.grad_norm?.[i] ?? null })) + .filter((p): p is TrainingSeriesPoint => p.value != null); + }, [status?.metric_history]); const onUpload = useCallback(async () => { const files = Array.from(fileInputRef.current?.files ?? []); @@ -784,7 +791,17 @@ export function DiffusionTrainPanel({ <>
- {status.status} + {/* A finished run should be unmistakable at a glance, so completed swaps + the plain status word for a celebratory line in the success color. */} + + {status.status === "completed" ? "Training complete \u{1F389}" : status.status} + {status.total_steps > 0 ? `${status.step}/${status.total_steps} steps` : ""} @@ -821,7 +838,11 @@ export function DiffusionTrainPanel({ )}
- + {completed && (
From 2c5955bda82ac679656fbe95f1f33e6874b8930e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 4 Jul 2026 04:31:07 +0000 Subject: [PATCH 6/7] Coerce num_epochs in normalized() and use utf-8 for run records num_epochs was only int-coerced for the range check, so a string value from a dict-built config would reach resolve_train_steps' arithmetic; normalized() now stores the coerced int. Run record reads/writes pass encoding utf-8 explicitly so non-ASCII prompts survive on Windows --- studio/backend/core/training/diffusion_train_common.py | 1 + studio/backend/core/training/diffusion_training_service.py | 6 +++--- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index e291c8c0fc..fc46898eb7 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -407,6 +407,7 @@ class DiffusionLoraConfig: lora_target_modules = targets, max_grad_norm = float(self.max_grad_norm), hf_token = token or None, + num_epochs = int(self.num_epochs), cache_variants = int(self.cache_variants), compile_transformer = compile_transformer, base_precision = base_precision, diff --git a/studio/backend/core/training/diffusion_training_service.py b/studio/backend/core/training/diffusion_training_service.py index 23c0cd2bbd..43946acefe 100644 --- a/studio/backend/core/training/diffusion_training_service.py +++ b/studio/backend/core/training/diffusion_training_service.py @@ -99,7 +99,7 @@ def list_diffusion_runs(limit: int = 20) -> list[dict]: out: list[dict] = [] for p in files[: max(0, int(limit))]: try: - rec = json.loads(p.read_text()) + rec = json.loads(p.read_text(encoding = "utf-8")) except Exception: # noqa: BLE001 -- a corrupt record never breaks the listing continue # A valid-JSON file with the wrong shape (an old or hand-edited record that is not a @@ -124,7 +124,7 @@ def get_diffusion_run(job_id: str) -> Optional[dict]: return None p = _runs_dir() / f"{job_id}.json" try: - return json.loads(p.read_text()) + return json.loads(p.read_text(encoding = "utf-8")) except Exception: # noqa: BLE001 -- missing/corrupt record return None @@ -403,7 +403,7 @@ class DiffusionTrainingService: }, } path = _runs_dir() / f"{s['job_id']}.json" - path.write_text(json.dumps(record)) + path.write_text(json.dumps(record), encoding = "utf-8") except Exception: # noqa: BLE001 -- persisting history must never break the run pass From ff59b435338080fac476b8a73afff8f20ae226db Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 4 Jul 2026 04:31:09 +0000 Subject: [PATCH 7/7] Fail clearly when a local Krea 2 dir lacks model_index.json Falling through to hf_hub_download with a filesystem path as the repo id raised an opaque HFValidationError; a local dir without the file now raises FileNotFoundError naming the directory --- studio/backend/core/inference/diffusion_krea2.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/studio/backend/core/inference/diffusion_krea2.py b/studio/backend/core/inference/diffusion_krea2.py index 4475559cdb..86a1d0413c 100644 --- a/studio/backend/core/inference/diffusion_krea2.py +++ b/studio/backend/core/inference/diffusion_krea2.py @@ -85,12 +85,20 @@ def load_krea2_text_encoder( def _load_model_index(repo_id: str, hf_token: Optional[str] = None) -> dict[str, Any]: """model_index.json as a dict, from a local path or the Hub cache.""" + is_local_dir = False try: - local = Path(repo_id).expanduser() / "model_index.json" + root = Path(repo_id).expanduser() + is_local_dir = root.is_dir() + local = root / "model_index.json" if local.is_file(): return json.loads(local.read_text()) except OSError: pass + if is_local_dir: + # A local checkpoint dir without the file must fail clearly here: falling through + # to hf_hub_download with a filesystem path as the repo id would die with an + # opaque HFValidationError instead. + raise FileNotFoundError(f"model_index.json not found in local model dir {repo_id}") from huggingface_hub import hf_hub_download path = hf_hub_download(repo_id, "model_index.json", token = hf_token or None)