From 1762dab12fcd736f3f11246a0afd04ec1817229b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 12 Jul 2026 11:46:23 +0000 Subject: [PATCH] Tighten comments across the remaining image stack files --- .../core/training/diffusion_lora_trainer.py | 120 +++++++----------- .../training/diffusion_training_service.py | 77 ++++------- 2 files changed, 69 insertions(+), 128 deletions(-) diff --git a/studio/backend/core/training/diffusion_lora_trainer.py b/studio/backend/core/training/diffusion_lora_trainer.py index 0e45a25a7c..fb8d911d3a 100644 --- a/studio/backend/core/training/diffusion_lora_trainer.py +++ b/studio/backend/core/training/diffusion_lora_trainer.py @@ -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()} ) diff --git a/studio/backend/core/training/diffusion_training_service.py b/studio/backend/core/training/diffusion_training_service.py index c02445c4f1..f55636f75a 100644 --- a/studio/backend/core/training/diffusion_training_service.py +++ b/studio/backend/core/training/diffusion_training_service.py @@ -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,