From bfbb902610116ea4cbd77dad22b608ba41266772 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 4 Jul 2026 03:23:05 +0000 Subject: [PATCH 1/3] 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 From 89d99e31af6958258053e22de6396b2647a76824 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 4 Jul 2026 03:23:16 +0000 Subject: [PATCH 2/3] Gate int8 and fp8 on a functional torchao import, not find_spec The Windows ROCm torchao import stub satisfies find_spec and even lets from torchao.quantization import quantize_ succeed, but its quantize_ is a no-op: auto would pick int8, leave the transformer dense, and disable compile as if it were quantized. has_functional_torchao imports the exact symbols the int8 path uses and rejects the stub via its sentinel; both the auto picker and the /info advertised modes now use it --- .../core/training/diffusion_dit_trainer.py | 12 +- .../core/training/diffusion_train_common.py | 53 ++++++++- .../tests/test_diffusion_base_precision.py | 108 +++++++++++++++--- 3 files changed, 146 insertions(+), 27 deletions(-) diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index d6288ce48c..43d6ba2a69 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -47,6 +47,7 @@ from core.training.diffusion_train_common import ( _publish_to_lora_catalog, _restore_perf_flags, discover_image_caption_pairs, + has_functional_torchao, repo_is_prequantized, ) @@ -333,11 +334,12 @@ def _resolve_base_precision(cfg, spec, device) -> str: free_gb = None capability = None has_fp8 = False - # int8 quantization has no runtime fallback, so gate the auto pick on torchao being - # importable (find_spec avoids the cost/side-effects of an actual import). - import importlib.util - - has_torchao = importlib.util.find_spec("torchao") is not None + # int8 quantization has no runtime fallback, so gate the auto pick on a FUNCTIONAL + # torchao: a plain find_spec("torchao") is satisfied by the Windows-ROCm import stub, + # whose quantize_ is a no-op that would leave the transformer dense while compile is + # disabled as if it were int8. has_functional_torchao imports the exact symbols + # _int8_quantize_base uses and rejects the stub. + has_torchao = has_functional_torchao() if device == "cuda": try: import torch diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index 8d168203c4..d4eb98877e 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -132,18 +132,63 @@ def repo_is_prequantized(base_model: str) -> bool: return "bnb-4bit" in name or "-4bit" in name or "int4" in name or "nf4" in name +def _module_is_torchao_stub(module: Any) -> bool: + """True iff ``module`` is the Unsloth Windows-ROCm torchao import stub rather than the + real package. The stub (core/_torchao_stub.py) satisfies find_spec and even lets + ``from torchao.quantization import quantize_`` succeed -- but the imported symbols are + no-op stub types, so the quantization never happens. Every stub module carries the + ``_unsloth_stub`` sentinel, so match on it (comparing against the stub module's own + sentinel object, not identity of a re-created one).""" + if module is None: + return False + sentinel = getattr(module, "_unsloth_stub", None) + if sentinel is None: + return False + try: + from core._torchao_stub import _STUB_SENTINEL + except Exception: # noqa: BLE001 -- stub module absent -> nothing to compare against + return False + return sentinel is _STUB_SENTINEL + + +def has_functional_torchao() -> bool: + """True iff the real torchao quantization API is importable (not the Windows-ROCm stub). + + ``_int8_quantize_base`` needs ``Int8WeightOnlyConfig`` + ``quantize_`` from + ``torchao.quantization`` and has no runtime fallback, so gate both the auto int8 pick + and the advertised int8 mode on a FUNCTIONAL import: a plain ``find_spec("torchao")`` + is satisfied by the stub, whose quantize_ is a no-op that leaves the transformer dense + while compile is disabled as if it were int8. Import the exact symbols the int8 path + uses and reject the stub module. Never raises.""" + try: + import importlib + + quant = importlib.import_module("torchao.quantization") + if _module_is_torchao_stub(quant): + return False + # The symbols the int8 path actually imports must exist on the real module. + return hasattr(quant, "Int8WeightOnlyConfig") and hasattr(quant, "quantize_") + except Exception: # noqa: BLE001 -- torchao absent / broken build -> treat as unavailable + return False + + def train_precision_modes() -> tuple[list[str], str]: """(supported base_precision modes, recommended pick) for the current machine: nf4 - always works; bf16/int8/auto need CUDA; fp8 needs an fp8-capable GPU (sm89+). Used by - the /info endpoint so the UI can gate the precision selector. Never raises.""" + always works; bf16/auto need CUDA; int8/fp8 additionally need a FUNCTIONAL torchao + (their explicit paths import torchao with no fallback, and the Windows-ROCm stub only + looks installed). fp8 also needs an fp8-capable GPU (sm89+). Used by the /info endpoint + so the UI can gate the precision selector. Never raises.""" modes = ["nf4"] recommended = "nf4" try: import torch if torch.cuda.is_available(): - modes += ["bf16", "int8"] + modes.append("bf16") + torchao_ok = has_functional_torchao() + if torchao_ok: + modes.append("int8") major, minor = torch.cuda.get_device_capability() - if (major, minor) >= (8, 9) and hasattr(torch, "float8_e4m3fn"): + if torchao_ok and (major, minor) >= (8, 9) and hasattr(torch, "float8_e4m3fn"): modes.append("fp8") modes.append("auto") recommended = "auto" diff --git a/studio/backend/tests/test_diffusion_base_precision.py b/studio/backend/tests/test_diffusion_base_precision.py index 1f8ca9318f..e9e28018b6 100644 --- a/studio/backend/tests/test_diffusion_base_precision.py +++ b/studio/backend/tests/test_diffusion_base_precision.py @@ -174,25 +174,14 @@ def test_resolve_auto_requires_bf16_compute(): def test_resolve_auto_int8_band_gates_on_torchao(monkeypatch): - # The int8 auto band needs torchao at runtime; when torchao is not importable - # _resolve_base_precision must fall to nf4 instead of picking an int8 that would crash - # in _int8_quantize_base. Drive the probe into the int8 band and toggle torchao. - import importlib.util as _ilu + # The int8 auto band needs a FUNCTIONAL torchao at runtime; when torchao is not + # importable _resolve_base_precision must fall to nf4 instead of picking an int8 that + # would crash in _int8_quantize_base. Drive the probe into the int8 band and toggle the + # functional-torchao probe (shared with train_precision_modes, imported into the trainer). import torch spec = dit._SPECS["flux.1"] # dense_bf16_gb = 23.8 cfg = _cfg(base_precision = "auto", mixed_precision = "bf16") - real_find_spec = _ilu.find_spec - - def _no_torchao(name, *args, **kwargs): - if name == "torchao": - return None # simulate torchao not installed - return real_find_spec(name, *args, **kwargs) - - def _has_torchao(name, *args, **kwargs): - if name == "torchao": - return object() # simulate torchao installed - return real_find_spec(name, *args, **kwargs) class _FakeCuda: # Free VRAM in the int8 band (30 > 23.8 * 1.15) but below the bf16 band. @@ -206,14 +195,75 @@ def test_resolve_auto_int8_band_gates_on_torchao(monkeypatch): monkeypatch.setattr(torch, "cuda", _FakeCuda) - monkeypatch.setattr(_ilu, "find_spec", _no_torchao) + monkeypatch.setattr(dit, "has_functional_torchao", lambda: False) # torchao absent / stub assert dit._resolve_base_precision(cfg, spec, "cuda") == "nf4" - # With torchao importable the same band picks int8. - monkeypatch.setattr(_ilu, "find_spec", _has_torchao) + # With a functional torchao the same band picks int8. + monkeypatch.setattr(dit, "has_functional_torchao", lambda: True) assert dit._resolve_base_precision(cfg, spec, "cuda") == "int8" +def test_resolve_auto_int8_band_treats_stub_as_absent(monkeypatch): + # Simulate the Windows-ROCm torchao STUB: has_functional_torchao returns False (the + # stub satisfies find_spec but its quantize_ is a no-op), so the int8 band must fall to + # nf4 rather than pick an int8 whose quantization silently does nothing. + import torch + + spec = dit._SPECS["flux.1"] + cfg = _cfg(base_precision = "auto", mixed_precision = "bf16") + + class _FakeCuda: + @staticmethod + def mem_get_info(): + return (int(30 * 1e9), int(80 * 1e9)) + + @staticmethod + def get_device_capability(): + return (10, 0) + + monkeypatch.setattr(torch, "cuda", _FakeCuda) + # The stub scenario: the probe reports no functional torchao. + monkeypatch.setattr(dit, "has_functional_torchao", lambda: False) + assert dit._resolve_base_precision(cfg, spec, "cuda") == "nf4" + + +def test_has_functional_torchao_rejects_stub(monkeypatch): + # has_functional_torchao must reject the Unsloth import stub: even though + # `from torchao.quantization import quantize_` would succeed against the stub, the + # symbols are no-op stub types. Simulate a stub torchao.quantization module carrying the + # stub sentinel and assert the probe returns False. + import importlib + import types + + from core._torchao_stub import _STUB_SENTINEL + + real_import_module = importlib.import_module + + stub_quant = types.ModuleType("torchao.quantization") + stub_quant._unsloth_stub = _STUB_SENTINEL + + def _fake_import(name, *args, **kwargs): + if name == "torchao.quantization": + return stub_quant + return real_import_module(name, *args, **kwargs) + + monkeypatch.setattr(importlib, "import_module", _fake_import) + assert common.has_functional_torchao() is False + + # A real module exposing the int8 symbols (no stub sentinel) probes True. + real_like = types.ModuleType("torchao.quantization") + real_like.Int8WeightOnlyConfig = object + real_like.quantize_ = lambda *a, **k: None + + def _fake_import_real(name, *args, **kwargs): + if name == "torchao.quantization": + return real_like + return real_import_module(name, *args, **kwargs) + + monkeypatch.setattr(importlib, "import_module", _fake_import_real) + assert common.has_functional_torchao() is True + + # ── _fp8_module_filter ──────────────────────────────────────────────────────── def test_fp8_module_filter(): lin = nn.Linear(64, 64) @@ -250,6 +300,28 @@ def test_train_precision_modes_no_cuda(monkeypatch): assert train_precision_modes() == (["nf4"], "nf4") +def test_train_precision_modes_gates_int8_fp8_on_torchao(monkeypatch): + # int8/fp8 are only advertised when torchao is FUNCTIONAL: on a CUDA host WITHOUT a real + # torchao (or with only the Windows-ROCm stub) /info must not offer int8/fp8, since their + # explicit paths import torchao with no fallback. bf16 + auto stay advertised. + import torch + + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: (10, 0)) + + # No functional torchao (absent or stub): bf16 + auto only, int8/fp8 dropped. + monkeypatch.setattr(common, "has_functional_torchao", lambda: False) + modes, recommended = train_precision_modes() + assert modes == ["nf4", "bf16", "auto"] + assert "int8" not in modes and "fp8" not in modes + assert recommended == "auto" + + # With a functional torchao on an fp8-capable GPU, int8 + fp8 are advertised again. + monkeypatch.setattr(common, "has_functional_torchao", lambda: True) + modes2, _ = train_precision_modes() + assert "int8" in modes2 and "fp8" in modes2 + + # ── family_train_infos precision fields ─────────────────────────────────────── def test_family_train_infos_carries_precision_fields(monkeypatch): # Pin the machine probe so the DiT families carry a deterministic mode list, while SDXL From 51de9da488b2ba19ed1c68b90e0ecf46f6fbec1b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 4 Jul 2026 03:23:20 +0000 Subject: [PATCH 3/3] Fix review findings: run history robustness, epoch-mode sentinel, seed 0, history refresh race - list_diffusion_runs skips wrong-shape records and the runs route tolerates per-record ValidationError so one bad file never breaks the panel - max_steps: 0 epoch-mode sentinel no longer trips train_steps validation before epochs are resolved - numberField keeps an explicit 0 (Seed, LR warmup) instead of falling back - previous-runs list refetches once more shortly after a terminal status so the just-finished run appears even if the record write races the fetch --- .../core/training/diffusion_train_common.py | 11 +++ .../training/diffusion_training_service.py | 8 ++ studio/backend/routes/training.py | 15 +++- .../backend/tests/test_diffusion_training.py | 84 +++++++++++++++++++ .../images/train/diffusion-train-panel.tsx | 27 ++++-- 5 files changed, 136 insertions(+), 9 deletions(-) diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index e0936a1dd4..bdbedbf5a0 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -651,6 +651,17 @@ def _config_from_dict(config: dict) -> DiffusionLoraConfig: for k, v in config.items(): if k in valid: kwargs[k] = v + # Epoch-mode payloads from the generic Studio UI carry max_steps: 0 as the "use epochs" + # sentinel, which the max_steps -> train_steps alias copies as train_steps: 0. Since + # normalized() rejects train_steps < 1 before resolve_train_steps() can apply num_epochs, + # drop a falsy/0 train_steps when num_epochs > 0 so the dataclass default stands in until + # epoch resolution replaces it. + try: + _num_epochs = int(kwargs.get("num_epochs") or 0) + except (TypeError, ValueError): + _num_epochs = 0 + if _num_epochs > 0 and not kwargs.get("train_steps"): + kwargs.pop("train_steps", None) if kwargs.get("lora_target_modules"): kwargs["lora_target_modules"] = tuple(kwargs["lora_target_modules"]) if "gradient_checkpointing" in kwargs: diff --git a/studio/backend/core/training/diffusion_training_service.py b/studio/backend/core/training/diffusion_training_service.py index 67f64f14ad..23c0cd2bbd 100644 --- a/studio/backend/core/training/diffusion_training_service.py +++ b/studio/backend/core/training/diffusion_training_service.py @@ -102,6 +102,14 @@ def list_diffusion_runs(limit: int = 20) -> list[dict]: rec = json.loads(p.read_text()) 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 + # dict, or is missing the required string job_id / status) would later blow up the + # route's DiffusionTrainingRunSummary(**r); skip it here so one bad record can never + # take down the whole Previous runs panel. + if not isinstance(rec, dict): + continue + if not (isinstance(rec.get("job_id"), str) and isinstance(rec.get("status"), str)): + continue rec.pop("metric_history", None) rec.pop("config", None) out.append(rec) diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 1a85089439..f6f2a0509a 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -80,6 +80,7 @@ from models.training import ( ) from models.responses import TrainingStopResponse, TrainingMetricsResponse from pydantic import BaseModel as PydanticBaseModel +from pydantic import ValidationError class TrainingStopRequest(PydanticBaseModel): @@ -1285,9 +1286,17 @@ async def list_diffusion_training_runs( """Previous diffusion training runs (terminal), newest first, from the persisted per-run records. Summaries only; fetch one run for its config + metric logs.""" from core.training.diffusion_training_service import list_diffusion_runs - return DiffusionTrainingRunsResponse( - runs = [DiffusionTrainingRunSummary(**r) for r in list_diffusion_runs(limit = limit)] - ) + + summaries: list[DiffusionTrainingRunSummary] = [] + for r in list_diffusion_runs(limit = limit): + # list_diffusion_runs already skips non-dict / missing-id records, but a record with + # a wrong-typed field (e.g. a non-numeric avg_loss) would still raise here; catch it + # per record so one bad file never breaks the whole Previous runs panel. + try: + summaries.append(DiffusionTrainingRunSummary(**r)) + except ValidationError: + continue + return DiffusionTrainingRunsResponse(runs = summaries) @router.get("/diffusion/runs/{job_id}", response_model = DiffusionTrainingRunDetail) diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index e4bd456f09..84446cc15c 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -381,6 +381,44 @@ def test_request_model_num_epochs_bounds(): DiffusionTrainingStartRequest(**base, num_epochs = bad) +def test_config_from_dict_epoch_mode_drops_max_steps_sentinel(): + # The generic Studio epoch-mode payload sends max_steps: 0 as the "use epochs" sentinel. + # The max_steps -> train_steps alias would copy that 0 and normalized() would reject + # train_steps < 1 before epochs are resolved; _config_from_dict must drop the falsy + # value so the default train_steps stands in until resolve_train_steps applies num_epochs. + from core.training.diffusion_train_common import DiffusionLoraConfig, _config_from_dict + + cfg = _config_from_dict( + { + "base_model": "stabilityai/stable-diffusion-xl-base-1.0", + "data_dir": "d", + "output_dir": "o", + "max_steps": 0, + "num_epochs": 2, + } + ) + # 0 was dropped: the dataclass default train_steps stands in and num_epochs carries over. + assert cfg.train_steps == DiffusionLoraConfig.train_steps + assert cfg.num_epochs == 2 + # normalized() no longer raises on the epoch-mode payload. + norm = cfg.normalized() + assert norm.num_epochs == 2 + + # An explicit non-zero max_steps in epochs mode is still honored (only the 0 sentinel is + # dropped), and a plain steps payload (no num_epochs) keeps max_steps: 0 -> train_steps 0 + # so normalized() surfaces the invalid value as before. + cfg_explicit = _config_from_dict( + { + "base_model": "stabilityai/stable-diffusion-xl-base-1.0", + "data_dir": "d", + "output_dir": "o", + "max_steps": 25, + "num_epochs": 2, + } + ) + assert cfg_explicit.train_steps == 25 + + def test_route_start_rejects_uncontained_paths(client): # An absolute path outside the Studio dataset roots is a 400, not silently accepted. r = client.post("/api/train/diffusion/start", json = {**_BODY, "data_dir": "/etc"}) @@ -804,3 +842,49 @@ def test_runs_endpoints_list_and_detail(client, _isolated_runs_dir): # Unknown and malformed ids 404 (malformed also covers path traversal). assert client.get(f"/api/train/diffusion/runs/{'c' * 32}").status_code == 404 assert client.get("/api/train/diffusion/runs/not-a-job-id").status_code == 404 + + +def test_list_diffusion_runs_skips_wrong_shape_records(_isolated_runs_dir): + # A valid-JSON file with the wrong shape (non-dict, or missing the required string + # job_id / status) must be skipped by list_diffusion_runs so it never reaches the route's + # DiffusionTrainingRunSummary(**r) and takes down the whole Previous runs panel. + import json + + from core.training.diffusion_training_service import list_diffusion_runs + + good = {"job_id": "a" * 32, "status": "completed", "adapter": "good", "saved": True} + (_isolated_runs_dir / "good.json").write_text(json.dumps(good)) + # A JSON list (not a dict). + (_isolated_runs_dir / "not_a_dict.json").write_text(json.dumps([1, 2, 3])) + # A dict missing the required job_id / status. + (_isolated_runs_dir / "no_ids.json").write_text(json.dumps({"adapter": "orphan"})) + # A dict whose job_id / status are the wrong type. + (_isolated_runs_dir / "bad_types.json").write_text( + json.dumps({"job_id": 123, "status": None, "adapter": "typed"}) + ) + + runs = list_diffusion_runs() + adapters = [r.get("adapter") for r in runs] + assert adapters == ["good"] # only the well-shaped record survives + + +def test_runs_route_tolerates_bad_field_record(client, _isolated_runs_dir): + # A record that passes the service's shape check but has a wrong-typed field (a + # non-numeric avg_loss) would raise pydantic ValidationError in the route; the route must + # catch it per record so one bad file never breaks the panel and the good runs still list. + import json + + good = {"job_id": "a" * 32, "status": "completed", "adapter": "good", "saved": True} + bad = { + "job_id": "b" * 32, + "status": "completed", + "adapter": "bad", + "avg_loss": "not-a-number", # str where the summary expects Optional[float] + } + (_isolated_runs_dir / f"{good['job_id']}.json").write_text(json.dumps(good)) + (_isolated_runs_dir / f"{bad['job_id']}.json").write_text(json.dumps(bad)) + + r = client.get("/api/train/diffusion/runs") + assert r.status_code == 200, r.text + adapters = [x["adapter"] for x in r.json()["runs"]] + assert adapters == ["good"] # the bad-field record was skipped, the good one remained 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 813a6c1fd2..463ca86dc4 100644 --- a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx +++ b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx @@ -520,13 +520,25 @@ export function DiffusionTrainPanel({ if (!active) return; if (status?.status === "running") return; let cancelled = false; - listDiffusionTrainingRuns() - .then((r) => { - if (!cancelled) setPrevRuns(r.runs); - }) - .catch(() => {}); + const refetch = () => { + listDiffusionTrainingRuns() + .then((r) => { + if (!cancelled) setPrevRuns(r.runs); + }) + .catch(() => {}); + }; + refetch(); + // The service exposes a terminal status before the pump has necessarily finished + // writing the run's JSON record, so the one-shot refetch above can win that race and + // miss the just-finished run. A short delayed second refetch after a terminal + // transition lets the record land so the newest run reliably appears. + let delayed: ReturnType | undefined; + if (status?.status === "completed" || status?.status === "stopped" || status?.status === "error") { + delayed = setTimeout(refetch, 1500); + } return () => { cancelled = true; + if (delayed !== undefined) clearTimeout(delayed); }; }, [active, status?.status]); @@ -745,7 +757,10 @@ export function DiffusionTrainPanel({ value={value} onChange={(e) => { settingsDirty.current = true; - set(Number(e.target.value) || fallback); + // Only fall back when the input parses to NaN (empty/invalid); a real 0 is a + // legal value for zero-legal fields (Seed, LR warmup steps) and must be kept. + const parsed = Number(e.target.value); + set(Number.isNaN(parsed) ? fallback : parsed); }} className="h-8 text-xs" />