From 07f27b23e8fd27dc4464d3d08248cba7c80c6343 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:14 +0000 Subject: [PATCH 1/4] [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 a7a1006e25..a7ef2240f4 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -430,8 +430,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 285c8fbd2083b79cfe001e02ed93a141cfbc535e 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 04:40:16 +0000 Subject: [PATCH 2/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/training/diffusion_train_common.py | 4 +++- .../backend/core/training/diffusion_training_service.py | 8 +++++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index ee8dd85c6f..ce31b1652f 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -360,7 +360,9 @@ def discover_image_caption_pairs( # 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(img.relative_to(root).as_posix()) + 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 b8d8aa27bf..da28f2d588 100644 --- a/studio/backend/core/training/diffusion_training_service.py +++ b/studio/backend/core/training/diffusion_training_service.py @@ -103,7 +103,13 @@ def _idle_state() -> dict[str, Any]: } -def _append_metric(state: dict[str, Any], step: Any, loss: Any, lr: Any, grad_norm: Any = None) -> None: +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``. From 1d3aa53d1f48069cf6224826db0bee9bb465c06e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 4 Jul 2026 05:01:58 +0000 Subject: [PATCH 3/4] Validate the training config before importing diffusers and pin the arbiter test's device The fp16-on-bf16-family refusal in run_dit_lora_training now fires before the heavy imports, so a host without diffusers gets the real validation error instead of ModuleNotFoundError. test_in_progress_returns_409_after_validation_passes pins the resolved device to cuda because the load route only takes the GPU arbiter for non-CPU loads, which made the ownership assert host-dependent. --- .../core/training/diffusion_dit_trainer.py | 29 ++++++++++--------- studio/backend/tests/test_diffusion_routes.py | 11 +++++++ 2 files changed, 26 insertions(+), 14 deletions(-) diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index 4f112b0a27..50bb0f5d3b 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -487,6 +487,21 @@ def run_dit_lora_training( should_stop: Optional[StopCb] = None, ) -> str: """Train a flow-matching DiT LoRA (FLUX.1-dev / Qwen-Image / Z-Image) and export it.""" + cfg = config.normalized() + spec = _SPECS.get(cfg.resolved_family) + if spec is None: + raise ValueError(f"No DiT trainer for family {cfg.resolved_family!r}") + + # DiT families train in bf16 (Z-Image/Qwen require it; FLUX prefers it). A caller that + # explicitly asks for fp16 on a bf16-only family is refused rather than silently + # upgraded, so the choice is never misrepresented. Validation runs before the heavy + # imports so a host without diffusers still sees the real error. + if cfg.mixed_precision == "fp16" and spec.force_bf16: + raise ValueError( + f"{spec.family} LoRA training requires bf16: fp16 overflows its fp32 RoPE / " + f"embedder internals. Set mixed precision to bf16." + ) + import torch import torch.nn.functional as F from diffusers import FlowMatchEulerDiscreteScheduler @@ -494,11 +509,6 @@ def run_dit_lora_training( from peft import LoraConfig from peft.utils import get_peft_model_state_dict - cfg = config.normalized() - spec = _SPECS.get(cfg.resolved_family) - if spec is None: - raise ValueError(f"No DiT trainer for family {cfg.resolved_family!r}") - rng = random.Random(cfg.seed) torch.manual_seed(cfg.seed) @@ -515,15 +525,6 @@ def run_dit_lora_training( save_on_stop = False return True - # DiT families train in bf16 (Z-Image/Qwen require it; FLUX prefers it). A caller that - # explicitly asks for fp16 on a bf16-only family is refused rather than silently - # upgraded, so the choice is never misrepresented. - if cfg.mixed_precision == "fp16" and spec.force_bf16: - raise ValueError( - f"{spec.family} LoRA training requires bf16: fp16 overflows its fp32 RoPE / " - f"embedder internals. Set mixed precision to bf16." - ) - device = "cuda" if torch.cuda.is_available() else "cpu" # The flow-matching + 4-bit path is bf16 throughout (fp32 on a CPU-only box, which is # unsupported for real runs but keeps import/unit tests architecture-agnostic). diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index f6a9ea2766..567cc39c26 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -659,6 +659,17 @@ def test_in_progress_returns_409_after_validation_passes(client, monkeypatch): backend = _FakeBackend() backend.begin_load = _busy monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend) + # Pin the resolved device to cuda: the route only takes the arbiter for non-CPU + # loads, so on a CPU-only host the ownership assert below would never hold. + import types as _types + + import core.inference.diffusion_device as devmod + + monkeypatch.setattr( + devmod, + "resolve_diffusion_device_target", + lambda: _types.SimpleNamespace(device = "cuda"), + ) resp = client.post( "/api/inference/images/load", json = {"model_path": "unsloth/Z-Image-Turbo-GGUF", "gguf_filename": "q.gguf"}, From f1007fb4662636678fdc10516812cfe08f95b3e6 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 4 Jul 2026 05:05:54 +0000 Subject: [PATCH 4/4] Skip the sigma-gather test when diffusers is not installed CI runs the backend suite without diffusers; the test checks our index math against the scheduler's own gather, so it skips rather than fails there. --- studio/backend/tests/test_diffusion_train_perf.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/studio/backend/tests/test_diffusion_train_perf.py b/studio/backend/tests/test_diffusion_train_perf.py index 4339fbb0e7..713f7e4dd7 100644 --- a/studio/backend/tests/test_diffusion_train_perf.py +++ b/studio/backend/tests/test_diffusion_train_perf.py @@ -130,6 +130,9 @@ def test_zimage_collate_list(): # ── index-based sigma gather ────────────────────────────────────────────────── def test_gather_sigmas_matches_search_based_gather(): + # CI installs the backend test deps without diffusers; the scheduler math is what we + # are checking, so skip rather than fail there. + pytest.importorskip("diffusers") from diffusers import FlowMatchEulerDiscreteScheduler torch.manual_seed(0)