Wan2.2 video fixes, pipeline cache, smarter offload, model-load tasks

Wan2.2 A14B (dual-expert) generation fixes:
- Fuse the Lightning distill LoRA into BOTH experts (transformer +
  transformer_2); diffusers' fuse_lora defaults to ["transformer"] only, which
  left the low-noise expert undistilled → 4-step clips collapsed to a solid
  colour. Also load per-request fighter/env LoRAs into both experts.
- Pre-configure the wan22_lightning_4step preset with the local high/low-noise
  LoRAs (lora_high/lora_low), used when acceleration is enabled, ignored when
  not; surfaced in the Acceleration UI.
- Safety net: only apply the preset's low step count when the distill LoRA
  actually fused, else fall back to safe steps.
- Skip bitsandbytes/quanto quant for the VAE (conv-only → "no linear modules").

VRAM / offload:
- Strategy auto-selection actually fires now ('auto' is normalised, not passed
  through as a no-op) and no longer double-counts the runtime/accel reserve.
- Graceful OOM degrade ladder: full-GPU → balanced @ configured% → 80 → 60 →
  40 → sequential → disk, respecting the model's balanced_gpu_percent as the
  starting cap. Expose 'balanced' as a selectable offload strategy.

Pipeline disk cache (--pipeline-cache / --rebuild-pipeline-cache):
- Cache the quantized base pipeline to disk and reload it on later starts,
  skipping re-download/re-quantization; accel LoRA re-fused per load. Fail-safe
  with self-healing invalidate-and-rebuild.

Tasks / misc:
- Show model loading as a (non-cancellable, non-pausable) Tasks entry.
- Filter the Tasks-page pollers from the access log unless --debug-web.
- Township gen script: per-image keyframe progress (no longer all-or-nothing).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

codai/admin/routes.py

@@ -2052,7 +2052,7 @@ async def api_tasks(username: str = Depends(require_admin)):
         seen.add(t["id"])
         t = dict(t)
         t["cancellable"] = bool(t.get("cancellable", True) and t.get("active", False))
-        t["pausable"] = (t.get("status") == "running")
+        t["pausable"] = bool(t.get("pausable", True) and t.get("status") == "running")
         t["restartable"] = False
         tasks.append(t)
 

codai/admin/templates/models.html

