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

Tighten comments across the remaining image stack files

Browse source
This commit is contained in:
Daniel Han 2026-07-12 11:46:23 +00:00
commit 1762dab12f
2 changed files with 69 additions and 128 deletions
Download patch file Download diff file Expand all files Collapse all files

120
studio/backend/core/training/diffusion_lora_trainer.py
View file

@ -92,8 +92,7 @@ def _load_image_tensor(
import torch
from PIL import Image, ImageOps
# Honour EXIF orientation before any geometry, or rotated camera/phone photos would
# train in their stored (sideways) orientation, mismatched to their captions.
# Honour EXIF orientation before any geometry, or rotated photos train sideways.
img = ImageOps.exif_transpose(Image.open(path)).convert("RGB")
original_w, original_h = img.size
scale = resolution / min(original_w, original_h)
@ -109,8 +108,7 @@ def _load_image_tensor(
crop_left = left
if random_flip and rng.random() < 0.5:
img = img.transpose(Image.FLIP_LEFT_RIGHT)
# A horizontal flip mirrors the crop's left origin, so report the mirrored offset
# (diffusers' SDXL training scripts do the same) to keep the conditioning honest.
# A flip mirrors the crop's left origin, so report the mirrored offset (as diffusers does).
crop_left = max(0, resized_w - resolution - left)
arr = np.asarray(img, dtype = np.float32) / 255.0
tensor = torch.from_numpy(arr).permute(2, 0, 1) * 2.0 - 1.0
@ -146,8 +144,7 @@ def _load_image_tensor_planned(
crop_left = left
if flip:
img = img.transpose(Image.FLIP_LEFT_RIGHT)
# Mirror the crop's left origin so the conditioning matches the flipped pixels, the
# same mirroring ``_load_image_tensor`` applies on a random flip.
# Mirror the crop's left origin (same as _load_image_tensor's random flip).
crop_left = max(0, resized_w - resolution - left)
arr = np.asarray(img, dtype = np.float32) / 255.0
tensor = torch.from_numpy(arr).permute(2, 0, 1) * 2.0 - 1.0
@ -175,7 +172,7 @@ def _encode_sdxl_prompts(
).input_ids.to(device)
with torch.no_grad():
out = text_encoder(tokens, output_hidden_states = True)
# The pooled embed always comes from the second (bigG) text encoder's [0] output.
# Pooled embed always comes from the second (bigG) encoder's [0] output.
pooled = out[0]
embeds_list.append(out.hidden_states[-2])
prompt_embeds = torch.concat(embeds_list, dim = -1)
@ -226,10 +223,8 @@ def _build_sdxl_latent_cache(
a = _hold(dist.mean * vae_scale)
b = _hold(dist.std * vae_scale)
if not forced and not gated:
# Size-gate the automatic cache off the first REAL encoded variant, before
# building the rest: thousands of images x variants of two fp32 tensors can
# exhaust host/pinned RAM with no fallback. Over budget we bail with the VAE
# still resident so the loop encodes latents per step instead.
# Size-gate the auto cache off the first real variant, before building the rest
# (it can exhaust host/pinned RAM). Over budget: bail with the VAE still resident.
per_variant = a.numel() * a.element_size() + b.numel() * b.element_size()
if _latent_cache_over_budget(per_variant, total_variants):
_emit(
@ -303,8 +298,7 @@ def run_diffusion_lora_training(
rng = random.Random(cfg.seed)
torch.manual_seed(cfg.seed)
# A stop signal may be a bare truthy value or a dict carrying save=False (cancel without
# saving a partial adapter). ``save_on_stop`` records that decision for the export step.
# A stop signal may be truthy or a dict with save=False; save_on_stop records that for export.
save_on_stop = True
def _check_stop() -> bool:
@ -321,16 +315,13 @@ def run_diffusion_lora_training(
device = "cuda" if torch.cuda.is_available() else "cpu"
precision = cfg.mixed_precision if device == "cuda" else "no"
if precision == "bf16" and device == "cuda" and not native_bf16_supported():
# The default is bf16, but pre-Ampere GPUs (T4 / V100 / RTX 20xx) have no
# NATIVE bf16 compute; torch.cuda.is_bf16_supported() counts emulated support there,
# so use the compute-capability probe (matches the DiT trainer) and fall back to fp16
# instead of failing at load/forward.
# Pre-Ampere GPUs (T4/V100/RTX 20xx) have no native bf16, and is_bf16_supported() counts
# emulation; use the compute-capability probe and fall back to fp16.
precision = "fp16"
weight_dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "no": torch.float32}[precision]
# TF32 / cudnn.benchmark for the run, restored on the way out (the trainer subprocess is
# disposable, but restoring keeps in-process callers -- tests, notebooks -- clean). Wraps
# the whole body so every return (early stop and normal) restores the backend flags.
# TF32 / cudnn.benchmark for the run, restored on the way out (keeps in-process callers
# clean). Wraps the whole body so every return restores the backend flags.
snap = _apply_perf_flags(cfg, device)
try:
# Preflight the base model against the same trust gate as inference, before any fetch.
@ -339,14 +330,12 @@ def run_diffusion_lora_training(
pairs = discover_image_caption_pairs(
cfg.data_dir, instance_prompt = cfg.instance_prompt, caption_column = cfg.caption_column
)
# Resolve num_epochs -> a concrete train_steps now that the dataset size is known, and
# rebind cfg so every downstream read (scheduler length, the loop range, progress
# total_steps, steps_run) sees the same resolved value.
# Resolve num_epochs -> a concrete train_steps now the dataset size is known, and rebind
# cfg so every downstream read sees the same value.
cfg = replace(cfg, train_steps = resolve_train_steps(cfg, len(pairs)), num_epochs = 0)
_emit(on_event, "model_load_started", num_images = len(pairs))
# Honour a stop requested before the (potentially large / slow) base model loads, the
# same way the LLM training worker checks its stop thread around model load.
# Honour a stop requested before the (slow) base model load.
if _check_stop():
out_dir = Path(cfg.output_dir).expanduser()
_emit(
@ -389,10 +378,8 @@ def run_diffusion_lora_training(
if weight_dtype != torch.float32:
cast_training_params(unet, dtype = torch.float32)
# Regionally torch.compile the U-Net's repeated BasicTransformerBlocks through the
# DiT trainer's never-fatal wrapper (a wrap/compile failure falls back to eager
# with a warning event). The U-Net is a dense bf16 base here, the combination that
# wrapper compiles under "auto".
# Regionally torch.compile the U-Net's repeated blocks via the DiT trainer's never-fatal
# wrapper (failure falls back to eager with a warning). Dense bf16 base compiles under "auto".
from core.training.diffusion_dit_trainer import _maybe_compile_transformer
compiled = _maybe_compile_transformer(
@ -401,10 +388,8 @@ def run_diffusion_lora_training(
lora_params = [p for p in unet.parameters() if p.requires_grad]
optimizer = _make_lora_optimizer(lora_params, cfg.learning_rate)
# The scheduler advances once per optimizer update: lr_sched.step() runs a single
# time per outer opt_step (after the accumulation inner loop), for cfg.train_steps
# total. Count warmup/decay in those optimizer steps -- multiplying by the
# accumulation factor would stretch warmup past the run and never reach the decay.
# The scheduler advances once per optimizer update (per opt_step, for cfg.train_steps
# total). Count warmup/decay in optimizer steps; the accumulation factor would stretch warmup.
lr_sched = get_scheduler(
cfg.lr_scheduler,
optimizer = optimizer,
@ -415,11 +400,9 @@ def run_diffusion_lora_training(
vae_scale = vae.config.scaling_factor
prediction_type = noise_scheduler.config.prediction_type
# Precompute text embeddings once per unique caption, then free the CLIP text encoders.
# SDXL re-encoded captions every step (pure waste: captions are constant) and kept both
# text encoders (~1.5 GB) resident. Embeddings are deterministic and this consumes no
# torch RNG, so the training math is bit-identical to in-loop encoding -- only faster and
# lighter. The env toggle exists purely so the accuracy guard can A/B the two paths.
# Precompute text embeddings once per unique caption, then free the ~1.5 GB CLIP encoders.
# Deterministic and RNG-free, so the training math is bit-identical to in-loop encoding.
# The env toggle lets the accuracy guard A/B the two paths.
precompute = os.environ.get("UNSLOTH_DIFFUSION_NO_PRECOMPUTE", "") not in ("1", "true")
caption_embeds: dict[str, tuple] = {}
if precompute:
@ -433,9 +416,8 @@ def run_diffusion_lora_training(
if device == "cuda":
torch.cuda.empty_cache()
# Precompute the VAE latent cache, then free the VAE: the cache holds the posterior
# affine pair (mean/std, scale folded in) so per-step sampling noise is preserved. The
# env toggle lets the accuracy guard A/B the cached vs in-loop encode paths.
# Precompute the VAE latent cache, then free the VAE: it holds the posterior affine pair
# so per-step sampling noise is preserved. The env toggle A/Bs cached vs in-loop encode.
use_cache = cfg.cache_latents and os.environ.get(
"UNSLOTH_DIFFUSION_NO_LATENT_CACHE", ""
) not in ("1", "true")
@ -452,8 +434,7 @@ def run_diffusion_lora_training(
_check_stop,
)
if latent_cache is LATENT_CACHE_OVER_BUDGET:
# The estimated cache exceeded the host-memory budget; keep the VAE resident
# and fall through to the in-loop encode path (latent_cache stays None).
# Over the host-memory budget; keep the VAE resident and encode in-loop.
latent_cache = None
elif latent_cache is None: # stopped during the cache build; nothing trained yet
out_dir = Path(cfg.output_dir).expanduser()
@ -476,25 +457,19 @@ def run_diffusion_lora_training(
gc.collect()
if device == "cuda":
torch.cuda.empty_cache()
# Variant picks use their own stream so the loop's index/noise draws stay on the same
# seed-deterministic sequence whether or not the cache is enabled.
# Variant picks use their own stream so the loop's index/noise draws stay seed-deterministic
# whether or not the cache is enabled.
variant_rng = random.Random(cfg.seed + 1)
_emit(on_event, "model_load_completed", compiled = compiled)
# Permutation-cycle index sampler (shared with the DiT trainer): each dataset image is
# visited once per cycle before any repeat, so a short run does not leave part of a
# small dataset unseen. Draws from the loop's own rng so the sequence stays
# seed-deterministic.
# Permutation-cycle index sampler: each image is visited once per cycle before any repeat,
# so a short run doesn't leave a small dataset partly unseen.
index_sampler = PermutationBatchSampler(len(pairs), rng)
def _next_batch() -> tuple[list[int], list[str], list[str]]:
# Draw the full configured batch, not min(batch, n): PermutationBatchSampler refills
# across permutation cycles so a dataset smaller than train_batch_size still yields
# exactly train_batch_size indices (the DiT trainer calls next_batch the same way).
# Clamping to len(pairs) would silently train a tiny dataset at a smaller effective
# batch than configured while the scheduler and samples-per-second still assume the
# full batch.
# Draw the full configured batch, not min(batch, n): the sampler refills across cycles
# so a dataset smaller than train_batch_size still yields exactly that many indices.
idx = index_sampler.next_batch(cfg.train_batch_size)
chosen = [pairs[i] for i in idx]
return idx, [c[0] for c in chosen], [c[1] for c in chosen]
@ -513,8 +488,7 @@ def run_diffusion_lora_training(
for _ in range(cfg.gradient_accumulation_steps):
idx, img_paths, captions = _next_batch()
if latent_cache is not None:
# Scale is folded into the cache; the sampler draws in fp32 and casts the
# result to weight_dtype (matching the in-loop path below).
# Scale folded into the cache; the sampler draws fp32 and casts to weight_dtype.
latents, batch_time_ids = _sample_sdxl_cached_latents(
latent_cache, idx, variant_rng, device, weight_dtype
)
@ -578,11 +552,11 @@ def run_diffusion_lora_training(
step_loss += float(loss.detach()) / cfg.gradient_accumulation_steps
micro += 1
# max_grad_norm <= 0 means "disable clipping" (the Studio payload sends 0.0 for that);
# passing 0.0 to clip_grad_norm_ would scale every gradient to zero (no learning).
# max_grad_norm <= 0 disables clipping (Studio sends 0.0); passing 0.0 to
# clip_grad_norm_ would zero every gradient (no learning).
grad_norm = None
if cfg.max_grad_norm and cfg.max_grad_norm > 0:
# The returned value is the total PRE-clip norm, reported to the UI chart.
# Returned value is the total PRE-clip norm, reported to the UI chart.
grad_norm = float(torch.nn.utils.clip_grad_norm_(lora_params, cfg.max_grad_norm))
optimizer.step()
lr_sched.step()
@ -591,14 +565,11 @@ def run_diffusion_lora_training(
done = opt_step + 1
now = time.time()
if done == 1:
# Step 1 pays the one-time costs (cudnn autotune, torch.compile warmup), so
# the reported rate starts after it and reflects the steady state (the DiT
# trainer does the same).
# Step 1 pays the one-time costs (cudnn autotune, compile warmup), so the rate
# starts after it and reflects steady state.
t_steady = now
if done % cfg.log_every == 0 or done == cfg.train_steps:
# ``learning_rate`` (not ``lr``) is the field the Studio training pump reads, so
# these progress events are directly consumable by the existing training
# status/SSE machinery when the diffusion trainer is wired into the worker.
# ``learning_rate`` (not ``lr``) is the field the Studio training pump reads.
if device == "cuda":
peak_gb = round(torch.cuda.max_memory_allocated() / 1e9, 2)
per_step = cfg.train_batch_size * cfg.gradient_accumulation_steps
@ -623,8 +594,7 @@ def run_diffusion_lora_training(
stopped = True
break
# Export the trained LoRA in diffusers format (loadable via load_lora_weights), unless
# the run was cancelled with save disabled -- then leave no partial adapter behind.
# Export the LoRA in diffusers format, unless cancelled with save disabled.
out_dir = Path(cfg.output_dir).expanduser()
lora_path: Optional[str] = None
catalog_path: Optional[str] = None
@ -638,8 +608,7 @@ def run_diffusion_lora_training(
weight_name = DEFAULT_LORA_FILENAME,
)
lora_path = str(out_dir / DEFAULT_LORA_FILENAME)
# Mirror into the Studio diffusion LoRA directory so the Images picker discovers it
# (its scan lists only files directly under loras/diffusion, not subdirectories).
# Mirror into loras/diffusion so the Images picker discovers it (its scan skips subdirs).
catalog_path = _publish_to_lora_catalog(lora_path, cfg)
_emit(
on_event,
@ -690,8 +659,7 @@ def run_diffusion_training_process(*, event_queue: Any, stop_queue: Any, config:
event_queue.put(ev)
def should_stop() -> Any:
# Drain the queue and return the last stop message (bool True, or a dict that may
# carry save=False for cancel-without-save); False when nothing was requested.
# Drain the queue and return the last stop message (True or a dict with save=False); else False.
got: Any = None
saw = False
try:
@ -703,15 +671,13 @@ def run_diffusion_training_process(*, event_queue: Any, stop_queue: Any, config:
return got if saw else False
try:
# normalized() resolves + validates the family; dispatch through the registry so a
# DiT family runs its own trainer while SDXL keeps this module's loop.
# normalized() resolves + validates the family; dispatch through the registry so a DiT
# family runs its own trainer while SDXL keeps this loop.
cfg = _config_from_dict(config).normalized()
trainer = get_trainer(cfg.resolved_family)
trainer(cfg, on_event = on_event, should_stop = should_stop)
except Exception as exc: # noqa: BLE001 -- surfaced to the parent as an error event
# Emit both keys: the diffusion service reads ``message``, but the generic Studio
# training worker reads ``error``; carrying both keeps the real failure visible on
# either path instead of surfacing as "Unknown error".
# Emit both keys: the diffusion service reads ``message``, the generic worker reads ``error``.
event_queue.put(
{"type": "error", "message": str(exc), "error": str(exc), "ts": time.time()}
)

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

@ -27,8 +27,7 @@ import uuid
from pathlib import Path
from typing import Any, Callable, Optional
# Spawn (not fork): a fresh interpreter, matching the LLM training worker, so CUDA/torch
# state from the parent never leaks into the trainer.
# Spawn (not fork): a fresh interpreter so parent CUDA/torch state never leaks into the trainer.
_CTX = mp.get_context("spawn")
# Terminal event types after which the pump stops.
@ -57,30 +56,22 @@ def _run_diffusion_child(*, event_queue: Any, stop_queue: Any, config: dict) ->
def _default_target(*, event_queue: Any, stop_queue: Any, config: dict) -> None:
# First thing in the spawned child (before torch is imported): bind to the parent's
# death on Linux and scrub the native path lease secret, exactly like the inference /
# export / LLM-training workers. multiprocessing children cannot be given a
# parent-set preexec_fn, so the child must self-bind; otherwise a Studio crash or
# kill leaves this trainer holding the GPU. Tests inject their own target, so this
# binding only runs for the real production spawn.
# First thing in the child (before torch): self-bind to parent death and scrub the native
# path secret, like the other workers (multiprocessing children can't be given a preexec_fn).
from utils.native_path_leases import run_without_native_path_secret
run_without_native_path_secret(
_run_diffusion_child, event_queue = event_queue, stop_queue = stop_queue, config = config
)
# Cap on retained metric points. When exceeded, the arrays are decimated (every other
# point dropped) so a long run stays bounded in memory while the live loss chart keeps a
# faithful shape. 4000 points comfortably covers a typical run at full resolution.
# Cap on retained metric points; over it, arrays are decimated (every other point) so a
# long run stays bounded while the loss chart keeps its shape.
_METRIC_CAP = 4000
# ── persisted run history ──────────────────────────────────────────────────────
# Every terminal run (completed / stopped / error) is recorded as one JSON file --
# summary + scrubbed config + the full bounded metric logs -- so the Train tab can show
# previous runs like the LLM trainer's history. JSON files (not the LLM sqlite tables)
# keep diffusion runs out of the LLM Runs page, whose resume/inspect actions assume an
# LLM-shaped run.
# Every terminal run is recorded as one JSON file (summary + scrubbed config + metric logs)
# so the Train tab can show history. JSON (not LLM sqlite) keeps diffusion runs off the LLM Runs page.
def _runs_dir() -> Path:
from utils.paths.storage_roots import studio_root
@ -102,10 +93,8 @@ def list_diffusion_runs(limit: int = 20) -> list[dict]:
rec = json.loads(p.read_text(encoding = "utf-8"))
except Exception: # noqa: BLE001 -- a corrupt record never breaks the listing
continue
# A valid-JSON file with the wrong shape (an old or hand-edited record that is not a
# 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.
# Skip a wrong-shape record (not a dict, or missing string job_id/status) so one bad
# file can't blow up the route's DiffusionTrainingRunSummary(**r) or the whole panel.
if not isinstance(rec, dict):
continue
if not (isinstance(rec.get("job_id"), str) and isinstance(rec.get("status"), str)):
@ -118,8 +107,7 @@ def list_diffusion_runs(limit: int = 20) -> list[dict]:
def get_diffusion_run(job_id: str) -> Optional[dict]:
"""The full persisted record for one run (summary + config + metric logs)."""
# Records are keyed by the uuid4 hex job id; reject anything else so a crafted id
# can never traverse out of the runs directory.
# Keyed by uuid4 hex; reject anything else so a crafted id can't traverse out of the dir.
if not re.fullmatch(r"[0-9a-f]{32}", str(job_id or "")):
return None
p = _runs_dir() / f"{job_id}.json"
@ -182,11 +170,9 @@ def _append_metric(
if istep <= 0 or loss is None:
return
floss = _finite_or_none(loss)
if floss is None: # non-numeric or non-finite (NaN/Inf): skip, keep the curve JSON-safe
if floss is None: # non-numeric or non-finite: skip, keep the curve JSON-safe
return
# lr / grad_norm may be None (sparse series) or non-finite; non-finite values are
# nulled, not dropped, so a bad point never taints the (loss-driven) history while the
# arrays stay index-aligned with steps.
# lr / grad_norm may be None or non-finite; non-finite is nulled (not dropped) to stay index-aligned.
flr = _finite_or_none(lr)
fgn = _finite_or_none(grad_norm)
steps = state["metric_steps"]
@ -222,10 +208,8 @@ class DiffusionTrainingService:
self._ctx = ctx if ctx is not None else _CTX
self._target = target if target is not None else _default_target
self._lock = threading.Lock()
# Set True by reserve() while a start is in flight, BEFORE the route frees resident GPU
# models, so the image/video load guards (which read is_active) refuse a concurrent load
# during the free-then-spawn window rather than double-allocate the GPU. Cleared by
# unreserve() once the proc is live (or the start failed).
# Set by reserve() while a start is in flight (before the route frees GPU models) so the
# load guards refuse a concurrent load during the free-then-spawn window. Cleared by unreserve().
self._reserved = False
self._proc: Any = None
self._stop_queue: Any = None
@ -274,10 +258,8 @@ class DiffusionTrainingService:
_config_from_dict(config).normalized()
# Join a finished job's pump OUTSIDE the lock: its final state writes take this
# lock (via _apply_event / the exit handler), so joining under it would stall
# the start for the whole timeout and then let the stale pump overwrite the new
# job's state once the lock was released.
# Join a finished job's pump OUTSIDE the lock: its final state writes take this lock, so
# joining under it would stall the start and let the stale pump overwrite the new state.
with self._lock:
if self._proc is not None and self._proc.is_alive():
raise RuntimeError("A diffusion training job is already running.")
@ -337,8 +319,7 @@ class DiffusionTrainingService:
if self._proc is None or not self._proc.is_alive() or self._stop_queue is None:
return False
try:
# Bare True keeps the wire format older trainers expect; the dict form
# carries the no-save cancel flag the trainer's _check_stop understands.
# Bare True = older wire format; the dict form carries the no-save cancel flag.
self._stop_queue.put(True if save else {"save": False})
except Exception: # noqa: BLE001
return False
@ -459,9 +440,8 @@ class DiffusionTrainingService:
elif etype == "model_load_completed":
s.update(in_model_load = False, message = "Training...")
elif etype == "preparing":
# A long precompute phase (e.g. the VAE latent cache) between model load and
# the first step; surfaced so the UI shows visible progress instead of a
# silent "Loading base model..." stall.
# A long precompute phase (e.g. VAE latent cache) before the first step; surfaced
# so the UI shows progress instead of a silent "Loading base model..." stall.
done, total = ev.get("done"), ev.get("total")
stage = str(ev.get("stage", "prepare")).replace("_", " ")
s.update(
@ -474,13 +454,11 @@ class DiffusionTrainingService:
),
)
elif etype == "warning":
# Non-fatal trainer notes (e.g. torch.compile falling back to eager); keep
# training state, surface the text.
# Non-fatal trainer notes; keep training state, surface the text.
s["message"] = str(ev.get("message", "warning"))
elif etype == "progress":
# Null any non-finite float (NaN/Inf from a divergent step or an inf grad
# norm) so the JSON status stays strict-parseable; a missing key keeps the
# last value, a present-but-non-finite one becomes None.
# Null any non-finite float so the JSON stays strict-parseable; a missing key
# keeps the last value, a present-but-non-finite one becomes None.
loss = _finite_or_none(ev["loss"]) if "loss" in ev else s["loss"]
avg_loss = _finite_or_none(ev["avg_loss"]) if "avg_loss" in ev else s["avg_loss"]
learning_rate = (
@ -501,8 +479,7 @@ class DiffusionTrainingService:
grad_norm = grad_norm,
message = "Training...",
)
# Fold optional perf fields (emitted by the trainers) so the UI can show
# throughput + peak VRAM without a separate channel.
# Fold optional perf fields so the UI shows throughput + peak VRAM.
if ev.get("samples_per_second") is not None:
s["samples_per_second"] = ev.get("samples_per_second")
if ev.get("peak_memory_gb") is not None:
@ -516,9 +493,8 @@ class DiffusionTrainingService:
ev.get("grad_norm"),
)
elif etype == "complete":
# Reset in_model_load: a stop during model load emits complete without a
# preceding model_load_completed, which would otherwise leave a stale
# loading indicator after the job ended.
# Reset in_model_load: a stop during model load emits complete with no preceding
# model_load_completed, which would otherwise leave a stale loading indicator.
s.update(
active = False,
in_model_load = False,
@ -540,8 +516,7 @@ class DiffusionTrainingService:
if ev.get("base_model") is not None:
s["base_model"] = ev.get("base_model")
elif etype == "error":
# Reset in_model_load too: an error raised during model loading has no
# model_load_completed, so the terminal state must clear it explicitly.
# Reset in_model_load too: an error during model loading has no model_load_completed.
s.update(
active = False,
in_model_load = False,