From bfbb902610116ea4cbd77dad22b608ba41266772 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 4 Jul 2026 03:23:05 +0000 Subject: [PATCH] Fix review findings: engine unload before training, caption precedence, dataset caption counts, local diffusers tagging, family LoRA targets - Unload the ACTIVE image engine (sd_cpp or diffusers) before diffusion training starts, not just the diffusers singleton - Count metadata.jsonl captions in dataset summaries so metadata-captioned datasets are not reported as uncaptioned - Sidecar captions now override metadata rows everywhere (grid edits win); trainer and dataset API agree - Tag local diffusers image checkpoints with text-to-image so they appear in the Images picker - Family LoRA targets (_FLUX_TARGETS etc) apply when the config carries the generic defaults; explicit overrides still win --- .../core/training/diffusion_dit_trainer.py | 17 ++++- .../core/training/diffusion_train_common.py | 24 ++++--- studio/backend/routes/models.py | 64 ++++++++++++++----- studio/backend/routes/training.py | 55 ++++++++++------ .../tests/test_diffusion_dataset_api.py | 38 ++++++++--- .../tests/test_diffusion_dit_trainer.py | 31 ++++++++- .../tests/test_diffusion_lora_trainer.py | 18 +++++- .../backend/tests/test_diffusion_training.py | 26 ++++++++ .../backend/tests/test_local_model_format.py | 45 +++++++++++++ 9 files changed, 259 insertions(+), 59 deletions(-) diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index 82b65ee5ad..d4ddf9d6d4 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -32,6 +32,7 @@ from typing import Any, Callable, Optional from core.training.diffusion_train_common import ( DEFAULT_LORA_FILENAME, + DEFAULT_LORA_TARGETS, DiffusionLoraConfig, EventCb, StopCb, @@ -58,6 +59,20 @@ _QWEN_TARGETS = _FLUX_TARGETS _ZIMAGE_TARGETS = ("to_q", "to_k", "to_v", "to_out.0") +def _select_lora_targets( + cfg_targets: tuple[str, ...], spec_targets: tuple[str, ...] +) -> tuple[str, ...]: + """Pick the LoRA target modules for a DiT run. + + ``normalized()`` always fills ``lora_target_modules`` with the generic + ``DEFAULT_LORA_TARGETS`` when a caller does not set it, so that value means "unset" + here: prefer the family's ``spec.lora_targets`` (which add the DiT-specific + projections). Any OTHER explicit tuple is a deliberate override and still wins.""" + if tuple(cfg_targets) == DEFAULT_LORA_TARGETS: + return tuple(spec_targets) + return tuple(cfg_targets) + + @dataclass class _FamilySpec: """Everything the shared loop needs that differs by family.""" @@ -513,7 +528,7 @@ def run_dit_lora_training( # 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). weight_dtype = torch.bfloat16 if device == "cuda" else torch.float32 - use_lora_targets = tuple(cfg.lora_target_modules) or spec.lora_targets + use_lora_targets = _select_lora_targets(cfg.lora_target_modules, spec.lora_targets) _assert_trusted_base_model(cfg.base_model) _assert_gated_access(cfg.base_model, cfg.hf_token) diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index 9458c4e5c9..fe077020ac 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -292,11 +292,15 @@ def discover_image_caption_pairs( """Resolve ``(image_path, caption)`` pairs from a dataset directory. Caption sources, in priority order per image: - 1. a ``metadata.jsonl`` / ``captions.jsonl`` row keyed by ``file_name`` (or ``image``) + 1. a per-image sidecar ``.txt`` / ``.caption``, + 2. a ``metadata.jsonl`` / ``captions.jsonl`` row keyed by ``file_name`` (or ``image``) carrying the caption in ``caption_column`` (default ``text``), - 2. a per-image sidecar ``.txt`` / ``.caption``, 3. ``instance_prompt`` (dreambooth) for any remaining image. + A sidecar wins over the metadata row because it is the user's explicit per-image edit + (the labeling grid writes a .txt sidecar), which must override the bulk metadata file. + Must agree with ``routes.training._image_record``, which resolves captions the same way. + Images with no caption from any source are skipped. Pure filesystem + JSON, so it is unit-testable without torch. Raises FileNotFoundError for a missing dir and ValueError when nothing is captionable. @@ -328,15 +332,15 @@ def discover_image_caption_pairs( pairs: list[tuple[str, str]] = [] for img in images: caption: Optional[str] = None - # 1. metadata row keyed by file name (basename or the name as written). - caption = meta_caption.get(img.name) or meta_caption.get(str(img.relative_to(root))) - # 2. per-image sidecar caption file. + # 1. per-image sidecar caption file (the user's explicit edit; wins over metadata). + for ext in _CAPTION_EXTS: + sidecar = img.with_suffix(ext) + 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). if caption is None: - for ext in _CAPTION_EXTS: - sidecar = img.with_suffix(ext) - if sidecar.is_file(): - caption = sidecar.read_text(encoding = "utf-8").strip() - break + caption = meta_caption.get(img.name) or meta_caption.get(str(img.relative_to(root))) # 3. dreambooth instance prompt. if caption is None and instance_prompt: caption = instance_prompt diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 65d64a135d..54bc490052 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -909,9 +909,10 @@ async def list_local_models( try: models = collect_local_models(models_root) - # Tag each GGUF with its task so the Images picker can filter to diffusion. + # 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.path, m.model_format)}) for m in models + m.model_copy(update = {"task": _local_model_task(m)}) for m in models ] return LocalModelListResponse( @@ -3229,24 +3230,55 @@ def _repo_gguf_task(repo_info) -> Optional[str]: return None -def _local_model_task(path: str, model_format: Optional[str]) -> Optional[str]: - """Same classification for a local model: read its GGUF architecture. The - path may be the .gguf file itself or a folder containing one.""" - if model_format != "gguf": +def _local_model_task(model: "LocalModelInfo") -> Optional[str]: + """Classify a local model into an HF pipeline task so the Images picker can filter. + + For a GGUF, read its architecture (the path may be the .gguf file itself or a folder + containing one). For a local non-GGUF image checkpoint (a diffusers pipeline dir or a + single-file safetensors), fall through to the diffusers detection so on-device image + models get the 'text-to-image' tag instead of being dropped as task=null; the load + path accepts these as a local pipeline.""" + path = model.path + if model.model_format == "gguf": + try: + p = Path(path) + if p.suffix.lower() == ".gguf" and p.is_file(): + return _arch_to_task(_gguf_architecture(str(p))) + for f in _iter_gguf_paths(p): + if _is_mmproj_filename(f.name): + continue + task = _arch_to_task(_gguf_architecture(str(f))) + if task is not None: + return task + except Exception: + pass return None + if _local_is_diffusers(model): + return "text-to-image" + return None + + +def _local_is_diffusers(model: "LocalModelInfo") -> bool: + """True for a local diffusers image checkpoint, mirroring the cached-repo + ``_repo_is_diffusers`` heuristics: a full pipeline carries a top-level + ``model_index.json``, while single-file / safetensors image checkpoints ship none, so + fall back to the model id resolving to a known diffusion family (the same resolver the + Images backend loads from). Family detection uses the clean model id / name, not the + on-disk path, so a parent directory keyword can't spuriously match.""" try: - p = Path(path) - if p.suffix.lower() == ".gguf" and p.is_file(): - return _arch_to_task(_gguf_architecture(str(p))) - for f in _iter_gguf_paths(p): - if _is_mmproj_filename(f.name): - continue - task = _arch_to_task(_gguf_architecture(str(f))) - if task is not None: - return task + p = Path(model.path) + if p.is_dir() and (p / "model_index.json").is_file(): + return True except Exception: pass - return None + try: + from core.inference.diffusion_families import detect_family + for needle in (model.model_id, model.display_name, model.id): + if needle and detect_family(needle) is not None: + return True + except Exception: + pass + return False @router.get("/cached-gguf") diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index a5b4c3b06e..e61042f1f6 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -1111,9 +1111,13 @@ def _free_gpu_for_diffusion_training() -> None: try: from core.inference import gpu_arbiter - from core.inference.diffusion import get_diffusion_backend + from core.inference.diffusion_engine_router import get_active_diffusion_engine - diffusion = get_diffusion_backend() + # The ACTIVE engine, not the diffusers singleton: on a native (sd_cpp) + # selection the diffusers backend reports unloaded while the resident + # sd-server still holds the GPU, so unloading only the singleton is a no-op. + # Mirrors the LLM training start path. + diffusion = get_active_diffusion_engine() if diffusion.is_loaded: logger.info("Unloading resident Images pipeline to free GPU memory for training") diffusion.unload() # no-op when nothing is loaded; also preempts an in-flight load @@ -1274,14 +1278,19 @@ _DIFFUSION_DATASET_TEXT_EXTS = {".txt", ".caption", ".jsonl"} def _diffusion_dataset_summary(folder: Path) -> DiffusionDatasetSummary: + # Count an image as captioned when a metadata/captions.jsonl row or a per-image + # sidecar (.txt / .caption) resolves a caption for it -- the same sources the + # trainer reads. Counting metadata-only captions here keeps a metadata-captioned + # dataset from reporting caption_count=0 and being treated as uncaptioned. + meta_captions = _load_metadata_captions(folder) images = captions = 0 for f in folder.iterdir(): - if not f.is_file(): + if not f.is_file() or f.suffix.lower() not in _DIFFUSION_DATASET_IMAGE_EXTS: continue - ext = f.suffix.lower() - if ext in _DIFFUSION_DATASET_IMAGE_EXTS: - images += 1 - elif ext in (".txt", ".caption"): + images += 1 + if f.name in meta_captions or any( + f.with_suffix(ext).is_file() for ext in (".txt", ".caption") + ): captions += 1 return DiffusionDatasetSummary( name = folder.name, path = str(folder), image_count = images, caption_count = captions @@ -1478,20 +1487,26 @@ def _load_metadata_captions(folder: Path) -> dict[str, str]: def _image_record( folder: Path, image_path: Path, meta_captions: dict[str, str] ) -> DiffusionDatasetImageRecord: - """Build one image record, resolving its caption with metadata > sidecar precedence - (the same order the trainer uses).""" - caption: Optional[str] = meta_captions.get(image_path.name) - source = "metadata" if caption is not None else "none" + """Build one image record, resolving its caption with sidecar > metadata precedence + (the same order the trainer uses). A per-image .txt / .caption sidecar wins because + it is the user's explicit edit from the labeling grid, which must override a + metadata.jsonl / captions.jsonl row for the image.""" + caption: Optional[str] = None + source = "none" + for ext in (".txt", ".caption"): + sidecar = image_path.with_suffix(ext) + if sidecar.is_file(): + try: + caption = sidecar.read_text(encoding = "utf-8").strip() + source = "sidecar" + except OSError: + caption = None + break if caption is None: - for ext in (".txt", ".caption"): - sidecar = image_path.with_suffix(ext) - if sidecar.is_file(): - try: - caption = sidecar.read_text(encoding = "utf-8").strip() - source = "sidecar" - except OSError: - caption = None - break + meta = meta_captions.get(image_path.name) + if meta is not None: + caption = meta + source = "metadata" try: size_bytes = image_path.stat().st_size except OSError: diff --git a/studio/backend/tests/test_diffusion_dataset_api.py b/studio/backend/tests/test_diffusion_dataset_api.py index b7a895281a..e1c2c0a186 100644 --- a/studio/backend/tests/test_diffusion_dataset_api.py +++ b/studio/backend/tests/test_diffusion_dataset_api.py @@ -61,12 +61,14 @@ def test_list_images_caption_precedence(client, ds_root): _write_png(folder / "a.png") _write_png(folder / "b.png") _write_png(folder / "c.png") - # a.png -> metadata (beats a stray sidecar), b.png -> sidecar, c.png -> none. + # 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", encoding = "utf-8" + 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("stray sidecar", encoding = "utf-8") - (folder / "b.txt").write_text("from sidecar", encoding = "utf-8") + (folder / "a.txt").write_text("edited sidecar", encoding = "utf-8") r = client.get("/api/train/diffusion/dataset/styleset/images") assert r.status_code == 200, r.text @@ -74,10 +76,11 @@ def test_list_images_caption_precedence(client, ds_root): assert body["name"] == "styleset" recs = {rec["filename"]: rec for rec in body["images"]} assert set(recs) == {"a.png", "b.png", "c.png"} - assert recs["a.png"]["caption"] == "from metadata" - assert recs["a.png"]["caption_source"] == "metadata" - assert recs["b.png"]["caption"] == "from sidecar" - assert recs["b.png"]["caption_source"] == "sidecar" + # A sidecar edit overrides the metadata row for the same image. + assert recs["a.png"]["caption"] == "edited sidecar" + assert recs["a.png"]["caption_source"] == "sidecar" + assert recs["b.png"]["caption"] == "from metadata" + assert recs["b.png"]["caption_source"] == "metadata" assert recs["c.png"]["caption"] is None assert recs["c.png"]["caption_source"] == "none" assert recs["a.png"]["width"] == 8 and recs["a.png"]["height"] == 8 @@ -133,6 +136,25 @@ def test_put_caption_roundtrip_and_clear(client, ds_root): assert not (folder / "x.txt").exists() +def test_put_caption_overrides_metadata_row(client, ds_root): + # Editing a caption for an image that already has a metadata.jsonl row must take + # effect: the sidecar edit wins over the metadata caption in the response (and in + # the data the trainer reads), not the other way round. + folder = ds_root / "cap" + folder.mkdir() + _write_png(folder / "x.png") + (folder / "metadata.jsonl").write_text( + json.dumps({"file_name": "x.png", "text": "from metadata"}) + "\n", encoding = "utf-8" + ) + + r = client.put( + "/api/train/diffusion/dataset/cap/caption/x.png", json = {"caption": "edited caption"} + ) + assert r.status_code == 200, r.text + assert r.json()["caption"] == "edited caption" + assert r.json()["caption_source"] == "sidecar" + + def test_put_caption_missing_image_404(client, ds_root): (ds_root / "cap").mkdir() r = client.put("/api/train/diffusion/dataset/cap/caption/ghost.png", json = {"caption": "hi"}) diff --git a/studio/backend/tests/test_diffusion_dit_trainer.py b/studio/backend/tests/test_diffusion_dit_trainer.py index ba54480003..e1aa7868a3 100644 --- a/studio/backend/tests/test_diffusion_dit_trainer.py +++ b/studio/backend/tests/test_diffusion_dit_trainer.py @@ -12,13 +12,21 @@ from __future__ import annotations import pytest from core.training.diffusion_dit_trainer import ( + _FLUX_TARGETS, _GATED_TRAIN_REPOS, + _QWEN_TARGETS, _SPECS, + _ZIMAGE_TARGETS, _assert_gated_access, _repo_is_prequantized, + _select_lora_targets, run_dit_lora_training, ) -from core.training.diffusion_train_common import DiffusionLoraConfig, family_train_infos +from core.training.diffusion_train_common import ( + DEFAULT_LORA_TARGETS, + DiffusionLoraConfig, + family_train_infos, +) def test_specs_cover_the_three_dit_families(): @@ -32,6 +40,27 @@ def test_specs_cover_the_three_dit_families(): assert _SPECS["qwen-image"].force_bf16 is True +def test_select_lora_targets_uses_family_default_for_generic_config(): + # normalized() fills lora_target_modules with the generic DEFAULT_LORA_TARGETS when a + # caller doesn't set it, so that value must resolve to the family's targets (which add + # the DiT-specific projections), not stay stuck on the generic SDXL list. + assert _select_lora_targets(DEFAULT_LORA_TARGETS, _FLUX_TARGETS) == _FLUX_TARGETS + assert _select_lora_targets(DEFAULT_LORA_TARGETS, _QWEN_TARGETS) == _QWEN_TARGETS + assert _select_lora_targets(DEFAULT_LORA_TARGETS, _ZIMAGE_TARGETS) == _ZIMAGE_TARGETS + + +def test_select_lora_targets_explicit_override_wins(): + # Any OTHER explicit tuple is a deliberate override and must win over the family spec. + override = ("to_q", "to_k") + assert _select_lora_targets(override, _FLUX_TARGETS) == override + # The default request path (config carrying the generic default) reaches the spec. + cfg = DiffusionLoraConfig( + 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 + + @pytest.mark.parametrize( "repo, expected", [ diff --git a/studio/backend/tests/test_diffusion_lora_trainer.py b/studio/backend/tests/test_diffusion_lora_trainer.py index ff8f3ec89e..75046165d0 100644 --- a/studio/backend/tests/test_diffusion_lora_trainer.py +++ b/studio/backend/tests/test_diffusion_lora_trainer.py @@ -27,15 +27,15 @@ def _touch(p): p.write_bytes(b"") -def test_discover_prefers_metadata_then_sidecar_then_instance(tmp_path): +def test_discover_prefers_sidecar_then_metadata_then_instance(tmp_path): _touch(tmp_path / "a.png") _touch(tmp_path / "b.jpg") _touch(tmp_path / "c.webp") - # a.png captioned via metadata.jsonl + # a.png captioned via metadata.jsonl only (tmp_path / "metadata.jsonl").write_text( json.dumps({"file_name": "a.png", "text": "from metadata"}) + "\n", encoding = "utf-8" ) - # b.jpg captioned via sidecar + # b.jpg captioned via sidecar only (tmp_path / "b.txt").write_text("from sidecar", encoding = "utf-8") # c.webp falls back to the instance prompt pairs = dict(discover_image_caption_pairs(tmp_path, instance_prompt = "from instance")) @@ -44,6 +44,18 @@ def test_discover_prefers_metadata_then_sidecar_then_instance(tmp_path): assert pairs[str(tmp_path / "c.webp")] == "from instance" +def test_discover_sidecar_overrides_metadata_row(tmp_path): + # A per-image sidecar is the user's explicit edit and must win over a metadata row + # for the same image (the labeling grid writes sidecars). + _touch(tmp_path / "a.png") + (tmp_path / "metadata.jsonl").write_text( + json.dumps({"file_name": "a.png", "text": "from metadata"}) + "\n", encoding = "utf-8" + ) + (tmp_path / "a.txt").write_text("edited sidecar", encoding = "utf-8") + pairs = dict(discover_image_caption_pairs(tmp_path)) + assert pairs[str(tmp_path / "a.png")] == "edited sidecar" + + def test_discover_skips_uncaptioned_without_instance_prompt(tmp_path): _touch(tmp_path / "cap.png") _touch(tmp_path / "nocap.png") diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index 886a3e5d4f..a7a1006e25 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -11,6 +11,7 @@ service, so wiring / validation / error mapping are covered without a GPU. from __future__ import annotations +import json import queue as _queue import threading import time @@ -417,6 +418,31 @@ def test_diffusion_info_lists_image_dataset_folders(client, dataset_roots): assert body["datasets"][0]["caption_count"] == 1 +def test_diffusion_info_counts_metadata_captions(client, dataset_roots): + # A dataset captioned via metadata.jsonl must not report caption_count=0 (which the + # Train UI treats as uncaptioned); metadata rows count like sidecars, without + # double-counting an image that has both. + ds_root, _ = dataset_roots + folder = ds_root / "meta-captioned" + folder.mkdir() + (folder / "a.png").write_bytes(b"x") + (folder / "b.png").write_bytes(b"x") + (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", + encoding = "utf-8", + ) + (folder / "a.txt").write_text("edited a", encoding = "utf-8") + + r = client.get("/api/train/diffusion/info") + assert r.status_code == 200, r.text + summary = next(d for d in r.json()["datasets"] if d["name"] == "meta-captioned") + assert summary["image_count"] == 3 + assert summary["caption_count"] == 2 + + def test_diffusion_dataset_upload_accumulates(client, dataset_roots): ds_root, _ = dataset_roots files = [ diff --git a/studio/backend/tests/test_local_model_format.py b/studio/backend/tests/test_local_model_format.py index b569163cc8..dd888eb50d 100644 --- a/studio/backend/tests/test_local_model_format.py +++ b/studio/backend/tests/test_local_model_format.py @@ -115,3 +115,48 @@ def test_scan_models_dir_classifies_root_gguf_with_config(tmp_path): assert row.path == str(root) assert row.model_format == "gguf" + + +# ── Images picker task tag for local (non-GGUF) diffusers models ────────────── +from models.models import LocalModelInfo # noqa: E402 + + +def _local(path, *, model_format = None, model_id = None, display_name = "m", id = "m"): + return LocalModelInfo( + id = id, + display_name = display_name, + path = str(path), + source = "models_dir", + model_id = model_id, + model_format = model_format, + ) + + +def test_local_task_tags_diffusers_pipeline_dir(tmp_path): + # A local diffusers pipeline (top-level model_index.json) is an image model even + # though its model_format is not "gguf": tag it so the Images picker keeps it. + d = tmp_path / "my-local-pipeline" + _touch(d / "model_index.json") + _touch(d / "unet" / "diffusion_pytorch_model.safetensors") + assert models_route._local_model_task(_local(d)) == "text-to-image" + + +def test_local_task_tags_diffusers_by_family_id(tmp_path): + # A single-file / safetensors image checkpoint ships no model_index.json; fall back + # to the model id resolving to a known diffusion family. + 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") + ) + == "text-to-image" + ) + + +def test_local_task_none_for_plain_llm(tmp_path): + # A plain non-GGUF LLM checkpoint (no pipeline, no image family) stays untagged. + d = tmp_path / "llama" + _touch(d / "config.json") + _touch(d / "model.safetensors") + assert models_route._local_model_task(_local(d, model_id = "meta-llama/Llama-3.1-8B")) is None