@@ -671,13 +671,15 @@ window.__DEFAULT_WHISPER_SERVER_PATH__ = {{ default_whisper_server_path|tojson }
         <div class="form-row" style="margin:0">
           <label class="form-label">Strategy</label>
           <select id="cfg-offload-strategy" class="form-input">
-            <option value="auto">Auto</option>
-            <option value="model">CPU offload (model)</option>
-            <option value="sequential">CPU offload (sequential)</option>
+            <option value="auto">Auto (pick from free VRAM)</option>
+            <option value="none">None (GPU only)</option>
+            <option value="balanced">Balanced (fill GPU, spill to CPU → disk)</option>
+            <option value="model">CPU offload (model — module-by-module)</option>
+            <option value="sequential">CPU offload (sequential — most aggressive)</option>
             <option value="cpu">CPU RAM (legacy)</option>
             <option value="disk">Disk</option>
-            <option value="none">None (GPU only)</option>
           </select>
+          <span class="form-hint">Auto picks full-GPU when the weights fit, else Balanced. Pick <b>Balanced</b> + lower the GPU % below for a model that's just over VRAM; <b>model</b>/<b>sequential</b> keep less on GPU (slower, but the safest fit).</span>
         </div>
         <div class="form-row" style="margin:0">
           <label class="form-label">Offload directory</label>
@@ -729,6 +731,12 @@ window.__DEFAULT_WHISPER_SERVER_PATH__ = {{ default_whisper_server_path|tojson }
             <label class="form-label">Distill LoRA <span class="muted">(path or HF repo, optionally repo:weight_name.safetensors; blank for turbo full-models)</span></label>
             <input type="text" id="cfg-accel-lora" class="form-input" placeholder="e.g. ByteDance/SDXL-Lightning:sdxl_lightning_4step_lora.safetensors">
           </div>
+          <div class="form-row" style="max-width:560px">
+            <label class="form-label">Distill LoRA — high/low noise <span class="muted">(Wan2.2 A14B two-expert only; overrides the single LoRA per expert)</span></label>
+            <input type="text" id="cfg-accel-lora-high" class="form-input" placeholder="high-noise → transformer  (e.g. repo:..._high_noise_..._4step.safetensors)">
+            <input type="text" id="cfg-accel-lora-low" class="form-input" style="margin-top:.4rem" placeholder="low-noise → transformer_2  (e.g. repo:..._low_noise_..._4step.safetensors)">
+            <span class="form-hint">Wan2.2 A14B has two experts; the distill LoRA must be fused into <b>both</b> or the clip collapses to a solid colour at 4 steps. Leave blank to apply the single LoRA above to both.</span>
+          </div>
           <div style="display:flex;gap:1rem;flex-wrap:wrap">
             <div class="form-row" style="max-width:130px">
               <label class="form-label">LoRA weight</label>
@@ -2801,6 +2809,8 @@ async function _populateAccel(a){
   const sel = document.getElementById('cfg-accel-preset');
   sel.value = [...sel.options].some(o=>o.value===(a.preset||'')) ? a.preset : 'custom';
   document.getElementById('cfg-accel-lora').value = a.lora || '';
+  document.getElementById('cfg-accel-lora-high').value = a.lora_high || '';
+  document.getElementById('cfg-accel-lora-low').value = a.lora_low || '';
   document.getElementById('cfg-accel-weight').value = a.lora_weight != null ? a.lora_weight : '';
   document.getElementById('cfg-accel-steps').value = a.steps != null ? a.steps : '';
   document.getElementById('cfg-accel-guidance').value = a.guidance_scale != null ? a.guidance_scale : '';
@@ -2815,6 +2825,8 @@ function _collectAccel(){
     enabled: true,
     preset: document.getElementById('cfg-accel-preset').value || 'custom',
     lora: document.getElementById('cfg-accel-lora').value.trim() || null,
+    lora_high: document.getElementById('cfg-accel-lora-high').value.trim() || null,
+    lora_low: document.getElementById('cfg-accel-lora-low').value.trim() || null,
     lora_weight: num('cfg-accel-weight'),
     steps: num('cfg-accel-steps'),
     guidance_scale: num('cfg-accel-guidance'),

codai/admin/templates/tasks.html

@@ -69,7 +69,7 @@ function fmtTime(s) {
   } catch { return ''; }
 }
 
-const KIND_LABEL = {training:'Training', image:'Image', video:'Video', audio:'Audio', text:'Text', pipeline:'Pipeline', request:'Request'};
+const KIND_LABEL = {training:'Training', image:'Image', video:'Video', audio:'Audio', text:'Text', pipeline:'Pipeline', request:'Request', loading:'Loading'};
 const STATUS_BADGE = {
   running:'badge-admin', queued:'badge-user', done:'badge-ok', error:'badge-err',
   cancelled:'badge-user', interrupted:'badge-warn'

codai/api/audio_gen.py

@@ -266,13 +266,15 @@ async def audio_generate(request: AudioGenerationRequest, http_request: Request
         device = _derive_device()
         model_type = _detect_audio_gen_type(model_name)
         _ag_cfg = model_info.get('config') or {}
+        from codai.tasks import loading_task
         try:
-            if model_type in ('musicgen', 'audiogen'):
-                pipe = await asyncio.get_event_loop().run_in_executor(
-                    None, _load_musicgen, model_name, device)
-            else:
-                pipe = await asyncio.get_event_loop().run_in_executor(
-                    None, _load_audioldm, model_name, device, _ag_cfg)
+            with loading_task(model_name, model_type="audio"):
+                if model_type in ('musicgen', 'audiogen'):
+                    pipe = await asyncio.get_event_loop().run_in_executor(
+                        None, _load_musicgen, model_name, device)
+                else:
+                    pipe = await asyncio.get_event_loop().run_in_executor(
+                        None, _load_audioldm, model_name, device, _ag_cfg)
         except Exception as e:
             raise HTTPException(status_code=500, detail=f"Failed to load audio gen model: {e}")
         multi_model_manager.models[model_key] = pipe

codai/api/embeddings.py

@@ -121,9 +121,11 @@ async def create_embeddings(request: EmbeddingsRequest, http_request: Request =
 
     if model_obj is None:
         device = _derive_device()
+        from codai.tasks import loading_task
         try:
-            model_obj = await asyncio.get_event_loop().run_in_executor(
-                None, _load_embedding_model, model_name, device, _emb_cfg)
+            with loading_task(model_name, model_type="embedding"):
+                model_obj = await asyncio.get_event_loop().run_in_executor(
+                    None, _load_embedding_model, model_name, device, _emb_cfg)
         except Exception as e:
             raise HTTPException(status_code=500, detail=f"Failed to load embedding model: {e}")
         multi_model_manager.models[model_key] = model_obj

codai/api/images.py

@@ -1370,7 +1370,14 @@ async def create_image_generation(request: ImageGenerationRequest, http_request:
         is_gguf = _is_gguf_model(model_name)
         diffusers_error = None
         sdcpp_error = None
-        
+
+        # Show the load as a (non-cancellable) Tasks-page entry spanning both
+        # backend attempts; finished done on success, error only if all fail.
+        from codai.tasks import task_registry as _treg
+        _ltid = _treg.register("loading", title=f"Loading {model_name}",
+                               model=model_key, status="running",
+                               cancellable=False, pausable=False)
+
         # Try diffusers first (for non-GGUF models)
         if not is_gguf:
             try:
@@ -1391,6 +1398,7 @@ async def create_image_generation(request: ImageGenerationRequest, http_request:
                         pass
                     print(f"Loaded diffusers model: {model_name}")
 
+                    _treg.finish(_ltid, "done")
                     return await _generate_with_diffusers(pipeline, request, global_args, http_request)
                     
             except ImportError as e:
@@ -1426,7 +1434,8 @@ async def create_image_generation(request: ImageGenerationRequest, http_request:
                     except Exception:
                         pass
                     print(f"Loaded sd.cpp model: {model_name}")
-                    
+
+                    _treg.finish(_ltid, "done")
                     return await _generate_with_sdcpp(sd_model, request, global_args,
                                                       http_request, model_config=cfg)
             else:
@@ -1449,6 +1458,7 @@ async def create_image_generation(request: ImageGenerationRequest, http_request:
         if sdcpp_error:
             error_details.append(f"sd.cpp: {sdcpp_error}")
         
+        _treg.finish(_ltid, "error", "; ".join(error_details) or "no compatible backend")
         raise HTTPException(
             status_code=400,
             detail=f"Failed to load image model '{model_name}'. Errors: {'; '.join(error_details) if error_details else 'No compatible backend found'}"

codai/api/video.py

@@ -724,6 +724,66 @@ def _free_pipeline_vram(pipe) -> None:
         pass
 
 
+def _lora_file_size_gb(ref) -> float:
+    """Size (GB) of a distill LoRA ref ('repo:weight' / local path) from cache."""
+    if not ref:
+        return 0.0
+    try:
+        import os
+        from codai.models.acceleration import _split_lora_ref
+        if os.path.isfile(ref):
+            return os.path.getsize(ref) / 1e9
+        repo, weight = _split_lora_ref(ref)
+        if repo and os.path.isfile(repo):
+            return os.path.getsize(repo) / 1e9
+        if repo and weight:
+            from huggingface_hub import try_to_load_from_cache
+            p = try_to_load_from_cache(repo, weight)
+            if isinstance(p, str) and os.path.isfile(p):
+                return os.path.getsize(p) / 1e9
+    except Exception:
+        pass
+    return 0.0
+
+
+def _accel_vram_gb(model_cfg: dict) -> float:
+    """VRAM the fused acceleration/distill LoRA(s) add. Sums the actual cached
+    file sizes when known (both Wan2.2 experts), else a conservative reserve when
+    acceleration is enabled but the size can't be resolved."""
+    try:
+        from codai.models.acceleration import resolve_acceleration
+        a = resolve_acceleration(model_cfg)
+        if not a:
+            return 0.0
+        refs = [r for r in (a.get('lora_high'), a.get('lora_low'), a.get('lora')) if r]
+        refs = list(dict.fromkeys(refs))  # a single LoRA on both experts counts once
+        total, known = 0.0, False
+        for r in refs:
+            sz = _lora_file_size_gb(r)
+            if sz > 0:
+                total += sz
+                known = True
+        if known:
+            return total
+        return 2.5 if refs else 0.0  # accel on but size unknown → reserve
+    except Exception:
+        return 0.0
+
+
+def _video_runtime_reserve_gb(request) -> float:
+    """Rough VRAM headroom for the denoise activations + VAE-decode spike, which
+    scales with frame count × resolution. Keeps the auto-offload decision from
+    being too optimistic for long/high-res clips."""
+    try:
+        nf = int(getattr(request, 'num_frames', None) or 16)
+        w = int(getattr(request, 'width', None) or 512)
+        h = int(getattr(request, 'height', None) or 512)
+        base = 3.0 * (nf / 16.0) * ((w * h) / (512.0 * 512.0))
+        return max(2.0, min(12.0, base))
+    except Exception:
+        return 3.0
+
+
 def _load_video_pipeline(model_name: str, device: str, mode: str, offload: str = None, model_cfg: dict = None):
     # GGUF models go through stable-diffusion.cpp, not diffusers
     from codai.api.images import _is_gguf_model
@@ -740,6 +800,30 @@ def _load_video_pipeline(model_name: str, device: str, mode: str, offload: str =
     dtype_map = {'bf16': torch.bfloat16, 'f16': torch.float16, 'f32': torch.float32}
     torch_dtype = dtype_map.get(_model_precision, torch.bfloat16)
 
+    # ── Pipeline disk cache (--pipeline-cache) ───────────────────────────────
+    # When a previously-built, quantized pipeline is cached on disk, load the
+    # pre-quantized weights from there (no re-download / re-quantization). The
+    # cache holds the BASE pipeline only — the acceleration LoRA is re-fused by
+    # the caller per load — so the source just swaps from the HF id to the cache
+    # dir and the quantization config is skipped (the saved components carry it).
+    _orig_model_name = model_name
+    _pc_save_path = None
+    _loading_from_cache = False
+    try:
+        from codai.models import pipeline_cache as _pcache
+        if _pcache.enabled():
+            _pc_path = _pcache.path(model_name, model_cfg)
+            if _pcache.valid(_pc_path):
+                print(f"  [pipeline-cache] HIT — loading pre-quantized pipeline "
+                      f"from {_pc_path}")
+                model_name = _pc_path
+                _loading_from_cache = True
+            else:
+                _pc_save_path = _pc_path
+                print(f"  [pipeline-cache] MISS — will build, then cache for next start")
+    except Exception as _pce:
+        print(f"  [pipeline-cache] unavailable ({_pce})")
+
     # ── Quantization (from per-model config) ─────────────────────────────────
     # Per-component overrides ('component_quantization' map) win; otherwise the
     # global load_in_4bit/8bit flag is applied to every heavy component (the
@@ -747,16 +831,22 @@ def _load_video_pipeline(model_name: str, device: str, mode: str, offload: str =
     # UMT5 text_encoder is ~11 GB in bf16 and must be quantized to fit on-GPU).
     from codai.models.hf_loading import (
         build_pipeline_quant_config, build_gguf_pipeline_components)
-    _quant_config, _quant_desc = build_pipeline_quant_config(
-        model_name, model_cfg, torch_dtype)
-    if _quant_config is not None:
-        print(f"  Video quantization: {_quant_desc}")
-    # GGUF-quantized components (Q5_K/Q6_K etc.) are loaded from their .gguf
-    # files and injected into the pipeline as pre-built components.
-    _gguf_components, _gguf_desc = build_gguf_pipeline_components(
-        model_name, model_cfg, torch_dtype)
-    if _gguf_components:
-        print(f"  Video GGUF components: {_gguf_desc}")
+    if _loading_from_cache:
+        # Cached components already carry their quantization config; don't rebuild
+        # it (and don't re-inject GGUF components — they were baked into the cache).
+        _quant_config, _quant_desc = None, ''
+        _gguf_components, _gguf_desc = {}, ''
+    else:
+        _quant_config, _quant_desc = build_pipeline_quant_config(
+            model_name, model_cfg, torch_dtype)
+        if _quant_config is not None:
+            print(f"  Video quantization: {_quant_desc}")
+        # GGUF-quantized components (Q5_K/Q6_K etc.) are loaded from their .gguf
+        # files and injected into the pipeline as pre-built components.
+        _gguf_components, _gguf_desc = build_gguf_pipeline_components(
+            model_name, model_cfg, torch_dtype)
+        if _gguf_components:
+            print(f"  Video GGUF components: {_gguf_desc}")
 
     def _with_quant(kw: dict) -> dict:
         """Inject quantization_config + GGUF components into from_pretrained kwargs."""
@@ -952,6 +1042,16 @@ def _load_video_pipeline(model_name: str, device: str, mode: str, offload: str =
         _report_device_map(pipe)
         _report_offload_dir_size()
         _mem_snapshot("after load")
+        # Persist the freshly-built quantized pipeline to the disk cache so the
+        # next start can skip the rebuild. Only on a cache MISS (we didn't load
+        # from it) and when --pipeline-cache is on. Best-effort; never fatal.
+        if _pc_save_path and not _loading_from_cache:
+            try:
+                from codai.models import pipeline_cache as _pcache
+                _pcache.save(pipe, _pc_save_path,
+                             model_name=_orig_model_name, model_cfg=model_cfg)
+            except Exception as _se:
+                print(f"  [pipeline-cache] save skipped ({_se})")
 
     # NOTE: we deliberately do NOT lower sys.setswitchinterval here.  A previous
     # version set it to 0.001s to keep the asyncio loop responsive during the
@@ -968,22 +1068,25 @@ def _load_video_pipeline(model_name: str, device: str, mode: str, offload: str =
         # RAM (and disk if needed). This is the preferred strategy when the
         # model won't fit entirely in VRAM but should maximise GPU utilisation.
         # `gpu_percent` (0–100) controls what fraction of FREE VRAM to occupy.
+        # Configured GPU cap for balanced (per-model balanced_gpu_percent, else 80).
+        # This is the STARTING cap for the balanced chain; on OOM it steps down.
         _gpu_pct = float((model_cfg or {}).get('balanced_gpu_percent') or 80)
-        try:
-            if torch.cuda.is_available():
-                _free_v, _ = torch.cuda.mem_get_info()
-                _avail_gpu_gb = (_free_v / 1e9) * (_gpu_pct / 100.0)
-            else:
-                _avail_gpu_gb = 0.0
-        except Exception:
-            _avail_gpu_gb = 0.0
-        _cpu_avail_gb = min(48, max(4, int(
-            _psutil.virtual_memory().available * 0.60 / 1e9)))
 
-        if offload == 'balanced':
-            _mem_snapshot("before balanced GPU+CPU load")
+        def _load_balanced(gpu_pct: float):
+            """device_map='balanced' capping GPU at gpu_pct% of free VRAM, spilling
+            to CPU then disk. Assigns the outer `pipe` so a failed attempt's VRAM
+            is reclaimable via _clear_mem. Raises (RuntimeError/MemoryError) on OOM."""
+            nonlocal pipe
+            try:
+                _free_v2, _ = torch.cuda.mem_get_info() if torch.cuda.is_available() else (0, 0)
+                _avail_gpu_gb = (_free_v2 / 1e9) * (gpu_pct / 100.0)
+            except Exception:
+                _avail_gpu_gb = 0.0
+            _cpu_avail_gb = min(48, max(4, int(
+                _psutil.virtual_memory().available * 0.60 / 1e9)))
+            _mem_snapshot(f"before balanced {gpu_pct:.0f}% GPU+CPU load")
             print(f"  Video load strategy: balanced GPU+CPU "
-                  f"({_gpu_pct:.0f}% GPU → {_avail_gpu_gb:.1f} GiB / "
+                  f"({gpu_pct:.0f}% GPU → {_avail_gpu_gb:.1f} GiB / "
                   f"CPU {_cpu_avail_gb} GiB / overflow → {_offload_dir})")
             os.makedirs(_offload_dir, exist_ok=True)
             try:
@@ -998,13 +1101,78 @@ def _load_video_pipeline(model_name: str, device: str, mode: str, offload: str =
                     low_cpu_mem_usage=True,
                 )))
             except (TypeError, ValueError):
+                # Pipeline doesn't accept device_map — fall to model CPU offload.
                 pipe = PClass.from_pretrained(**_with_quant(dict(
                     pretrained_model_name_or_path=model_name,
                     torch_dtype=torch_dtype, low_cpu_mem_usage=True)))
                 pipe.enable_model_cpu_offload()
-            _report_loaded(pipe, f"balanced {_gpu_pct:.0f}%GPU+CPU")
+            _report_loaded(pipe, f"balanced {gpu_pct:.0f}%GPU+CPU")
+            return pipe
+
+        def _try_balanced_chain():
+            """Try balanced starting at the configured GPU %, stepping down through
+            60% then 40% on OOM. Returns the pipe, or None if every step OOM'd
+            (caller then falls through to model/sequential/disk offload)."""
+            # Start at the configured cap, then step down through the standard
+            # checkpoints (80/60/40) that sit below it. So 90% → 90,80,60,40;
+            # 80% → 80,60,40; 70% → 70,60,40; 50% → 50,40.
+            _pcts = sorted({_gpu_pct} | {p for p in (80.0, 60.0, 40.0)
+                                         if p < _gpu_pct}, reverse=True)
+            for _i, _pct in enumerate(_pcts):
+                try:
+                    return _load_balanced(_pct)
+                except (RuntimeError, MemoryError) as _e:
+                    if not _is_oom(_e):
+                        raise
+                    _nxt = _pcts[_i + 1] if _i + 1 < len(_pcts) else None
+                    if _nxt is not None:
+                        print(f"  Video: balanced {_pct:.0f}% GPU OOM — "
+                              f"retrying at {_nxt:.0f}% GPU…")
+                    else:
+                        print(f"  Video: balanced {_pct:.0f}% GPU OOM — "
+                              f"falling back to sequential CPU offload…")
+                    _clear_mem()
+            return None
+
+        def _load_sequential():
+            """Most aggressive fit: stream each submodule GPU↔CPU during the
+            forward pass (slowest, lowest VRAM). Assigns the outer `pipe`. Raises
+            on OOM."""
+            nonlocal pipe
+            _mem_snapshot("before sequential CPU offload load")
+            print("  Video load strategy: sequential CPU offload "
+                  "(each submodule GPU↔CPU during forward; slowest, lowest VRAM)")
+            pipe = PClass.from_pretrained(**_with_quant(dict(
+                pretrained_model_name_or_path=model_name,
+                torch_dtype=torch_dtype, low_cpu_mem_usage=True)))
+            pipe.enable_sequential_cpu_offload()
+            _report_loaded(pipe, "sequential CPU offload")
             return pipe
 
+        def _try_balanced_then_sequential():
+            """Balanced chain (configured% → 60 → 40), then sequential CPU offload
+            if all balanced steps OOM. Returns the pipe, or None if even sequential
+            OOMs (caller falls through to the disk-offload attempts)."""
+            p = _try_balanced_chain()
+            if p is not None:
+                return p
+            _clear_mem()
+            try:
+                return _load_sequential()
+            except (RuntimeError, MemoryError) as _e:
+                if not _is_oom(_e):
+                    raise
+                print(f"  Video: sequential CPU offload OOM ({_e}) — "
+                      f"trying disk offload…")
+                _clear_mem()
+            return None
+
+        if offload == 'balanced':
+            pipe = _try_balanced_then_sequential()
+            if pipe is not None:
+                return pipe
+            # Even sequential OOM'd → continue to the disk-offload attempts.
+
         # ── Attempt 0: full GPU ──────────────────────────────────────────────
         if offload not in ('model', 'sequential', 'disk', 'balanced'):
             _mem_snapshot("before full-GPU load")
@@ -1030,7 +1198,15 @@ def _load_video_pipeline(model_name: str, device: str, mode: str, offload: str =
             except (RuntimeError, MemoryError) as e:
                 if not _is_oom(e):
                     raise
-                print(f"  Video: full-GPU OOM ({e}) — trying model CPU offload…")
+                print(f"  Video: full-GPU OOM ({e}) — falling back to balanced "
+                      f"GPU+CPU (starting at {_gpu_pct:.0f}% GPU)…")
+                _clear_mem()
+                # Graceful degrade: balanced at the configured %, then 60%, then
+                # 40%, then sequential CPU offload, before the slower disk paths.
+                pipe = _try_balanced_then_sequential()
+                if pipe is not None:
+                    return pipe
+                print("  Video: balanced + sequential all OOM — trying disk offload…")
                 _clear_mem()
 
         # ── Attempt 1: model CPU offload ─────────────────────────────────────
@@ -1323,12 +1499,23 @@ def _sync_video_loras(pipe, loras) -> None:
                       or getattr(pipe, 'unet', None) or pipe).__name__
     loaded = []  # (name, weight) that actually registered on the model
     before = _present_adapters(pipe)
+    # Wan2.2 A14B is a two-expert MoE — a LoRA loaded only into `transformer`
+    # leaves the low-noise expert (transformer_2) un-adapted, so the concept
+    # fades out as denoising hands over to it. Load into both when present.
+    _has_t2 = getattr(pipe, 'transformer_2', None) is not None
     for model, name, w in desired:
         try:
             pipe.load_lora_weights(model, adapter_name=name)
         except Exception as e:
             print(f"  [video][lora] failed to load '{name}': {e}")
             continue
+        if _has_t2:
+            try:
+                pipe.load_lora_weights(model, adapter_name=name,
+                                       load_into_transformer_2=True)
+            except Exception as e:
+                print(f"  [video][lora] '{name}' not loaded into transformer_2 "
+                      f"(low-noise expert): {e}")
         now = _present_adapters(pipe)
         if name in now and name not in before:
             loaded.append((name, w))
@@ -1488,10 +1675,19 @@ def _generate_video(pipe, request: VideoGenerationRequest):
     # _build_call_kwargs only populates them when the request specified them, so
     # setdefault below correctly leaves an explicit request value untouched.
     _accel = getattr(pipe, '_coderai_accel', None)
-    if _accel:
+    # Only trust the preset's low step-count / guidance when the distill LoRA
+    # actually fused. If it didn't (e.g. the Wan2.2 low-noise expert never got its
+    # LoRA, or the ref failed to load), running 4 steps un-distilled collapses the
+    # video to a solid colour — so fall back to a safe step count instead.
+    _accel_fused = getattr(pipe, '_coderai_accel_fused', None)
+    if _accel and _accel_fused is not False:
         from codai.models.acceleration import accel_call_defaults
         for _k, _v in accel_call_defaults(_accel).items():
             kw.setdefault(_k, _v)
+    elif _accel and _accel_fused is False:
+        print("  [video][accel] distill LoRA not fused — ignoring the preset's "
+              "low step count and using safe defaults (25 steps) to avoid a "
+              "collapsed/solid-colour result")
     kw.setdefault('num_inference_steps', 25)
     kw.setdefault('guidance_scale', 7.5)
     kw.setdefault('num_frames', 16)
@@ -2104,6 +2300,12 @@ async def video_generations(request: VideoGenerationRequest,
 
     if pipe is None:
         _offload = _model_cfg.get('offload_strategy') or None
+        # 'auto' (the default) means "let coderai pick from available VRAM" — it is
+        # NOT a diffusers strategy, and passing it through lands on the full-GPU
+        # path that then disk-thrashes. Normalise it to None so the VRAM check
+        # below decides between full-GPU and balanced GPU+CPU.
+        if _offload == 'auto':
+            _offload = None
         # Auto-select "balanced" strategy when the model (including runtime
         # reserve: KV/activation spike, VAE decode) exceeds available VRAM even
         # after eviction. Going straight to "balanced" (GPU-first + CPU spill)
@@ -2117,22 +2319,40 @@ async def video_generations(request: VideoGenerationRequest,
                 if _t.cuda.is_available():
                     _free_v, _ = _t.cuda.mem_get_info()
                     _free_gb = _free_v / 1e9
-                    _need_gb = multi_model_manager._get_model_used_vram_gb(
+                    # `_get_model_used_vram_gb` is the *measured total* footprint —
+                    # it already includes the runtime/activation reserve AND the
+                    # fused acceleration LoRA (it's measured after fusion). So do
+                    # NOT re-add those (that over-counts and wrongly forces the
+                    # slow balanced+disk path). Per-request LoRAs are extra.
+                    _base_gb = multi_model_manager._get_model_used_vram_gb(
                         model_key, model_name)
-                    if _need_gb > 0 and _free_gb < _need_gb:
+                    # full-GPU only needs the WEIGHTS to fit at load time (the
+                    # bundled ~runtime reserve is a gen-time allowance, and the
+                    # full-GPU path has its own OOM→offload fallback). Keep a
+                    # headroom margin so a model that *marginally* fits uses the
+                    # much faster full-GPU strategy rather than balanced+disk.
+                    _need_gb = _base_gb + _lora_extra_gb
+                    _margin = 2.5  # ≈ the bundled runtime reserve
+                    if _base_gb > 0 and _free_gb < (_need_gb - _margin):
                         _gpu_pct = float(_model_cfg.get('balanced_gpu_percent') or 80)
-                        print(f"  VRAM insufficient for full-GPU load "
-                              f"({_need_gb:.1f} GB needed, {_free_gb:.1f} GB free) "
-                              f"— auto-selecting balanced strategy "
-                              f"({_gpu_pct:.0f}% GPU + CPU spill)")
+                        print(f"  VRAM well short for full-GPU load "
+                              f"({_need_gb:.1f} GB measured need + LoRA; "
+                              f"{_free_gb:.1f} GB free) — auto-selecting balanced "
+                              f"strategy ({_gpu_pct:.0f}% GPU + CPU spill)")
                         _offload = 'balanced'
+                    else:
+                        print(f"  Full-GPU load looks viable "
+                              f"({_need_gb:.1f} GB measured need, {_free_gb:.1f} GB "
+                              f"free) — using full GPU (it falls back to offload on OOM)")
             except Exception:
                 pass
         # Snapshot free VRAM so we can record the real footprint after load.
         _vram_before = multi_model_manager.vram_before_load()
+        from codai.tasks import loading_task
         try:
-            pipe = await asyncio.get_event_loop().run_in_executor(
-                None, _load_video_pipeline, model_name, device, request.mode, _offload, _model_cfg)
+            with loading_task(model_name, model_type="video"):
+                pipe = await asyncio.get_event_loop().run_in_executor(
+                    None, _load_video_pipeline, model_name, device, request.mode, _offload, _model_cfg)
         except Exception as e:
             multi_model_manager._mark_cuda_poisoned_if_fatal(e)
             if getattr(multi_model_manager, 'cuda_context_poisoned', False):
@@ -2140,7 +2360,24 @@ async def video_generations(request: VideoGenerationRequest,
                     "CUDA context corrupted (device-side assert) while loading the "
                     "video model. Restart coderai to recover. "
                     f"Original error: {str(e).splitlines()[0]}"))
-            raise HTTPException(status_code=500, detail=f"Failed to load video model: {e}")
+            # Self-heal: a failed load from a (possibly stale/corrupt) pipeline
+            # cache should drop the cache and rebuild once rather than wedging.
+            _retried_fresh = False
+            try:
+                from codai.models import pipeline_cache as _pcache
+                if _pcache.enabled() and _pcache.valid(_pcache.path(model_name, _model_cfg)):
+                    print(f"  [pipeline-cache] load failed ({str(e).splitlines()[0]}) "
+                          f"— invalidating cache and rebuilding")
+                    _pcache.invalidate(model_name, _model_cfg)
+                    with loading_task(model_name, model_type="video"):
+                        pipe = await asyncio.get_event_loop().run_in_executor(
+                            None, _load_video_pipeline, model_name, device,
+                            request.mode, _offload, _model_cfg)
+                    _retried_fresh = True
+            except Exception:
+                _retried_fresh = False
+            if not _retried_fresh:
+                raise HTTPException(status_code=500, detail=f"Failed to load video model: {e}")
         # Fuse any configured acceleration/distillation LoRA (Lightning / Lightx2v /
         # LCM) into the freshly loaded pipeline. Done once at load; cached pipes keep
         # it. No-op for sd.cpp pipes and when no acceleration is configured.

codai/cli.py

@@ -254,4 +254,18 @@ configuration directory (--config DIR, default: OS-specific CoderAI directory).
              "Mid-flight jobs are marked 'cancelled' (checkpoints are kept, so they "
              "can still be restarted manually from the Tasks page).",
     )
+    parser.add_argument(
+        "--pipeline-cache",
+        action="store_true",
+        help="Cache quantized diffusers pipelines to disk after the first build "
+             "and reload them from that cache on later starts — skipping the "
+             "expensive re-download/re-quantization (e.g. the Wan2.2 A14B). The "
+             "fast acceleration LoRA fuse is re-applied per load. Uses extra disk.",
+    )
+    parser.add_argument(
+        "--rebuild-pipeline-cache",
+        action="store_true",
+        help="Ignore any existing pipeline cache and rebuild it from scratch this "
+             "run (use after changing a model's quantization/precision config).",
+    )
     return parser.parse_args()

codai/main.py

@@ -878,6 +878,14 @@ def main():
     if not _resume_jobs:
         print("LoRA job recovery: DISABLED (interrupted training will be cancelled on restart)")
 
+    if getattr(args, "pipeline_cache", False):
+        try:
+            from codai.models.pipeline_cache import cache_root
+            _pc_extra = " (rebuilding this run)" if getattr(args, "rebuild_pipeline_cache", False) else ""
+            print(f"Pipeline cache: ENABLED{_pc_extra} — quantized pipelines cached at {cache_root()}")
+        except Exception:
+            print("Pipeline cache: ENABLED")
+
     # Set environment profiles module global args
     from codai.api.environments import set_global_args as set_envs_global_args
     set_envs_global_args(global_args)
@@ -964,13 +972,18 @@ def main():
     if not _debug_web:
         class _AccessNoiseFilter(logging.Filter):
             # uvicorn.access record args: (client_addr, method, full_path, http_ver, status)
-            _NOISY = ("/v1/loras/progress",)
+            _NOISY_PREFIX = ("/v1/loras/progress",)
+            # Exact-match only, so the live Tasks-page pollers are dropped but the
+            # user-initiated action endpoints (/admin/api/tasks/{id}/pause, …) still log.
+            _NOISY_EXACT = ("/admin/api/tasks", "/admin/api/system-stats")
             def filter(self, record):
                 try:
                     args = record.args
                     if isinstance(args, (tuple, list)) and len(args) >= 3:
                         path = str(args[2]).split("?", 1)[0]
-                        if any(path == p or path.startswith(p) for p in self._NOISY):
+                        if path in self._NOISY_EXACT:
+                            return False
+                        if any(path == p or path.startswith(p) for p in self._NOISY_PREFIX):
                             return False
                 except Exception:
                     pass

codai/models/acceleration.py

@@ -52,7 +52,15 @@ ACCEL_PRESETS: dict = {
         "label": "Wan2.2 Lightning (4-step DMD)",
         "family": "wan",
         "applies_to": ["video"],
-        "lora": "lightx2v/Wan2.2-Lightning",
+        # Wan2.2 A14B is a two-expert MoE: the distill LoRA must be fused into BOTH
+        # the high-noise (transformer) and low-noise (transformer_2) experts, or the
+        # clip collapses to a solid colour at 4 steps. These default to the locally
+        # installed lightx2v/Wan2.2-Lightning weights (resolved from cache — not a
+        # download); override per model in the Acceleration config for T2V or a
+        # different rank/version. `lora` stays None because the two experts differ.
+        "lora": None,
+        "lora_high": "lightx2v/Wan2.2-Lightning:Wan2.2-I2V-A14B-4steps-lora-rank64-Seko-V1/high_noise_model.safetensors",
+        "lora_low": "lightx2v/Wan2.2-Lightning:Wan2.2-I2V-A14B-4steps-lora-rank64-Seko-V1/low_noise_model.safetensors",
         "lora_weight": 1.0,
         "steps": 4,
         "guidance_scale": 1.0,
@@ -175,6 +183,12 @@ def resolve_acceleration(model_cfg: Optional[dict]) -> Optional[dict]:
     out = {
         "preset": preset_key or "custom",
         "lora": _pick("lora"),
+        # Wan2.2 A14B is a two-expert MoE: the distill LoRA differs for the
+        # high-noise (transformer) and low-noise (transformer_2) experts. When
+        # these are set they take precedence over the single `lora` per expert;
+        # otherwise the single `lora` is applied to BOTH experts.
+        "lora_high": _pick("lora_high"),
+        "lora_low": _pick("lora_low"),
         "lora_weight": _pick("lora_weight", 1.0),
         "steps": _pick("steps"),
         "guidance_scale": _pick("guidance_scale"),
@@ -255,35 +269,94 @@ def apply_accel_to_pipeline(pipe, accel: Optional[dict]) -> None:
             log.warning("[accel] flow_shift apply failed: %s", e)
 
     # 3. Fuse the distill LoRA (when one is configured — turbo has none).
-    lora_ref = accel.get("lora")
-    if not lora_ref:
+    #    `_coderai_accel_fused` records whether a distill LoRA actually baked in,
+    #    so the generator only drops to the preset's low step count when the model
+    #    is genuinely distilled (running 4 steps un-distilled collapses the video
+    #    to a solid colour — exactly the Wan2.2 dual-expert failure mode).
+    try:
+        pipe._coderai_accel_fused = False
+    except Exception:
+        pass
+
+    has_t2 = getattr(pipe, "transformer_2", None) is not None
+    lora_high = accel.get("lora_high") or accel.get("lora")
+    lora_low = accel.get("lora_low") or accel.get("lora")
+    if not lora_high and not lora_low:
+        # No LoRA (e.g. a full distilled model like SDXL-Turbo) — treat as distilled.
+        try:
+            pipe._coderai_accel_fused = True
+        except Exception:
+            pass
         return
     if not hasattr(pipe, "load_lora_weights"):
         log.warning("[accel] pipeline %s has no load_lora_weights — cannot fuse "
                     "acceleration LoRA", type(pipe).__name__)
         return
-    repo, weight_name = _split_lora_ref(lora_ref)
+
     weight = float(accel.get("lora_weight") or 1.0)
-    try:
-        load_kwargs = {"adapter_name": "__accel__"}
+
+    def _load_one(ref, into_t2: bool, adapter: str) -> bool:
+        repo, weight_name = _split_lora_ref(ref)
+        kw = {"adapter_name": adapter}
         if weight_name:
-            load_kwargs["weight_name"] = weight_name
-        pipe.load_lora_weights(repo, **load_kwargs)
+            kw["weight_name"] = weight_name
+        if into_t2:
+            kw["load_into_transformer_2"] = True
+        pipe.load_lora_weights(repo, **kw)
+        return True
+
+    loaded_adapters = []
+    try:
+        # High-noise expert (transformer) — always.
+        if lora_high and _load_one(lora_high, False, "__accel__"):
+            loaded_adapters.append("__accel__")
+        # Low-noise expert (transformer_2) — only on dual-expert Wan2.2 models.
+        if has_t2 and lora_low:
+            try:
+                if _load_one(lora_low, True, "__accel_2__"):
+                    loaded_adapters.append("__accel_2__")
+            except Exception as e2:
+                log.warning("[accel] could not load distill LoRA into transformer_2 "
+                            "(%s) — low-noise expert stays un-distilled", e2)
+        elif has_t2 and not lora_low:
+            log.warning("[accel] model has a second expert (transformer_2) but no "
+                        "low-noise distill LoRA — set acceleration.lora_low")
+        if not loaded_adapters:
+            raise RuntimeError("no distill adapter registered on the pipeline")
         try:
-            pipe.set_adapters(["__accel__"], [weight])
+            pipe.set_adapters(loaded_adapters, [weight] * len(loaded_adapters))
         except Exception:
             pass
-        # Bake it in, then drop the adapter handle so per-request LoRAs are clean.
-        pipe.fuse_lora(lora_scale=weight)
+        # Bake them in, then drop the adapter handles so per-request LoRAs are clean.
+        # CRITICAL: diffusers' Wan fuse_lora defaults to components=["transformer"],
+        # so without naming transformer_2 the low-noise expert's distill adapter is
+        # never fused — and the subsequent unload strips it off, leaving that expert
+        # undistilled. At 4 steps that collapses the clip to a solid colour. Fuse
+        # BOTH experts explicitly.
+        _fuse_components = ["transformer"]
+        if has_t2:
+            _fuse_components.append("transformer_2")
+        try:
+            pipe.fuse_lora(components=_fuse_components, lora_scale=weight)
+        except TypeError:
+            # Older diffusers without the `components` kwarg — best effort.
+            pipe.fuse_lora(lora_scale=weight)
         try:
             pipe.unload_lora_weights()
         except Exception:
             pass
-        log.info("[accel] fused distillation LoRA %s (weight=%s) into %s",
-                 repo, weight, type(pipe).__name__)
+        try:
+            pipe._coderai_accel_fused = True
+        except Exception:
+            pass
+        log.info("[accel] fused distillation LoRA(s) %s (weight=%s) into %s%s",
+                 loaded_adapters, weight, type(pipe).__name__,
+                 " (both experts)" if len(loaded_adapters) > 1 else "")
     except Exception as e:
-        log.warning("[accel] failed to fuse acceleration LoRA %s: %s — generating "
-                    "without acceleration", lora_ref, e)
+        log.warning("[accel] failed to fuse acceleration LoRA (high=%s low=%s): %s "
+                    "— generating WITHOUT acceleration (step count will fall back to "
+                    "a safe default, not the preset's distilled count)",
+                    lora_high, lora_low, e)
 
 
 def accel_call_defaults(accel: Optional[dict]) -> dict:

codai/models/hf_loading.py

@@ -182,6 +182,15 @@ def build_pipeline_quant_config(model_name: str, cfg: Optional[Dict[str, Any]],
                        bnb_4bit_compute_dtype=dtype, bnb_4bit_use_double_quant=True)
         return BnB(load_in_8bit=True)
 
+    def _bnb_incompatible(name: str) -> bool:
+        # bitsandbytes (4/8-bit) and optimum-quanto (2-bit) only quantize
+        # nn.Linear. A fully-convolutional component (the VAE) has no Linear
+        # layers, so applying them triggers a hard "no linear modules were found"
+        # error. Such components must stay full precision (a smaller VAE comes
+        # from a GGUF VAE instead, handled separately).
+        n = (name or '').lower()
+        return n == 'vae' or n.endswith('_vae') or n.startswith('vae')
+
     quant_mapping: Dict[str, Any] = {}
     descs = []
     if comp_q:
@@ -189,6 +198,11 @@ def build_pipeline_quant_config(model_name: str, cfg: Optional[Dict[str, Any]],
             mode = _normalize_quant_mode(raw_mode)  # GGUF/none → None here
             if mode is None:
                 continue
+            if _bnb_incompatible(name):
+                print(f"  Skipping {mode} for '{name}': it has no Linear layers "
+                      f"(conv-only) — bitsandbytes/quanto cannot quantize it; "
+                      f"leaving full precision (use a GGUF VAE to shrink the VAE).")
+                continue
             cfg_obj = _mk(comp_lib.get(name, 'diffusers'), mode)
             if cfg_obj is not None:
                 quant_mapping[name] = cfg_obj
@@ -198,6 +212,8 @@ def build_pipeline_quant_config(model_name: str, cfg: Optional[Dict[str, Any]],
         targets = [n for n in comp_lib if _is_heavy(n)] or \
             ['transformer', 'transformer_2', 'text_encoder', 'unet']
         for name in targets:
+            if _bnb_incompatible(name):
+                continue
             cfg_obj = _mk(comp_lib.get(name, 'diffusers'), mode)
             if cfg_obj is not None:
                 quant_mapping[name] = cfg_obj

codai/models/manager.py

@@ -817,7 +817,9 @@ class MultiModelManager:
                 print(f"Loading default model on demand: {self.default_model}")
                 _snap = self.vram_before_load()
                 kwargs['expected_vram_gb'] = self._get_model_used_vram_gb(self.default_model)
-                model_manager.load_model(self.default_model, backend_type=backend_type, **kwargs)
+                from codai.tasks import loading_task
+                with loading_task(self.default_model, model_type="text"):
+                    model_manager.load_model(self.default_model, backend_type=backend_type, **kwargs)
                 self.add_model(self.default_model, model_manager)
                 self.record_vram_delta(self.default_model, _snap)
                 self.current_model_key = self.default_model
@@ -916,7 +918,9 @@ class MultiModelManager:
                 # it can decide whether Flash-Attention-2 is safe (FA2 requires the
                 # whole model on GPU; it device-side-asserts when layers offload).
                 kwargs['expected_vram_gb'] = self._get_model_used_vram_gb(model_name)
-                model_manager.load_model(model_name, backend_type=backend_type, **kwargs)
+                from codai.tasks import loading_task
+                with loading_task(model_name, model_type="text"):
+                    model_manager.load_model(model_name, backend_type=backend_type, **kwargs)
                 self.add_model(model_name, model_manager)
                 self.record_vram_delta(model_name, _snap)
                 self.current_model_key = model_name

codai/models/pipeline_cache.py

+# CoderAI - OpenAI-compatible API server
+# Copyright (C) 2026 Stefy Lanza <stefy@nexlab.net>
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program. If not, see <https://www.gnu.org/licenses/>.
+
+"""On-disk cache of *built* diffusers pipelines.
+
+Building a large quantized video pipeline (e.g. Wan2.2 A14B at 4-bit) is slow:
+download + bitsandbytes quantization of ~28B parameters. The weights don't change
+between restarts, so once built we ``save_pretrained`` the pipeline to a local
+cache keyed by ``(model, quantization, precision)``. A later start with
+``--pipeline-cache`` reloads from there with a plain ``from_pretrained`` of the
+already-quantized weights — no re-download, no re-quantization.
+
+Scope: only the *base* pipeline is cached. The acceleration/distillation LoRA is
+NOT baked into the cache — it is re-fused on every load (a fast operation), so the
+cache stays independent of the (cheap to re-apply) ``acceleration`` config and we
+avoid the fragile round-trip of serialising a fused + quantized model.
+
+Everything here is best-effort: any failure (save or load) is swallowed and the
+caller falls back to a normal build, so the cache can never break generation.
+"""
+
+import hashlib
+import json
+import os
+import shutil
+import time
+from typing import Optional
+
+# Bump when the cache layout / marker format changes so stale caches are ignored.
+_CACHE_VERSION = 1
+
+
+def _global_args():
+    try:
+        from codai.api.state import get_global_args
+        return get_global_args()
+    except Exception:
+        return None
+
+
+def enabled() -> bool:
+    """True when --pipeline-cache was passed."""
+    ga = _global_args()
+    return bool(ga is not None and getattr(ga, "pipeline_cache", False))
+
+
+def _force_rebuild() -> bool:
+    ga = _global_args()
+    return bool(ga is not None and getattr(ga, "rebuild_pipeline_cache", False))
+
+
+def cache_root() -> str:
+    """Root dir for cached pipelines. Sits next to the offload dir by default."""
+    ga = _global_args()
+    offload_dir = getattr(ga, "offload_dir", None) if ga else None
+    if offload_dir:
+        root = os.path.join(os.path.dirname(os.path.abspath(os.path.expanduser(offload_dir))),
+                            "pipeline_cache")
+    else:
+        root = os.path.join(os.path.expanduser("~"), ".cache", "coderai", "pipeline_cache")
+    return root
+
+
+def _signature(model_name: str, model_cfg: Optional[dict]) -> str:
+    """Stable hash of everything that changes the *built* (quantized) weights:
+    the model id, the quantization choices, and the precision. NOT acceleration
+    (re-applied per load) and NOT offload (a runtime placement decision)."""
+    c = model_cfg or {}
+    payload = {
+        "v": _CACHE_VERSION,
+        "model": model_name,
+        "precision": c.get("precision") or "bf16",
+        "load_in_4bit": bool(c.get("load_in_4bit", False)),
+        "load_in_8bit": bool(c.get("load_in_8bit", False)),
+        "component_quantization": c.get("component_quantization") or {},
+    }
+    blob = json.dumps(payload, sort_keys=True, default=str)
+    return hashlib.sha256(blob.encode()).hexdigest()[:16]
+
+
+def _safe_name(model_name: str) -> str:
+    return "".join(ch if ch.isalnum() or ch in "-._" else "_" for ch in model_name)[:80]
+
+
+def path(model_name: str, model_cfg: Optional[dict]) -> str:
+    """Absolute cache directory for this model + quant/precision signature."""
+    return os.path.join(cache_root(),
+                        f"{_safe_name(model_name)}__{_signature(model_name, model_cfg)}")
+
+
+def _marker(p: str) -> str:
+    return os.path.join(p, ".coderai_pipeline_cache.json")
+
+
+def valid(p: str) -> bool:
+    """True if a complete, current cache exists at ``p`` and rebuild wasn't forced."""
+    if not p or _force_rebuild():
+        return False
+    try:
+        if not os.path.isfile(os.path.join(p, "model_index.json")):
+            return False
+        with open(_marker(p)) as f:
+            meta = json.load(f)
+        return meta.get("version") == _CACHE_VERSION and meta.get("complete") is True
+    except Exception:
+        return False
+
+
+def invalidate(model_name: str, model_cfg: Optional[dict]) -> None:
+    """Delete a model's cache dir (e.g. after a failed cache load) so the next
+    build rewrites it. Best-effort."""
+    try:
+        p = path(model_name, model_cfg)
+        if p and os.path.isdir(p):
+            shutil.rmtree(p, ignore_errors=True)
+            print(f"  [pipeline-cache] invalidated {p}")
+    except Exception:
+        pass
+
+
+def save(pipe, p: str, *, model_name: str = "", model_cfg: Optional[dict] = None) -> bool:
+    """Serialize ``pipe`` to the cache dir ``p`` (atomic via a temp dir).
+
+    Returns True on success. Any failure is logged and returns False — the caller
+    keeps the freshly built in-memory pipeline regardless."""
+    if not p:
+        return False
+    tmp = p + ".building"
+    try:
+        os.makedirs(cache_root(), exist_ok=True)
+        if os.path.exists(tmp):
+            shutil.rmtree(tmp, ignore_errors=True)
+        print(f"  [pipeline-cache] saving quantized pipeline → {p}")
+        t0 = time.time()
+        pipe.save_pretrained(tmp)
+        with open(_marker(tmp), "w") as f:
+            json.dump({
+                "version": _CACHE_VERSION, "complete": True,
+                "model": model_name, "saved_at": time.time(),
+                "signature": _signature(model_name, model_cfg),
+            }, f)
+        if os.path.exists(p):
+            shutil.rmtree(p, ignore_errors=True)
+        os.replace(tmp, p)
+        print(f"  [pipeline-cache] saved in {time.time() - t0:.0f}s")
+        return True
+    except Exception as e:
+        print(f"  [pipeline-cache] save failed ({e}) — continuing without a cache")
+        try:
+            shutil.rmtree(tmp, ignore_errors=True)
+        except Exception:
+            pass
+        return False

codai/tasks/__init__.py

@@ -23,6 +23,7 @@ from codai.tasks.registry import (
     task_registry,
     raise_if_cancelled,
     wait_if_paused,
+    loading_task,
 )
 
 __all__ = [
@@ -32,4 +33,5 @@ __all__ = [
     "task_registry",
     "raise_if_cancelled",
     "wait_if_paused",
+    "loading_task",
 ]

codai/tasks/registry.py

@@ -32,6 +32,7 @@ a task with a ``job_id`` links the two.
 import threading
 import time
 import uuid
+from contextlib import contextmanager
 from dataclasses import asdict, dataclass, field
 from typing import Dict, List, Optional
 
@@ -60,6 +61,7 @@ class Task:
     ended_at: Optional[float] = None
     cancellable: bool = True
     restartable: bool = False
+    pausable: bool = True
     paused: bool = False
 
     def to_dict(self) -> dict:
@@ -83,13 +85,14 @@ class TaskRegistry:
     def register(self, kind: str, *, title: str = "", model: str = "",
                  total: int = 0, job_id: Optional[str] = None,
                  status: str = "queued", cancellable: bool = True,
-                 restartable: bool = False, task_id: Optional[str] = None) -> str:
+                 restartable: bool = False, pausable: bool = True,
+                 task_id: Optional[str] = None) -> str:
         tid = task_id or f"task-{uuid.uuid4().hex[:12]}"
         with self._lock:
             self._tasks[tid] = Task(
                 id=tid, kind=kind, title=title, model=model, total=total,
                 job_id=job_id, status=status, cancellable=cancellable,
-                restartable=restartable,
+                restartable=restartable, pausable=pausable,
             )
             self._events[tid] = threading.Event()
             self._pause_events[tid] = threading.Event()
@@ -256,3 +259,25 @@ def wait_if_paused(task_id: Optional[str]) -> None:
     Returns immediately when not paused; raises :class:`TaskCancelled` if the
     task is cancelled while paused. A falsy ``task_id`` is a no-op."""
     task_registry.wait_if_paused(task_id)
+
+
+@contextmanager
+def loading_task(model: str, *, model_type: str = "model", title: Optional[str] = None):
+    """Context manager that shows a model load as a Tasks-page entry.
+
+    Model loading can't be paused or cancelled (it's a single blocking
+    ``from_pretrained`` / ``Llama(...)`` call), so the task is registered
+    non-cancellable and non-pausable — the Tasks UI shows it with no action
+    buttons. The task finishes ``done`` on success or ``error`` on exception.
+    Re-entrant guard: a nested load of the same model_key reuses no task; each
+    call is independent (loads don't nest in practice)."""
+    label = title or f"Loading {model}"
+    tid = task_registry.register(
+        "loading", title=label, model=model or "", status="running",
+        cancellable=False, restartable=False, pausable=False)
+    try:
+        yield tid
+        task_registry.finish(tid, "done")
+    except BaseException as e:  # noqa: BLE001 — record then re-raise
+        task_registry.finish(tid, "error", str(e)[:200] or e.__class__.__name__)
+        raise

tools/gen_township_fighters.py

@@ -1709,10 +1709,21 @@ def _generate_keyframes(client: CoderAIClient, image_model: str, keyframe_dir: P
                         fight_plan: list, outcome_plan: list, consistency: set,
                         lora_map: dict, char_strength: float, keyframe_steps: int,
                         keyframe_size: str, lora_weight: float,
-                        env_lora_map: dict = None, env_lora_weight: float = 0.8):
+                        env_lora_map: dict = None, env_lora_weight: float = 0.8,
+                        kf_cb=None):
     """Generate one keyframe still per clip (image model). Saved as PNG keyed by
     the clip's output stem so the render phase can pick them up as init images.
-    Resumable: existing PNGs are kept."""
+    Resumable: existing PNGs are kept.
+
+    kf_cb(stem, phase, ok) — optional; fired so callers (the web match-render job)
+    can show per-image progress. phase is "start" (this keyframe begins) or "end"
+    (finished, ok=True/False); a reused/existing PNG fires "end" with ok=True."""
+    def _kf(stem, phase, ok=None):
+        if kf_cb:
+            try:
+                kf_cb(stem, phase, ok)
+            except Exception:
+                pass
     keyframe_dir.mkdir(parents=True, exist_ok=True)
     use_ip = "ipadapter" in consistency or "keyframe" in consistency
     use_lora = "lora" in consistency
@@ -1751,7 +1762,9 @@ def _generate_keyframes(client: CoderAIClient, image_model: str, keyframe_dir: P
         out_png = keyframe_dir / f"{stem}.png"
         if out_png.exists() and out_png.stat().st_size > 0:
             skipped += 1
+            _kf(stem, "end", True)   # already present — show it as done
             continue
+        _kf(stem, "start")
         profiles = list(fighters) if use_ip else None
         loras = None
         if use_lora:
@@ -1771,8 +1784,10 @@ def _generate_keyframes(client: CoderAIClient, image_model: str, keyframe_dir: P
             )
             out_png.write_bytes(img)
             made += 1
+            _kf(stem, "end", True)
         except Exception as e:
             failed += 1
+            _kf(stem, "end", False)
             _log(f"    ✗ keyframe {stem} failed: {e}")
     _log(f"  ── Keyframes: {made} new, {skipped} reused, {failed} failed ──")
 
@@ -2968,13 +2983,30 @@ def launch_web_ui(default_args):
                     _done("no missing keyframes — all present")
                     return
                 _set_items([f"keyframe {s}" for s in work])
-                for i, s in enumerate(work):
-                    _item(i, "start")
-                    if not missing_only:
+                _kf_idx = {s: i for i, s in enumerate(work)}
+                # Delete the targeted PNGs first so they're actually regenerated
+                # (missing-only keeps existing ones → reported done by the callback).
+                if not missing_only:
+                    for s in work:
                         try:
                             (kdir / f"{s}.png").unlink()
                         except Exception:
                             pass
+                # Per-image progress: _generate_keyframes fires kf_cb as each
+                # keyframe starts/finishes, so the bars advance image-by-image
+                # instead of all flipping at the end.
+                _kf_done = [0]
+
+                def _kf_cb(stem, phase, ok=None):
+                    i = _kf_idx.get(stem)
+                    if i is None:
+                        return
+                    _item(i, phase, ok)
+                    if phase == "end":
+                        _kf_done[0] += 1
+                        _prog(10 + int(88 * _kf_done[0] / max(1, len(work))),
+                              f"keyframe {_kf_done[0]}/{len(work)} done")
+
                 _prog(10, ("filling in {n} missing keyframe(s)…" if missing_only
                            else "regenerating {n} keyframe(s)…").format(n=len(work)))
                 try:
@@ -2984,11 +3016,13 @@ def launch_web_ui(default_args):
                         float(getattr(default_args, "character_strength", 0.7)),
                         int(getattr(default_args, "keyframe_steps", 28)),
                         getattr(default_args, "keyframe_size", "512x512"), lw,
-                        env_lora_map=env_lora_map, env_lora_weight=elw)
+                        env_lora_map=env_lora_map, env_lora_weight=elw,
+                        kf_cb=_kf_cb)
                 except Exception as e:
                     _fail(f"keyframe regeneration failed: {e}")
                     return
-                # Mark each item done/failed by whether its PNG now exists.
+                # Safety net: resolve any item the callback didn't (e.g. a stem
+                # _generate_keyframes never visited) by whether its PNG now exists.
                 for i, s in enumerate(work):
                     _item(i, "end", (kdir / f"{s}.png").exists())
                 made = sum(1 for s in work if (kdir / f"{s}.png").exists())