video: OOM fallback LADDER (no 5xx until exhausted) + monolithic cache finalize

From the latest render log: the per-component cache + same-model reuse worked
(clip 02 ran on clip 01's resident pipe, LoRA-synced, no reload), but BIG clips
(46-50 frames) OOM'd in the dual-expert forward under 'model' CPU offload, and the
single retry then crashed two ways:
  - NameError: loading_task was imported only in the load-only block, so on a
    REUSED pipe the retry's `with loading_task(...)` was unbound.
  - enable_sequential_cpu_offload on injected bnb-4bit components hits
    "Params4bit.__new__() got an unexpected keyword argument '_is_hf_initialized'".
So every big clip returned 500 instead of recovering.

Fixes:
1. Generation fallback LADDER. On a recoverable failure (OOM / device-mismatch) we
   no longer return 5xx — we free the pipe (outside the except, so the failed
   forward's traceback can't pin its VRAM) and reload with the next strategy, then
   retry generation. Return on the FIRST success, or 500 only after EVERY rung
   fails. Rungs: [already-loaded] -> balanced -> disk, all PLAIN device_map builds
   (incremental=False) so device_map places every component itself — sidestepping
   the injected-component failure modes ('model' OOM, 'sequential' Params4bit,
   device-map-can't-place-injected). loading_task imported up-front.
2. Monolithic cache finalize. After the per-component heavy build, complete the
   cache DIR (save light components + model_index.json + marker) so later loads use
   the monolithic HIT path: from_pretrained(dir, device_map=balanced), NO injection
   — big-clip generation fits AND loads pre-quantized from cache (fast, no
   re-quantize). Conservative: only marked complete when every component is present.

Net: first clip builds+finalizes the cache; thereafter loads are balanced-from-cache
(fits any clip size, no injection bugs), same-model clips reuse, and a transient OOM
walks the ladder instead of failing the clip.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RdMufYvtTbtGDWsiZVoXce

codai/api/video.py

@@ -3378,34 +3378,70 @@ async def video_generations(request: VideoGenerationRequest,
     except ImportError:
         pass
 
-    _gen_exc = None
-    _gen_err_str = ""
-    try:
-        if _is_sdcpp_video:
-            frames, fps = await asyncio.get_event_loop().run_in_executor(
-                None, _generate_sdcpp_video, pipe, request, _model_cfg)
-        else:
-            frames, fps = await asyncio.get_event_loop().run_in_executor(
-                None, _generate_video, pipe, request)
-    except Exception as e:
-        _vid_progress_done()
-        import traceback as _tb
-        _gen_exc = e
-        _gen_err_str = str(e).lower()
-        print(f"  [video] generation FAILED for '{model_name}': "
-              f"{str(e).splitlines()[0] if str(e) else type(e).__name__}\n"
-              + _tb.format_exc())
-        # CRITICAL: drop the traceback NOW. It pins the failed forward pass's GPU
-        # activations (and the resident expert) via this except's sys.exc_info(),
-        # so _free_pipeline_vram()/empty_cache() below would reclaim nothing and a
-        # retry reload would see the card still ~20 GB full and OOM again (the leak
-        # seen in the logs). We free + retry OUTSIDE this except block, below.
-        e.__traceback__ = None
-
-    if _gen_exc is not None:
-        # Outside the except block: sys.exc_info() no longer pins the failed
-        # forward's tensors, so the pipeline + activations are now collectable.
-        import gc as _gc, traceback as _tb
+    # ── Generation with an offload FALLBACK LADDER ───────────────────────────
+    # The model is already loaded (rung 0 = the pipe loaded above). On a recoverable
+    # failure (OOM / device-mismatch) we DON'T return to the client — we free the
+    # pipe (outside the except, so the failed forward's traceback can't pin its VRAM)
+    # and reload with the next, safer strategy, then retry generation. We return only
+    # on the FIRST success, or 500 after EVERY rung is exhausted.
+    #
+    # Fallback rungs are PLAIN device_map builds (incremental=False): no injected
+    # components, so device_map places everything itself — which sidesteps all the
+    # injected-component failure modes (hook-based 'model' OOMs big forwards;
+    # 'sequential' hits a bnb Params4bit bug; device_map can't place injected
+    # components). With the monolithic cache finalized after the first build, these
+    # load the pre-quantized weights from cache (fast), not a fresh re-quantize.
+    import gc as _gc, traceback as _tb
+    from codai.tasks import loading_task as _loading_task
+    _loop = asyncio.get_event_loop()
+    _rungs = [None, ('balanced', False), ('disk', False)]  # None = already-loaded pipe
+    frames = fps = None
+    _last_err = None
+    for _ri, _rung in enumerate(_rungs):
+        # (Re)load for fallback rungs; rung 0 reuses the already-loaded pipe.
+        if _rung is not None:
+            _strat, _incr = _rung
+            print(f"  [video] fallback {_ri}/{len(_rungs)-1}: reloading "
+                  f"'{model_name}' with strategy='{_strat}' (incremental={_incr}) — "
+                  f"NOT returning to client…")
+            try:
+                with _loading_task(model_name, model_type="video"):
+                    pipe = await _loop.run_in_executor(
+                        None, _load_video_pipeline, model_name, device, request.mode,
+                        _strat, _model_cfg, _incr)
+                multi_model_manager.models[model_key] = pipe
+                multi_model_manager.current_model_key = model_key
+            except Exception as _le:
+                print(f"  [video] fallback {_ri} load FAILED: "
+                      f"{str(_le).splitlines()[0] if str(_le) else type(_le).__name__}\n"
+                      + _tb.format_exc())
+                _last_err = _le
+                _le.__traceback__ = None
+                multi_model_manager.models.pop(model_key, None)
+                multi_model_manager.model_pools.pop(model_key, None)
+                _free_pipeline_vram(pipe)
+                pipe = None
+                continue  # try the next rung
+        # Generate on the current pipe.
+        try:
+            if _is_sdcpp_video:
+                frames, fps = await _loop.run_in_executor(
+                    None, _generate_sdcpp_video, pipe, request, _model_cfg)
+            else:
+                frames, fps = await _loop.run_in_executor(
+                    None, _generate_video, pipe, request)
+            _last_err = None
+            break  # success → return to client
+        except Exception as e:
+            _vid_progress_done()
+            _err_str = str(e).lower()
+            _which = "initial" if _rung is None else f"fallback {_ri}"
+            print(f"  [video] generation ({_which}) FAILED for '{model_name}': "
+                  f"{str(e).splitlines()[0] if str(e) else type(e).__name__}\n"
+                  + _tb.format_exc())
+            _last_err = e
+            e.__traceback__ = None  # drop NOW so the free below can reclaim VRAM
+        # Failure handling OUTSIDE the except: free the pipe + activations.
         multi_model_manager.models.pop(model_key, None)
         multi_model_manager.model_pools.pop(model_key, None)
         _free_pipeline_vram(pipe)
@@ -3419,46 +3455,18 @@ async def video_generations(request: VideoGenerationRequest,
                 _torch.cuda.empty_cache()
         except Exception:
             pass
+        _recoverable = ("out of memory" in _err_str
+                        or ("cuda" in _err_str and "memory" in _err_str)
+                        or "same device" in _err_str or "index_select" in _err_str
+                        or "cuda:0" in _err_str)
+        # sd.cpp pipes don't offload-retry; a non-recoverable error won't be helped
+        # by a different strategy — stop laddering and report.
+        if _is_sdcpp_video or not _recoverable:
+            break
 
-        _is_oom = ("out of memory" in _gen_err_str
-                   or ("cuda" in _gen_err_str and "memory" in _gen_err_str))
-        _dev_mismatch = ("same device" in _gen_err_str or "index_select" in _gen_err_str
-                         or "cuda:0" in _gen_err_str)
-        # Retry ONCE for OOM or device-mismatch (sd.cpp pipes don't offload-retry).
-        if (_is_oom or _dev_mismatch) and not _is_sdcpp_video:
-            # OOM  -> sequential offload (minimal VRAM, reliably fits the forward).
-            # mismatch -> drop incremental injection; the plain device_map build
-            # places every component itself, so there's nothing un-hooked.
-            _retry_strategy = 'sequential' if _is_oom else _offload
-            _retry_incremental = not _dev_mismatch
-            print(f"  [video] retrying generation once "
-                  f"(strategy={_retry_strategy or 'auto'}, "
-                  f"incremental={'on' if _retry_incremental else 'off'}) "
-                  f"after {'OOM' if _is_oom else 'device-mismatch'}…")
-            try:
-                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,
-                        _retry_strategy, _model_cfg, _retry_incremental)
-                multi_model_manager.models[model_key] = pipe
-                multi_model_manager.current_model_key = model_key
-                frames, fps = await asyncio.get_event_loop().run_in_executor(
-                    None, _generate_video, pipe, request)
-                _gen_exc = None  # recovered
-            except Exception as _e2:
-                print(f"  [video] generation retry ALSO failed for '{model_name}': "
-                      f"{str(_e2).splitlines()[0] if str(_e2) else type(_e2).__name__}\n"
-                      + _tb.format_exc())
-                multi_model_manager.models.pop(model_key, None)
-                multi_model_manager.model_pools.pop(model_key, None)
-                _free_pipeline_vram(pipe)
-                pipe = None
-                _e2.__traceback__ = None
-                raise HTTPException(status_code=500,
-                                    detail=f"Video generation failed (retry): {_e2}")
-        if _gen_exc is not None:
-            raise HTTPException(status_code=500,
-                                detail=f"Video generation failed: {_gen_exc}")
+    if _last_err is not None:
+        raise HTTPException(status_code=500,
+                            detail=f"Video generation failed: {_last_err}")
 
     # Encode raw frames to MP4 (per-model output quality via CRF when configured).
     try:

codai/models/hf_loading.py

@@ -437,9 +437,90 @@ def build_cached_components(model_name: str, cfg: Optional[Dict[str, Any]], dtyp
             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")
+
+    # If every heavy component is now cached, complete the cache DIR into a
+    # normally-loadable pipeline (light components + model_index.json + marker), so
+    # subsequent loads use the monolithic HIT path — from_pretrained(dir,
+    # device_map=…) with NO injection. That lets big-clip generation fit via
+    # device_map (the injected-component path is stuck on hook-based offload, where
+    # 'model' OOMs big forwards and 'sequential' hits a bnb Params4bit bug).
+    if components and len(components) == len(targets):
+        try:
+            finalize_monolithic_cache(model_name, cfg, dtype, set(components.keys()))
+        except Exception as _fe:
+            print(f"  [pipeline-cache] monolithic finalize skipped ({_fe})")
+
     return components, ', '.join(descs)
 
 
+def finalize_monolithic_cache(model_name: str, cfg: Optional[Dict[str, Any]],
+                              dtype, built_names) -> bool:
+    """Complete a per-component cache DIR into a pipeline from_pretrained() can load
+    directly: save the LIGHT components the heavy build didn't cover (vae /
+    scheduler / tokenizer / …), copy model_index.json, then mark it complete.
+
+    Best-effort and conservative: it marks the cache complete ONLY if every
+    component named in model_index.json is present, so a partial dir is left
+    UNMARKED (valid() stays False) and the injection path keeps working. Returns
+    True only when the monolithic cache is fully written + marked."""
+    import os
+    import json
+    try:
+        from codai.models import pipeline_cache as _pcache
+        import diffusers
+        import transformers
+    except Exception as e:
+        print(f"  [pipeline-cache] finalize unavailable: {e}")
+        return False
+    cdir = _pcache.path(model_name, cfg)
+    try:
+        idx = diffusers.DiffusionPipeline.load_config(model_name)
+    except Exception as e:
+        print(f"  [pipeline-cache] finalize: can't read model_index ({e})")
+        return False
+    comps = {k: v for k, v in idx.items()
+             if not k.startswith('_') and isinstance(v, (list, tuple))
+             and len(v) >= 2 and v[1]}
+    for name, spec in comps.items():
+        sub = os.path.join(cdir, name)
+        # Already cached? (heavy component subdir, or a light one from a prior run)
+        if (os.path.isfile(os.path.join(sub, 'config.json'))
+                or os.path.isfile(os.path.join(sub, 'tokenizer_config.json'))
+                or os.path.isfile(os.path.join(sub, 'scheduler_config.json'))):
+            continue
+        lib, cls_name = spec[0], spec[1]
+        mod = transformers if lib == 'transformers' else diffusers
+        cls = getattr(mod, cls_name, None)
+        if cls is None or not hasattr(cls, 'from_pretrained'):
+            print(f"  [pipeline-cache] finalize: no loadable class for '{name}' "
+                  f"({spec}) — leaving monolithic cache incomplete")
+            return False
+        try:
+            print(f"  [pipeline-cache] finalize: caching light component '{name}'")
+            try:
+                obj = cls.from_pretrained(model_name, subfolder=name, torch_dtype=dtype)
+            except TypeError:
+                # tokenizer / scheduler don't accept torch_dtype
+                obj = cls.from_pretrained(model_name, subfolder=name)
+            obj.save_pretrained(sub)
+            del obj
+        except Exception as e:
+            print(f"  [pipeline-cache] finalize: couldn't cache '{name}' ({e}) — "
+                  f"leaving monolithic cache incomplete")
+            return False
+    try:
+        with open(os.path.join(cdir, 'model_index.json'), 'w') as f:
+            json.dump(idx, f)
+    except Exception as e:
+        print(f"  [pipeline-cache] finalize: couldn't write model_index.json ({e})")
+        return False
+    ok = _pcache.mark_monolithic_complete(model_name, cfg)
+    if ok:
+        print(f"  [pipeline-cache] monolithic cache complete → next load uses "
+              f"device_map from cache (no injection)")
+    return ok
+
+
 def build_from_pretrained_kwargs(
     cfg: Optional[Dict[str, Any]],
     *,

codai/models/pipeline_cache.py

@@ -133,6 +133,29 @@ def valid(p: str) -> bool:
 # 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 mark_monolithic_complete(model_name: str, model_cfg: Optional[dict]) -> bool:
+    """Write the monolithic completion marker for a cache dir that was populated
+    component-by-component (heavy components via save_component, light components +
+    model_index.json by the finalizer). After this, valid() returns True and the
+    dir loads as a normal pipeline via from_pretrained(dir, device_map=…) — no
+    injection — so device_map can place every component and big-clip generation
+    fits. Only call once model_index.json and every component subdir are present."""
+    try:
+        p = path(model_name, model_cfg)
+        if not os.path.isfile(os.path.join(p, "model_index.json")):
+            return False
+        with open(_marker(p), "w") as f:
+            json.dump({
+                "version": _CACHE_VERSION, "complete": True,
+                "model": model_name, "saved_at": time.time(),
+                "signature": _signature(model_name, model_cfg),
+                "built": "incremental",
+            }, f)
+        return True
+    except Exception:
+        return False
+
+
 def _component_marker(cdir: str) -> str:
     return os.path.join(cdir, ".coderai_component.json")