ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge diffusion-krea2: Raw training default + stacked review fixes

Browse source 
# Conflicts:
#	studio/backend/core/training/diffusion_train_common.py
#	studio/backend/tests/test_diffusion_dit_trainer.py
#	studio/backend/tests/test_diffusion_training.py
This commit is contained in:
Daniel Han 2026-07-04 03:28:32 +00:00
commit 95f783ae1e
16 changed files with 577 additions and 100 deletions
Download patch file Download diff file Expand all files Collapse all files

5
studio/backend/core/inference/diffusion.py
View file

@ -204,9 +204,12 @@ _TRUSTED_NON_GGUF_REPOS = frozenset(
"black-forest-labs/flux.1-dev",
"tongyi-mai/z-image-turbo",
"qwen/qwen-image",
# Krea 2 Turbo: official vendor repo, safetensors-only, no remote code. Loaded
# Krea 2: official vendor repos, safetensors-only, no remote code. Loaded
# per-component via core/inference/diffusion_krea2.py (no GGUF variant yet).
# Turbo is the inference model; Raw is the undistilled base Krea recommends
# training LoRAs on (train on Raw, run adapters on Turbo).
"krea/krea-2-turbo",
"krea/krea-2-raw",
}
)

7
studio/backend/core/inference/diffusion_families.py
View file

@ -294,9 +294,12 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
base_repo = "krea/Krea-2-Turbo",
aliases = ("krea2",),
# LoRA training via the DiT trainer (no prequant repo yet, so nf4 quantizes the
# 12B transformer on the fly under the default precision).
# 12B transformer on the fly under the default precision). Krea's release guidance
# is explicit: train LoRAs on the undistilled Raw checkpoint and apply them on
# Turbo for inference, so Raw is the default training base and Turbo stays the
# inference/base repo.
trainable = True,
train_base_repos = ("krea/Krea-2-Turbo",),
train_base_repos = ("krea/Krea-2-Raw", "krea/Krea-2-Turbo"),
# The checkpoint is exported bf16-only (the model card pins bfloat16); fp16 is
# unvalidated upstream, so keep the fp16 fallback off like z-image.
fp16_incompatible = True,

9
studio/backend/core/inference/diffusion_krea2.py
View file

@ -113,6 +113,15 @@ def load_krea2_pipeline(
"""
import diffusers
# diffusers gained Krea2Pipeline in 0.39; on an older install the getattr chain below
# would die with a bare AttributeError mid-load, so fail first with the actionable fix.
if not hasattr(diffusers, "Krea2Pipeline"):
raise RuntimeError(
f"Krea 2 needs diffusers >= 0.39.0 (Krea2Pipeline); this environment has "
f"diffusers {getattr(diffusers, '__version__', 'unknown')}. "
f"Upgrade with: pip install -U diffusers"
)
token = hf_token or None
tokenizer = load_krea2_tokenizer(repo_id, hf_token = token)
text_encoder = load_krea2_text_encoder(repo_id, dtype, hf_token = token)

29
studio/backend/core/training/diffusion_dit_trainer.py
View file

@ -37,6 +37,7 @@ from typing import Any, Callable, Optional
from core.training.diffusion_train_common import (
DEFAULT_LORA_FILENAME,
DEFAULT_LORA_TARGETS,
DiffusionLoraConfig,
EventCb,
StopCb,
@ -47,6 +48,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,
resolve_train_steps,
)
@ -88,6 +90,20 @@ _KREA2_TARGETS = (
)
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."""
@ -407,11 +423,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
@ -1205,7 +1222,7 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto
from peft import LoraConfig
from peft.utils import get_peft_model_state_dict
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)
out_dir = Path(cfg.output_dir).expanduser()
# Phase 1: conditioning only. The pipeline loads WITHOUT its transformer, so the text

93
studio/backend/core/training/diffusion_train_common.py
View file

