video: incremental per-component pipeline cache (load/build/cache one component at a time)

A Wan2.2 A14B pipeline that only loads via offload could never be cached — diffusers
can't save_pretrained an offloaded pipe, so every load re-quantized from bf16 shards
(~14 min) and risked OOM. Cache it per-component instead:

- pipeline_cache: per-component subdirs each with their own completion marker +
  signature (component_dir / component_valid / save_component), atomic via .building
  + sweep_stale(). A load evicted before all components are cached keeps the finished
  ones; the next load mixes cached + freshly-built (now-cached) ones, converging to a
  full cache. (Plus the earlier sweep_stale / _unsavable_reason robustness.)

- hf_loading.build_cached_components: for each bnb-quantizable heavy component, load
  from the per-component cache if present, else build fresh from the model — each ONE
  AT A TIME on the GPU (where bnb quantizes it to uniform 4-bit), saved, then moved to
  CPU + empty_cache. Peak VRAM is a single component even when the whole model doesn't
  fit; uniform 4-bit by construction (no bf16/4-bit mix). Verified the bnb round-trip
  (load 4-bit on GPU -> save_pretrained -> to('cpu') -> reload).

- video._load_video_pipeline: inject the cached components like GGUF components so the
  existing offload ladder assembles them. Kill-switch CODERAI_INCREMENTAL_CACHE=0.
  Fully fallback-safe: on any load failure the request handler retries once with
  use_incremental_cache=False (plain build = today's behaviour).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RdMufYvtTbtGDWsiZVoXce

codai/api/video.py

@@ -843,7 +843,7 @@ def _video_runtime_reserve_gb(request) -> float:
         return 3.0
 
 
-def _load_video_pipeline(model_name: str, device: str, mode: str, offload: str = None, model_cfg: dict = None):
+def _load_video_pipeline(model_name: str, device: str, mode: str, offload: str = None, model_cfg: dict = None, use_incremental_cache: bool = True):
     # GGUF models go through stable-diffusion.cpp, not diffusers
     from codai.api.images import _is_gguf_model
     if _is_gguf_model(model_name):
@@ -906,6 +906,22 @@ def _load_video_pipeline(model_name: str, device: str, mode: str, offload: str =
             model_name, model_cfg, torch_dtype)
         if _gguf_components:
             print(f"  Video GGUF components: {_gguf_desc}")
+        # Incremental per-component cache: load each bnb-quantizable heavy component
+        # from the per-component cache (or build+cache it one-at-a-time on the GPU),
+        # then inject like GGUF components so the offload ladder assembles them. This
+        # lets a model that only loads via offload still get cached (a wholesale
+        # save_pretrained of an offloaded pipe is impossible). Fully fallback-safe;
+        # the caller retries with use_incremental_cache=False if assembly fails.
+        if use_incremental_cache:
+            try:
+                from codai.models.hf_loading import build_cached_components
+                _cached_comps, _cc_desc = build_cached_components(
+                    model_name, model_cfg, torch_dtype)
+                if _cached_comps:
+                    print(f"  Video cached components: {_cc_desc}")
+                    _gguf_components = {**_gguf_components, **_cached_comps}
+            except Exception as _cce:
+                print(f"  [pipeline-cache] incremental components skipped ({_cce})")
 
     def _with_quant(kw: dict) -> dict:
         """Inject quantization_config + GGUF components into from_pretrained kwargs."""
@@ -3264,21 +3280,30 @@ 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]}"))
-            # Self-heal: a failed load from a (possibly stale/corrupt) pipeline
-            # cache should drop the cache and rebuild once rather than wedging.
+            # Self-heal: a failed load should retry once WITHOUT the incremental
+            # per-component cache + with any monolithic cache dropped — i.e. fall
+            # back to a plain from_pretrained build (today's behaviour) rather than
+            # wedging. This catches both a stale/corrupt cache and any breakage in
+            # the component-injection assembly path.
             _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")
+                          f"— invalidating monolithic cache")
                     _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:
+                print(f"  [pipeline-cache] retrying load without incremental "
+                      f"component cache (plain build)")
+                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, False)
+                _retried_fresh = True
+            except Exception as _e2:
+                import traceback as _tb2
+                print(f"  [video] fallback (no-incremental) load ALSO failed for "
+                      f"'{model_name}': {str(_e2).splitlines()[0] if str(_e2) else type(_e2).__name__}\n"
+                      + _tb2.format_exc())
                 _retried_fresh = False
             if not _retried_fresh:
                 raise HTTPException(status_code=500, detail=f"Failed to load video model: {e}")

codai/models/hf_loading.py

@@ -316,6 +316,130 @@ def build_gguf_pipeline_components(model_name: str, cfg: Optional[Dict[str, Any]
     return components, ', '.join(descs)
 
 
+def build_cached_components(model_name: str, cfg: Optional[Dict[str, Any]], dtype):
+    """Incremental per-component cache for big bnb-quantized video pipelines.
+
+    A Wan2.2 A14B pipeline often won't fit on the GPU and can only load with
+    offload — and an offloaded pipeline can't be save_pretrained wholesale. But
+    each component IS small enough to pass through the GPU alone, where bnb
+    quantizes it to a uniform 4-bit. So for every bnb-quantizable heavy component
+    we: load it from the per-component cache if present, else build it fresh from
+    the model (quantizing on the GPU) and cache it — then move it to CPU to free
+    the GPU slot for the next one. Peak VRAM is ONE component even when the whole
+    model doesn't fit; the result is injected into the pipeline's from_pretrained
+    exactly like GGUF components, and the offload ladder places them at runtime.
+
+    Caching is incremental and resumable: a load evicted before all components are
+    cached leaves the finished ones valid, and the next load mixes cached with
+    freshly-built (now cached) ones until the set is complete.
+
+    Returns ``(components_dict, description)``. Empty when disabled (kill-switch
+    CODERAI_INCREMENTAL_CACHE=0, --pipeline-cache off, or no bnb quant). Fully
+    fallback-safe: any per-component failure skips that component (the normal
+    pipeline load handles it) and is logged."""
+    import os
+    if os.environ.get("CODERAI_INCREMENTAL_CACHE", "1") == "0":
+        return {}, ''
+    c = _norm(cfg)
+    comp_q = c.get('component_quantization') or {}
+    global_4 = bool(c.get('load_in_4bit', False))
+    global_8 = bool(c.get('load_in_8bit', False))
+    if not comp_q and not (global_4 or global_8):
+        return {}, ''
+    try:
+        from codai.models import pipeline_cache as _pcache
+        if not _pcache.enabled():
+            return {}, ''
+    except Exception:
+        return {}, ''
+    try:
+        import diffusers
+        import transformers
+        from diffusers import BitsAndBytesConfig as DiffBnb
+        from transformers import BitsAndBytesConfig as TfBnb
+    except Exception as e:
+        print(f"  [pipeline-cache] incremental build unavailable: {e}")
+        return {}, ''
+
+    specs = _discover_components(model_name)  # name -> [library, class_name]
+
+    def _is_heavy(name: str) -> bool:
+        return (name.startswith('transformer') or name == 'unet'
+                or name.startswith('text_encoder'))
+
+    def _bnb_incompatible(name: str) -> bool:
+        n = (name or '').lower()
+        return n == 'vae' or n.endswith('_vae') or n.startswith('vae')
+
+    # Which components to handle, and at what bnb width.
+    targets: Dict[str, str] = {}
+    if comp_q:
+        for name, raw in comp_q.items():
+            mode = _normalize_quant_mode(raw)  # gguf/none/2bit -> not bnb here
+            if mode in ('4bit', '8bit') and not _bnb_incompatible(name) and name in specs:
+                targets[name] = mode
+    else:
+        mode = '4bit' if global_4 else '8bit'
+        for name in specs:
+            if _is_heavy(name) and not _bnb_incompatible(name):
+                targets[name] = mode
+    if not targets:
+        return {}, ''
+
+    def _mk(lib: str, mode: str):
+        BnB = TfBnb if lib == 'transformers' else DiffBnb
+        if mode == '4bit':
+            return BnB(load_in_4bit=True, bnb_4bit_quant_type='nf4',
+                       bnb_4bit_compute_dtype=dtype, bnb_4bit_use_double_quant=True)
+        return BnB(load_in_8bit=True)
+
+    def _cls(name):
+        spec = specs.get(name) or []
+        lib = spec[0] if spec else 'diffusers'
+        cls_name = spec[1] if len(spec) > 1 else None
+        mod = transformers if lib == 'transformers' else diffusers
+        return (getattr(mod, cls_name, None) if cls_name else None), lib
+
+    components: Dict[str, Any] = {}
+    descs = []
+    for name, mode in targets.items():
+        try:
+            cls, lib = _cls(name)
+            if cls is None or not hasattr(cls, 'from_pretrained'):
+                continue
+            cached = _pcache.component_valid(model_name, cfg, name)
+            if cached:
+                cdir = _pcache.component_dir(model_name, cfg, name)
+                print(f"  [pipeline-cache] component '{name}' HIT — loading 4-bit from cache")
+                comp = cls.from_pretrained(
+                    cdir, quantization_config=_mk(lib, mode),
+                    torch_dtype=dtype, device_map={'': 0})
+            else:
+                print(f"  [pipeline-cache] component '{name}' MISS — building 4-bit "
+                      f"(one component on GPU, then cached)")
+                comp = cls.from_pretrained(
+                    model_name, subfolder=name, quantization_config=_mk(lib, mode),
+                    torch_dtype=dtype, device_map={'': 0})
+                _pcache.save_component(comp, model_name, cfg, name)
+            # Free the single GPU slot for the next component; the offload ladder
+            # re-places it at runtime. bnb 4-bit survives the move to CPU (verified).
+            try:
+                comp = comp.to('cpu')
+                import torch as _t
+                if _t.cuda.is_available():
+                    _t.cuda.synchronize()
+                    _t.cuda.empty_cache()
+            except Exception:
+                pass
+            components[name] = comp
+            descs.append(f"{name}:{mode}{'(cached)' if cached else '(built)'}")
+        except Exception as e:
+            print(f"  [pipeline-cache] component '{name}' incremental build failed "
+                  f"({str(e).splitlines()[0] if str(e) else type(e).__name__}) — "
+                  f"leaving it to the normal pipeline load")
+    return components, ', '.join(descs)
+
+
 def build_from_pretrained_kwargs(
     cfg: Optional[Dict[str, Any]],
     *,

codai/models/pipeline_cache.py

@@ -120,6 +120,99 @@ def valid(p: str) -> bool:
         return False
 
 
+# ── Per-component (incremental) cache ────────────────────────────────────────
+# A big quantized video pipeline (Wan2.2 A14B) often won't fit on the GPU and can
+# only be loaded with offload — and an offloaded pipeline can't be save_pretrained
+# wholesale. But we CAN cache it one component at a time: each component is small
+# enough to pass through the GPU alone (where bitsandbytes quantizes it to a
+# uniform 4-bit), so we save each as it's built. The cache is then a set of
+# per-component subdirs under the same model+signature dir as the monolithic
+# cache, each with its OWN completion marker. A load that's evicted before every
+# component is cached still leaves the finished ones valid; the next load mixes
+# cached components with freshly-built (and now-cached) ones, converging to a full
+# cache over a couple of runs. Each component carries the same model+quant+
+# precision signature, so a config change invalidates it like the monolithic one.
+
+def _component_marker(cdir: str) -> str:
+    return os.path.join(cdir, ".coderai_component.json")
+
+
+def component_dir(model_name: str, model_cfg: Optional[dict], comp: str) -> str:
+    """Cache subdir for a single pipeline component (e.g. 'transformer')."""
+    return os.path.join(path(model_name, model_cfg), comp)
+
+
+def component_valid(model_name: str, model_cfg: Optional[dict], comp: str) -> bool:
+    """True if a complete, signature-matching cache for this component exists."""
+    if _force_rebuild():
+        return False
+    try:
+        cdir = component_dir(model_name, model_cfg, comp)
+        if not os.path.isfile(os.path.join(cdir, "config.json")):
+            return False
+        with open(_component_marker(cdir)) as f:
+            meta = json.load(f)
+        return (meta.get("version") == _CACHE_VERSION
+                and meta.get("complete") is True
+                and meta.get("signature") == _signature(model_name, model_cfg))
+    except Exception:
+        return False
+
+
+def save_component(comp_obj, model_name: str, model_cfg: Optional[dict], comp: str) -> bool:
+    """Atomically save one freshly-built (4-bit) component to its cache subdir.
+
+    Writes to ``<comp>.building`` then renames, with the completion marker written
+    last, so an interrupted save never leaves a half-written component that
+    component_valid() would accept. Best-effort: any failure returns False and the
+    caller keeps the in-memory component regardless."""
+    cdir = component_dir(model_name, model_cfg, comp)
+    if not cdir:
+        return False
+    tmp = cdir + ".building"
+    try:
+        os.makedirs(os.path.dirname(cdir), exist_ok=True)
+        sweep_stale()
+        if os.path.exists(tmp):
+            shutil.rmtree(tmp, ignore_errors=True)
+        print(f"  [pipeline-cache] caching component '{comp}' → {cdir}")
+        t0 = time.time()
+        comp_obj.save_pretrained(tmp)
+        with open(_component_marker(tmp), "w") as f:
+            json.dump({
+                "version": _CACHE_VERSION, "complete": True,
+                "model": model_name, "component": comp, "saved_at": time.time(),
+                "signature": _signature(model_name, model_cfg),
+                "bytes": _dir_size(tmp),
+            }, f)
+        if os.path.exists(cdir):
+            shutil.rmtree(cdir, ignore_errors=True)
+        os.replace(tmp, cdir)
+        print(f"  [pipeline-cache] cached '{comp}' in {time.time() - t0:.0f}s")
+        return True
+    except Exception as e:
+        print(f"  [pipeline-cache] component '{comp}' save failed ({e}) — skipping")
+        try:
+            shutil.rmtree(tmp, ignore_errors=True)
+        except Exception:
+            pass
+        return False
+
+
+def _dir_size(d: str) -> int:
+    total = 0
+    try:
+        for root, _, files in os.walk(d):
+            for fn in files:
+                try:
+                    total += os.path.getsize(os.path.join(root, fn))
+                except Exception:
+                    pass
+    except Exception:
+        pass
+    return total
+
+
 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."""