@ -133,9 +133,51 @@ 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+); mxfp8
always works; bf16/auto need CUDA; int8/fp8/mxfp8 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+); mxfp8
(block-scaled fp8 compute) needs the Blackwell tensor cores (sm100+) its cuBLAS
kernels target. Used by the /info endpoint so the UI can gate the precision
selector. Never raises."""
@ -144,11 +186,14 @@ def train_precision_modes() -> tuple[list[str], str]:
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")
if (major, minor) >= (10, 0):
if torchao_ok and (major, minor) >= (10, 0):
modes.append("mxfp8")
modes.append("auto")
recommended = "auto"
@ -206,7 +251,10 @@ _FAMILY_VRAM_NOTES = {
),
"qwen-image": "20B model, QLoRA (nf4) by default (~24 GB+). The heaviest option.",
"z-image": "6B model, QLoRA (nf4) by default (~12 GB+). bf16 only.",
"krea-2": "12B model, QLoRA (nf4) by default (~18 GB+). bf16 only.",
"krea-2": (
"12B model, QLoRA (nf4) by default (~18 GB+). bf16 only. Trains on the "
"undistilled Krea-2-Raw (Krea's guidance: train on Raw, run adapters on Turbo)."
),
}
@ -400,11 +448,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 ``<stem>.txt`` / ``<stem>.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 ``<stem>.txt`` / ``<stem>.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.
@ -436,15 +488,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
@ -661,6 +713,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:

8
studio/backend/core/training/diffusion_training_service.py
View file

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

64
studio/backend/routes/models.py
View file

@ -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")

70
studio/backend/routes/training.py
View file

@ -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):
@ -1115,9 +1116,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
@ -1285,9 +1290,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)
@ -1311,14 +1324,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
@ -1515,20 +1533,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:

108
studio/backend/tests/test_diffusion_base_precision.py
View file

@ -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

38
studio/backend/tests/test_diffusion_dataset_api.py
View file

@ -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"})

26
studio/backend/tests/test_diffusion_dit_trainer.py
View file

@ -14,16 +14,21 @@ import sys
import pytest
from core.training.diffusion_dit_trainer import (
_FLUX_TARGETS,
_GATED_TRAIN_REPOS,
_QWEN_TARGETS,
_SPECS,
_ZIMAGE_TARGETS,
_apply_mxfp8_training,
_assert_gated_access,
_mx_module_filter,
_repo_is_prequantized,
_select_lora_targets,
_should_compile,
run_dit_lora_training,
)
from core.training.diffusion_train_common import (
DEFAULT_LORA_TARGETS,
DiffusionLoraConfig,
family_train_infos,
train_precision_modes,
@ -44,6 +49,27 @@ def test_specs_cover_the_dit_families():
assert _SPECS["krea-2"].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",
[

20
studio/backend/tests/test_diffusion_krea2.py
View file

@ -120,6 +120,17 @@ def test_load_krea2_pipeline_threads_init_config(monkeypatch, tmp_path):
# ── registry / trust / int8 exclusion wiring ─────────────────────────────────
def test_load_krea2_pipeline_requires_krea_capable_diffusers(monkeypatch):
# On diffusers < 0.39 (no Krea2Pipeline) the loader must fail fast with the upgrade
# hint instead of dying with a bare AttributeError mid-load.
import pytest
fake = SimpleNamespace(__version__ = "0.38.0")
monkeypatch.setitem(sys.modules, "diffusers", fake)
with pytest.raises(RuntimeError, match = "0.39"):
load_krea2_pipeline("krea/Krea-2-Turbo", "bf16")
def test_krea2_family_wiring():
from core.inference.diffusion import _is_trusted_diffusion_repo
from core.inference.diffusion_families import detect_family, family_sd_cpp_supported
@ -127,8 +138,10 @@ def test_krea2_family_wiring():
fam = detect_family("krea/Krea-2-Turbo")
assert fam is not None and fam.name == KREA2_FAMILY_NAME
# The vendor repo is non-GGUF allowlisted; no sd.cpp mapping -> diffusers fallback.
# Both vendor repos are non-GGUF allowlisted (Turbo for inference, Raw for training);
# no sd.cpp mapping -> diffusers fallback.
assert _is_trusted_diffusion_repo("krea/Krea-2-Turbo")
assert _is_trusted_diffusion_repo("krea/Krea-2-Raw")
assert not family_sd_cpp_supported(fam)
# Krea2TimestepEmbedding runs at M = batch; int8 (torch._int_mm, M > 16) must skip it.
assert "time_embed" in exclude_tokens_for_scheme(TQ_INT8)
@ -155,7 +168,10 @@ def test_krea2_training_registry():
"resolution": 512,
}
info = {i["name"]: i for i in family_train_infos()}["krea-2"]
assert info["default_base"] == "krea/Krea-2-Turbo"
# Krea's guidance: train LoRAs on the undistilled Raw model, run them on Turbo, so
# Raw leads the training bases while Turbo stays available.
assert info["default_base"] == "krea/Krea-2-Raw"
assert info["base_repos"] == ["krea/Krea-2-Raw", "krea/Krea-2-Turbo"]
assert info["supports_compile"] is True

18
studio/backend/tests/test_diffusion_lora_trainer.py
View file

@ -28,15 +28,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"))
@ -45,6 +45,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")

110
studio/backend/tests/test_diffusion_training.py
View file

@ -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
@ -395,6 +396,44 @@ def test_request_model_base_precision_accepts_mxfp8():
DiffusionTrainingStartRequest(**base, base_precision = "bogus")
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"})
@ -510,6 +549,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 = [
@ -818,3 +882,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

45
studio/backend/tests/test_local_model_format.py
View file

@ -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

27
studio/frontend/src/features/images/train/diffusion-train-panel.tsx
View file

@ -525,13 +525,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<typeof setTimeout> | 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]);
@ -750,7 +762,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"
/>